{-# LANGUAGE BangPatterns, NamedFieldPuns, RecordWildCards,
             ScopedTypeVariables #-}

module Data.SearchEngine.Autosuggest (

    -- * Query auto-completion \/ auto-suggestion
    queryAutosuggest,
    ResultsFilter(..),

    queryAutosuggestPredicate,
    queryAutosuggestMatchingDocuments

  ) where

import Data.SearchEngine.Types
import Data.SearchEngine.Query (ResultsFilter(..))
import qualified Data.SearchEngine.Query       as Query
import qualified Data.SearchEngine.SearchIndex as SI
import qualified Data.SearchEngine.DocIdSet    as DocIdSet
import qualified Data.SearchEngine.DocTermIds  as DocTermIds
import qualified Data.SearchEngine.BM25F       as BM25F

import Data.Ix
import Data.Ord
import Data.List
import Data.Maybe
import qualified Data.Map as Map
import qualified Data.IntSet as IntSet
import qualified Data.Vector.Unboxed as Vec


-- | Execute an \"auto-suggest\" query. This is where one of the search terms
-- is an incomplete prefix and we are looking for possible completions of that
-- search term, and result documents to go with the possible completions.
--
-- An auto-suggest query only gives useful results when the 'SearchEngine' is
-- configured to use a non-term feature score. That is, when we can give
-- documents an importance score independent of what terms we are looking for.
-- This is because an auto-suggest query is backwards from a normal query: we
-- are asking for relevant terms occurring in important or popular documents
-- so we need some notion of important or popular. Without this we would just
-- be ranking based on term frequency which while it makes sense for normal
-- \"forward\" queries is pretty meaningless for auto-suggest \"reverse\"
-- queries. Indeed for single-term auto-suggest queries the ranking function
-- we use will assign 0 for all documents and completions if there is no 
-- non-term feature scores.
--
queryAutosuggest :: (Ix field, Bounded field, Ix feature, Bounded feature) =>
                    SearchEngine doc key field feature ->
                    ResultsFilter key ->
                    [Term] -> Term -> ([(Term, Float)], [(key, Float)])
queryAutosuggest :: forall field feature doc key.
(Ix field, Bounded field, Ix feature, Bounded feature) =>
SearchEngine doc key field feature
-> ResultsFilter key
-> [Term]
-> Term
-> ([(Term, Float)], [(key, Float)])
queryAutosuggest SearchEngine doc key field feature
se ResultsFilter key
resultsFilter [Term]
precedingTerms Term
partialTerm =

     ([(TermId, Float)], [(DocId, Float)])
-> ([(Term, Float)], [(key, Float)])
forall {v}.
([(TermId, v)], [(DocId, v)]) -> ([(Term, v)], [(key, v)])
step_external
   (([(TermId, Float)], [(DocId, Float)])
 -> ([(Term, Float)], [(key, Float)]))
-> (AutosuggestQuery -> ([(TermId, Float)], [(DocId, Float)]))
-> AutosuggestQuery
-> ([(Term, Float)], [(key, Float)])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([(TermId, Float)], [(DocId, Float)])
-> ([(TermId, Float)], [(DocId, Float)])
forall {a} {b}.
([(a, Float)], [(b, Float)]) -> ([(a, Float)], [(b, Float)])
step_rank
   (([(TermId, Float)], [(DocId, Float)])
 -> ([(TermId, Float)], [(DocId, Float)]))
-> (AutosuggestQuery -> ([(TermId, Float)], [(DocId, Float)]))
-> AutosuggestQuery
-> ([(TermId, Float)], [(DocId, Float)])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Map DocId (Float, Map TermId Float), [(TermId, DocIdSet)],
 [(TermId, Float)])
-> ([(TermId, Float)], [(DocId, Float)])
step_scoreDs
   ((Map DocId (Float, Map TermId Float), [(TermId, DocIdSet)],
  [(TermId, Float)])
 -> ([(TermId, Float)], [(DocId, Float)]))
-> (AutosuggestQuery
    -> (Map DocId (Float, Map TermId Float), [(TermId, DocIdSet)],
        [(TermId, Float)]))
-> AutosuggestQuery
-> ([(TermId, Float)], [(DocId, Float)])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Map DocId (Float, Map TermId Float), [(TermId, DocIdSet)])
-> (Map DocId (Float, Map TermId Float), [(TermId, DocIdSet)],
    [(TermId, Float)])
step_scoreTs
   ((Map DocId (Float, Map TermId Float), [(TermId, DocIdSet)])
 -> (Map DocId (Float, Map TermId Float), [(TermId, DocIdSet)],
     [(TermId, Float)]))
-> (AutosuggestQuery
    -> (Map DocId (Float, Map TermId Float), [(TermId, DocIdSet)]))
-> AutosuggestQuery
-> (Map DocId (Float, Map TermId Float), [(TermId, DocIdSet)],
    [(TermId, Float)])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([TermId], [(DocId, (key, DocTermIds field, DocFeatVals feature))],
 [(TermId, DocIdSet)], [TermId])
-> (Map DocId (Float, Map TermId Float), [(TermId, DocIdSet)])
forall {b}.
([TermId], [(DocId, (key, DocTermIds field, DocFeatVals feature))],
 b, [TermId])
-> (Map DocId (Float, Map TermId Float), b)
step_cache
   (([TermId],
  [(DocId, (key, DocTermIds field, DocFeatVals feature))],
  [(TermId, DocIdSet)], [TermId])
 -> (Map DocId (Float, Map TermId Float), [(TermId, DocIdSet)]))
-> (AutosuggestQuery
    -> ([TermId],
        [(DocId, (key, DocTermIds field, DocFeatVals feature))],
        [(TermId, DocIdSet)], [TermId]))
-> AutosuggestQuery
-> (Map DocId (Float, Map TermId Float), [(TermId, DocIdSet)])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([TermId], [(DocId, (key, DocTermIds field, DocFeatVals feature))],
 [(TermId, DocIdSet)], [TermId])
-> ([TermId],
    [(DocId, (key, DocTermIds field, DocFeatVals feature))],
    [(TermId, DocIdSet)], [TermId])
forall {a} {a} {a} {a}.
([a], [a], [a], [a]) -> ([a], [a], [a], [a])
step_postfilterlimit
   (([TermId],
  [(DocId, (key, DocTermIds field, DocFeatVals feature))],
  [(TermId, DocIdSet)], [TermId])
 -> ([TermId],
     [(DocId, (key, DocTermIds field, DocFeatVals feature))],
     [(TermId, DocIdSet)], [TermId]))
-> (AutosuggestQuery
    -> ([TermId],
        [(DocId, (key, DocTermIds field, DocFeatVals feature))],
        [(TermId, DocIdSet)], [TermId]))
-> AutosuggestQuery
-> ([TermId],
    [(DocId, (key, DocTermIds field, DocFeatVals feature))],
    [(TermId, DocIdSet)], [TermId])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([TermId], DocIdSet, [(TermId, DocIdSet)], [TermId])
-> ([TermId],
    [(DocId, (key, DocTermIds field, DocFeatVals feature))],
    [(TermId, DocIdSet)], [TermId])
forall {a} {c} {d}.
(a, DocIdSet, c, d)
-> (a, [(DocId, (key, DocTermIds field, DocFeatVals feature))], c,
    d)
step_filter
   (([TermId], DocIdSet, [(TermId, DocIdSet)], [TermId])
 -> ([TermId],
     [(DocId, (key, DocTermIds field, DocFeatVals feature))],
     [(TermId, DocIdSet)], [TermId]))
-> (AutosuggestQuery
    -> ([TermId], DocIdSet, [(TermId, DocIdSet)], [TermId]))
-> AutosuggestQuery
-> ([TermId],
    [(DocId, (key, DocTermIds field, DocFeatVals feature))],
    [(TermId, DocIdSet)], [TermId])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([TermId], DocIdSet, [(TermId, DocIdSet)], [TermId])
-> ([TermId], DocIdSet, [(TermId, DocIdSet)], [TermId])
forall {a} {a} {a}.
([a], DocIdSet, [a], [a]) -> ([a], DocIdSet, [a], [a])
step_prefilterlimit
   (([TermId], DocIdSet, [(TermId, DocIdSet)], [TermId])
 -> ([TermId], DocIdSet, [(TermId, DocIdSet)], [TermId]))
-> (AutosuggestQuery
    -> ([TermId], DocIdSet, [(TermId, DocIdSet)], [TermId]))
-> AutosuggestQuery
-> ([TermId], DocIdSet, [(TermId, DocIdSet)], [TermId])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AutosuggestQuery
-> ([TermId], DocIdSet, [(TermId, DocIdSet)], [TermId])
step_process
   (AutosuggestQuery -> ([(Term, Float)], [(key, Float)]))
-> AutosuggestQuery -> ([(Term, Float)], [(key, Float)])
forall a b. (a -> b) -> a -> b
$ [Term] -> Term -> AutosuggestQuery
step_prep
       [Term]
precedingTerms Term
partialTerm

  where
    -- Construct the auto-suggest query from the query terms
    step_prep :: [Term] -> Term -> AutosuggestQuery
step_prep [Term]
pre_ts Term
t = SearchEngine doc key field feature
-> [Term] -> Term -> AutosuggestQuery
forall field doc key feature.
(Ix field, Bounded field) =>
SearchEngine doc key field feature
-> [Term] -> Term -> AutosuggestQuery
mkAutosuggestQuery SearchEngine doc key field feature
se [Term]
pre_ts Term
t

    -- Find the appropriate subset of ts and ds
    -- and an intermediate result that will be useful later:
    -- { (t, ds ∩ ds_t) | t ∈ ts, ds ∩ ds_t ≠ ∅ }
    step_process :: AutosuggestQuery
-> ([TermId], DocIdSet, [(TermId, DocIdSet)], [TermId])
step_process (Map TermId DocIdSet
ts, Maybe DocIdSet
ds, [TermId]
pre_ts) = ([TermId]
ts', DocIdSet
ds', [(TermId, DocIdSet)]
tdss', [TermId]
pre_ts)
      where
        ([(TermId, DocIdSet)]
tdss', [TermId]
ts', DocIdSet
ds') = SearchEngine doc key field feature
-> AutosuggestQuery -> ([(TermId, DocIdSet)], [TermId], DocIdSet)
forall doc key field feature.
SearchEngine doc key field feature
-> AutosuggestQuery -> ([(TermId, DocIdSet)], [TermId], DocIdSet)
processAutosuggestQuery SearchEngine doc key field feature
se (Map TermId DocIdSet
ts, Maybe DocIdSet
ds, [TermId]
pre_ts)

    -- If the number of docs results is huge then we may not want to bother
    -- and just return no results. Even the filtering of a huge number of
    -- docs can be expensive.
    step_prefilterlimit :: ([a], DocIdSet, [a], [a]) -> ([a], DocIdSet, [a], [a])
step_prefilterlimit args :: ([a], DocIdSet, [a], [a])
args@([a]
_, DocIdSet
ds, [a]
_, [a]
_)
      | SearchEngine doc key field feature -> DocIdSet -> Bool
forall doc key field feature.
SearchEngine doc key field feature -> DocIdSet -> Bool
withinPrefilterLimit SearchEngine doc key field feature
se DocIdSet
ds = ([a], DocIdSet, [a], [a])
args
      | Bool
otherwise                  = ([], DocIdSet
DocIdSet.empty, [], [])

    -- Filter ds to those that are visible for this query
    -- and at the same time, do the docid -> docinfo lookup
    -- (needed at this step anyway to do the filter)
    step_filter :: (a, DocIdSet, c, d)
-> (a, [(DocId, (key, DocTermIds field, DocFeatVals feature))], c,
    d)
step_filter (a
ts, DocIdSet
ds, c
tdss, d
pre_ts) = (a
ts, [(DocId, (key, DocTermIds field, DocFeatVals feature))]
ds_info, c
tdss, d
pre_ts)
      where
        ds_info :: [(DocId, (key, DocTermIds field, DocFeatVals feature))]
ds_info = SearchEngine doc key field feature
-> ResultsFilter key
-> DocIdSet
-> [(DocId, (key, DocTermIds field, DocFeatVals feature))]
forall doc key field feature.
SearchEngine doc key field feature
-> ResultsFilter key
-> DocIdSet
-> [(DocId, (key, DocTermIds field, DocFeatVals feature))]
filterAutosuggestQuery SearchEngine doc key field feature
se ResultsFilter key
resultsFilter DocIdSet
ds

    -- If the number of docs results is huge then we may not want to bother
    -- and just return no results. Scoring a large number of docs is expensive.
    step_postfilterlimit :: ([a], [a], [a], [a]) -> ([a], [a], [a], [a])
step_postfilterlimit args :: ([a], [a], [a], [a])
args@([a]
_, [a]
ds_info, [a]
_, [a]
_)
      | SearchEngine doc key field feature -> [a] -> Bool
forall doc key field feature a.
SearchEngine doc key field feature -> [a] -> Bool
withinPostfilterLimit SearchEngine doc key field feature
se [a]
ds_info = ([a], [a], [a], [a])
args
      | Bool
otherwise                        = ([], [], [], [])

    -- For all ds, calculate and cache a couple bits of info needed
    -- later for scoring completion terms and doc results
    step_cache :: ([TermId], [(DocId, (key, DocTermIds field, DocFeatVals feature))],
 b, [TermId])
-> (Map DocId (Float, Map TermId Float), b)
step_cache ([TermId]
ts, [(DocId, (key, DocTermIds field, DocFeatVals feature))]
ds_info, b
tdss, [TermId]
pre_ts) = (Map DocId (Float, Map TermId Float)
ds_info', b
tdss)
      where
        ds_info' :: Map DocId (Float, Map TermId Float)
ds_info' = SearchEngine doc key field feature
-> [TermId]
-> [(DocId, (key, DocTermIds field, DocFeatVals feature))]
-> [TermId]
-> Map DocId (Float, Map TermId Float)
forall field feature doc key.
(Ix field, Bounded field, Ix feature, Bounded feature) =>
SearchEngine doc key field feature
-> [TermId]
-> [(DocId, (key, DocTermIds field, DocFeatVals feature))]
-> [TermId]
-> Map DocId (Float, Map TermId Float)
cacheDocScoringInfo SearchEngine doc key field feature
se [TermId]
ts [(DocId, (key, DocTermIds field, DocFeatVals feature))]
ds_info [TermId]
pre_ts

    -- Score the completion terms
    step_scoreTs :: (Map DocId (Float, Map TermId Float), [(TermId, DocIdSet)])
-> (Map DocId (Float, Map TermId Float), [(TermId, DocIdSet)],
    [(TermId, Float)])
step_scoreTs (Map DocId (Float, Map TermId Float)
ds_info, [(TermId, DocIdSet)]
tdss) = (Map DocId (Float, Map TermId Float)
ds_info, [(TermId, DocIdSet)]
tdss, [(TermId, Float)]
ts_scored)
      where
        ts_scored :: [(TermId, Float)]
ts_scored = [(TermId, DocIdSet)]
-> Map DocId (Float, Map TermId Float) -> [(TermId, Float)]
scoreAutosuggestQueryCompletions [(TermId, DocIdSet)]
tdss Map DocId (Float, Map TermId Float)
ds_info

    -- Score the doc results (making use of the completion scores)
    step_scoreDs :: (Map DocId (Float, Map TermId Float), [(TermId, DocIdSet)],
 [(TermId, Float)])
-> ([(TermId, Float)], [(DocId, Float)])
step_scoreDs (Map DocId (Float, Map TermId Float)
ds_info, [(TermId, DocIdSet)]
tdss, [(TermId, Float)]
ts_scored) = ([(TermId, Float)]
ts_scored, [(DocId, Float)]
ds_scored)
      where
        ds_scored :: [(DocId, Float)]
ds_scored = [(TermId, DocIdSet)]
-> Map DocId (Float, Map TermId Float)
-> [(TermId, Float)]
-> [(DocId, Float)]
scoreAutosuggestQueryResults [(TermId, DocIdSet)]
tdss Map DocId (Float, Map TermId Float)
ds_info [(TermId, Float)]
ts_scored

    -- Rank the completions and results based on their scores
    step_rank :: ([(a, Float)], [(b, Float)]) -> ([(a, Float)], [(b, Float)])
step_rank = ([(a, Float)], [(b, Float)]) -> ([(a, Float)], [(b, Float)])
forall av bv a b.
(Ord av, Ord bv) =>
([(a, av)], [(b, bv)]) -> ([(a, av)], [(b, bv)])
sortResults

    -- Convert from internal Ids into external forms: Term and doc key
    step_external :: ([(TermId, v)], [(DocId, v)]) -> ([(Term, v)], [(key, v)])
step_external = SearchEngine doc key field feature
-> ([(TermId, v)], [(DocId, v)]) -> ([(Term, v)], [(key, v)])
forall doc key field feature v.
SearchEngine doc key field feature
-> ([(TermId, v)], [(DocId, v)]) -> ([(Term, v)], [(key, v)])
convertIdsToExternal SearchEngine doc key field feature
se


-- | Given an incomplete prefix query, find the set of documents that match
-- possible completions of that query.  This should be less computationally
-- expensive than 'queryAutosuggest' as it does not do any ranking of documents.
-- However, it does not apply the pre-filter or post-filter limits, and the list
-- may be large when the query terms occur in many documents.  The order of
-- returned keys is unspecified.
queryAutosuggestMatchingDocuments :: (Ix field, Bounded field, Ord key) =>
                                     SearchEngine doc key field feature ->
                                     [Term] -> Term -> [key]
queryAutosuggestMatchingDocuments :: forall field key doc feature.
(Ix field, Bounded field, Ord key) =>
SearchEngine doc key field feature -> [Term] -> Term -> [key]
queryAutosuggestMatchingDocuments se :: SearchEngine doc key field feature
se@SearchEngine{SearchIndex key field feature
searchIndex :: SearchIndex key field feature
searchIndex :: forall doc key field feature.
SearchEngine doc key field feature -> SearchIndex key field feature
searchIndex} [Term]
precedingTerms Term
partialTerm =
    let ([(TermId, DocIdSet)]
_, [TermId]
_, DocIdSet
ds) = SearchEngine doc key field feature
-> AutosuggestQuery -> ([(TermId, DocIdSet)], [TermId], DocIdSet)
forall doc key field feature.
SearchEngine doc key field feature
-> AutosuggestQuery -> ([(TermId, DocIdSet)], [TermId], DocIdSet)
processAutosuggestQuery SearchEngine doc key field feature
se (SearchEngine doc key field feature
-> [Term] -> Term -> AutosuggestQuery
forall field doc key feature.
(Ix field, Bounded field) =>
SearchEngine doc key field feature
-> [Term] -> Term -> AutosuggestQuery
mkAutosuggestQuery SearchEngine doc key field feature
se [Term]
precedingTerms Term
partialTerm)
    in (DocId -> key) -> [DocId] -> [key]
forall a b. (a -> b) -> [a] -> [b]
map (SearchIndex key field feature -> DocId -> key
forall key field feature.
SearchIndex key field feature -> DocId -> key
SI.getDocKey SearchIndex key field feature
searchIndex) (DocIdSet -> [DocId]
DocIdSet.toList DocIdSet
ds)

-- | Given an incomplete prefix query, return a predicate that indicates whether
-- a key is in the set of documents that match possible completions of that
-- query.  This is equivalent to calling 'queryAutosuggestMatchingDocuments' and
-- testing whether the key is in the list, but should be more efficient.
--
-- This does not apply the pre-filter or post-filter limits.
queryAutosuggestPredicate :: (Ix field, Bounded field, Ord key) =>
                             SearchEngine doc key field feature ->
                             [Term] -> Term -> (key -> Bool)
queryAutosuggestPredicate :: forall field key doc feature.
(Ix field, Bounded field, Ord key) =>
SearchEngine doc key field feature -> [Term] -> Term -> key -> Bool
queryAutosuggestPredicate se :: SearchEngine doc key field feature
se@SearchEngine{SearchIndex key field feature
searchIndex :: forall doc key field feature.
SearchEngine doc key field feature -> SearchIndex key field feature
searchIndex :: SearchIndex key field feature
searchIndex} [Term]
precedingTerms Term
partialTerm =
    let ([(TermId, DocIdSet)]
_, [TermId]
_, DocIdSet
ds) = SearchEngine doc key field feature
-> AutosuggestQuery -> ([(TermId, DocIdSet)], [TermId], DocIdSet)
forall doc key field feature.
SearchEngine doc key field feature
-> AutosuggestQuery -> ([(TermId, DocIdSet)], [TermId], DocIdSet)
processAutosuggestQuery SearchEngine doc key field feature
se (SearchEngine doc key field feature
-> [Term] -> Term -> AutosuggestQuery
forall field doc key feature.
(Ix field, Bounded field) =>
SearchEngine doc key field feature
-> [Term] -> Term -> AutosuggestQuery
mkAutosuggestQuery SearchEngine doc key field feature
se [Term]
precedingTerms Term
partialTerm)
    in (\ key
key -> Bool -> (DocId -> Bool) -> Maybe DocId -> Bool
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Bool
False ((DocId -> DocIdSet -> Bool) -> DocIdSet -> DocId -> Bool
forall a b c. (a -> b -> c) -> b -> a -> c
flip DocId -> DocIdSet -> Bool
DocIdSet.member DocIdSet
ds) (SearchIndex key field feature -> key -> Maybe DocId
forall key field feature.
Ord key =>
SearchIndex key field feature -> key -> Maybe DocId
SI.lookupDocKeyDocId SearchIndex key field feature
searchIndex key
key))


-- We apply hard limits both before and after filtering.
-- The post-filter limit is to avoid scoring 1000s of documents.
-- The pre-filter limit is to avoid filtering 1000s of docs (which in some
-- apps may be expensive itself)

withinPrefilterLimit :: SearchEngine doc key field feature ->
                        DocIdSet -> Bool
withinPrefilterLimit :: forall doc key field feature.
SearchEngine doc key field feature -> DocIdSet -> Bool
withinPrefilterLimit SearchEngine{SearchRankParameters field feature
searchRankParams :: SearchRankParameters field feature
searchRankParams :: forall doc key field feature.
SearchEngine doc key field feature
-> SearchRankParameters field feature
searchRankParams} DocIdSet
ds =
    DocIdSet -> Int
DocIdSet.size DocIdSet
ds Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= SearchRankParameters field feature -> Int
forall field feature. SearchRankParameters field feature -> Int
paramAutosuggestPrefilterLimit SearchRankParameters field feature
searchRankParams

withinPostfilterLimit :: SearchEngine doc key field feature ->
                         [a] -> Bool
withinPostfilterLimit :: forall doc key field feature a.
SearchEngine doc key field feature -> [a] -> Bool
withinPostfilterLimit SearchEngine{SearchRankParameters field feature
searchRankParams :: forall doc key field feature.
SearchEngine doc key field feature
-> SearchRankParameters field feature
searchRankParams :: SearchRankParameters field feature
searchRankParams} [a]
ds_info =
    [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
ds_info Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= SearchRankParameters field feature -> Int
forall field feature. SearchRankParameters field feature -> Int
paramAutosuggestPostfilterLimit SearchRankParameters field feature
searchRankParams


sortResults :: (Ord av, Ord bv) => ([(a,av)], [(b,bv)]) -> ([(a,av)], [(b,bv)])
sortResults :: forall av bv a b.
(Ord av, Ord bv) =>
([(a, av)], [(b, bv)]) -> ([(a, av)], [(b, bv)])
sortResults ([(a, av)]
xs, [(b, bv)]
ys) =
    ( [(a, av)] -> [(a, av)]
forall v x. Ord v => [(x, v)] -> [(x, v)]
sortBySndDescending [(a, av)]
xs
    , [(b, bv)] -> [(b, bv)]
forall v x. Ord v => [(x, v)] -> [(x, v)]
sortBySndDescending [(b, bv)]
ys )
  where
    sortBySndDescending :: Ord v => [(x,v)] -> [(x,v)]
    sortBySndDescending :: forall v x. Ord v => [(x, v)] -> [(x, v)]
sortBySndDescending = ((x, v) -> (x, v) -> Ordering) -> [(x, v)] -> [(x, v)]
forall a. (a -> a -> Ordering) -> [a] -> [a]
sortBy (((x, v) -> (x, v) -> Ordering) -> (x, v) -> (x, v) -> Ordering
forall a b c. (a -> b -> c) -> b -> a -> c
flip (((x, v) -> v) -> (x, v) -> (x, v) -> Ordering
forall a b. Ord a => (b -> a) -> b -> b -> Ordering
comparing (x, v) -> v
forall a b. (a, b) -> b
snd))

convertIdsToExternal :: SearchEngine doc key field feature ->
                        ([(TermId, v)], [(DocId, v)]) -> ([(Term, v)], [(key, v)])
convertIdsToExternal :: forall doc key field feature v.
SearchEngine doc key field feature
-> ([(TermId, v)], [(DocId, v)]) -> ([(Term, v)], [(key, v)])
convertIdsToExternal SearchEngine{SearchIndex key field feature
searchIndex :: forall doc key field feature.
SearchEngine doc key field feature -> SearchIndex key field feature
searchIndex :: SearchIndex key field feature
searchIndex} ([(TermId, v)]
termids, [(DocId, v)]
docids) =
    ( [ (SearchIndex key field feature -> TermId -> Term
forall key field feature.
SearchIndex key field feature -> TermId -> Term
SI.getTerm   SearchIndex key field feature
searchIndex TermId
termid, v
s) | (TermId
termid, v
s) <- [(TermId, v)]
termids ]
    , [ (SearchIndex key field feature -> DocId -> key
forall key field feature.
SearchIndex key field feature -> DocId -> key
SI.getDocKey SearchIndex key field feature
searchIndex DocId
docid,  v
s) | (DocId
docid,  v
s) <- [(DocId, v)]
docids  ]
    )


-- From Bast and Weber:
--
--   An autocompletion query is a pair (T, D), where T is a range of terms
--   (all possible completions of the last term which the user has started
--   typing) and D is a set of documents (the hits for the preceding part of
--   the query).
--
-- We augment this with the preceding terms because we will need these to
-- score the set of documents D.
--
-- Note that the set D will be the entire collection in the case that the
-- preceding part of the query is empty. For efficiency we represent that
-- case specially with Maybe.

type AutosuggestQuery = (Map.Map TermId DocIdSet, Maybe DocIdSet, [TermId])

mkAutosuggestQuery :: (Ix field, Bounded field) =>
                      SearchEngine doc key field feature ->
                      [Term] -> Term -> AutosuggestQuery
mkAutosuggestQuery :: forall field doc key feature.
(Ix field, Bounded field) =>
SearchEngine doc key field feature
-> [Term] -> Term -> AutosuggestQuery
mkAutosuggestQuery se :: SearchEngine doc key field feature
se@SearchEngine{ SearchIndex key field feature
searchIndex :: forall doc key field feature.
SearchEngine doc key field feature -> SearchIndex key field feature
searchIndex :: SearchIndex key field feature
searchIndex }
                   [Term]
precedingTerms Term
partialTerm =
    (Map TermId DocIdSet
completionTerms, Maybe DocIdSet
precedingDocHits, [TermId]
precedingTerms')
  where
    completionTerms :: Map TermId DocIdSet
completionTerms =
      [Map TermId DocIdSet] -> Map TermId DocIdSet
forall (f :: * -> *) k a.
(Foldable f, Ord k) =>
f (Map k a) -> Map k a
Map.unions
        [ [(TermId, DocIdSet)] -> Map TermId DocIdSet
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList (SearchIndex key field feature -> Term -> [(TermId, DocIdSet)]
forall key field feature.
SearchIndex key field feature -> Term -> [(TermId, DocIdSet)]
SI.lookupTermsByPrefix SearchIndex key field feature
searchIndex Term
partialTerm')
        | Term
partialTerm' <- SearchEngine doc key field feature -> Term -> [Term]
forall field doc key feature.
(Ix field, Bounded field) =>
SearchEngine doc key field feature -> Term -> [Term]
Query.expandTransformedQueryTerm SearchEngine doc key field feature
se Term
partialTerm
        ]

    ([TermId]
precedingTerms', Maybe DocIdSet
precedingDocHits)
      | [Term] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Term]
precedingTerms = ([], Maybe DocIdSet
forall a. Maybe a
Nothing)
      | Bool
otherwise           = ([DocIdSet] -> Maybe DocIdSet)
-> ([TermId], [DocIdSet]) -> ([TermId], Maybe DocIdSet)
forall a b. (a -> b) -> ([TermId], a) -> ([TermId], b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [DocIdSet] -> Maybe DocIdSet
carefulUnions
                                   ([Term] -> ([TermId], [DocIdSet])
lookupRawResults [Term]
precedingTerms)

    -- For the preceding terms, we compute the union of the sets of documents in
    -- which they appear.  This means that a query like "Apple Blackberry C"
    -- will look for documents containing "Apple" or "Blackberry", then later
    -- intersect that set with documents containing completions of "C".
    --
    -- In general we want to use union rather than intersection here, because
    -- the preceding terms might contain some useful and some missing terms, and
    -- if we took the intersection we would end up with no results; thus we rely
    -- on scoring to rank the best matches highest.
    --
    -- However, this leads to an issue: if some of the terms are extremely
    -- common, we might end up taking unions of very large document sets, which
    -- is a performance disaster.  We address this by unioning only sets smaller
    -- than the pre-filter limit (but falling back on the whole collection if
    -- all sets are too large).  This means that:
    --
    --  * A query containing a mixture of common and uncommon preceding terms
    --    will be completed/ranked solely based on the uncommon terms.  For
    --    example, "Apple Blackberry C" will be equivalent to "Blackberry C" if
    --    there are many apples.
    --
    --  * A query containing only common preceding terms will be
    --    completed/ranked as if only the final term was present.  For example,
    --    "Apple Blackberry C" will be equivalent to "C" if there are many
    --    apples and blackberries.
    --
    carefulUnions :: [DocIdSet] -> Maybe DocIdSet
    carefulUnions :: [DocIdSet] -> Maybe DocIdSet
carefulUnions [DocIdSet]
dss
      | [DocIdSet] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [DocIdSet]
dss  = DocIdSet -> Maybe DocIdSet
forall a. a -> Maybe a
Just DocIdSet
DocIdSet.empty
      | [DocIdSet] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [DocIdSet]
dss' = Maybe DocIdSet
forall a. Maybe a
Nothing
      | Bool
otherwise = DocIdSet -> Maybe DocIdSet
forall a. a -> Maybe a
Just ([DocIdSet] -> DocIdSet
DocIdSet.unions [DocIdSet]
dss')
      where
        dss' :: [DocIdSet]
dss' = (DocIdSet -> Bool) -> [DocIdSet] -> [DocIdSet]
forall a. (a -> Bool) -> [a] -> [a]
filter (SearchEngine doc key field feature -> DocIdSet -> Bool
forall doc key field feature.
SearchEngine doc key field feature -> DocIdSet -> Bool
withinPrefilterLimit SearchEngine doc key field feature
se) [DocIdSet]
dss

    lookupRawResults :: [Term] -> ([TermId], [DocIdSet])
    lookupRawResults :: [Term] -> ([TermId], [DocIdSet])
lookupRawResults [Term]
ts =
      [(TermId, DocIdSet)] -> ([TermId], [DocIdSet])
forall a b. [(a, b)] -> ([a], [b])
unzip ([(TermId, DocIdSet)] -> ([TermId], [DocIdSet]))
-> [(TermId, DocIdSet)] -> ([TermId], [DocIdSet])
forall a b. (a -> b) -> a -> b
$ [Maybe (TermId, DocIdSet)] -> [(TermId, DocIdSet)]
forall a. [Maybe a] -> [a]
catMaybes
        [ SearchIndex key field feature -> Term -> Maybe (TermId, DocIdSet)
forall key field feature.
SearchIndex key field feature -> Term -> Maybe (TermId, DocIdSet)
SI.lookupTerm SearchIndex key field feature
searchIndex Term
t'
        | Term
t  <- [Term]
ts
        , Term
t' <- SearchEngine doc key field feature -> Term -> [Term]
forall field doc key feature.
(Ix field, Bounded field) =>
SearchEngine doc key field feature -> Term -> [Term]
Query.expandTransformedQueryTerm SearchEngine doc key field feature
se Term
t
        ]



-- From Bast and Weber:
--
--   To process the query means to compute the subset T' ⊆ T of terms that
--   occur in at least one document from D, as well as the subset D' ⊆ D of
--   documents that contain at least one of these words.
--
--   The obvious way to use an inverted index to process an autocompletion
--   query (T, D) is to compute, for each t ∈ T, the intersections D ∩ Dt.
--   Then, T' is simply the set of all t for which the intersection was
--   non-empty, and D' is the union of all (non-empty) intersections.
--
-- We will do this but additionally we will return all the non-empty
-- intersections because they will be useful when scoring.

processAutosuggestQuery :: SearchEngine doc key field feature ->
                           AutosuggestQuery ->
                           ([(TermId, DocIdSet)], [TermId], DocIdSet)
processAutosuggestQuery :: forall doc key field feature.
SearchEngine doc key field feature
-> AutosuggestQuery -> ([(TermId, DocIdSet)], [TermId], DocIdSet)
processAutosuggestQuery SearchEngine doc key field feature
se (Map TermId DocIdSet
completionTerms, Maybe DocIdSet
precedingDocHits, [TermId]
_)
  -- Check all the individual document sets are smaller than the pre-filter
  -- limit.  If any are larger, their union must also be too large, so we return
  -- no results now rather than having to compute the union (which may be
  -- expensive) only for it to inevitably hit the limit.
  | (DocIdSet -> Bool) -> [DocIdSet] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (SearchEngine doc key field feature -> DocIdSet -> Bool
forall doc key field feature.
SearchEngine doc key field feature -> DocIdSet -> Bool
withinPrefilterLimit SearchEngine doc key field feature
se) [DocIdSet]
docSets =
    ( [(TermId, DocIdSet)]
completionTermAndDocSets
    , [TermId]
completionTerms'
    , DocIdSet
allTermDocSet
    )
  | Bool
otherwise = ([], [], DocIdSet
DocIdSet.empty)
  where
    -- We look up each candidate completion to find the set of documents
    -- it appears in, and filtering (intersecting) down to just those
    -- appearing in the existing partial query results (if any).
    -- Candidate completions not appearing at all within the existing
    -- partial query results are excluded at this stage.
    --
    -- We have to keep these doc sets for the whole process, so we keep
    -- them as the compact DocIdSet type.
    --
    completionTermAndDocSets :: [(TermId, DocIdSet)]
    completionTermAndDocSets :: [(TermId, DocIdSet)]
completionTermAndDocSets =
      [ (TermId
t, DocIdSet
ds_t')
      | (TermId
t, DocIdSet
ds_t) <- Map TermId DocIdSet -> [(TermId, DocIdSet)]
forall k a. Map k a -> [(k, a)]
Map.toList Map TermId DocIdSet
completionTerms
      , let ds_t' :: DocIdSet
ds_t' = case Maybe DocIdSet
precedingDocHits of
                      Just DocIdSet
ds -> DocIdSet
ds DocIdSet -> DocIdSet -> DocIdSet
`DocIdSet.intersection` DocIdSet
ds_t
                      Maybe DocIdSet
Nothing -> DocIdSet
ds_t
      , Bool -> Bool
not (DocIdSet -> Bool
DocIdSet.null DocIdSet
ds_t')
      ]

    -- The remaining candidate completions
    completionTerms' :: [TermId]
    docSets :: [DocIdSet]
    ([TermId]
completionTerms', [DocIdSet]
docSets) = [(TermId, DocIdSet)] -> ([TermId], [DocIdSet])
forall a b. [(a, b)] -> ([a], [b])
unzip [(TermId, DocIdSet)]
completionTermAndDocSets

    -- The union of all these is this set of documents that form the results.
    allTermDocSet :: DocIdSet
    allTermDocSet :: DocIdSet
allTermDocSet = [DocIdSet] -> DocIdSet
DocIdSet.unions [DocIdSet]
docSets


filterAutosuggestQuery :: SearchEngine doc key field feature ->
                          ResultsFilter key ->
                          DocIdSet ->
                          [(DocId, (key, DocTermIds field, DocFeatVals feature))]
filterAutosuggestQuery :: forall doc key field feature.
SearchEngine doc key field feature
-> ResultsFilter key
-> DocIdSet
-> [(DocId, (key, DocTermIds field, DocFeatVals feature))]
filterAutosuggestQuery SearchEngine{ SearchIndex key field feature
searchIndex :: forall doc key field feature.
SearchEngine doc key field feature -> SearchIndex key field feature
searchIndex :: SearchIndex key field feature
searchIndex } ResultsFilter key
resultsFilter DocIdSet
ds =
    case ResultsFilter key
resultsFilter of
      ResultsFilter key
NoFilter ->
        [ (DocId
docid, (key, DocTermIds field, DocFeatVals feature)
doc)
        | DocId
docid <- DocIdSet -> [DocId]
DocIdSet.toList DocIdSet
ds
        , let doc :: (key, DocTermIds field, DocFeatVals feature)
doc = SearchIndex key field feature
-> DocId -> (key, DocTermIds field, DocFeatVals feature)
forall key field feature.
SearchIndex key field feature
-> DocId -> (key, DocTermIds field, DocFeatVals feature)
SI.lookupDocId SearchIndex key field feature
searchIndex DocId
docid ]

      FilterPredicate key -> Bool
predicate ->
        [ (DocId
docid, (key, DocTermIds field, DocFeatVals feature)
doc)
        | DocId
docid <- DocIdSet -> [DocId]
DocIdSet.toList DocIdSet
ds
        , let doc :: (key, DocTermIds field, DocFeatVals feature)
doc@(key
k,DocTermIds field
_,DocFeatVals feature
_) = SearchIndex key field feature
-> DocId -> (key, DocTermIds field, DocFeatVals feature)
forall key field feature.
SearchIndex key field feature
-> DocId -> (key, DocTermIds field, DocFeatVals feature)
SI.lookupDocId SearchIndex key field feature
searchIndex DocId
docid
        , key -> Bool
predicate key
k ]

      FilterBulkPredicate [key] -> [Bool]
bulkPredicate ->
        [ (DocId
docid, (key, DocTermIds field, DocFeatVals feature)
doc)
        | let docids :: [DocId]
docids = DocIdSet -> [DocId]
DocIdSet.toList DocIdSet
ds
              docinf :: [(key, DocTermIds field, DocFeatVals feature)]
docinf = (DocId -> (key, DocTermIds field, DocFeatVals feature))
-> [DocId] -> [(key, DocTermIds field, DocFeatVals feature)]
forall a b. (a -> b) -> [a] -> [b]
map (SearchIndex key field feature
-> DocId -> (key, DocTermIds field, DocFeatVals feature)
forall key field feature.
SearchIndex key field feature
-> DocId -> (key, DocTermIds field, DocFeatVals feature)
SI.lookupDocId SearchIndex key field feature
searchIndex) [DocId]
docids
              keep :: [Bool]
keep   = [key] -> [Bool]
bulkPredicate [ key
k | (key
k,DocTermIds field
_,DocFeatVals feature
_) <- [(key, DocTermIds field, DocFeatVals feature)]
docinf ]
        , (DocId
docid, (key, DocTermIds field, DocFeatVals feature)
doc, Bool
True) <- [DocId]
-> [(key, DocTermIds field, DocFeatVals feature)]
-> [Bool]
-> [(DocId, (key, DocTermIds field, DocFeatVals feature), Bool)]
forall a b c. [a] -> [b] -> [c] -> [(a, b, c)]
zip3 [DocId]
docids [(key, DocTermIds field, DocFeatVals feature)]
docinf [Bool]
keep ]


-- Scoring
-------------
--
-- From Bast and Weber:
--   In practice, only a selection of items from these lists can and will be
--   presented to the user, and it is of course crucial that the most relevant
--   completions and hits are selected.
--
--   A standard approach for this task in ad-hoc retrieval is to have a
--   precomputed score for each term-in-document pair, and when a query is
--   being processed, to aggregate these scores for each candidate document,
--   and return documents with the highest such aggregated scores.
--
--   Both INV and HYB can be easily adapted to implement any such scoring and
--   aggregation scheme: store by each term-in-document pair its precomputed
--   score, and when intersecting, aggregate the scores. A decision has to be
--   made on how to reconcile scores from different completions within the
--   same document. We suggest the following: when merging the intersections
--   (which gives the set D' according to Definition 1), compute for each
--   document in D' the maximal score achieved for some completion in T'
--   contained in that document, and compute for each completion in T' the
--   maximal score achieved for a hit from D' achieved for this completion.
--
-- So firstly let us explore what this means and then discuss why it does not
-- work for BM25.
--
-- The "precomputed score for each term-in-document pair" refers to the bm25
-- score for this term in this document (and obviously doesn't have to be
-- precomputed, though that'd be faster).
--
-- So the score for a document d ∈ D' is:
--   maximum of score for d ∈ D ∩ Dt, for any t ∈ T'
--
-- While the score for a completion t ∈ T' is:
--   maximum of score for d ∈ D ∩ Dt
--
-- So for documents we say their score is their best score for any of the
-- completion terms they contain. While for completions we say their score
-- is their best score for any of the documents they appear in.
--
-- For a scoring function like BM25 this appears to be not a good method, both
-- in principle and in practice. Consider what terms get high BM25 scores:
-- very rare ones. So this means we're going to score highly documents that
-- contain the least frequent terms, and completions that are themselves very
-- rare. This is silly.
--
-- Another important thing to note is that if we use this scoring method then
-- we are using the BM25 score in a way that makes no sense. The BM25 score
-- for different documents for the /same/ set of terms are comparable. The
-- score for the same for different document with different terms are simply
-- not comparable.
--
-- This also makes sense if you consider what question the BM25 score is
-- answering: "what is the likelihood that this document is relevant given that
-- I judge these terms to be relevant". However an auto-suggest query is
-- different: "what is the likelihood that this term is relevant given the
-- importance/popularity of the documents (and any preceding terms I've judged
-- to be relevant)". They are both conditional likelihood questions but with
-- different premises.
--
-- More generally, term frequency information simply isn't useful for
-- auto-suggest queries. We don't want results that have the most obscure terms
-- nor the most common terms, not even something in-between. Term frequency
-- just doesn't tell us anything unless we've already judged terms to be
-- relevant, and in an auto-suggest query we've not done that yet.
--
-- What we really need is information on the importance/popularity of the
-- documents. We can actually do something with that.
--
-- So, instead we follow a different strategy. We require that we have
-- importance/popularity info for the documents.
--
-- A first approximation would be to rank result documents by their importance
-- and completion terms by the sum of the importance of the documents each
-- term appears in.
--
-- Score for a document d ∈ D'
--   importance score for d
--
-- Score for a completion t ∈ T'
--   sum of importance score for d ∈ D ∩ Dt
--
-- The only problem with this is that just because a term appears in an
-- important document, doesn't mean that term is /about/ that document, or to
-- put it another way, that term may not be relevant for that document. For
-- example common words like "the" likely appear in all important documents
-- but this doesn't really tell us anything because "the" isn't an important
-- keyword.
--
-- So what we want to do is to weight the document importance by the relevance
-- of the keyword to the document. So now if we have an important document and
-- a relevant keyword for that document then we get a high score, but an
-- irrelevant term like "the" would get a very low weighting and so would not
-- contribute much to the score, even for very important documents.
--
-- The intuition is that we will score term completions by taking the
-- document importance weighted by the relevance of that term to that document
-- and summing over all the documents where the term occurs.
--
-- We define document importance (for the set D') to be the BM25F score for
-- the documents with any preceding terms. So this includes the non-term
-- feature score for the importance/popularity, and also takes account of
-- preceding terms if there were any.
--
-- We define term relevance (for terms in documents) to be the BM25F score for
-- that term in that document as a fraction of the total BM25F score for all
-- terms in the document. Thus the relevance of all terms in a document sums
-- to 1.
--
-- Now we can re-weight the document importance by the term relevance:
--
-- Score for a completion t ∈ T'
--   sum (for d ∈ D ∩ Dt) of ( importance for d * relevance for t in d )
--
-- And now for document result scores. We don't want to just stick with the
-- innate document importance. We want to re-weight by the completion term
-- scores:
--
-- Score for a document d ∈ D'
--   sum (for t ∈ T' ∩ d) (importance score for d * score for completion t)
--
-- Clear as mud?

type DocImportance = Float
type TermRelevanceBreakdown = Map.Map TermId Float

-- | Precompute the document importance and the term relevance breakdown for
-- all the documents. This will be used in scoring the term completions
-- and the result documents. They will all be used and some used many
-- times so it's better to compute up-front and share.
--
-- This is actually the expensive bit (which is why we've filtered already).
--
cacheDocScoringInfo :: (Ix field, Bounded field, Ix feature, Bounded feature) =>
                       SearchEngine doc key field feature ->
                       [TermId] ->
                       [(DocId, (key, DocTermIds field, DocFeatVals feature))] ->
                       [TermId] ->
                       Map.Map DocId (DocImportance, TermRelevanceBreakdown)
cacheDocScoringInfo :: forall field feature doc key.
(Ix field, Bounded field, Ix feature, Bounded feature) =>
SearchEngine doc key field feature
-> [TermId]
-> [(DocId, (key, DocTermIds field, DocFeatVals feature))]
-> [TermId]
-> Map DocId (Float, Map TermId Float)
cacheDocScoringInfo SearchEngine doc key field feature
se [TermId]
completionTerms [(DocId, (key, DocTermIds field, DocFeatVals feature))]
allTermDocInfo [TermId]
precedingTerms =
    [(DocId, (Float, Map TermId Float))]
-> Map DocId (Float, Map TermId Float)
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList
      [ (DocId
docid, (Float
docImportance, Map TermId Float
termRelevances))
      | (DocId
docid, (key
_dockey, DocTermIds field
doctermids, DocFeatVals feature
docfeatvals)) <- [(DocId, (key, DocTermIds field, DocFeatVals feature))]
allTermDocInfo
      , let docImportance :: Float
docImportance  = SearchEngine doc key field feature
-> [TermId] -> DocTermIds field -> DocFeatVals feature -> Float
forall field feature doc key.
(Ix field, Bounded field, Ix feature, Bounded feature) =>
SearchEngine doc key field feature
-> [TermId] -> DocTermIds field -> DocFeatVals feature -> Float
Query.relevanceScore SearchEngine doc key field feature
se [TermId]
precedingTerms
                                                  DocTermIds field
doctermids DocFeatVals feature
docfeatvals
            termRelevances :: Map TermId Float
termRelevances = SearchEngine doc key field feature
-> DocTermIds field
-> DocFeatVals feature
-> [TermId]
-> Map TermId Float
forall doc key field feature.
(Ix field, Bounded field, Ix feature, Bounded feature) =>
SearchEngine doc key field feature
-> DocTermIds field
-> DocFeatVals feature
-> [TermId]
-> Map TermId Float
relevanceBreakdown SearchEngine doc key field feature
se DocTermIds field
doctermids DocFeatVals feature
docfeatvals
                                                [TermId]
completionTerms
      ]

-- | Calculate the relevance of each of a given set of terms to the given
-- document.
--
-- We define the \"relevance\" of each term in a document to be its
-- term-in-document score as a fraction of the total of the scores for all
-- terms in the document. Thus the sum of all the relevance values in the
-- document is 1.
--
-- Note: we have to calculate the relevance for all terms in the document
-- but we only keep the relevance value for the terms of interest.
--
relevanceBreakdown :: forall doc key field feature.
                      (Ix field, Bounded field, Ix feature, Bounded feature) =>
                      SearchEngine doc key field feature ->
                      DocTermIds field -> DocFeatVals feature ->
                      [TermId] -> TermRelevanceBreakdown
relevanceBreakdown :: forall doc key field feature.
(Ix field, Bounded field, Ix feature, Bounded feature) =>
SearchEngine doc key field feature
-> DocTermIds field
-> DocFeatVals feature
-> [TermId]
-> Map TermId Float
relevanceBreakdown SearchEngine{ Context TermId field feature
bm25Context :: Context TermId field feature
bm25Context :: forall doc key field feature.
SearchEngine doc key field feature -> Context TermId field feature
bm25Context } DocTermIds field
doctermids DocFeatVals feature
docfeatvals [TermId]
ts =
    let -- We'll calculate the bm25 score for each term in this document
        bm25Doc :: Doc TermId field feature
bm25Doc     = DocTermIds field -> DocFeatVals feature -> Doc TermId field feature
forall field feature.
(Ix field, Bounded field, Ix feature, Bounded feature) =>
DocTermIds field -> DocFeatVals feature -> Doc TermId field feature
Query.indexDocToBM25Doc DocTermIds field
doctermids DocFeatVals feature
docfeatvals

        -- Cache the info that depends only on this doc, not the terms
        termScore   :: (TermId -> (field -> Int) -> Float)
        termScore :: TermId -> (field -> Int) -> Float
termScore   = Context TermId field feature
-> Doc TermId field feature -> TermId -> (field -> Int) -> Float
forall field term feature.
(Ix field, Bounded field) =>
Context term field feature
-> Doc term field feature -> term -> (field -> Int) -> Float
BM25F.scoreTermsBulk Context TermId field feature
bm25Context Doc TermId field feature
bm25Doc

        -- The DocTermIds has the info we need to do bulk scoring, but it's
        -- a sparse representation, so we first convert it to a dense table
        term        :: Int -> TermId
        count       :: Int -> field -> Int
        (!Int
numTerms, Int -> TermId
term, Int -> field -> Int
count) = DocTermIds field -> (Int, Int -> TermId, Int -> field -> Int)
forall field.
(Ix field, Bounded field) =>
DocTermIds field -> (Int, Int -> TermId, Int -> field -> Int)
DocTermIds.denseTable DocTermIds field
doctermids

        -- We generate the vector of scores for all terms, based on looking up
        -- the termid and the per-field counts in the dense table
        termScores  :: Vec.Vector Float
        !termScores :: Vector Float
termScores = Int -> (Int -> Float) -> Vector Float
forall a. Unbox a => Int -> (Int -> a) -> Vector a
Vec.generate Int
numTerms ((Int -> Float) -> Vector Float) -> (Int -> Float) -> Vector Float
forall a b. (a -> b) -> a -> b
$ \Int
i ->
                       TermId -> (field -> Int) -> Float
termScore (Int -> TermId
term Int
i) (\field
f -> Int -> field -> Int
count Int
i field
f)

        -- We keep only the values for the terms we're interested in
        -- and normalise so we get the relevence fraction
        !scoreSum :: Float
scoreSum   = Vector Float -> Float
forall a. (Unbox a, Num a) => Vector a -> a
Vec.sum Vector Float
termScores
        !tset :: IntSet
tset       = [Int] -> IntSet
IntSet.fromList ((TermId -> Int) -> [TermId] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map TermId -> Int
forall a. Enum a => a -> Int
fromEnum [TermId]
ts)
     in [(TermId, Float)] -> Map TermId Float
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList
      ([(TermId, Float)] -> Map TermId Float)
-> (Vector Float -> [(TermId, Float)])
-> Vector Float
-> Map TermId Float
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector (TermId, Float) -> [(TermId, Float)]
forall a. Unbox a => Vector a -> [a]
Vec.toList
      (Vector (TermId, Float) -> [(TermId, Float)])
-> (Vector Float -> Vector (TermId, Float))
-> Vector Float
-> [(TermId, Float)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((TermId, Float) -> (TermId, Float))
-> Vector (TermId, Float) -> Vector (TermId, Float)
forall a b. (Unbox a, Unbox b) => (a -> b) -> Vector a -> Vector b
Vec.map    (\(TermId
t,Float
s) -> (TermId
t, Float
sFloat -> Float -> Float
forall a. Fractional a => a -> a -> a
/Float
scoreSum))
      (Vector (TermId, Float) -> Vector (TermId, Float))
-> (Vector Float -> Vector (TermId, Float))
-> Vector Float
-> Vector (TermId, Float)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((TermId, Float) -> Bool)
-> Vector (TermId, Float) -> Vector (TermId, Float)
forall a. Unbox a => (a -> Bool) -> Vector a -> Vector a
Vec.filter (\(TermId
t,Float
_) -> TermId -> Int
forall a. Enum a => a -> Int
fromEnum TermId
t Int -> IntSet -> Bool
`IntSet.member` IntSet
tset)
      (Vector (TermId, Float) -> Vector (TermId, Float))
-> (Vector Float -> Vector (TermId, Float))
-> Vector Float
-> Vector (TermId, Float)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int -> Float -> (TermId, Float))
-> Vector Float -> Vector (TermId, Float)
forall a b.
(Unbox a, Unbox b) =>
(Int -> a -> b) -> Vector a -> Vector b
Vec.imap   (\Int
i Float
s   -> (Int -> TermId
term Int
i, Float
s))
      (Vector Float -> Map TermId Float)
-> Vector Float -> Map TermId Float
forall a b. (a -> b) -> a -> b
$ Vector Float
termScores


scoreAutosuggestQueryCompletions :: [(TermId, DocIdSet)]
                                 -> Map.Map DocId (Float, Map.Map TermId Float)
                                 -> [(TermId, Float)]
scoreAutosuggestQueryCompletions :: [(TermId, DocIdSet)]
-> Map DocId (Float, Map TermId Float) -> [(TermId, Float)]
scoreAutosuggestQueryCompletions [(TermId, DocIdSet)]
completionTermAndDocSets Map DocId (Float, Map TermId Float)
allTermDocInfo =
    [ (TermId
t, TermId -> DocIdSet -> Float
candidateScore TermId
t DocIdSet
ds_t)
    | (TermId
t, DocIdSet
ds_t) <- [(TermId, DocIdSet)]
completionTermAndDocSets ]
  where
    -- The score for a completion is the sum of the importance of the
    -- documents in which that completion occurs, weighted by the relevance
    -- of the term to each document. For example we can have a very
    -- important document and our completion term is highly relevant to it
    -- or we could have a large number of moderately important documents
    -- that our term is quite relevant to. In either example the completion
    -- term would score highly.
    candidateScore :: TermId -> DocIdSet -> Float
    candidateScore :: TermId -> DocIdSet -> Float
candidateScore TermId
t DocIdSet
ds_t =
      [Float] -> Float
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [ Float
docImportance Float -> Float -> Float
forall a. Num a => a -> a -> a
* Float
termRelevance
          | Just (Float
docImportance, Map TermId Float
termRelevances) <-
               (DocId -> Maybe (Float, Map TermId Float))
-> [DocId] -> [Maybe (Float, Map TermId Float)]
forall a b. (a -> b) -> [a] -> [b]
map (DocId
-> Map DocId (Float, Map TermId Float)
-> Maybe (Float, Map TermId Float)
forall k a. Ord k => k -> Map k a -> Maybe a
`Map.lookup` Map DocId (Float, Map TermId Float)
allTermDocInfo) (DocIdSet -> [DocId]
DocIdSet.toList DocIdSet
ds_t)
          , let termRelevance :: Float
termRelevance = Map TermId Float
termRelevances Map TermId Float -> TermId -> Float
forall k a. Ord k => Map k a -> k -> a
Map.! TermId
t
          ]


scoreAutosuggestQueryResults :: [(TermId, DocIdSet)] ->
                                Map.Map DocId (Float, Map.Map TermId Float) ->
                                [(TermId, Float)] ->
                                [(DocId, Float)]
scoreAutosuggestQueryResults :: [(TermId, DocIdSet)]
-> Map DocId (Float, Map TermId Float)
-> [(TermId, Float)]
-> [(DocId, Float)]
scoreAutosuggestQueryResults [(TermId, DocIdSet)]
completionTermAndDocSets Map DocId (Float, Map TermId Float)
allTermDocInfo
                             [(TermId, Float)]
scoredCandidates =
  Map DocId Float -> [(DocId, Float)]
forall k a. Map k a -> [(k, a)]
Map.toList (Map DocId Float -> [(DocId, Float)])
-> Map DocId Float -> [(DocId, Float)]
forall a b. (a -> b) -> a -> b
$ (Float -> Float -> Float) -> [(DocId, Float)] -> Map DocId Float
forall k a. Ord k => (a -> a -> a) -> [(k, a)] -> Map k a
Map.fromListWith Float -> Float -> Float
forall a. Num a => a -> a -> a
(+)
    [ (DocId
docid, Float
docImportance Float -> Float -> Float
forall a. Num a => a -> a -> a
* Float
score_t)
    | ((TermId
_, DocIdSet
ds_t), (TermId
_, Float
score_t)) <- [(TermId, DocIdSet)]
-> [(TermId, Float)] -> [((TermId, DocIdSet), (TermId, Float))]
forall a b. [a] -> [b] -> [(a, b)]
zip [(TermId, DocIdSet)]
completionTermAndDocSets [(TermId, Float)]
scoredCandidates
    , let docids :: [DocId]
docids  = DocIdSet -> [DocId]
DocIdSet.toList DocIdSet
ds_t
          docinfo :: [Maybe (Float, Map TermId Float)]
docinfo = (DocId -> Maybe (Float, Map TermId Float))
-> [DocId] -> [Maybe (Float, Map TermId Float)]
forall a b. (a -> b) -> [a] -> [b]
map (DocId
-> Map DocId (Float, Map TermId Float)
-> Maybe (Float, Map TermId Float)
forall k a. Ord k => k -> Map k a -> Maybe a
`Map.lookup` Map DocId (Float, Map TermId Float)
allTermDocInfo) [DocId]
docids
    , (DocId
docid, Just (Float
docImportance, Map TermId Float
_)) <- [DocId]
-> [Maybe (Float, Map TermId Float)]
-> [(DocId, Maybe (Float, Map TermId Float))]
forall a b. [a] -> [b] -> [(a, b)]
zip [DocId]
docids [Maybe (Float, Map TermId Float)]
docinfo
    ]