packages feed

full-text-search 0.2.2.2 → 0.2.2.3

raw patch · 26 files changed

+2471/−2445 lines, 26 filesdep ~Cabaldep ~arraydep ~basenew-uploaderPVP ok

version bump matches the API change (PVP)

Dependency ranges changed: Cabal, array, base, bytestring, containers, directory, filepath, mtl, split, tar, text, time, tokenize, transformers, vector

API changes (from Hackage documentation)

Files

− Data/SearchEngine.hs 
@@ -1,44 +0,0 @@-module Data.SearchEngine (--    -- * Basic interface--    -- ** Querying-    Term,-    query,--    -- *** Query auto-completion \/ auto-suggestion-    queryAutosuggest,-    ResultsFilter(..),-    queryAutosuggestPredicate,-    queryAutosuggestMatchingDocuments,--    -- ** Making a search engine instance-    initSearchEngine,-    SearchEngine,-    SearchConfig(..),-    SearchRankParameters(..),-    FeatureFunction(..),--    -- ** Helper type for non-term features-    NoFeatures,-    noFeatures,--    -- ** Managing documents to be searched-    insertDoc,-    insertDocs,-    deleteDoc,--    -- * Explain mode for query result rankings-    queryExplain,-    Explanation(..),-    setRankParams,--    -- * Internal sanity check-    invariant,-  ) where--import Data.SearchEngine.Types-import Data.SearchEngine.Update-import Data.SearchEngine.Query-import Data.SearchEngine.Autosuggest-
− Data/SearchEngine/Autosuggest.hs 
@@ -1,573 +0,0 @@-{-# 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 se resultsFilter precedingTerms partialTerm =--     step_external-   . step_rank-   . step_scoreDs-   . step_scoreTs-   . step_cache-   . step_postfilterlimit-   . step_filter-   . step_prefilterlimit-   . step_process-   $ step_prep-       precedingTerms partialTerm--  where-    -- Construct the auto-suggest query from the query terms-    step_prep pre_ts t = mkAutosuggestQuery se pre_ts 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 (ts, ds, pre_ts) = (ts', ds', tdss', pre_ts)-      where-        (tdss', ts', ds') = processAutosuggestQuery se (ts, ds, 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 args@(_, ds, _, _)-      | withinPrefilterLimit se ds = args-      | otherwise                  = ([], 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 (ts, ds, tdss, pre_ts) = (ts, ds_info, tdss, pre_ts)-      where-        ds_info = filterAutosuggestQuery se resultsFilter 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 args@(_, ds_info, _, _)-      | withinPostfilterLimit se ds_info = args-      | otherwise                        = ([], [], [], [])--    -- For all ds, calculate and cache a couple bits of info needed-    -- later for scoring completion terms and doc results-    step_cache (ts, ds_info, tdss, pre_ts) = (ds_info', tdss)-      where-        ds_info' = cacheDocScoringInfo se ts ds_info pre_ts--    -- Score the completion terms-    step_scoreTs (ds_info, tdss) = (ds_info, tdss, ts_scored)-      where-        ts_scored = scoreAutosuggestQueryCompletions tdss ds_info--    -- Score the doc results (making use of the completion scores)-    step_scoreDs (ds_info, tdss, ts_scored) = (ts_scored, ds_scored)-      where-        ds_scored = scoreAutosuggestQueryResults tdss ds_info ts_scored--    -- Rank the completions and results based on their scores-    step_rank = sortResults--    -- Convert from internal Ids into external forms: Term and doc key-    step_external = convertIdsToExternal 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 se@SearchEngine{searchIndex} precedingTerms partialTerm =-    let (_, _, ds) = processAutosuggestQuery se (mkAutosuggestQuery se precedingTerms partialTerm)-    in map (SI.getDocKey searchIndex) (DocIdSet.toList 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 se@SearchEngine{searchIndex} precedingTerms partialTerm =-    let (_, _, ds) = processAutosuggestQuery se (mkAutosuggestQuery se precedingTerms partialTerm)-    in (\ key -> maybe False (flip DocIdSet.member ds) (SI.lookupDocKeyDocId searchIndex 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 SearchEngine{searchRankParams} ds =-    DocIdSet.size ds <= paramAutosuggestPrefilterLimit searchRankParams--withinPostfilterLimit :: SearchEngine doc key field feature ->-                         [a] -> Bool-withinPostfilterLimit SearchEngine{searchRankParams} ds_info =-    length ds_info <= paramAutosuggestPostfilterLimit searchRankParams---sortResults :: (Ord av, Ord bv) => ([(a,av)], [(b,bv)]) -> ([(a,av)], [(b,bv)])-sortResults (xs, ys) =-    ( sortBySndDescending xs-    , sortBySndDescending ys )-  where-    sortBySndDescending :: Ord v => [(x,v)] -> [(x,v)]-    sortBySndDescending = sortBy (flip (comparing snd))--convertIdsToExternal :: SearchEngine doc key field feature ->-                        ([(TermId, v)], [(DocId, v)]) -> ([(Term, v)], [(key, v)])-convertIdsToExternal SearchEngine{searchIndex} (termids, docids) =-    ( [ (SI.getTerm   searchIndex termid, s) | (termid, s) <- termids ]-    , [ (SI.getDocKey searchIndex docid,  s) | (docid,  s) <- 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 se@SearchEngine{ searchIndex }-                   precedingTerms partialTerm =-    (completionTerms, precedingDocHits, precedingTerms')-  where-    completionTerms =-      Map.unions-        [ Map.fromList (SI.lookupTermsByPrefix searchIndex partialTerm')-        | partialTerm' <- Query.expandTransformedQueryTerm se partialTerm-        ]--    (precedingTerms', precedingDocHits)-      | null precedingTerms = ([], Nothing)-      | otherwise           = fmap carefulUnions-                                   (lookupRawResults 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 dss-      | null dss  = Just DocIdSet.empty-      | null dss' = Nothing-      | otherwise = Just (DocIdSet.unions dss')-      where-        dss' = filter (withinPrefilterLimit se) dss--    lookupRawResults :: [Term] -> ([TermId], [DocIdSet])-    lookupRawResults ts =-      unzip $ catMaybes-        [ SI.lookupTerm searchIndex t'-        | t  <- ts-        , t' <- Query.expandTransformedQueryTerm se 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 se (completionTerms, precedingDocHits, _)-  -- 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.-  | all (withinPrefilterLimit se) docSets =-    ( completionTermAndDocSets-    , completionTerms'-    , allTermDocSet-    )-  | otherwise = ([], [], 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 =-      [ (t, ds_t')-      | (t, ds_t) <- Map.toList completionTerms-      , let ds_t' = case precedingDocHits of-                      Just ds -> ds `DocIdSet.intersection` ds_t-                      Nothing -> ds_t-      , not (DocIdSet.null ds_t')-      ]--    -- The remaining candidate completions-    completionTerms' :: [TermId]-    docSets :: [DocIdSet]-    (completionTerms', docSets) = unzip completionTermAndDocSets--    -- The union of all these is this set of documents that form the results.-    allTermDocSet :: DocIdSet-    allTermDocSet = DocIdSet.unions docSets---filterAutosuggestQuery :: SearchEngine doc key field feature ->-                          ResultsFilter key ->-                          DocIdSet ->-                          [(DocId, (key, DocTermIds field, DocFeatVals feature))]-filterAutosuggestQuery SearchEngine{ searchIndex } resultsFilter ds =-    case resultsFilter of-      NoFilter ->-        [ (docid, doc)-        | docid <- DocIdSet.toList ds-        , let doc = SI.lookupDocId searchIndex docid ]--      FilterPredicate predicate ->-        [ (docid, doc)-        | docid <- DocIdSet.toList ds-        , let doc@(k,_,_) = SI.lookupDocId searchIndex docid-        , predicate k ]--      FilterBulkPredicate bulkPredicate ->-        [ (docid, doc)-        | let docids = DocIdSet.toList ds-              docinf = map (SI.lookupDocId searchIndex) docids-              keep   = bulkPredicate [ k | (k,_,_) <- docinf ]-        , (docid, doc, True) <- zip3 docids docinf 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 se completionTerms allTermDocInfo precedingTerms =-    Map.fromList-      [ (docid, (docImportance, termRelevances))-      | (docid, (_dockey, doctermids, docfeatvals)) <- allTermDocInfo-      , let docImportance  = Query.relevanceScore se precedingTerms-                                                  doctermids docfeatvals-            termRelevances = relevanceBreakdown se doctermids docfeatvals-                                                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 SearchEngine{ bm25Context } doctermids docfeatvals ts =-    let -- We'll calculate the bm25 score for each term in this document-        bm25Doc     = Query.indexDocToBM25Doc doctermids docfeatvals--        -- Cache the info that depends only on this doc, not the terms-        termScore   :: (TermId -> (field -> Int) -> Float)-        termScore   = BM25F.scoreTermsBulk bm25Context 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-        (!numTerms, term, count) = DocTermIds.denseTable 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 = Vec.generate numTerms $ \i ->-                       termScore (term i) (\f -> count i f)--        -- We keep only the values for the terms we're interested in-        -- and normalise so we get the relevence fraction-        !scoreSum   = Vec.sum termScores-        !tset       = IntSet.fromList (map fromEnum ts)-     in Map.fromList-      . Vec.toList-      . Vec.map    (\(t,s) -> (t, s/scoreSum))-      . Vec.filter (\(t,_) -> fromEnum t `IntSet.member` tset)-      . Vec.imap   (\i s   -> (term i, s))-      $ termScores---scoreAutosuggestQueryCompletions :: [(TermId, DocIdSet)]-                                 -> Map.Map DocId (Float, Map.Map TermId Float)-                                 -> [(TermId, Float)]-scoreAutosuggestQueryCompletions completionTermAndDocSets allTermDocInfo =-    [ (t, candidateScore t ds_t)-    | (t, ds_t) <- 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 t ds_t =-      sum [ docImportance * termRelevance-          | Just (docImportance, termRelevances) <--               map (`Map.lookup` allTermDocInfo) (DocIdSet.toList ds_t)-          , let termRelevance = termRelevances Map.! t-          ]---scoreAutosuggestQueryResults :: [(TermId, DocIdSet)] ->-                                Map.Map DocId (Float, Map.Map TermId Float) ->-                                [(TermId, Float)] ->-                                [(DocId, Float)]-scoreAutosuggestQueryResults completionTermAndDocSets allTermDocInfo-                             scoredCandidates =-  Map.toList $ Map.fromListWith (+)-    [ (docid, docImportance * score_t)-    | ((_, ds_t), (_, score_t)) <- zip completionTermAndDocSets scoredCandidates-    , let docids  = DocIdSet.toList ds_t-          docinfo = map (`Map.lookup` allTermDocInfo) docids-    , (docid, Just (docImportance, _)) <- zip docids docinfo-    ]-
− Data/SearchEngine/BM25F.hs 
@@ -1,253 +0,0 @@-{-# LANGUAGE RecordWildCards, BangPatterns, ScopedTypeVariables #-}---- | An implementation of BM25F ranking. See:------ * A quick overview: <http://en.wikipedia.org/wiki/Okapi_BM25>------ * /The Probabilistic Relevance Framework: BM25 and Beyond/---   <http://www.staff.city.ac.uk/~sbrp622/papers/foundations_bm25_review.pdf>------ * /An Introduction to Information Retrieval/---   <http://nlp.stanford.edu/IR-book/pdf/irbookonlinereading.pdf>----module Data.SearchEngine.BM25F (-    -- * The ranking function-    score,-    Context(..),-    FeatureFunction(..),-    Doc(..),-    -- ** Specialised variants-    scoreTermsBulk,--    -- * Explaining the score-    Explanation(..),-    explain,-  ) where--import Data.Ix-import Data.Array.Unboxed--data Context term field feature = Context {-         numDocsTotal     :: !Int,-         avgFieldLength   :: field -> Float,-         numDocsWithTerm  :: term -> Int,-         paramK1          :: !Float,-         paramB           :: field -> Float,-         -- consider minimum length to prevent massive B bonus?-         fieldWeight      :: field -> Float,-         featureWeight    :: feature -> Float,-         featureFunction  :: feature -> FeatureFunction-       }--data Doc term field feature = Doc {-         docFieldLength        :: field -> Int,-         docFieldTermFrequency :: field -> term -> Int,-         docFeatureValue       :: feature -> Float-       }----- | The BM25F score for a document for a given set of terms.----score :: (Ix field, Bounded field, Ix feature, Bounded feature) =>-         Context term field feature ->-         Doc term field feature -> [term] -> Float-score ctx doc terms =-    sum (map (weightedTermScore    ctx doc) terms)-  + sum (map (weightedNonTermScore ctx doc) features)--  where-    features = range (minBound, maxBound)---weightedTermScore :: (Ix field, Bounded field) =>-                     Context term field feature ->-                     Doc term field feature -> term -> Float-weightedTermScore ctx doc t =-    weightIDF ctx t *     tf'-                     / (k1 + tf')-  where-    tf' = weightedDocTermFrequency ctx doc t-    k1  = paramK1 ctx---weightIDF :: Context term field feature -> term -> Float-weightIDF ctx t =-    log ((n - n_t + 0.5) / (n_t + 0.5))-  where-    n   = fromIntegral (numDocsTotal ctx)-    n_t = fromIntegral (numDocsWithTerm ctx t)---weightedDocTermFrequency :: (Ix field, Bounded field) =>-                            Context term field feature ->-                            Doc term field feature -> term -> Float-weightedDocTermFrequency ctx doc t =-    sum [ w_f * tf_f / _B_f-        | field <- range (minBound, maxBound)-        , let w_f  = fieldWeight ctx field-              tf_f = fromIntegral (docFieldTermFrequency doc field t)-              _B_f = lengthNorm ctx doc field-        , not (isNaN _B_f)-        ]-    -- When the avgFieldLength is 0 we have a field which is empty for all-    -- documents. Unfortunately it leads to a NaN because the-    -- docFieldTermFrequency will also be 0 so we get 0/0. What we want to-    -- do in this situation is have that field contribute nothing to the-    -- score. The simplest way to achieve that is to skip if _B_f is NaN.-    -- So I think this is fine and not an ugly hack.--lengthNorm :: Context term field feature ->-              Doc term field feature -> field -> Float-lengthNorm ctx doc field =-    (1-b_f) + b_f * sl_f / avgsl_f-  where-    b_f     = paramB ctx field-    sl_f    = fromIntegral (docFieldLength doc field)-    avgsl_f = avgFieldLength ctx field---weightedNonTermScore :: (Ix feature, Bounded feature) =>-                        Context term field feature ->-                        Doc term field feature -> feature -> Float-weightedNonTermScore ctx doc feature =-    w_f * _V_f f_f-  where-    w_f  = featureWeight ctx feature-    _V_f = applyFeatureFunction (featureFunction ctx feature)-    f_f  = docFeatureValue doc feature---data FeatureFunction-   = LogarithmicFunction   Float -- ^ @log (\lambda_i + f_i)@-   | RationalFunction      Float -- ^ @f_i / (\lambda_i + f_i)@-   | SigmoidFunction Float Float -- ^ @1 / (\lambda + exp(-(\lambda' * f_i))@--applyFeatureFunction :: FeatureFunction -> (Float -> Float)-applyFeatureFunction (LogarithmicFunction p1) = \fi -> log (p1 + fi)-applyFeatureFunction (RationalFunction    p1) = \fi -> fi / (p1 + fi)-applyFeatureFunction (SigmoidFunction  p1 p2) = \fi -> 1 / (p1 + exp (-fi * p2))----------------------------------- Bulk scoring of many terms------- | Most of the time we want to score several different documents for the same--- set of terms, but sometimes we want to score one document for many terms--- and in that case we can save a bit of work by doing it in bulk. It lets us--- calculate once and share things that depend only on the document, and not--- the term.------ To take advantage of the sharing you must partially apply and name the--- per-doc score functon, e.g.------ > let score :: term -> (field -> Int) -> Float--- >     score = BM25.bulkScorer ctx doc--- >  in sum [ score t (\f -> counts ! (t, f)) | t <- ts ]----scoreTermsBulk :: forall field term feature. (Ix field, Bounded field) =>-                  Context term field feature ->-                  Doc term field feature ->-                  (term -> (field -> Int) -> Float)-scoreTermsBulk ctx doc = -    -- This is just a rearrangement of weightedTermScore and-    -- weightedDocTermFrequency above, with the doc-constant bits hoisted out.--    \t tFreq ->-    let !tf' = sum [ w!f * tf_f / _B!f-                   | f <- range (minBound, maxBound)-                   , let tf_f = fromIntegral (tFreq f)-                         _B_f = _B!f-                   , not (isNaN _B_f)-                   ]--     in weightIDF ctx t *     tf'-                         / (k1 + tf')-  where-    -- So long as the caller does the partial application thing then these-    -- values can all be shared between many calls with different terms.--    !k1 = paramK1 ctx-    w, _B :: UArray field Float-    !w  = array (minBound, maxBound)-                [ (field, fieldWeight ctx field)-                | field <- range (minBound, maxBound) ]-    !_B = array (minBound, maxBound)-                [ (field, lengthNorm ctx doc field)-                | field <- range (minBound, maxBound) ]------------------------ Explanation------- | A breakdown of the BM25F score, to explain somewhat how it relates to--- the inputs, and so you can compare the scores of different documents.----data Explanation field feature term = Explanation {-       -- | The overall score is the sum of the 'termScores', 'positionScore'-       -- and 'nonTermScore'-       overallScore  :: Float,--       -- | There is a score contribution from each query term. This is the-       -- score for the term across all fields in the document (but see-       -- 'termFieldScores').-       termScores    :: [(term, Float)],-{--       -- | There is a score contribution for positional information. Terms-       -- appearing in the document close together give a bonus.-       positionScore :: [(field, Float)],--}-       -- | The document can have an inate bonus score independent of the terms-       -- in the query. For example this might be a popularity score.-       nonTermScores :: [(feature, Float)],--       -- | This does /not/ contribute to the 'overallScore'. It is an-       -- indication of how the 'termScores' relates to per-field scores.-       -- Note however that the term score for all fields is /not/ simply-       -- sum of the per-field scores. The point of the BM25F scoring function-       -- is that a linear combination of per-field scores is wrong, and BM25F-       -- does a more cunning non-linear combination.-       ---       -- However, it is still useful as an indication to see scores for each-       -- field for a term, to see how the compare.-       ---       termFieldScores :: [(term, [(field, Float)])]-     }-  deriving Show--instance Functor (Explanation field feature) where-  fmap f e@Explanation{..} =-    e {-      termScores      = [ (f t, s)  | (t, s)  <- termScores ],-      termFieldScores = [ (f t, fs) | (t, fs) <- termFieldScores ]-    }--explain :: (Ix field, Bounded field, Ix feature, Bounded feature) =>-           Context term field feature ->-           Doc term field feature -> [term] -> Explanation field feature term-explain ctx doc ts =-    Explanation {..}-  where-    overallScore  = sum (map snd termScores)---                  + sum (map snd positionScore)-                  + sum (map snd nonTermScores)-    termScores    = [ (t, weightedTermScore ctx doc t) | t <- ts ]---    positionScore = [ (f, 0) | f <- range (minBound, maxBound) ]-    nonTermScores = [ (feature, weightedNonTermScore ctx doc feature)-                    | feature <- range (minBound, maxBound) ]--    termFieldScores =-      [ (t, fieldScores)-      | t <- ts-      , let fieldScores =-              [ (f, weightedTermScore ctx' doc t)-              | f <- range (minBound, maxBound)-              , let ctx' = ctx { fieldWeight = fieldWeightOnly f }-              ]-      ]-    fieldWeightOnly f f' | sameField f f' = fieldWeight ctx f'-                         | otherwise      = 0--    sameField f f' = index (minBound, maxBound) f-                  == index (minBound, maxBound) f'
− Data/SearchEngine/DocFeatVals.hs 
@@ -1,25 +0,0 @@-{-# LANGUAGE BangPatterns, GeneralizedNewtypeDeriving #-}-module Data.SearchEngine.DocFeatVals (-    DocFeatVals,-    featureValue,-    create,-  ) where--import Data.SearchEngine.DocTermIds (vecIndexIx, vecCreateIx)-import Data.Vector (Vector)-import Data.Ix (Ix)----- | Storage for the non-term feature values i a document.----newtype DocFeatVals feature = DocFeatVals (Vector Float)-  deriving (Show)--featureValue :: (Ix feature, Bounded feature) => DocFeatVals feature -> feature -> Float-featureValue (DocFeatVals featVec) = vecIndexIx featVec--create :: (Ix feature, Bounded feature) =>-          (feature -> Float) -> DocFeatVals feature-create docFeatVals =-    DocFeatVals (vecCreateIx docFeatVals)-
− Data/SearchEngine/DocIdSet.hs 
@@ -1,207 +0,0 @@-{-# LANGUAGE BangPatterns, GeneralizedNewtypeDeriving, MultiParamTypeClasses,-             TypeFamilies #-}-module Data.SearchEngine.DocIdSet (-    DocId(DocId),-    DocIdSet(..),-    null,-    size,-    empty,-    singleton,-    fromList,-    toList,-    insert,-    delete,-    member,-    union,-    unions,-    intersection,-    invariant,-  ) where--import Data.Word-import qualified Data.Vector.Unboxed         as Vec-import qualified Data.Vector.Unboxed.Mutable as MVec-import qualified Data.Vector.Generic         as GVec-import qualified Data.Vector.Generic.Mutable as GMVec-import Control.Monad.ST-import Control.Monad (liftM)-import qualified Data.Set as Set-import Data.List (foldl', sortBy)-import Data.Function (on)--import Prelude hiding (null)---newtype DocId = DocId { unDocId :: Word32 }-  deriving (Eq, Ord, Show, Enum, Bounded)--newtype DocIdSet = DocIdSet (Vec.Vector DocId)-  deriving (Eq, Show)---- represented as a sorted sequence of ids-invariant :: DocIdSet -> Bool-invariant (DocIdSet vec) =-    strictlyAscending (Vec.toList vec)-  where-    strictlyAscending (a:xs@(b:_)) = a < b && strictlyAscending xs-    strictlyAscending _  = True---size :: DocIdSet -> Int-size (DocIdSet vec) = Vec.length vec--null :: DocIdSet -> Bool-null (DocIdSet vec) = Vec.null vec--empty :: DocIdSet-empty = DocIdSet Vec.empty--singleton :: DocId -> DocIdSet-singleton = DocIdSet . Vec.singleton--fromList :: [DocId] -> DocIdSet-fromList = DocIdSet . Vec.fromList . Set.toAscList . Set.fromList--toList ::  DocIdSet -> [DocId]-toList (DocIdSet vec) = Vec.toList vec--insert :: DocId -> DocIdSet -> DocIdSet-insert x (DocIdSet vec) =-    case binarySearch vec 0 (Vec.length vec - 1) x of-      (_, True)  -> DocIdSet vec-      (i, False) -> case Vec.splitAt i vec of-                      (before, after) ->-                        DocIdSet (Vec.concat [before, Vec.singleton x, after])--delete :: DocId -> DocIdSet -> DocIdSet-delete x (DocIdSet vec) =-    case binarySearch vec 0 (Vec.length vec - 1) x of-      (_, False) -> DocIdSet vec-      (i, True)  -> case Vec.splitAt i vec of-                      (before, after) ->-                        DocIdSet (before Vec.++ Vec.tail after)--member :: DocId -> DocIdSet -> Bool-member x (DocIdSet vec) = snd (binarySearch vec 0 (Vec.length vec - 1) x)--binarySearch :: Vec.Vector DocId -> Int -> Int -> DocId -> (Int, Bool)-binarySearch vec !a !b !key-  | a > b     = (a, False)-  | otherwise =-    let mid = (a + b) `div` 2-     in case compare key (vec Vec.! mid) of-          LT -> binarySearch vec a (mid-1) key-          EQ -> (mid, True)-          GT -> binarySearch vec (mid+1) b key--unions :: [DocIdSet] -> DocIdSet-unions = foldl' union empty-         -- a bit more effecient if we merge small ones first-       . sortBy (compare `on` size)--union :: DocIdSet -> DocIdSet -> DocIdSet-union x y | null x = y-          | null y = x-union (DocIdSet xs) (DocIdSet ys) =-    DocIdSet (Vec.create (MVec.new sizeBound >>= writeMergedUnion xs ys))-  where-    sizeBound = Vec.length xs + Vec.length ys--writeMergedUnion :: Vec.Vector DocId -> Vec.Vector DocId ->-                    MVec.MVector s DocId -> ST s (MVec.MVector s DocId)-writeMergedUnion xs0 ys0 !out = do-    i <- go xs0 ys0 0-    return $! MVec.take i out-  where-    go !xs !ys !i-      | Vec.null xs = do Vec.copy (MVec.slice i (Vec.length ys) out) ys-                         return (i + Vec.length ys)-      | Vec.null ys = do Vec.copy (MVec.slice i (Vec.length xs) out) xs-                         return (i + Vec.length xs)-      | otherwise   = let x = Vec.head xs; y = Vec.head ys-                      in case compare x y of-                          GT -> do MVec.write out i y-                                   go           xs  (Vec.tail ys) (i+1)-                          EQ -> do MVec.write out i x-                                   go (Vec.tail xs) (Vec.tail ys) (i+1)-                          LT -> do MVec.write out i x-                                   go (Vec.tail xs)           ys  (i+1)--intersection :: DocIdSet -> DocIdSet -> DocIdSet-intersection x y | null x = empty-                 | null y = empty-intersection (DocIdSet xs) (DocIdSet ys) =-    DocIdSet (Vec.create (MVec.new sizeBound >>= writeMergedIntersection xs ys))-  where-    sizeBound = max (Vec.length xs) (Vec.length ys)--writeMergedIntersection :: Vec.Vector DocId -> Vec.Vector DocId ->-                           MVec.MVector s DocId -> ST s (MVec.MVector s DocId)-writeMergedIntersection xs0 ys0 !out = do-    i <- go xs0 ys0 0-    return $! MVec.take i out-  where-    go !xs !ys !i-      | Vec.null xs = return i-      | Vec.null ys = return i-      | otherwise   = let x = Vec.head xs; y = Vec.head ys-                      in case compare x y of-                          GT ->    go           xs  (Vec.tail ys)  i-                          EQ -> do MVec.write out i x-                                   go (Vec.tail xs) (Vec.tail ys) (i+1)-                          LT ->    go (Vec.tail xs)           ys   i----------------------------------------------------------------------------------- verbose Unbox instances-----instance MVec.Unbox DocId--newtype instance MVec.MVector s DocId = MV_DocId (MVec.MVector s Word32)--instance GMVec.MVector MVec.MVector DocId where-    basicLength          (MV_DocId v) = GMVec.basicLength v-    basicUnsafeSlice i l (MV_DocId v) = MV_DocId (GMVec.basicUnsafeSlice i l v)-    basicUnsafeNew     l              = MV_DocId `liftM` GMVec.basicUnsafeNew l-    basicInitialize      (MV_DocId v) = GMVec.basicInitialize v-    basicUnsafeReplicate l x          = MV_DocId `liftM` GMVec.basicUnsafeReplicate l (unDocId x)-    basicUnsafeRead  (MV_DocId v) i   = DocId `liftM`    GMVec.basicUnsafeRead v i-    basicUnsafeWrite (MV_DocId v) i x = GMVec.basicUnsafeWrite v i (unDocId x)-    basicClear       (MV_DocId v)     = GMVec.basicClear v-    basicSet         (MV_DocId v) x   = GMVec.basicSet v (unDocId x)-    basicUnsafeGrow  (MV_DocId v) l   = MV_DocId `liftM` GMVec.basicUnsafeGrow v l-    basicUnsafeCopy  (MV_DocId v) (MV_DocId v') = GMVec.basicUnsafeCopy v v'-    basicUnsafeMove  (MV_DocId v) (MV_DocId v') = GMVec.basicUnsafeMove v v'-    basicOverlaps    (MV_DocId v) (MV_DocId v') = GMVec.basicOverlaps   v v'-    {-# INLINE basicLength #-}-    {-# INLINE basicUnsafeSlice #-}-    {-# INLINE basicOverlaps #-}-    {-# INLINE basicUnsafeNew #-}-    {-# INLINE basicInitialize #-}-    {-# INLINE basicUnsafeReplicate #-}-    {-# INLINE basicUnsafeRead #-}-    {-# INLINE basicUnsafeWrite #-}-    {-# INLINE basicClear #-}-    {-# INLINE basicSet #-}-    {-# INLINE basicUnsafeCopy #-}-    {-# INLINE basicUnsafeMove #-}-    {-# INLINE basicUnsafeGrow #-}--newtype instance Vec.Vector DocId = V_DocId (Vec.Vector Word32)--instance GVec.Vector Vec.Vector DocId where-    basicUnsafeFreeze (MV_DocId mv)  = V_DocId  `liftM` GVec.basicUnsafeFreeze mv-    basicUnsafeThaw   (V_DocId  v)   = MV_DocId `liftM` GVec.basicUnsafeThaw v-    basicLength       (V_DocId  v)   = GVec.basicLength v-    basicUnsafeSlice i l (V_DocId v) = V_DocId (GVec.basicUnsafeSlice i l v)-    basicUnsafeIndexM (V_DocId  v) i = DocId `liftM` GVec.basicUnsafeIndexM v i-    basicUnsafeCopy   (MV_DocId mv)-                      (V_DocId  v)   = GVec.basicUnsafeCopy mv v-    elemseq           (V_DocId  v) x = GVec.elemseq v (unDocId x)-    {-# INLINE basicUnsafeFreeze #-}-    {-# INLINE basicUnsafeThaw #-}-    {-# INLINE basicLength #-}-    {-# INLINE basicUnsafeSlice #-}-    {-# INLINE basicUnsafeIndexM #-}-    {-# INLINE basicUnsafeCopy #-}-    {-# INLINE elemseq #-}
− Data/SearchEngine/DocTermIds.hs 
@@ -1,81 +0,0 @@-{-# LANGUAGE BangPatterns, GeneralizedNewtypeDeriving #-}-module Data.SearchEngine.DocTermIds (-    DocTermIds,-    TermId,-    fieldLength,-    fieldTermCount,-    fieldElems,-    create,-    denseTable,-    vecIndexIx,-    vecCreateIx,-  ) where--import Data.SearchEngine.TermBag (TermBag, TermId)-import qualified Data.SearchEngine.TermBag as TermBag--import Data.Vector (Vector, (!))-import qualified Data.Vector as Vec-import qualified Data.Vector.Unboxed as UVec-import Data.Ix (Ix)-import qualified Data.Ix as Ix----- | The 'TermId's for the 'Term's that occur in a document. Documents may have--- multiple fields and the 'DocTerms' type holds them separately for each field.----newtype DocTermIds field = DocTermIds (Vector TermBag)-  deriving (Show)--getField :: (Ix field, Bounded field) => DocTermIds field -> field -> TermBag-getField (DocTermIds fieldVec) = vecIndexIx fieldVec--create :: (Ix field, Bounded field) =>-          (field -> [TermId]) -> DocTermIds field-create docTermIds =-    DocTermIds (vecCreateIx (TermBag.fromList . docTermIds))---- | The number of terms in a field within the document.-fieldLength :: (Ix field, Bounded field) => DocTermIds field -> field -> Int-fieldLength docterms field =-    TermBag.size (getField docterms field)---- | /O(log n)/ The frequency of a particular term in a field within the document.----fieldTermCount :: (Ix field, Bounded field) =>-                  DocTermIds field -> field -> TermId -> Int-fieldTermCount docterms field termid =-    fromIntegral (TermBag.termCount (getField docterms field) termid)--fieldElems :: (Ix field, Bounded field) => DocTermIds field -> field -> [TermId]-fieldElems docterms field =-    TermBag.elems (getField docterms field)---- | The 'DocTermIds' is really a sparse 2d array, and doing lookups with--- 'fieldTermCount' has a O(log n) cost. This function converts to a dense--- tabular representation which then enables linear scans.----denseTable :: (Ix field, Bounded field) => DocTermIds field ->-              (Int, Int -> TermId, Int -> field -> Int)-denseTable (DocTermIds fieldVec) =-    let (!termids, !termcounts) = TermBag.denseTable (Vec.toList fieldVec)-        !numTerms = UVec.length termids-     in ( numTerms-        , \i    -> termids UVec.! i-        , \i ix -> let j = Ix.index (minBound, maxBound) ix-                    in fromIntegral (termcounts UVec.! (j * numTerms + i))-        )-------------------------------------- Vector indexed by Ix Bounded-----vecIndexIx  :: (Ix ix, Bounded ix) => Vector a -> ix -> a-vecIndexIx vec ix = vec ! Ix.index (minBound, maxBound) ix--vecCreateIx :: (Ix ix, Bounded ix) => (ix -> a) -> Vector a-vecCreateIx f = Vec.fromListN (Ix.rangeSize bounds)-                  [ y | ix <- Ix.range bounds, let !y = f ix ]-  where-    bounds = (minBound, maxBound)-
− Data/SearchEngine/Query.hs 
@@ -1,257 +0,0 @@-{-# LANGUAGE BangPatterns, NamedFieldPuns, RecordWildCards #-}--module Data.SearchEngine.Query (--    -- * Querying-    query,-    ResultsFilter(..),--    -- * Explain mode for query result rankings-    queryExplain,-    BM25F.Explanation(..),-    setRankParams,--    -- ** Utils used by autosuggest-    relevanceScore,-    indexDocToBM25Doc,-    expandTransformedQueryTerm,-  ) where--import Data.SearchEngine.Types-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.DocFeatVals as DocFeatVals-import qualified Data.SearchEngine.BM25F       as BM25F--import Data.Ix-import Data.List-import Data.Function-import Data.Maybe----- | Execute a normal query. Find the documents in which one or more of--- the search terms appear and return them in ranked order.------ The number of documents returned is limited by the 'paramResultsetSoftLimit'--- and 'paramResultsetHardLimit' paramaters. This also limits the cost of the--- query (which is primarily the cost of scoring each document).------ The given terms are all assumed to be complete (as opposed to prefixes--- like with 'queryAutosuggest').----query :: (Ix field, Bounded field, Ix feature, Bounded feature) =>-         SearchEngine doc key field feature ->-         [Term] -> [key]-query se@SearchEngine{ searchIndex,-                       searchRankParams = SearchRankParameters{..} }-      terms =--  let -- Start by transforming/normalising all the query terms.-      -- This can be done differently for each field we search by.-      lookupTerms :: [Term]-      lookupTerms = concatMap (expandTransformedQueryTerm se) terms--      -- Then we look up all the normalised terms in the index.-      rawresults :: [Maybe (TermId, DocIdSet)]-      rawresults = map (SI.lookupTerm searchIndex) lookupTerms--      -- For the terms that occur in the index, this gives us the term's id-      -- and the set of documents that the term occurs in.-      termids   :: [TermId]-      docidsets :: [DocIdSet]-      (termids, docidsets) = unzip (catMaybes rawresults)--      -- We looked up the documents that *any* of the term occur in (not all)-      -- so this could be rather a lot of docs if the user uses a few common-      -- terms. Scoring these result docs is a non-trivial cost so we want to-      -- limit the number that we have to score. The standard trick is to-      -- consider the doc sets in the order of size, smallest to biggest. Once-      -- we have gone over a certain threshold of docs then don't bother with-      -- the doc sets for the remaining terms. This tends to work because the-      -- scoring gives lower weight to terms that occur in many documents.-      unrankedResults :: DocIdSet-      unrankedResults = pruneRelevantResults-                          paramResultsetSoftLimit-                          paramResultsetHardLimit-                          docidsets--      --TODO: technically this isn't quite correct. Because each field can-      -- be normalised differently, we can end up with different termids for-      -- the same original search term, and then we score those as if they-      -- were different terms, which makes a difference when the term appears-      -- in multiple fields (exactly the case BM25F is supposed to deal with).-      -- What we ought to have instead is an Array (Int, field) TermId, and-      -- make the scoring use the appropriate termid for each field, but to-      -- consider them the "same" term.-   in rankResults se termids (DocIdSet.toList unrankedResults)---- | Before looking up a term in the main index we need to normalise it--- using the 'transformQueryTerm'. Of course the transform can be different--- for different fields, so we have to collect all the forms (eliminating--- duplicates).----expandTransformedQueryTerm :: (Ix field, Bounded field) =>-                              SearchEngine doc key field feature ->-                              Term -> [Term]-expandTransformedQueryTerm SearchEngine{searchConfig} term =-    nub [ transformForField field-        | let transformForField = transformQueryTerm searchConfig term-        , field <- range (minBound, maxBound) ]---rankResults :: (Ix field, Bounded field, Ix feature, Bounded feature) =>-               SearchEngine doc key field feature ->-               [TermId] -> [DocId] -> [key]-rankResults se@SearchEngine{searchIndex} queryTerms docids =-    map snd-  $ sortBy (flip compare `on` fst)-      [ (relevanceScore se queryTerms doctermids docfeatvals, dockey)-      | docid <- docids-      , let (dockey, doctermids, docfeatvals) = SI.lookupDocId searchIndex docid ]--relevanceScore :: (Ix field, Bounded field, Ix feature, Bounded feature) =>-                  SearchEngine doc key field feature ->-                  [TermId] -> DocTermIds field -> DocFeatVals feature -> Float-relevanceScore SearchEngine{bm25Context} queryTerms doctermids docfeatvals =-    BM25F.score bm25Context doc queryTerms-  where-    doc = indexDocToBM25Doc doctermids docfeatvals--indexDocToBM25Doc :: (Ix field, Bounded field, Ix feature, Bounded feature) =>-                     DocTermIds field ->-                     DocFeatVals feature ->-                     BM25F.Doc TermId field feature-indexDocToBM25Doc doctermids docfeatvals =-    BM25F.Doc {-      BM25F.docFieldLength        = DocTermIds.fieldLength    doctermids,-      BM25F.docFieldTermFrequency = DocTermIds.fieldTermCount doctermids,-      BM25F.docFeatureValue       = DocFeatVals.featureValue docfeatvals-    }--pruneRelevantResults :: Int -> Int -> [DocIdSet] -> DocIdSet-pruneRelevantResults softLimit hardLimit =-    -- Look at the docsets starting with the smallest ones. Smaller docsets-    -- correspond to the rarer terms, which are the ones that score most highly.-    go DocIdSet.empty . sortBy (compare `on` DocIdSet.size)-  where-    go !acc [] = acc-    go !acc (d:ds)-        -- If this is the first one, we add it anyway, otherwise we're in-        -- danger of returning no results at all.-      | DocIdSet.null acc = go d ds-        -- We consider the size our docset would be if we add this extra one...-        -- If it puts us over the hard limit then stop.-      | size > hardLimit  = acc-        -- If it puts us over soft limit then we add it and stop-      | size > softLimit  = DocIdSet.union acc d-        -- Otherwise we can add it and carry on to consider the remainder-      | otherwise         = go (DocIdSet.union acc d) ds-      where-        size = DocIdSet.size acc + DocIdSet.size d-------------------------------------- Normal query with explanation-----queryExplain :: (Ix field, Bounded field, Ix feature, Bounded feature) =>-                SearchEngine doc key field feature ->-                [Term] -> [(BM25F.Explanation field feature Term, key)]-queryExplain se@SearchEngine{ searchIndex,-                              searchConfig     = SearchConfig{transformQueryTerm},-                              searchRankParams = SearchRankParameters{..} }-      terms =--  -- See 'query' above for explanation. Really we ought to combine them.-  let lookupTerms :: [Term]-      lookupTerms = [ term'-                    | term  <- terms-                    , let transformForField = transformQueryTerm term-                    , term' <- nub [ transformForField field-                                   | field <- range (minBound, maxBound) ]-                    ]--      rawresults :: [Maybe (TermId, DocIdSet)]-      rawresults = map (SI.lookupTerm searchIndex) lookupTerms--      termids   :: [TermId]-      docidsets :: [DocIdSet]-      (termids, docidsets) = unzip (catMaybes rawresults)--      unrankedResults :: DocIdSet-      unrankedResults = pruneRelevantResults-                          paramResultsetSoftLimit-                          paramResultsetHardLimit-                          docidsets--   in rankExplainResults se termids (DocIdSet.toList unrankedResults)--rankExplainResults :: (Ix field, Bounded field, Ix feature, Bounded feature) =>-                      SearchEngine doc key field feature ->-                      [TermId] ->-                      [DocId] ->-                      [(BM25F.Explanation field feature Term, key)]-rankExplainResults se@SearchEngine{searchIndex} queryTerms docids =-    sortBy (flip compare `on` (BM25F.overallScore . fst))-      [ (explainRelevanceScore se queryTerms doctermids docfeatvals, dockey)-      | docid <- docids-      , let (dockey, doctermids, docfeatvals) = SI.lookupDocId searchIndex docid ]---explainRelevanceScore :: (Ix field, Bounded field, Ix feature, Bounded feature) =>-                         SearchEngine doc key field feature ->-                         [TermId] ->-                         DocTermIds field ->-                         DocFeatVals feature ->-                         BM25F.Explanation field feature Term-explainRelevanceScore SearchEngine{bm25Context, searchIndex}-                      queryTerms doctermids docfeatvals =-    fmap (SI.getTerm searchIndex) (BM25F.explain bm25Context doc queryTerms)-  where-    doc = indexDocToBM25Doc doctermids docfeatvals---setRankParams :: SearchRankParameters field feature ->-                 SearchEngine doc key field feature ->-                 SearchEngine doc key field feature-setRankParams params@SearchRankParameters{..} se =-    se {-      searchRankParams = params,-      bm25Context      = (bm25Context se) {-        BM25F.paramK1         = paramK1,-        BM25F.paramB          = paramB,-        BM25F.fieldWeight     = paramFieldWeights,-        BM25F.featureWeight   = paramFeatureWeights,-        BM25F.featureFunction = paramFeatureFunctions-      }-    }-------------------------------------- Results filter------- | In some applications it is necessary to enforce some security or--- visibility rule about the query results (e.g. in a typical DB-based--- application different users can see different data items). Typically--- it would be too expensive to build different search indexes for the--- different contexts and so the strategy is to use one index containing--- everything and filter for visibility in the results. This means the--- filter condition is different for different queries (e.g. performed--- on behalf of different users).------ Filtering the results after a query is possible but not the most efficient--- thing to do because we've had to score all the not-visible documents.--- The better thing to do is to filter as part of the query, this way we can--- filter before the expensive scoring.------ We provide one further optimisation: bulk predicates. In some applications--- it can be quicker to check the security\/visibility of a whole bunch of--- results all in one go.----data ResultsFilter key = NoFilter-                       | FilterPredicate     (key -> Bool)-                       | FilterBulkPredicate ([key] -> [Bool])---TODO: allow filtering & non-feature score lookup in one bulk op-
− Data/SearchEngine/SearchIndex.hs 
@@ -1,456 +0,0 @@-{-# LANGUAGE BangPatterns, NamedFieldPuns #-}--module Data.SearchEngine.SearchIndex (-    SearchIndex,-    Term,-    TermId,-    DocId,--    emptySearchIndex,-    insertDoc,-    deleteDoc,--    docCount,-    lookupTerm,-    lookupTermsByPrefix,-    lookupTermId,-    lookupDocId,-    lookupDocKey,-    lookupDocKeyDocId,--    getTerm,-    getDocKey,--    invariant,-  ) where--import Data.SearchEngine.DocIdSet (DocIdSet, DocId)-import qualified Data.SearchEngine.DocIdSet as DocIdSet-import Data.SearchEngine.DocTermIds (DocTermIds, TermId, vecIndexIx, vecCreateIx)-import qualified Data.SearchEngine.DocTermIds as DocTermIds-import Data.SearchEngine.DocFeatVals (DocFeatVals)-import qualified Data.SearchEngine.DocFeatVals as DocFeatVals--import Data.Ix (Ix)-import qualified Data.Ix as Ix-import Data.Map (Map)-import qualified Data.Map as Map-import Data.IntMap (IntMap)-import qualified Data.IntMap as IntMap-import qualified Data.Set as Set-import Data.Text (Text)-import qualified Data.Text as T-import Data.List (foldl')--import Control.Exception (assert)---- | Terms are short strings, usually whole words.----type Term = Text---- | The search index is essentially a many-to-many mapping between documents--- and terms. Each document contains many terms and each term occurs in many--- documents. It is a bidirectional mapping as we need to support lookups in--- both directions.------ Documents are identified by a key (in Ord) while terms are text values.--- Inside the index however we assign compact numeric ids to both documents and--- terms. The advantage of this is a much more compact in-memory representation--- and the disadvantage is greater complexity. In particular it means we have--- to manage bidirectional mappings between document keys and ids, and between--- terms and term ids.------ So the mappings we maintain can be depicted as:------ >  Term   <-- 1:1 -->   TermId--- >          \              ^--- >           \             |--- >           1:many    many:many--- >                \        |--- >                 \->     v--- > DocKey  <-- 1:1 -->   DocId------ For efficiency, these details are exposed in the interface. In particular--- the mapping from TermId to many DocIds is exposed via a 'DocIdSet',--- and the mapping from DocIds to TermIds is exposed via 'DocTermIds'.------ The main reason we need to keep the DocId -> TermId is to allow for--- efficient incremental updates.----data SearchIndex key field feature = SearchIndex {-       -- the indexes-       termMap           :: !(Map Term TermInfo),-       termIdMap         :: !(IntMap TermIdInfo),-       docIdMap          :: !(IntMap (DocInfo key field feature)),-       docKeyMap         :: !(Map key DocId),--       -- auto-increment key counters-       nextTermId        :: TermId,-       nextDocId         :: DocId-     }-  deriving Show--data TermInfo = TermInfo !TermId !DocIdSet-  deriving Show--data TermIdInfo = TermIdInfo !Term !DocIdSet-  deriving (Show, Eq)--data DocInfo key field feature = DocInfo !key !(DocTermIds field)-                                              !(DocFeatVals feature)-  deriving Show----------------------------- SearchIndex basics-----emptySearchIndex :: SearchIndex key field feature-emptySearchIndex =-    SearchIndex-      Map.empty-      IntMap.empty-      IntMap.empty-      Map.empty-      minBound-      minBound--checkInvariant :: (Ord key, Ix field, Bounded field) =>-                  SearchIndex key field feature -> SearchIndex key field feature-checkInvariant si = assert (invariant si) si--invariant :: (Ord key, Ix field, Bounded field) =>-             SearchIndex key field feature -> Bool-invariant SearchIndex{termMap, termIdMap, docKeyMap, docIdMap} =-      and [ IntMap.lookup (fromEnum termId) termIdMap-            == Just (TermIdInfo term docidset)-          | (term, (TermInfo termId docidset)) <- Map.assocs termMap ]-  &&  and [ case Map.lookup term termMap of-              Just (TermInfo termId' docidset') -> toEnum termId == termId'-                                                   && docidset == docidset'-              Nothing                           -> False-          | (termId, (TermIdInfo term docidset)) <- IntMap.assocs termIdMap ]-  &&  and [ case IntMap.lookup (fromEnum docId) docIdMap of-              Just (DocInfo docKey' _ _) -> docKey == docKey'-              Nothing                  -> False-          | (docKey, docId) <- Map.assocs docKeyMap ]-  &&  and [ Map.lookup docKey docKeyMap == Just (toEnum docId)-          | (docId, DocInfo docKey _ _) <- IntMap.assocs docIdMap ]-  &&  and [ DocIdSet.invariant docIdSet-          | (_term, (TermInfo _ docIdSet)) <- Map.assocs termMap ]-  &&  and [ any (\field -> DocTermIds.fieldTermCount docterms field termId > 0) fields-          | (_term, (TermInfo termId docIdSet)) <- Map.assocs termMap-          , docId <- DocIdSet.toList docIdSet-          , let DocInfo _ docterms _ = docIdMap IntMap.! fromEnum docId ]-  &&  and [ IntMap.member (fromEnum termid) termIdMap-          | (_docId, DocInfo _ docTerms _) <- IntMap.assocs docIdMap-          , field <- fields-          , termid <- DocTermIds.fieldElems docTerms field ]-  where-    fields = Ix.range (minBound, maxBound)------------------------- Lookups-----docCount :: SearchIndex key field feature -> Int-docCount SearchIndex{docIdMap} = IntMap.size docIdMap--lookupTerm :: SearchIndex key field feature -> Term -> Maybe (TermId, DocIdSet)-lookupTerm SearchIndex{termMap} term =-    case Map.lookup term termMap of-      Nothing                         -> Nothing-      Just (TermInfo termid docidset) -> Just (termid, docidset)--lookupTermsByPrefix :: SearchIndex key field feature ->-                       Term -> [(TermId, DocIdSet)]-lookupTermsByPrefix SearchIndex{termMap} term =-    [ (termid, docidset)-    | (TermInfo termid docidset) <- lookupPrefix term termMap ]--lookupTermId :: SearchIndex key field feature -> TermId -> DocIdSet-lookupTermId SearchIndex{termIdMap} termid =-    case IntMap.lookup (fromEnum termid) termIdMap of-      Nothing -> error $ "lookupTermId: not found " ++ show termid-      Just (TermIdInfo _ docidset) -> docidset--lookupDocId :: SearchIndex key field feature ->-               DocId -> (key, DocTermIds field, DocFeatVals feature)-lookupDocId SearchIndex{docIdMap} docid =-    case IntMap.lookup (fromEnum docid) docIdMap of-      Nothing                                   -> errNotFound-      Just (DocInfo key doctermids docfeatvals) -> (key, doctermids, docfeatvals)-  where-    errNotFound = error $ "lookupDocId: not found " ++ show docid--lookupDocKey :: Ord key => SearchIndex key field feature ->-                key -> Maybe (DocTermIds field)-lookupDocKey SearchIndex{docKeyMap, docIdMap} key = do-    case Map.lookup key docKeyMap of-      Nothing    -> Nothing-      Just docid ->-        case IntMap.lookup (fromEnum docid) docIdMap of-          Nothing                          -> error "lookupDocKey: internal error"-          Just (DocInfo _key doctermids _) -> Just doctermids--lookupDocKeyDocId :: Ord key => SearchIndex key field feature -> key -> Maybe DocId-lookupDocKeyDocId SearchIndex{docKeyMap} key = Map.lookup key docKeyMap---getTerm :: SearchIndex key field feature -> TermId -> Term-getTerm SearchIndex{termIdMap} termId =-    case termIdMap IntMap.! fromEnum termId of TermIdInfo term _ -> term--getTermId :: SearchIndex key field feature -> Term -> TermId-getTermId SearchIndex{termMap} term =-    case termMap Map.! term of TermInfo termid _ -> termid--getDocKey :: SearchIndex key field feature -> DocId -> key-getDocKey SearchIndex{docIdMap} docid =-    case docIdMap IntMap.! fromEnum docid of-      DocInfo dockey _ _ -> dockey--getDocTermIds :: SearchIndex key field feature -> DocId -> DocTermIds field-getDocTermIds SearchIndex{docIdMap} docid =-    case docIdMap IntMap.! fromEnum docid of-      DocInfo _ doctermids _ -> doctermids------------------------- Insert & delete------- Procedure for adding a new doc...--- (key, field -> [Term])--- alloc docid for key--- add term occurences for docid (include rev map for termid)--- construct indexdoc now that we have all the term -> termid entries--- insert indexdoc---- Procedure for updating a doc...--- (key, field -> [Term])--- find docid for key--- lookup old terms for docid (using termid rev map)--- calc term occurrences to add, term occurrences to delete--- add new term occurrences, delete old term occurrences--- construct indexdoc now that we have all the term -> termid entries--- insert indexdoc---- Procedure for deleting a doc...--- (key, field -> [Term])--- find docid for key--- lookup old terms for docid (using termid rev map)--- delete old term occurrences--- delete indexdoc---- | This is the representation for documents to be added to the index.--- Documents may ----type DocTerms         field   = field   -> [Term]-type DocFeatureValues feature = feature -> Float--insertDoc :: (Ord key, Ix field, Bounded field, Ix feature, Bounded feature) =>-              key -> DocTerms field -> DocFeatureValues feature ->-              SearchIndex key field feature -> SearchIndex key field feature-insertDoc key userDocTerms userDocFeats si@SearchIndex{docKeyMap}-  | Just docid <- Map.lookup key docKeyMap-  = -- Some older version of the doc is already present in the index,-    -- So we keep its docid. Now have to update the doc itself-    -- and update the terms by removing old ones and adding new ones.-    let oldTermsIds   = getDocTermIds si docid-        userDocTerms' = memoiseDocTerms userDocTerms-        newTerms      = docTermSet userDocTerms'-        oldTerms      = docTermIdsTermSet si oldTermsIds-        -- We optimise for the typical case of significant overlap between-        -- the terms in the old and new versions of the document.-        delTerms      = oldTerms `Set.difference` newTerms-        addTerms      = newTerms `Set.difference` oldTerms--     -- Note: adding the doc relies on all the terms being in the termMap-     -- already, so we first add all the term occurences for the docid.-     in checkInvariant-      . insertDocIdToDocEntry docid key userDocTerms' userDocFeats-      . insertTermToDocIdEntries (Set.toList addTerms) docid-      . deleteTermToDocIdEntries (Set.toList delTerms) docid-      $ si--  | otherwise-  = -- We're dealing with a new doc, so allocate a docid for the key-    let (si', docid)  = allocFreshDocId si-        userDocTerms' = memoiseDocTerms userDocTerms-        addTerms      = docTermSet userDocTerms'--     -- Note: adding the doc relies on all the terms being in the termMap-     -- already, so we first add all the term occurences for the docid.-     in checkInvariant-      . insertDocIdToDocEntry docid key userDocTerms' userDocFeats-      . insertDocKeyToIdEntry key docid-      . insertTermToDocIdEntries (Set.toList addTerms) docid-      $ si'--deleteDoc :: (Ord key, Ix field, Bounded field) =>-             key ->-             SearchIndex key field feature -> SearchIndex key field feature-deleteDoc key si@SearchIndex{docKeyMap}-  | Just docid <- Map.lookup key docKeyMap-  = let oldTermsIds = getDocTermIds si docid-        oldTerms    = docTermIdsTermSet si oldTermsIds-     in checkInvariant-      . deleteDocEntry docid key-      . deleteTermToDocIdEntries (Set.toList oldTerms) docid-      $ si-  -  | otherwise = si---------------------------------------- Insert & delete support utils------memoiseDocTerms :: (Ix field, Bounded field) => DocTerms field -> DocTerms field-memoiseDocTerms docTermsFn =-    \field -> vecIndexIx vec field-  where-    vec = vecCreateIx docTermsFn--docTermSet :: (Bounded t, Ix t) => DocTerms t -> Set.Set Term-docTermSet docterms =-    Set.unions [ Set.fromList (docterms field)-               | field <- Ix.range (minBound, maxBound) ]--docTermIdsTermSet :: (Bounded field, Ix field) =>-                     SearchIndex key field feature ->-                     DocTermIds field -> Set.Set Term-docTermIdsTermSet si doctermids =-    Set.unions [ Set.fromList terms-               | field <- Ix.range (minBound, maxBound)-               , let termids = DocTermIds.fieldElems doctermids field-                     terms   = map (getTerm si) termids ]------- The Term <-> DocId mapping------- | Add an entry into the 'Term' to 'DocId' mapping.-insertTermToDocIdEntry :: Term -> DocId -> -                          SearchIndex key field feature ->-                          SearchIndex key field feature-insertTermToDocIdEntry term !docid si@SearchIndex{termMap, termIdMap, nextTermId} =-    case Map.lookup term termMap of-      Nothing ->-        let docIdSet'    = DocIdSet.singleton docid-            !termInfo'   = TermInfo nextTermId docIdSet'-            !termIdInfo' = TermIdInfo term     docIdSet'-         in si { termMap    = Map.insert term termInfo' termMap-               , termIdMap  = IntMap.insert (fromEnum nextTermId)-                                            termIdInfo' termIdMap-               , nextTermId = succ nextTermId }--      Just (TermInfo termId docIdSet) ->-        let docIdSet'    = DocIdSet.insert docid docIdSet-            !termInfo'   = TermInfo termId docIdSet'-            !termIdInfo' = TermIdInfo term docIdSet'-         in si { termMap   = Map.insert term termInfo' termMap-               , termIdMap = IntMap.insert (fromEnum termId)-                                           termIdInfo' termIdMap-               }---- | Add multiple entries into the 'Term' to 'DocId' mapping: many terms that--- map to the same document.-insertTermToDocIdEntries :: [Term] -> DocId ->-                            SearchIndex key field feature ->-                            SearchIndex key field feature-insertTermToDocIdEntries terms !docid si =-    foldl' (\si' term -> insertTermToDocIdEntry term docid si') si terms---- | Delete an entry from the 'Term' to 'DocId' mapping.-deleteTermToDocIdEntry :: Term -> DocId ->-                          SearchIndex key field feature ->-                          SearchIndex key field feature-deleteTermToDocIdEntry term !docid si@SearchIndex{termMap, termIdMap} =-    case  Map.lookup term termMap of-      Nothing -> si-      Just (TermInfo termId docIdSet) ->-        let docIdSet'    = DocIdSet.delete docid docIdSet-            !termInfo'   = TermInfo termId docIdSet'-            !termIdInfo' = TermIdInfo term docIdSet'-        in if DocIdSet.null docIdSet'-            then si { termMap = Map.delete term termMap-                    , termIdMap = IntMap.delete (fromEnum termId) termIdMap }-            else si { termMap   = Map.insert term termInfo' termMap-                    , termIdMap = IntMap.insert (fromEnum termId)-                                                termIdInfo' termIdMap-                    }---- | Delete multiple entries from the 'Term' to 'DocId' mapping: many terms--- that map to the same document.-deleteTermToDocIdEntries :: [Term] -> DocId ->-                            SearchIndex key field feature ->-                            SearchIndex key field feature-deleteTermToDocIdEntries terms !docid si =-    foldl' (\si' term -> deleteTermToDocIdEntry term docid si') si terms------- The DocId <-> Doc mapping-----allocFreshDocId :: SearchIndex key field feature ->-                  (SearchIndex key field feature, DocId)-allocFreshDocId si@SearchIndex{nextDocId} =-    let !si' = si { nextDocId = succ nextDocId }-     in (si', nextDocId)--insertDocKeyToIdEntry :: Ord key => key -> DocId ->-                         SearchIndex key field feature ->-                         SearchIndex key field feature-insertDocKeyToIdEntry dockey !docid si@SearchIndex{docKeyMap} =-    si { docKeyMap = Map.insert dockey docid docKeyMap }--insertDocIdToDocEntry :: (Ix field, Bounded field,-                          Ix feature, Bounded feature) =>-                         DocId -> key ->-                         DocTerms field ->-                         DocFeatureValues feature ->-                         SearchIndex key field feature ->-                         SearchIndex key field feature-insertDocIdToDocEntry !docid dockey userdocterms userdocfeats-                       si@SearchIndex{docIdMap} =-    let doctermids = DocTermIds.create (map (getTermId si) . userdocterms)-        docfeatvals= DocFeatVals.create userdocfeats-        !docinfo   = DocInfo dockey doctermids docfeatvals-     in si { docIdMap  = IntMap.insert (fromEnum docid) docinfo docIdMap }--deleteDocEntry :: Ord key => DocId -> key ->-                  SearchIndex key field feature -> SearchIndex key field feature-deleteDocEntry docid key si@SearchIndex{docIdMap, docKeyMap} =-     si { docIdMap  = IntMap.delete (fromEnum docid) docIdMap-        , docKeyMap = Map.delete key docKeyMap }------- Data.Map utils------- Data.Map does not support prefix lookups directly (unlike a trie)--- but we can implement it reasonably efficiently using split:---- | Lookup values for a range of keys (inclusive lower bound and exclusive--- upper bound)----lookupRange :: Ord k => (k, k) -> Map k v -> [v]-lookupRange (lb, ub) m =-  let (_, mv, gt)  = Map.splitLookup lb m-      (between, _) = Map.split       ub gt-   in case mv of-        Just v  -> v : Map.elems between-        Nothing ->     Map.elems between--lookupPrefix :: Text -> Map Text v -> [v]-lookupPrefix t _ | T.null t = []-lookupPrefix t m = lookupRange (t, prefixUpperBound t) m--prefixUpperBound :: Text -> Text-prefixUpperBound = succLast . T.dropWhileEnd (== maxBound)-  where-    succLast t = T.init t `T.snoc` succ (T.last t)-
− Data/SearchEngine/TermBag.hs 
@@ -1,263 +0,0 @@-{-# LANGUAGE BangPatterns, GeneralizedNewtypeDeriving, MultiParamTypeClasses,-             TypeFamilies #-}-module Data.SearchEngine.TermBag (-    TermId(TermId), TermCount,-    TermBag,-    size,-    fromList,-    toList,-    elems,-    termCount,-    denseTable,-    invariant-  ) where--import qualified Data.Vector.Unboxed         as Vec-import qualified Data.Vector.Unboxed.Mutable as MVec-import qualified Data.Vector.Generic         as GVec-import qualified Data.Vector.Generic.Mutable as GMVec-import Control.Monad.ST-import Control.Monad (liftM)-import qualified Data.Map as Map-import Data.Word (Word32, Word8)-import Data.Bits-import Data.List (sortBy, foldl')-import Data.Function (on)--newtype TermId = TermId { unTermId :: Word32 }-  deriving (Eq, Ord, Show, Enum)--instance Bounded TermId where-  minBound = TermId 0-  maxBound = TermId 0x00FFFFFF--data TermBag = TermBag !Int !(Vec.Vector TermIdAndCount)-  deriving Show---- We sneakily stuff both the TermId and the bag count into one 32bit word-type TermIdAndCount = Word32-type TermCount      = Word8---- Bottom 24 bits is the TermId, top 8 bits is the bag count-termIdAndCount :: TermId -> Int -> TermIdAndCount-termIdAndCount (TermId termid) freq =-      (min (fromIntegral freq) 255 `shiftL` 24)-  .|. (termid .&. 0x00FFFFFF)--getTermId :: TermIdAndCount -> TermId-getTermId word = TermId (word .&. 0x00FFFFFF)--getTermCount :: TermIdAndCount -> TermCount-getTermCount word = fromIntegral (word `shiftR` 24)--invariant :: TermBag -> Bool-invariant (TermBag _ vec) =-    strictlyAscending (Vec.toList vec)-  where-    strictlyAscending (a:xs@(b:_)) = getTermId a < getTermId b-                                  && strictlyAscending xs-    strictlyAscending _  = True--size :: TermBag -> Int-size (TermBag sz _) = sz--elems :: TermBag -> [TermId]-elems (TermBag _ vec) = map getTermId (Vec.toList vec)--toList :: TermBag -> [(TermId, TermCount)]-toList (TermBag _ vec) = [ (getTermId x, getTermCount x)-                         | x <- Vec.toList vec ]--termCount :: TermBag -> TermId -> TermCount-termCount (TermBag _ vec) =-    binarySearch 0 (Vec.length vec - 1)-  where-    binarySearch :: Int -> Int -> TermId -> TermCount-    binarySearch !a !b !key-      | a > b     = 0-      | otherwise =-        let mid         = (a + b) `div` 2-            tidAndCount = vec Vec.! mid-         in case compare key (getTermId tidAndCount) of-              LT -> binarySearch a (mid-1) key-              EQ -> getTermCount tidAndCount-              GT -> binarySearch (mid+1) b key--fromList :: [TermId] -> TermBag-fromList termids =-    let bag = Map.fromListWith (+) [ (t, 1) | t <- termids ]-        sz  = Map.foldl' (+) 0 bag-        vec = Vec.fromListN (Map.size bag)-                            [ termIdAndCount termid freq-                            | (termid, freq) <- Map.toAscList bag ]-     in TermBag sz vec---- | Given a bunch of term bags, merge them into a table for easier subsequent--- processing. This is bascially a sparse to dense conversion. Missing entries--- are filled in with 0. We represent the table as one vector for the--- term ids and a 2d array for the counts.------ Unfortunately vector does not directly support 2d arrays and array does--- not make it easy to trim arrays.----denseTable :: [TermBag] -> (Vec.Vector TermId, Vec.Vector TermCount)-denseTable termbags =-    (tids, tcts)-  where-    -- First merge the TermIds into one array-    -- then make a linear pass to create the counts array-    -- filling in 0s or the counts as we find them-    !numBags   = length termbags-    !tids      = unionsTermId termbags-    !numTerms  = Vec.length tids-    !numCounts = numTerms * numBags-    !tcts      = Vec.create (do-                   out <- MVec.new numCounts-                   sequence_-                     [ writeMergedTermCounts tids bag out i-                     | (n, TermBag _ bag) <- zip [0..] termbags-                     , let i = n * numTerms ]-                   return out-                 )--writeMergedTermCounts :: Vec.Vector TermId -> Vec.Vector TermIdAndCount ->-                         MVec.MVector s TermCount -> Int -> ST s ()-writeMergedTermCounts xs0 ys0 !out i0 =-    -- assume xs & ys are sorted, and ys contains a subset of xs-    go xs0 ys0 i0-  where-    go !xs !ys !i-      | Vec.null ys = MVec.set (MVec.slice i (Vec.length xs) out) 0-      | Vec.null xs = return ()-      | otherwise   = let x   = Vec.head xs-                          ytc = Vec.head ys-                          y   = getTermId ytc-                          c   = getTermCount ytc-                      in case x == y of-                           True  -> do MVec.write out i c-                                       go (Vec.tail xs) (Vec.tail ys) (i+1)-                           False -> do MVec.write out i 0-                                       go (Vec.tail xs)           ys  (i+1)---- | Given a set of term bags, form the set of TermIds----unionsTermId :: [TermBag] -> Vec.Vector TermId-unionsTermId tbs =-    case sortBy (compare `on` bagVecLength) tbs of-      []             -> Vec.empty-      [TermBag _ xs] -> (Vec.map getTermId xs)-      (x0:x1:xs)     -> foldl' union3 (union2 x0 x1) xs-  where-    bagVecLength (TermBag _ vec) = Vec.length vec--union2 :: TermBag -> TermBag -> Vec.Vector TermId-union2 (TermBag _ xs) (TermBag _ ys) =-    Vec.create (MVec.new sizeBound >>= writeMergedUnion2 xs ys)-  where-    sizeBound = Vec.length xs + Vec.length ys--writeMergedUnion2 :: Vec.Vector TermIdAndCount -> Vec.Vector TermIdAndCount ->-                     MVec.MVector s TermId -> ST s (MVec.MVector s TermId)-writeMergedUnion2 xs0 ys0 !out = do-    i <- go xs0 ys0 0-    return $! MVec.take i out-  where-    go !xs !ys !i-      | Vec.null xs = do Vec.copy (MVec.slice i (Vec.length ys) out)-                                  (Vec.map getTermId ys)-                         return (i + Vec.length ys)-      | Vec.null ys = do Vec.copy (MVec.slice i (Vec.length xs) out)-                                  (Vec.map getTermId xs)-                         return (i + Vec.length xs)-      | otherwise   = let x = getTermId (Vec.head xs)-                          y = getTermId (Vec.head ys)-                      in case compare x y of-                          GT -> do MVec.write out i y-                                   go           xs  (Vec.tail ys) (i+1)-                          EQ -> do MVec.write out i x-                                   go (Vec.tail xs) (Vec.tail ys) (i+1)-                          LT -> do MVec.write out i x-                                   go (Vec.tail xs)           ys  (i+1)--union3 :: Vec.Vector TermId -> TermBag -> Vec.Vector TermId-union3 xs (TermBag _ ys) =-    Vec.create (MVec.new sizeBound >>= writeMergedUnion3 xs ys)-  where-    sizeBound = Vec.length xs + Vec.length ys--writeMergedUnion3 :: Vec.Vector TermId -> Vec.Vector TermIdAndCount ->-                     MVec.MVector s TermId -> ST s (MVec.MVector s TermId)-writeMergedUnion3 xs0 ys0 !out = do-    i <- go xs0 ys0 0-    return $! MVec.take i out-  where-    go !xs !ys !i-      | Vec.null xs = do Vec.copy (MVec.slice i (Vec.length ys) out)-                                  (Vec.map getTermId ys)-                         return (i + Vec.length ys)-      | Vec.null ys = do Vec.copy (MVec.slice i (Vec.length xs) out) xs-                         return (i + Vec.length xs)-      | otherwise   = let x =            Vec.head xs-                          y = getTermId (Vec.head ys)-                      in case compare x y of-                          GT -> do MVec.write out i y-                                   go           xs  (Vec.tail ys) (i+1)-                          EQ -> do MVec.write out i x-                                   go (Vec.tail xs) (Vec.tail ys) (i+1)-                          LT -> do MVec.write out i x-                                   go (Vec.tail xs)           ys  (i+1)----------------------------------------------------------------------------------- verbose Unbox instances-----instance MVec.Unbox TermId--newtype instance MVec.MVector s TermId = MV_TermId (MVec.MVector s Word32)--instance GMVec.MVector MVec.MVector TermId where-    basicLength          (MV_TermId v) = GMVec.basicLength v-    basicUnsafeSlice i l (MV_TermId v) = MV_TermId (GMVec.basicUnsafeSlice i l v)-    basicUnsafeNew     l               = MV_TermId `liftM` GMVec.basicUnsafeNew l-    basicInitialize      (MV_TermId v) = GMVec.basicInitialize v-    basicUnsafeReplicate l x           = MV_TermId `liftM` GMVec.basicUnsafeReplicate l (unTermId x)-    basicUnsafeRead  (MV_TermId v) i   = TermId `liftM`    GMVec.basicUnsafeRead v i-    basicUnsafeWrite (MV_TermId v) i x = GMVec.basicUnsafeWrite v i (unTermId x)-    basicClear       (MV_TermId v)     = GMVec.basicClear v-    basicSet         (MV_TermId v) x   = GMVec.basicSet v (unTermId x)-    basicUnsafeGrow  (MV_TermId v) l   = MV_TermId `liftM` GMVec.basicUnsafeGrow v l-    basicUnsafeCopy  (MV_TermId v) (MV_TermId v') = GMVec.basicUnsafeCopy v v'-    basicUnsafeMove  (MV_TermId v) (MV_TermId v') = GMVec.basicUnsafeMove v v'-    basicOverlaps    (MV_TermId v) (MV_TermId v') = GMVec.basicOverlaps   v v'-    {-# INLINE basicLength #-}-    {-# INLINE basicUnsafeSlice #-}-    {-# INLINE basicOverlaps #-}-    {-# INLINE basicUnsafeNew #-}-    {-# INLINE basicInitialize #-}-    {-# INLINE basicUnsafeReplicate #-}-    {-# INLINE basicUnsafeRead #-}-    {-# INLINE basicUnsafeWrite #-}-    {-# INLINE basicClear #-}-    {-# INLINE basicSet #-}-    {-# INLINE basicUnsafeCopy #-}-    {-# INLINE basicUnsafeMove #-}-    {-# INLINE basicUnsafeGrow #-}--newtype instance Vec.Vector TermId = V_TermId (Vec.Vector Word32)--instance GVec.Vector Vec.Vector TermId where-    basicUnsafeFreeze (MV_TermId mv)  = V_TermId  `liftM` GVec.basicUnsafeFreeze mv-    basicUnsafeThaw   (V_TermId  v)   = MV_TermId `liftM` GVec.basicUnsafeThaw v-    basicLength       (V_TermId  v)   = GVec.basicLength v-    basicUnsafeSlice i l (V_TermId v) = V_TermId (GVec.basicUnsafeSlice i l v)-    basicUnsafeIndexM (V_TermId  v) i = TermId `liftM` GVec.basicUnsafeIndexM v i-    basicUnsafeCopy   (MV_TermId mv)-                      (V_TermId  v)   = GVec.basicUnsafeCopy mv v-    elemseq           (V_TermId  v) x = GVec.elemseq v (unTermId x)-    {-# INLINE basicUnsafeFreeze #-}-    {-# INLINE basicUnsafeThaw #-}-    {-# INLINE basicLength #-}-    {-# INLINE basicUnsafeSlice #-}-    {-# INLINE basicUnsafeIndexM #-}-    {-# INLINE basicUnsafeCopy #-}-    {-# INLINE elemseq #-}
− Data/SearchEngine/Types.hs 
@@ -1,124 +0,0 @@-{-# LANGUAGE NamedFieldPuns, RecordWildCards #-}--module Data.SearchEngine.Types (-    -- * Search engine types and helper functions-    SearchEngine(..),-    SearchConfig(..),-    SearchRankParameters(..),-    BM25F.FeatureFunction(..),-    initSearchEngine,-    cacheBM25Context,--    -- ** Helper type for non-term features-    NoFeatures,-    noFeatures,--    -- * Re-export SearchIndex and other types-    SearchIndex, Term, TermId,-    DocIdSet, DocId,-    DocTermIds, DocFeatVals,--    -- * Internal sanity check-    invariant,-  ) where--import Data.SearchEngine.SearchIndex (SearchIndex, Term, TermId)-import qualified Data.SearchEngine.SearchIndex as SI-import Data.SearchEngine.DocIdSet (DocIdSet, DocId)-import qualified Data.SearchEngine.DocIdSet as DocIdSet-import Data.SearchEngine.DocFeatVals (DocFeatVals)-import Data.SearchEngine.DocTermIds  (DocTermIds)-import qualified Data.SearchEngine.BM25F as BM25F--import Data.Ix-import Data.Array.Unboxed----data SearchConfig doc key field feature = SearchConfig {-       documentKey          :: doc -> key,-       extractDocumentTerms :: doc -> field -> [Term],-       transformQueryTerm   :: Term -> field -> Term,-       documentFeatureValue :: doc -> feature -> Float-     }--data SearchRankParameters field feature = SearchRankParameters {-       paramK1                 :: !Float,-       paramB                  :: field -> Float,-       paramFieldWeights       :: field -> Float,-       paramFeatureWeights     :: feature -> Float,-       paramFeatureFunctions   :: feature -> BM25F.FeatureFunction,--       paramResultsetSoftLimit   :: !Int,-       paramResultsetHardLimit   :: !Int,-       paramAutosuggestPrefilterLimit  :: !Int,-       paramAutosuggestPostfilterLimit :: !Int-     }--data SearchEngine doc key field feature = SearchEngine {-       searchIndex      :: !(SearchIndex      key field feature),-       searchConfig     :: !(SearchConfig doc key field feature),-       searchRankParams :: !(SearchRankParameters field feature),--       -- cached info-       sumFieldLengths :: !(UArray field Int),-       bm25Context     :: BM25F.Context TermId field feature-     }--invariant :: (Ord key, Ix field, Bounded field) =>-             SearchEngine doc key field feature -> Bool-invariant SearchEngine{searchIndex} =-    SI.invariant searchIndex--- && check caches--initSearchEngine :: (Ix field, Bounded field, Ix feature, Bounded feature) =>-                    SearchConfig doc key field feature ->-                    SearchRankParameters field feature ->-                    SearchEngine doc key field feature-initSearchEngine config params =-    cacheBM25Context-      SearchEngine {-        searchIndex      = SI.emptySearchIndex,-        searchConfig     = config,-        searchRankParams = params,-        sumFieldLengths  = listArray (minBound, maxBound) (repeat 0),-        bm25Context      = undefined-      }--cacheBM25Context :: Ix field =>-                    SearchEngine doc key field feature ->-                    SearchEngine doc key field feature-cacheBM25Context-    se@SearchEngine {-      searchRankParams = SearchRankParameters{..},-      searchIndex,-      sumFieldLengths-    }-  = se { bm25Context = bm25Context' }-  where-    bm25Context' = BM25F.Context {-      BM25F.numDocsTotal    = SI.docCount searchIndex,-      BM25F.avgFieldLength  = \f -> fromIntegral (sumFieldLengths ! f)-                                  / fromIntegral (SI.docCount searchIndex),-      BM25F.numDocsWithTerm = DocIdSet.size . SI.lookupTermId searchIndex,-      BM25F.paramK1         = paramK1,-      BM25F.paramB          = paramB,-      BM25F.fieldWeight     = paramFieldWeights,-      BM25F.featureWeight   = paramFeatureWeights,-      BM25F.featureFunction = paramFeatureFunctions-    }----------------------------------data NoFeatures = NoFeatures-  deriving (Eq, Ord, Bounded, Show)--instance Ix NoFeatures where-  range   _   = []-  inRange _ _ = False-  index   _ _ = -1--noFeatures :: NoFeatures -> a-noFeatures _ = error "noFeatures"-
− Data/SearchEngine/Update.hs 
@@ -1,90 +0,0 @@-{-# LANGUAGE BangPatterns, NamedFieldPuns, RecordWildCards #-}--module Data.SearchEngine.Update (--    -- * Managing documents to be searched-    insertDoc,-    insertDocs,-    deleteDoc,--  ) where--import Data.SearchEngine.Types-import qualified Data.SearchEngine.SearchIndex as SI-import qualified Data.SearchEngine.DocTermIds as DocTermIds--import Data.Ix-import Data.Array.Unboxed-import Data.List---insertDocs :: (Ord key, Ix field, Bounded field, Ix feature, Bounded feature) =>-              [doc] ->-              SearchEngine doc key field feature ->-              SearchEngine doc key field feature-insertDocs docs se = foldl' (\se' doc -> insertDoc doc se') se docs---insertDoc :: (Ord key, Ix field, Bounded field, Ix feature, Bounded feature) =>-             doc ->-             SearchEngine doc key field feature ->-             SearchEngine doc key field feature-insertDoc doc se@SearchEngine{ searchConfig = SearchConfig {-                                 documentKey,-                                 extractDocumentTerms,-                                 documentFeatureValue-                               }-                             , searchIndex } =-    let key = documentKey doc-        searchIndex' = SI.insertDoc key (extractDocumentTerms doc)-                                        (documentFeatureValue doc)-                                        searchIndex-        oldDoc       = SI.lookupDocKey searchIndex  key-        newDoc       = SI.lookupDocKey searchIndex' key--     in cacheBM25Context $-        updateCachedFieldLengths oldDoc newDoc $-          se { searchIndex = searchIndex' }---deleteDoc :: (Ord key, Ix field, Bounded field) =>-             key ->-             SearchEngine doc key field feature ->-             SearchEngine doc key field feature-deleteDoc key se@SearchEngine{searchIndex} =-    let searchIndex' = SI.deleteDoc key searchIndex-        oldDoc       = SI.lookupDocKey searchIndex key--     in cacheBM25Context $-        updateCachedFieldLengths oldDoc Nothing $-          se { searchIndex = searchIndex' }---updateCachedFieldLengths :: (Ix field, Bounded field) =>-                            Maybe (DocTermIds field) -> Maybe (DocTermIds field) ->-                            SearchEngine doc key field feature ->-                            SearchEngine doc key field feature-updateCachedFieldLengths Nothing (Just newDoc) se@SearchEngine{sumFieldLengths} =-    se {-      sumFieldLengths =-        array (bounds sumFieldLengths)-              [ (i, n + DocTermIds.fieldLength newDoc i)-              | (i, n) <- assocs sumFieldLengths ]-    }-updateCachedFieldLengths (Just oldDoc) (Just newDoc) se@SearchEngine{sumFieldLengths} =-    se {-      sumFieldLengths =-        array (bounds sumFieldLengths)-              [ (i, n - DocTermIds.fieldLength oldDoc i-                      + DocTermIds.fieldLength newDoc i)-              | (i, n) <- assocs sumFieldLengths ]-    }-updateCachedFieldLengths (Just oldDoc) Nothing se@SearchEngine{sumFieldLengths} =-    se {-      sumFieldLengths =-        array (bounds sumFieldLengths)-              [ (i, n - DocTermIds.fieldLength oldDoc i)-              | (i, n) <- assocs sumFieldLengths ]-    }-updateCachedFieldLengths Nothing Nothing se = se-
changelog view
@@ -1,3 +1,7 @@+0.2.2.3 Ben Gamari <ben@well-typed.com> January 2025+	* Introduce compatibility with GHCs up to 9.12+	* Drop compatibility with GHC versions pre-9.2+ 0.2.2.2 Adam Gundry <adam@well-typed.com> March 2023 	* Fix bug in 0.2.2.1 autosuggest patch
demo/ExtractNameTerms.hs view
@@ -93,8 +93,8 @@   case parts t of     []             -> return ()     ts             -> forEach (t:ts) >>= put-   + splitDot :: String -> [String] splitDot = split (dropBlanks $ dropDelims $ whenElt (=='.')) @@ -149,7 +149,7 @@     let mostFreq :: [String]         pkgs     :: [PackageDescription]         (mostFreq, pkgs) = read pkgsFile-        + --    wordsFile <- T.readFile "/usr/share/dict/words" --    let ws = Set.fromList (map T.toLower $ T.lines wordsFile)
full-text-search.cabal view
@@ -1,47 +1,52 @@-name:                full-text-search-version:             0.2.2.2-synopsis:            In-memory full text search engine-description:         An in-memory full text search engine library. It lets you-                     run full-text queries on a collection of your documents.-                     .-                     Features:-                     .-                     * Keyword queries and auto-complete\/auto-suggest queries.-                     .-                     * Can search over any type of \"document\".-                       (You explain how to extract search terms from them.)-                     .-                     * Supports documents with multiple fields-                       (e.g. title, body)-                     .-                     * Supports documents with non-term features-                       (e.g. quality score, page rank)-                     .-                     * Uses the state of the art BM25F ranking function-                     .-                     * Adjustable ranking parameters (including field weights-                       and non-term feature scores)-                     .-                     * In-memory but quite compact. It does not keep a copy of-                       your original documents.-                     .-                     * Quick incremental index updates, making it possible to-                       keep your text search in-sync with your data.-                     .-                     It is independent of the document type, so you have to-                     write the document-specific parts: extracting search terms-                     and any stop words, case-normalisation or stemming. This-                     is quite easy using libraries such as-                     <http://hackage.haskell.org/package/tokenize tokenize> and-                     <http://hackage.haskell.org/package/snowball snowball>.-                     .-                     The source package includes a demo to illustrate how to-                     use the library. The demo is a simplified version of how-                     the library is used in the-                     <http://hackage.haskell.org/package/hackage-server hackage-server>-                     where it provides the backend for the package search feature.+cabal-version: 3.0++name: full-text-search+version: 0.2.2.3+synopsis: In-memory full text search engine++description:+  An in-memory full text search engine library. It lets you+  run full-text queries on a collection of your documents.++  Features:++  * Keyword queries and auto-complete\/auto-suggest queries.++  * Can search over any type of \"document\".+    (You explain how to extract search terms from them.)++  * Supports documents with multiple fields+    (e.g. title, body)++  * Supports documents with non-term features+    (e.g. quality score, page rank)++  * Uses the state of the art BM25F ranking function++  * Adjustable ranking parameters (including field weights+    and non-term feature scores)++  * In-memory but quite compact. It does not keep a copy of+    your original documents.++  * Quick incremental index updates, making it possible to+    keep your text search in-sync with your data.++  It is independent of the document type, so you have to+  write the document-specific parts: extracting search terms+  and any stop words, case-normalisation or stemming. This+  is quite easy using libraries such as+  <https://hackage.haskell.org/package/tokenize tokenize> and+  <https://hackage.haskell.org/package/snowball snowball>.++  The source package includes a demo to illustrate how to+  use the library. The demo is a simplified version of how+  the library is used in the+  <https://hackage.haskell.org/package/hackage-server hackage-server>+  where it provides the backend for the package search feature.+ bug-reports:         https://github.com/well-typed/full-text-search/issues-license:             BSD3+license:             BSD-3-Clause license-file:        LICENSE author:              Duncan Coutts maintainer:          Duncan Coutts <duncan@well-typed.com>,@@ -50,10 +55,9 @@                      2014-2023 IRIS Connect Ltd. category:            Data, Text, Search build-type:          Simple-cabal-version:       >=1.10-extra-source-files:  changelog+extra-doc-files:     changelog -tested-with:         GHC ==8.10.7 || ==9.0.2 || ==9.2.7 || ==9.4.4 || ==9.6.1+tested-with:         GHC ==9.2.8 || ==9.4.8 || ==9.6.6 || ==9.8.2 || ==9.10.1 || ==9.12.1  source-repository head   type:              git@@ -64,7 +68,15 @@   description:       Build a little program illustrating the use of the library   manual:            True +common base-deps+  build-depends:       base       >=4.16 && <4.22,+                       array      >=0.4  && <0.6,+                       vector     >=0.11 && <0.14,+                       containers >=0.4  && <0.8,+                       text       >=0.11 && <2.2+ library+  import:              base-deps   exposed-modules:     Data.SearchEngine,                        Data.SearchEngine.BM25F   other-modules:       Data.SearchEngine.Types,@@ -76,16 +88,12 @@                        Data.SearchEngine.TermBag,                        Data.SearchEngine.DocTermIds,                        Data.SearchEngine.DocIdSet+  hs-source-dirs:      src   other-extensions:    BangPatterns,                        NamedFieldPuns,                        RecordWildCards,                        GeneralizedNewtypeDeriving,                        ScopedTypeVariables-  build-depends:       base       >=4.5  && <4.19,-                       array      >=0.4  && <0.6,-                       vector     >=0.11 && <0.14,-                       containers >=0.4  && <0.7,-                       text       >=0.11 && <2.1   default-language:    Haskell2010   ghc-options:         -Wall -funbox-strict-fields @@ -106,28 +114,32 @@     buildable:         False   else     build-depends:     full-text-search,-                       base, text, containers, array,-                       tokenize      >= 0.1,-                       snowball      == 1.0.*,-                       transformers,-                       split         >= 0.2,-                       Cabal         >= 1.14 && < 3,-                       bytestring, filepath, directory, tar, time, mtl-  build-tools:         alex, happy+                       base,+                       text,+                       containers,+                       array,+                       tokenize      >= 0.1  && <0.4,+                       snowball      >= 1.0  && <1.1,+                       transformers  >= 0.5  && <0.6,+                       split         >= 0.2  && <0.3,+                       Cabal         >= 1.14 && <3.15,+                       bytestring    >= 0.12 && <0.13,+                       filepath      >= 1.5  && <1.6,+                       directory     >= 1.3  && <1.5,+                       tar           >= 0.6  && <0.7,+                       time          >= 1.14 && <1.15,+                       mtl           >= 2.2  && <2.4+  build-tool-depends:  alex:alex, happy:happy   default-language:    Haskell2010   other-extensions:    GeneralizedNewtypeDeriving   ghc-options:         -Wall  test-suite qc-props+  import:              base-deps   type:                exitcode-stdio-1.0   main-is:             Main.hs-  hs-source-dirs:      . tests-  build-depends:       base,-                       array,-                       vector,-                       containers,-                       text,-                       QuickCheck       ==2.*,+  hs-source-dirs:      src, tests+  build-depends:       QuickCheck       ==2.*,                        tasty            >=0.8,                        tasty-quickcheck >=0.8   other-modules:       Test.Data.SearchEngine.TermBag,
+ src/Data/SearchEngine.hs view
@@ -0,0 +1,44 @@+module Data.SearchEngine (++    -- * Basic interface++    -- ** Querying+    Term,+    query,++    -- *** Query auto-completion \/ auto-suggestion+    queryAutosuggest,+    ResultsFilter(..),+    queryAutosuggestPredicate,+    queryAutosuggestMatchingDocuments,++    -- ** Making a search engine instance+    initSearchEngine,+    SearchEngine,+    SearchConfig(..),+    SearchRankParameters(..),+    FeatureFunction(..),++    -- ** Helper type for non-term features+    NoFeatures,+    noFeatures,++    -- ** Managing documents to be searched+    insertDoc,+    insertDocs,+    deleteDoc,++    -- * Explain mode for query result rankings+    queryExplain,+    Explanation(..),+    setRankParams,++    -- * Internal sanity check+    invariant,+  ) where++import Data.SearchEngine.Types+import Data.SearchEngine.Update+import Data.SearchEngine.Query+import Data.SearchEngine.Autosuggest+
+ src/Data/SearchEngine/Autosuggest.hs view
@@ -0,0 +1,573 @@+{-# 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 se resultsFilter precedingTerms partialTerm =++     step_external+   . step_rank+   . step_scoreDs+   . step_scoreTs+   . step_cache+   . step_postfilterlimit+   . step_filter+   . step_prefilterlimit+   . step_process+   $ step_prep+       precedingTerms partialTerm++  where+    -- Construct the auto-suggest query from the query terms+    step_prep pre_ts t = mkAutosuggestQuery se pre_ts 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 (ts, ds, pre_ts) = (ts', ds', tdss', pre_ts)+      where+        (tdss', ts', ds') = processAutosuggestQuery se (ts, ds, 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 args@(_, ds, _, _)+      | withinPrefilterLimit se ds = args+      | otherwise                  = ([], 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 (ts, ds, tdss, pre_ts) = (ts, ds_info, tdss, pre_ts)+      where+        ds_info = filterAutosuggestQuery se resultsFilter 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 args@(_, ds_info, _, _)+      | withinPostfilterLimit se ds_info = args+      | otherwise                        = ([], [], [], [])++    -- For all ds, calculate and cache a couple bits of info needed+    -- later for scoring completion terms and doc results+    step_cache (ts, ds_info, tdss, pre_ts) = (ds_info', tdss)+      where+        ds_info' = cacheDocScoringInfo se ts ds_info pre_ts++    -- Score the completion terms+    step_scoreTs (ds_info, tdss) = (ds_info, tdss, ts_scored)+      where+        ts_scored = scoreAutosuggestQueryCompletions tdss ds_info++    -- Score the doc results (making use of the completion scores)+    step_scoreDs (ds_info, tdss, ts_scored) = (ts_scored, ds_scored)+      where+        ds_scored = scoreAutosuggestQueryResults tdss ds_info ts_scored++    -- Rank the completions and results based on their scores+    step_rank = sortResults++    -- Convert from internal Ids into external forms: Term and doc key+    step_external = convertIdsToExternal 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 se@SearchEngine{searchIndex} precedingTerms partialTerm =+    let (_, _, ds) = processAutosuggestQuery se (mkAutosuggestQuery se precedingTerms partialTerm)+    in map (SI.getDocKey searchIndex) (DocIdSet.toList 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 se@SearchEngine{searchIndex} precedingTerms partialTerm =+    let (_, _, ds) = processAutosuggestQuery se (mkAutosuggestQuery se precedingTerms partialTerm)+    in (\ key -> maybe False (flip DocIdSet.member ds) (SI.lookupDocKeyDocId searchIndex 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 SearchEngine{searchRankParams} ds =+    DocIdSet.size ds <= paramAutosuggestPrefilterLimit searchRankParams++withinPostfilterLimit :: SearchEngine doc key field feature ->+                         [a] -> Bool+withinPostfilterLimit SearchEngine{searchRankParams} ds_info =+    length ds_info <= paramAutosuggestPostfilterLimit searchRankParams+++sortResults :: (Ord av, Ord bv) => ([(a,av)], [(b,bv)]) -> ([(a,av)], [(b,bv)])+sortResults (xs, ys) =+    ( sortBySndDescending xs+    , sortBySndDescending ys )+  where+    sortBySndDescending :: Ord v => [(x,v)] -> [(x,v)]+    sortBySndDescending = sortBy (flip (comparing snd))++convertIdsToExternal :: SearchEngine doc key field feature ->+                        ([(TermId, v)], [(DocId, v)]) -> ([(Term, v)], [(key, v)])+convertIdsToExternal SearchEngine{searchIndex} (termids, docids) =+    ( [ (SI.getTerm   searchIndex termid, s) | (termid, s) <- termids ]+    , [ (SI.getDocKey searchIndex docid,  s) | (docid,  s) <- 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 se@SearchEngine{ searchIndex }+                   precedingTerms partialTerm =+    (completionTerms, precedingDocHits, precedingTerms')+  where+    completionTerms =+      Map.unions+        [ Map.fromList (SI.lookupTermsByPrefix searchIndex partialTerm')+        | partialTerm' <- Query.expandTransformedQueryTerm se partialTerm+        ]++    (precedingTerms', precedingDocHits)+      | null precedingTerms = ([], Nothing)+      | otherwise           = fmap carefulUnions+                                   (lookupRawResults 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 dss+      | null dss  = Just DocIdSet.empty+      | null dss' = Nothing+      | otherwise = Just (DocIdSet.unions dss')+      where+        dss' = filter (withinPrefilterLimit se) dss++    lookupRawResults :: [Term] -> ([TermId], [DocIdSet])+    lookupRawResults ts =+      unzip $ catMaybes+        [ SI.lookupTerm searchIndex t'+        | t  <- ts+        , t' <- Query.expandTransformedQueryTerm se 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 se (completionTerms, precedingDocHits, _)+  -- 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.+  | all (withinPrefilterLimit se) docSets =+    ( completionTermAndDocSets+    , completionTerms'+    , allTermDocSet+    )+  | otherwise = ([], [], 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 =+      [ (t, ds_t')+      | (t, ds_t) <- Map.toList completionTerms+      , let ds_t' = case precedingDocHits of+                      Just ds -> ds `DocIdSet.intersection` ds_t+                      Nothing -> ds_t+      , not (DocIdSet.null ds_t')+      ]++    -- The remaining candidate completions+    completionTerms' :: [TermId]+    docSets :: [DocIdSet]+    (completionTerms', docSets) = unzip completionTermAndDocSets++    -- The union of all these is this set of documents that form the results.+    allTermDocSet :: DocIdSet+    allTermDocSet = DocIdSet.unions docSets+++filterAutosuggestQuery :: SearchEngine doc key field feature ->+                          ResultsFilter key ->+                          DocIdSet ->+                          [(DocId, (key, DocTermIds field, DocFeatVals feature))]+filterAutosuggestQuery SearchEngine{ searchIndex } resultsFilter ds =+    case resultsFilter of+      NoFilter ->+        [ (docid, doc)+        | docid <- DocIdSet.toList ds+        , let doc = SI.lookupDocId searchIndex docid ]++      FilterPredicate predicate ->+        [ (docid, doc)+        | docid <- DocIdSet.toList ds+        , let doc@(k,_,_) = SI.lookupDocId searchIndex docid+        , predicate k ]++      FilterBulkPredicate bulkPredicate ->+        [ (docid, doc)+        | let docids = DocIdSet.toList ds+              docinf = map (SI.lookupDocId searchIndex) docids+              keep   = bulkPredicate [ k | (k,_,_) <- docinf ]+        , (docid, doc, True) <- zip3 docids docinf 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 se completionTerms allTermDocInfo precedingTerms =+    Map.fromList+      [ (docid, (docImportance, termRelevances))+      | (docid, (_dockey, doctermids, docfeatvals)) <- allTermDocInfo+      , let docImportance  = Query.relevanceScore se precedingTerms+                                                  doctermids docfeatvals+            termRelevances = relevanceBreakdown se doctermids docfeatvals+                                                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 SearchEngine{ bm25Context } doctermids docfeatvals ts =+    let -- We'll calculate the bm25 score for each term in this document+        bm25Doc     = Query.indexDocToBM25Doc doctermids docfeatvals++        -- Cache the info that depends only on this doc, not the terms+        termScore   :: (TermId -> (field -> Int) -> Float)+        termScore   = BM25F.scoreTermsBulk bm25Context 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+        (!numTerms, term, count) = DocTermIds.denseTable 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 = Vec.generate numTerms $ \i ->+                       termScore (term i) (\f -> count i f)++        -- We keep only the values for the terms we're interested in+        -- and normalise so we get the relevence fraction+        !scoreSum   = Vec.sum termScores+        !tset       = IntSet.fromList (map fromEnum ts)+     in Map.fromList+      . Vec.toList+      . Vec.map    (\(t,s) -> (t, s/scoreSum))+      . Vec.filter (\(t,_) -> fromEnum t `IntSet.member` tset)+      . Vec.imap   (\i s   -> (term i, s))+      $ termScores+++scoreAutosuggestQueryCompletions :: [(TermId, DocIdSet)]+                                 -> Map.Map DocId (Float, Map.Map TermId Float)+                                 -> [(TermId, Float)]+scoreAutosuggestQueryCompletions completionTermAndDocSets allTermDocInfo =+    [ (t, candidateScore t ds_t)+    | (t, ds_t) <- 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 t ds_t =+      sum [ docImportance * termRelevance+          | Just (docImportance, termRelevances) <-+               map (`Map.lookup` allTermDocInfo) (DocIdSet.toList ds_t)+          , let termRelevance = termRelevances Map.! t+          ]+++scoreAutosuggestQueryResults :: [(TermId, DocIdSet)] ->+                                Map.Map DocId (Float, Map.Map TermId Float) ->+                                [(TermId, Float)] ->+                                [(DocId, Float)]+scoreAutosuggestQueryResults completionTermAndDocSets allTermDocInfo+                             scoredCandidates =+  Map.toList $ Map.fromListWith (+)+    [ (docid, docImportance * score_t)+    | ((_, ds_t), (_, score_t)) <- zip completionTermAndDocSets scoredCandidates+    , let docids  = DocIdSet.toList ds_t+          docinfo = map (`Map.lookup` allTermDocInfo) docids+    , (docid, Just (docImportance, _)) <- zip docids docinfo+    ]+
+ src/Data/SearchEngine/BM25F.hs view
@@ -0,0 +1,253 @@+{-# LANGUAGE RecordWildCards, BangPatterns, ScopedTypeVariables #-}++-- | An implementation of BM25F ranking. See:+--+-- * A quick overview: <http://en.wikipedia.org/wiki/Okapi_BM25>+--+-- * /The Probabilistic Relevance Framework: BM25 and Beyond/+--   <http://www.staff.city.ac.uk/~sbrp622/papers/foundations_bm25_review.pdf>+--+-- * /An Introduction to Information Retrieval/+--   <http://nlp.stanford.edu/IR-book/pdf/irbookonlinereading.pdf>+--+module Data.SearchEngine.BM25F (+    -- * The ranking function+    score,+    Context(..),+    FeatureFunction(..),+    Doc(..),+    -- ** Specialised variants+    scoreTermsBulk,++    -- * Explaining the score+    Explanation(..),+    explain,+  ) where++import Data.Ix+import Data.Array.Unboxed++data Context term field feature = Context {+         numDocsTotal     :: !Int,+         avgFieldLength   :: field -> Float,+         numDocsWithTerm  :: term -> Int,+         paramK1          :: !Float,+         paramB           :: field -> Float,+         -- consider minimum length to prevent massive B bonus?+         fieldWeight      :: field -> Float,+         featureWeight    :: feature -> Float,+         featureFunction  :: feature -> FeatureFunction+       }++data Doc term field feature = Doc {+         docFieldLength        :: field -> Int,+         docFieldTermFrequency :: field -> term -> Int,+         docFeatureValue       :: feature -> Float+       }+++-- | The BM25F score for a document for a given set of terms.+--+score :: (Ix field, Bounded field, Ix feature, Bounded feature) =>+         Context term field feature ->+         Doc term field feature -> [term] -> Float+score ctx doc terms =+    sum (map (weightedTermScore    ctx doc) terms)+  + sum (map (weightedNonTermScore ctx doc) features)++  where+    features = range (minBound, maxBound)+++weightedTermScore :: (Ix field, Bounded field) =>+                     Context term field feature ->+                     Doc term field feature -> term -> Float+weightedTermScore ctx doc t =+    weightIDF ctx t *     tf'+                     / (k1 + tf')+  where+    tf' = weightedDocTermFrequency ctx doc t+    k1  = paramK1 ctx+++weightIDF :: Context term field feature -> term -> Float+weightIDF ctx t =+    log ((n - n_t + 0.5) / (n_t + 0.5))+  where+    n   = fromIntegral (numDocsTotal ctx)+    n_t = fromIntegral (numDocsWithTerm ctx t)+++weightedDocTermFrequency :: (Ix field, Bounded field) =>+                            Context term field feature ->+                            Doc term field feature -> term -> Float+weightedDocTermFrequency ctx doc t =+    sum [ w_f * tf_f / _B_f+        | field <- range (minBound, maxBound)+        , let w_f  = fieldWeight ctx field+              tf_f = fromIntegral (docFieldTermFrequency doc field t)+              _B_f = lengthNorm ctx doc field+        , not (isNaN _B_f)+        ]+    -- When the avgFieldLength is 0 we have a field which is empty for all+    -- documents. Unfortunately it leads to a NaN because the+    -- docFieldTermFrequency will also be 0 so we get 0/0. What we want to+    -- do in this situation is have that field contribute nothing to the+    -- score. The simplest way to achieve that is to skip if _B_f is NaN.+    -- So I think this is fine and not an ugly hack.++lengthNorm :: Context term field feature ->+              Doc term field feature -> field -> Float+lengthNorm ctx doc field =+    (1-b_f) + b_f * sl_f / avgsl_f+  where+    b_f     = paramB ctx field+    sl_f    = fromIntegral (docFieldLength doc field)+    avgsl_f = avgFieldLength ctx field+++weightedNonTermScore :: (Ix feature, Bounded feature) =>+                        Context term field feature ->+                        Doc term field feature -> feature -> Float+weightedNonTermScore ctx doc feature =+    w_f * _V_f f_f+  where+    w_f  = featureWeight ctx feature+    _V_f = applyFeatureFunction (featureFunction ctx feature)+    f_f  = docFeatureValue doc feature+++data FeatureFunction+   = LogarithmicFunction   Float -- ^ @log (\lambda_i + f_i)@+   | RationalFunction      Float -- ^ @f_i / (\lambda_i + f_i)@+   | SigmoidFunction Float Float -- ^ @1 / (\lambda + exp(-(\lambda' * f_i))@++applyFeatureFunction :: FeatureFunction -> (Float -> Float)+applyFeatureFunction (LogarithmicFunction p1) = \fi -> log (p1 + fi)+applyFeatureFunction (RationalFunction    p1) = \fi -> fi / (p1 + fi)+applyFeatureFunction (SigmoidFunction  p1 p2) = \fi -> 1 / (p1 + exp (-fi * p2))+++-----------------------------+-- Bulk scoring of many terms+--++-- | Most of the time we want to score several different documents for the same+-- set of terms, but sometimes we want to score one document for many terms+-- and in that case we can save a bit of work by doing it in bulk. It lets us+-- calculate once and share things that depend only on the document, and not+-- the term.+--+-- To take advantage of the sharing you must partially apply and name the+-- per-doc score functon, e.g.+--+-- > let score :: term -> (field -> Int) -> Float+-- >     score = BM25.bulkScorer ctx doc+-- >  in sum [ score t (\f -> counts ! (t, f)) | t <- ts ]+--+scoreTermsBulk :: forall field term feature. (Ix field, Bounded field) =>+                  Context term field feature ->+                  Doc term field feature ->+                  (term -> (field -> Int) -> Float)+scoreTermsBulk ctx doc = +    -- This is just a rearrangement of weightedTermScore and+    -- weightedDocTermFrequency above, with the doc-constant bits hoisted out.++    \t tFreq ->+    let !tf' = sum [ w!f * tf_f / _B!f+                   | f <- range (minBound, maxBound)+                   , let tf_f = fromIntegral (tFreq f)+                         _B_f = _B!f+                   , not (isNaN _B_f)+                   ]++     in weightIDF ctx t *     tf'+                         / (k1 + tf')+  where+    -- So long as the caller does the partial application thing then these+    -- values can all be shared between many calls with different terms.++    !k1 = paramK1 ctx+    w, _B :: UArray field Float+    !w  = array (minBound, maxBound)+                [ (field, fieldWeight ctx field)+                | field <- range (minBound, maxBound) ]+    !_B = array (minBound, maxBound)+                [ (field, lengthNorm ctx doc field)+                | field <- range (minBound, maxBound) ]+++------------------+-- Explanation+--++-- | A breakdown of the BM25F score, to explain somewhat how it relates to+-- the inputs, and so you can compare the scores of different documents.+--+data Explanation field feature term = Explanation {+       -- | The overall score is the sum of the 'termScores', 'positionScore'+       -- and 'nonTermScore'+       overallScore  :: Float,++       -- | There is a score contribution from each query term. This is the+       -- score for the term across all fields in the document (but see+       -- 'termFieldScores').+       termScores    :: [(term, Float)],+{-+       -- | There is a score contribution for positional information. Terms+       -- appearing in the document close together give a bonus.+       positionScore :: [(field, Float)],+-}+       -- | The document can have an inate bonus score independent of the terms+       -- in the query. For example this might be a popularity score.+       nonTermScores :: [(feature, Float)],++       -- | This does /not/ contribute to the 'overallScore'. It is an+       -- indication of how the 'termScores' relates to per-field scores.+       -- Note however that the term score for all fields is /not/ simply+       -- sum of the per-field scores. The point of the BM25F scoring function+       -- is that a linear combination of per-field scores is wrong, and BM25F+       -- does a more cunning non-linear combination.+       --+       -- However, it is still useful as an indication to see scores for each+       -- field for a term, to see how the compare.+       --+       termFieldScores :: [(term, [(field, Float)])]+     }+  deriving Show++instance Functor (Explanation field feature) where+  fmap f e@Explanation{..} =+    e {+      termScores      = [ (f t, s)  | (t, s)  <- termScores ],+      termFieldScores = [ (f t, fs) | (t, fs) <- termFieldScores ]+    }++explain :: (Ix field, Bounded field, Ix feature, Bounded feature) =>+           Context term field feature ->+           Doc term field feature -> [term] -> Explanation field feature term+explain ctx doc ts =+    Explanation {..}+  where+    overallScore  = sum (map snd termScores)+--                  + sum (map snd positionScore)+                  + sum (map snd nonTermScores)+    termScores    = [ (t, weightedTermScore ctx doc t) | t <- ts ]+--    positionScore = [ (f, 0) | f <- range (minBound, maxBound) ]+    nonTermScores = [ (feature, weightedNonTermScore ctx doc feature)+                    | feature <- range (minBound, maxBound) ]++    termFieldScores =+      [ (t, fieldScores)+      | t <- ts+      , let fieldScores =+              [ (f, weightedTermScore ctx' doc t)+              | f <- range (minBound, maxBound)+              , let ctx' = ctx { fieldWeight = fieldWeightOnly f }+              ]+      ]+    fieldWeightOnly f f' | sameField f f' = fieldWeight ctx f'+                         | otherwise      = 0++    sameField f f' = index (minBound, maxBound) f+                  == index (minBound, maxBound) f'
+ src/Data/SearchEngine/DocFeatVals.hs view
@@ -0,0 +1,25 @@+{-# LANGUAGE BangPatterns, GeneralizedNewtypeDeriving #-}+module Data.SearchEngine.DocFeatVals (+    DocFeatVals,+    featureValue,+    create,+  ) where++import Data.SearchEngine.DocTermIds (vecIndexIx, vecCreateIx)+import Data.Vector (Vector)+import Data.Ix (Ix)+++-- | Storage for the non-term feature values i a document.+--+newtype DocFeatVals feature = DocFeatVals (Vector Float)+  deriving (Show)++featureValue :: (Ix feature, Bounded feature) => DocFeatVals feature -> feature -> Float+featureValue (DocFeatVals featVec) = vecIndexIx featVec++create :: (Ix feature, Bounded feature) =>+          (feature -> Float) -> DocFeatVals feature+create docFeatVals =+    DocFeatVals (vecCreateIx docFeatVals)+
+ src/Data/SearchEngine/DocIdSet.hs view
@@ -0,0 +1,207 @@+{-# LANGUAGE BangPatterns, GeneralizedNewtypeDeriving, MultiParamTypeClasses,+             TypeFamilies #-}+module Data.SearchEngine.DocIdSet (+    DocId(DocId),+    DocIdSet(..),+    null,+    size,+    empty,+    singleton,+    fromList,+    toList,+    insert,+    delete,+    member,+    union,+    unions,+    intersection,+    invariant,+  ) where++import Data.Word+import qualified Data.Vector.Unboxed         as Vec+import qualified Data.Vector.Unboxed.Mutable as MVec+import qualified Data.Vector.Generic         as GVec+import qualified Data.Vector.Generic.Mutable as GMVec+import Control.Monad.ST+import Control.Monad (liftM)+import qualified Data.Set as Set+import qualified Data.List as List+import Data.Function (on)++import Prelude hiding (null)+++newtype DocId = DocId { unDocId :: Word32 }+  deriving (Eq, Ord, Show, Enum, Bounded)++newtype DocIdSet = DocIdSet (Vec.Vector DocId)+  deriving (Eq, Show)++-- represented as a sorted sequence of ids+invariant :: DocIdSet -> Bool+invariant (DocIdSet vec) =+    strictlyAscending (Vec.toList vec)+  where+    strictlyAscending (a:xs@(b:_)) = a < b && strictlyAscending xs+    strictlyAscending _  = True+++size :: DocIdSet -> Int+size (DocIdSet vec) = Vec.length vec++null :: DocIdSet -> Bool+null (DocIdSet vec) = Vec.null vec++empty :: DocIdSet+empty = DocIdSet Vec.empty++singleton :: DocId -> DocIdSet+singleton = DocIdSet . Vec.singleton++fromList :: [DocId] -> DocIdSet+fromList = DocIdSet . Vec.fromList . Set.toAscList . Set.fromList++toList ::  DocIdSet -> [DocId]+toList (DocIdSet vec) = Vec.toList vec++insert :: DocId -> DocIdSet -> DocIdSet+insert x (DocIdSet vec) =+    case binarySearch vec 0 (Vec.length vec - 1) x of+      (_, True)  -> DocIdSet vec+      (i, False) -> case Vec.splitAt i vec of+                      (before, after) ->+                        DocIdSet (Vec.concat [before, Vec.singleton x, after])++delete :: DocId -> DocIdSet -> DocIdSet+delete x (DocIdSet vec) =+    case binarySearch vec 0 (Vec.length vec - 1) x of+      (_, False) -> DocIdSet vec+      (i, True)  -> case Vec.splitAt i vec of+                      (before, after) ->+                        DocIdSet (before Vec.++ Vec.tail after)++member :: DocId -> DocIdSet -> Bool+member x (DocIdSet vec) = snd (binarySearch vec 0 (Vec.length vec - 1) x)++binarySearch :: Vec.Vector DocId -> Int -> Int -> DocId -> (Int, Bool)+binarySearch vec !a !b !key+  | a > b     = (a, False)+  | otherwise =+    let mid = (a + b) `div` 2+     in case compare key (vec Vec.! mid) of+          LT -> binarySearch vec a (mid-1) key+          EQ -> (mid, True)+          GT -> binarySearch vec (mid+1) b key++unions :: [DocIdSet] -> DocIdSet+unions = List.foldl' union empty+         -- a bit more effecient if we merge small ones first+       . List.sortBy (compare `on` size)++union :: DocIdSet -> DocIdSet -> DocIdSet+union x y | null x = y+          | null y = x+union (DocIdSet xs) (DocIdSet ys) =+    DocIdSet (Vec.create (MVec.new sizeBound >>= writeMergedUnion xs ys))+  where+    sizeBound = Vec.length xs + Vec.length ys++writeMergedUnion :: Vec.Vector DocId -> Vec.Vector DocId ->+                    MVec.MVector s DocId -> ST s (MVec.MVector s DocId)+writeMergedUnion xs0 ys0 !out = do+    i <- go xs0 ys0 0+    return $! MVec.take i out+  where+    go !xs !ys !i+      | Vec.null xs = do Vec.copy (MVec.slice i (Vec.length ys) out) ys+                         return (i + Vec.length ys)+      | Vec.null ys = do Vec.copy (MVec.slice i (Vec.length xs) out) xs+                         return (i + Vec.length xs)+      | otherwise   = let x = Vec.head xs; y = Vec.head ys+                      in case compare x y of+                          GT -> do MVec.write out i y+                                   go           xs  (Vec.tail ys) (i+1)+                          EQ -> do MVec.write out i x+                                   go (Vec.tail xs) (Vec.tail ys) (i+1)+                          LT -> do MVec.write out i x+                                   go (Vec.tail xs)           ys  (i+1)++intersection :: DocIdSet -> DocIdSet -> DocIdSet+intersection x y | null x = empty+                 | null y = empty+intersection (DocIdSet xs) (DocIdSet ys) =+    DocIdSet (Vec.create (MVec.new sizeBound >>= writeMergedIntersection xs ys))+  where+    sizeBound = max (Vec.length xs) (Vec.length ys)++writeMergedIntersection :: Vec.Vector DocId -> Vec.Vector DocId ->+                           MVec.MVector s DocId -> ST s (MVec.MVector s DocId)+writeMergedIntersection xs0 ys0 !out = do+    i <- go xs0 ys0 0+    return $! MVec.take i out+  where+    go !xs !ys !i+      | Vec.null xs = return i+      | Vec.null ys = return i+      | otherwise   = let x = Vec.head xs; y = Vec.head ys+                      in case compare x y of+                          GT ->    go           xs  (Vec.tail ys)  i+                          EQ -> do MVec.write out i x+                                   go (Vec.tail xs) (Vec.tail ys) (i+1)+                          LT ->    go (Vec.tail xs)           ys   i++------------------------------------------------------------------------------+-- verbose Unbox instances+--++instance MVec.Unbox DocId++newtype instance MVec.MVector s DocId = MV_DocId (MVec.MVector s Word32)++instance GMVec.MVector MVec.MVector DocId where+    basicLength          (MV_DocId v) = GMVec.basicLength v+    basicUnsafeSlice i l (MV_DocId v) = MV_DocId (GMVec.basicUnsafeSlice i l v)+    basicUnsafeNew     l              = MV_DocId `liftM` GMVec.basicUnsafeNew l+    basicInitialize      (MV_DocId v) = GMVec.basicInitialize v+    basicUnsafeReplicate l x          = MV_DocId `liftM` GMVec.basicUnsafeReplicate l (unDocId x)+    basicUnsafeRead  (MV_DocId v) i   = DocId `liftM`    GMVec.basicUnsafeRead v i+    basicUnsafeWrite (MV_DocId v) i x = GMVec.basicUnsafeWrite v i (unDocId x)+    basicClear       (MV_DocId v)     = GMVec.basicClear v+    basicSet         (MV_DocId v) x   = GMVec.basicSet v (unDocId x)+    basicUnsafeGrow  (MV_DocId v) l   = MV_DocId `liftM` GMVec.basicUnsafeGrow v l+    basicUnsafeCopy  (MV_DocId v) (MV_DocId v') = GMVec.basicUnsafeCopy v v'+    basicUnsafeMove  (MV_DocId v) (MV_DocId v') = GMVec.basicUnsafeMove v v'+    basicOverlaps    (MV_DocId v) (MV_DocId v') = GMVec.basicOverlaps   v v'+    {-# INLINE basicLength #-}+    {-# INLINE basicUnsafeSlice #-}+    {-# INLINE basicOverlaps #-}+    {-# INLINE basicUnsafeNew #-}+    {-# INLINE basicInitialize #-}+    {-# INLINE basicUnsafeReplicate #-}+    {-# INLINE basicUnsafeRead #-}+    {-# INLINE basicUnsafeWrite #-}+    {-# INLINE basicClear #-}+    {-# INLINE basicSet #-}+    {-# INLINE basicUnsafeCopy #-}+    {-# INLINE basicUnsafeMove #-}+    {-# INLINE basicUnsafeGrow #-}++newtype instance Vec.Vector DocId = V_DocId (Vec.Vector Word32)++instance GVec.Vector Vec.Vector DocId where+    basicUnsafeFreeze (MV_DocId mv)  = V_DocId  `liftM` GVec.basicUnsafeFreeze mv+    basicUnsafeThaw   (V_DocId  v)   = MV_DocId `liftM` GVec.basicUnsafeThaw v+    basicLength       (V_DocId  v)   = GVec.basicLength v+    basicUnsafeSlice i l (V_DocId v) = V_DocId (GVec.basicUnsafeSlice i l v)+    basicUnsafeIndexM (V_DocId  v) i = DocId `liftM` GVec.basicUnsafeIndexM v i+    basicUnsafeCopy   (MV_DocId mv)+                      (V_DocId  v)   = GVec.basicUnsafeCopy mv v+    elemseq           (V_DocId  v) x = GVec.elemseq v (unDocId x)+    {-# INLINE basicUnsafeFreeze #-}+    {-# INLINE basicUnsafeThaw #-}+    {-# INLINE basicLength #-}+    {-# INLINE basicUnsafeSlice #-}+    {-# INLINE basicUnsafeIndexM #-}+    {-# INLINE basicUnsafeCopy #-}+    {-# INLINE elemseq #-}
+ src/Data/SearchEngine/DocTermIds.hs view
@@ -0,0 +1,81 @@+{-# LANGUAGE BangPatterns, GeneralizedNewtypeDeriving #-}+module Data.SearchEngine.DocTermIds (+    DocTermIds,+    TermId,+    fieldLength,+    fieldTermCount,+    fieldElems,+    create,+    denseTable,+    vecIndexIx,+    vecCreateIx,+  ) where++import Data.SearchEngine.TermBag (TermBag, TermId)+import qualified Data.SearchEngine.TermBag as TermBag++import Data.Vector (Vector, (!))+import qualified Data.Vector as Vec+import qualified Data.Vector.Unboxed as UVec+import Data.Ix (Ix)+import qualified Data.Ix as Ix+++-- | The 'TermId's for the 'Term's that occur in a document. Documents may have+-- multiple fields and the 'DocTerms' type holds them separately for each field.+--+newtype DocTermIds field = DocTermIds (Vector TermBag)+  deriving (Show)++getField :: (Ix field, Bounded field) => DocTermIds field -> field -> TermBag+getField (DocTermIds fieldVec) = vecIndexIx fieldVec++create :: (Ix field, Bounded field) =>+          (field -> [TermId]) -> DocTermIds field+create docTermIds =+    DocTermIds (vecCreateIx (TermBag.fromList . docTermIds))++-- | The number of terms in a field within the document.+fieldLength :: (Ix field, Bounded field) => DocTermIds field -> field -> Int+fieldLength docterms field =+    TermBag.size (getField docterms field)++-- | /O(log n)/ The frequency of a particular term in a field within the document.+--+fieldTermCount :: (Ix field, Bounded field) =>+                  DocTermIds field -> field -> TermId -> Int+fieldTermCount docterms field termid =+    fromIntegral (TermBag.termCount (getField docterms field) termid)++fieldElems :: (Ix field, Bounded field) => DocTermIds field -> field -> [TermId]+fieldElems docterms field =+    TermBag.elems (getField docterms field)++-- | The 'DocTermIds' is really a sparse 2d array, and doing lookups with+-- 'fieldTermCount' has a O(log n) cost. This function converts to a dense+-- tabular representation which then enables linear scans.+--+denseTable :: (Ix field, Bounded field) => DocTermIds field ->+              (Int, Int -> TermId, Int -> field -> Int)+denseTable (DocTermIds fieldVec) =+    let (!termids, !termcounts) = TermBag.denseTable (Vec.toList fieldVec)+        !numTerms = UVec.length termids+     in ( numTerms+        , \i    -> termids UVec.! i+        , \i ix -> let j = Ix.index (minBound, maxBound) ix+                    in fromIntegral (termcounts UVec.! (j * numTerms + i))+        )++---------------------------------+-- Vector indexed by Ix Bounded+--++vecIndexIx  :: (Ix ix, Bounded ix) => Vector a -> ix -> a+vecIndexIx vec ix = vec ! Ix.index (minBound, maxBound) ix++vecCreateIx :: (Ix ix, Bounded ix) => (ix -> a) -> Vector a+vecCreateIx f = Vec.fromListN (Ix.rangeSize bounds)+                  [ y | ix <- Ix.range bounds, let !y = f ix ]+  where+    bounds = (minBound, maxBound)+
+ src/Data/SearchEngine/Query.hs view
@@ -0,0 +1,257 @@+{-# LANGUAGE BangPatterns, NamedFieldPuns, RecordWildCards #-}++module Data.SearchEngine.Query (++    -- * Querying+    query,+    ResultsFilter(..),++    -- * Explain mode for query result rankings+    queryExplain,+    BM25F.Explanation(..),+    setRankParams,++    -- ** Utils used by autosuggest+    relevanceScore,+    indexDocToBM25Doc,+    expandTransformedQueryTerm,+  ) where++import Data.SearchEngine.Types+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.DocFeatVals as DocFeatVals+import qualified Data.SearchEngine.BM25F       as BM25F++import Data.Ix+import Data.List+import Data.Function+import Data.Maybe+++-- | Execute a normal query. Find the documents in which one or more of+-- the search terms appear and return them in ranked order.+--+-- The number of documents returned is limited by the 'paramResultsetSoftLimit'+-- and 'paramResultsetHardLimit' paramaters. This also limits the cost of the+-- query (which is primarily the cost of scoring each document).+--+-- The given terms are all assumed to be complete (as opposed to prefixes+-- like with 'queryAutosuggest').+--+query :: (Ix field, Bounded field, Ix feature, Bounded feature) =>+         SearchEngine doc key field feature ->+         [Term] -> [key]+query se@SearchEngine{ searchIndex,+                       searchRankParams = SearchRankParameters{..} }+      terms =++  let -- Start by transforming/normalising all the query terms.+      -- This can be done differently for each field we search by.+      lookupTerms :: [Term]+      lookupTerms = concatMap (expandTransformedQueryTerm se) terms++      -- Then we look up all the normalised terms in the index.+      rawresults :: [Maybe (TermId, DocIdSet)]+      rawresults = map (SI.lookupTerm searchIndex) lookupTerms++      -- For the terms that occur in the index, this gives us the term's id+      -- and the set of documents that the term occurs in.+      termids   :: [TermId]+      docidsets :: [DocIdSet]+      (termids, docidsets) = unzip (catMaybes rawresults)++      -- We looked up the documents that *any* of the term occur in (not all)+      -- so this could be rather a lot of docs if the user uses a few common+      -- terms. Scoring these result docs is a non-trivial cost so we want to+      -- limit the number that we have to score. The standard trick is to+      -- consider the doc sets in the order of size, smallest to biggest. Once+      -- we have gone over a certain threshold of docs then don't bother with+      -- the doc sets for the remaining terms. This tends to work because the+      -- scoring gives lower weight to terms that occur in many documents.+      unrankedResults :: DocIdSet+      unrankedResults = pruneRelevantResults+                          paramResultsetSoftLimit+                          paramResultsetHardLimit+                          docidsets++      --TODO: technically this isn't quite correct. Because each field can+      -- be normalised differently, we can end up with different termids for+      -- the same original search term, and then we score those as if they+      -- were different terms, which makes a difference when the term appears+      -- in multiple fields (exactly the case BM25F is supposed to deal with).+      -- What we ought to have instead is an Array (Int, field) TermId, and+      -- make the scoring use the appropriate termid for each field, but to+      -- consider them the "same" term.+   in rankResults se termids (DocIdSet.toList unrankedResults)++-- | Before looking up a term in the main index we need to normalise it+-- using the 'transformQueryTerm'. Of course the transform can be different+-- for different fields, so we have to collect all the forms (eliminating+-- duplicates).+--+expandTransformedQueryTerm :: (Ix field, Bounded field) =>+                              SearchEngine doc key field feature ->+                              Term -> [Term]+expandTransformedQueryTerm SearchEngine{searchConfig} term =+    nub [ transformForField field+        | let transformForField = transformQueryTerm searchConfig term+        , field <- range (minBound, maxBound) ]+++rankResults :: (Ix field, Bounded field, Ix feature, Bounded feature) =>+               SearchEngine doc key field feature ->+               [TermId] -> [DocId] -> [key]+rankResults se@SearchEngine{searchIndex} queryTerms docids =+    map snd+  $ sortBy (flip compare `on` fst)+      [ (relevanceScore se queryTerms doctermids docfeatvals, dockey)+      | docid <- docids+      , let (dockey, doctermids, docfeatvals) = SI.lookupDocId searchIndex docid ]++relevanceScore :: (Ix field, Bounded field, Ix feature, Bounded feature) =>+                  SearchEngine doc key field feature ->+                  [TermId] -> DocTermIds field -> DocFeatVals feature -> Float+relevanceScore SearchEngine{bm25Context} queryTerms doctermids docfeatvals =+    BM25F.score bm25Context doc queryTerms+  where+    doc = indexDocToBM25Doc doctermids docfeatvals++indexDocToBM25Doc :: (Ix field, Bounded field, Ix feature, Bounded feature) =>+                     DocTermIds field ->+                     DocFeatVals feature ->+                     BM25F.Doc TermId field feature+indexDocToBM25Doc doctermids docfeatvals =+    BM25F.Doc {+      BM25F.docFieldLength        = DocTermIds.fieldLength    doctermids,+      BM25F.docFieldTermFrequency = DocTermIds.fieldTermCount doctermids,+      BM25F.docFeatureValue       = DocFeatVals.featureValue docfeatvals+    }++pruneRelevantResults :: Int -> Int -> [DocIdSet] -> DocIdSet+pruneRelevantResults softLimit hardLimit =+    -- Look at the docsets starting with the smallest ones. Smaller docsets+    -- correspond to the rarer terms, which are the ones that score most highly.+    go DocIdSet.empty . sortBy (compare `on` DocIdSet.size)+  where+    go !acc [] = acc+    go !acc (d:ds)+        -- If this is the first one, we add it anyway, otherwise we're in+        -- danger of returning no results at all.+      | DocIdSet.null acc = go d ds+        -- We consider the size our docset would be if we add this extra one...+        -- If it puts us over the hard limit then stop.+      | size > hardLimit  = acc+        -- If it puts us over soft limit then we add it and stop+      | size > softLimit  = DocIdSet.union acc d+        -- Otherwise we can add it and carry on to consider the remainder+      | otherwise         = go (DocIdSet.union acc d) ds+      where+        size = DocIdSet.size acc + DocIdSet.size d+++--------------------------------+-- Normal query with explanation+--++queryExplain :: (Ix field, Bounded field, Ix feature, Bounded feature) =>+                SearchEngine doc key field feature ->+                [Term] -> [(BM25F.Explanation field feature Term, key)]+queryExplain se@SearchEngine{ searchIndex,+                              searchConfig     = SearchConfig{transformQueryTerm},+                              searchRankParams = SearchRankParameters{..} }+      terms =++  -- See 'query' above for explanation. Really we ought to combine them.+  let lookupTerms :: [Term]+      lookupTerms = [ term'+                    | term  <- terms+                    , let transformForField = transformQueryTerm term+                    , term' <- nub [ transformForField field+                                   | field <- range (minBound, maxBound) ]+                    ]++      rawresults :: [Maybe (TermId, DocIdSet)]+      rawresults = map (SI.lookupTerm searchIndex) lookupTerms++      termids   :: [TermId]+      docidsets :: [DocIdSet]+      (termids, docidsets) = unzip (catMaybes rawresults)++      unrankedResults :: DocIdSet+      unrankedResults = pruneRelevantResults+                          paramResultsetSoftLimit+                          paramResultsetHardLimit+                          docidsets++   in rankExplainResults se termids (DocIdSet.toList unrankedResults)++rankExplainResults :: (Ix field, Bounded field, Ix feature, Bounded feature) =>+                      SearchEngine doc key field feature ->+                      [TermId] ->+                      [DocId] ->+                      [(BM25F.Explanation field feature Term, key)]+rankExplainResults se@SearchEngine{searchIndex} queryTerms docids =+    sortBy (flip compare `on` (BM25F.overallScore . fst))+      [ (explainRelevanceScore se queryTerms doctermids docfeatvals, dockey)+      | docid <- docids+      , let (dockey, doctermids, docfeatvals) = SI.lookupDocId searchIndex docid ]+++explainRelevanceScore :: (Ix field, Bounded field, Ix feature, Bounded feature) =>+                         SearchEngine doc key field feature ->+                         [TermId] ->+                         DocTermIds field ->+                         DocFeatVals feature ->+                         BM25F.Explanation field feature Term+explainRelevanceScore SearchEngine{bm25Context, searchIndex}+                      queryTerms doctermids docfeatvals =+    fmap (SI.getTerm searchIndex) (BM25F.explain bm25Context doc queryTerms)+  where+    doc = indexDocToBM25Doc doctermids docfeatvals+++setRankParams :: SearchRankParameters field feature ->+                 SearchEngine doc key field feature ->+                 SearchEngine doc key field feature+setRankParams params@SearchRankParameters{..} se =+    se {+      searchRankParams = params,+      bm25Context      = (bm25Context se) {+        BM25F.paramK1         = paramK1,+        BM25F.paramB          = paramB,+        BM25F.fieldWeight     = paramFieldWeights,+        BM25F.featureWeight   = paramFeatureWeights,+        BM25F.featureFunction = paramFeatureFunctions+      }+    }+++--------------------------------+-- Results filter+--++-- | In some applications it is necessary to enforce some security or+-- visibility rule about the query results (e.g. in a typical DB-based+-- application different users can see different data items). Typically+-- it would be too expensive to build different search indexes for the+-- different contexts and so the strategy is to use one index containing+-- everything and filter for visibility in the results. This means the+-- filter condition is different for different queries (e.g. performed+-- on behalf of different users).+--+-- Filtering the results after a query is possible but not the most efficient+-- thing to do because we've had to score all the not-visible documents.+-- The better thing to do is to filter as part of the query, this way we can+-- filter before the expensive scoring.+--+-- We provide one further optimisation: bulk predicates. In some applications+-- it can be quicker to check the security\/visibility of a whole bunch of+-- results all in one go.+--+data ResultsFilter key = NoFilter+                       | FilterPredicate     (key -> Bool)+                       | FilterBulkPredicate ([key] -> [Bool])+--TODO: allow filtering & non-feature score lookup in one bulk op+
+ src/Data/SearchEngine/SearchIndex.hs view
@@ -0,0 +1,456 @@+{-# LANGUAGE BangPatterns, NamedFieldPuns #-}++module Data.SearchEngine.SearchIndex (+    SearchIndex,+    Term,+    TermId,+    DocId,++    emptySearchIndex,+    insertDoc,+    deleteDoc,++    docCount,+    lookupTerm,+    lookupTermsByPrefix,+    lookupTermId,+    lookupDocId,+    lookupDocKey,+    lookupDocKeyDocId,++    getTerm,+    getDocKey,++    invariant,+  ) where++import Data.SearchEngine.DocIdSet (DocIdSet, DocId)+import qualified Data.SearchEngine.DocIdSet as DocIdSet+import Data.SearchEngine.DocTermIds (DocTermIds, TermId, vecIndexIx, vecCreateIx)+import qualified Data.SearchEngine.DocTermIds as DocTermIds+import Data.SearchEngine.DocFeatVals (DocFeatVals)+import qualified Data.SearchEngine.DocFeatVals as DocFeatVals++import Data.Ix (Ix)+import qualified Data.Ix as Ix+import Data.Map (Map)+import qualified Data.Map as Map+import Data.IntMap (IntMap)+import qualified Data.IntMap as IntMap+import qualified Data.Set as Set+import Data.Text (Text)+import qualified Data.Text as T+import qualified Data.List as List++import Control.Exception (assert)++-- | Terms are short strings, usually whole words.+--+type Term = Text++-- | The search index is essentially a many-to-many mapping between documents+-- and terms. Each document contains many terms and each term occurs in many+-- documents. It is a bidirectional mapping as we need to support lookups in+-- both directions.+--+-- Documents are identified by a key (in Ord) while terms are text values.+-- Inside the index however we assign compact numeric ids to both documents and+-- terms. The advantage of this is a much more compact in-memory representation+-- and the disadvantage is greater complexity. In particular it means we have+-- to manage bidirectional mappings between document keys and ids, and between+-- terms and term ids.+--+-- So the mappings we maintain can be depicted as:+--+-- >  Term   <-- 1:1 -->   TermId+-- >          \              ^+-- >           \             |+-- >           1:many    many:many+-- >                \        |+-- >                 \->     v+-- > DocKey  <-- 1:1 -->   DocId+--+-- For efficiency, these details are exposed in the interface. In particular+-- the mapping from TermId to many DocIds is exposed via a 'DocIdSet',+-- and the mapping from DocIds to TermIds is exposed via 'DocTermIds'.+--+-- The main reason we need to keep the DocId -> TermId is to allow for+-- efficient incremental updates.+--+data SearchIndex key field feature = SearchIndex {+       -- the indexes+       termMap           :: !(Map Term TermInfo),+       termIdMap         :: !(IntMap TermIdInfo),+       docIdMap          :: !(IntMap (DocInfo key field feature)),+       docKeyMap         :: !(Map key DocId),++       -- auto-increment key counters+       nextTermId        :: TermId,+       nextDocId         :: DocId+     }+  deriving Show++data TermInfo = TermInfo !TermId !DocIdSet+  deriving Show++data TermIdInfo = TermIdInfo !Term !DocIdSet+  deriving (Show, Eq)++data DocInfo key field feature = DocInfo !key !(DocTermIds field)+                                              !(DocFeatVals feature)+  deriving Show+++-----------------------+-- SearchIndex basics+--++emptySearchIndex :: SearchIndex key field feature+emptySearchIndex =+    SearchIndex+      Map.empty+      IntMap.empty+      IntMap.empty+      Map.empty+      minBound+      minBound++checkInvariant :: (Ord key, Ix field, Bounded field) =>+                  SearchIndex key field feature -> SearchIndex key field feature+checkInvariant si = assert (invariant si) si++invariant :: (Ord key, Ix field, Bounded field) =>+             SearchIndex key field feature -> Bool+invariant SearchIndex{termMap, termIdMap, docKeyMap, docIdMap} =+      and [ IntMap.lookup (fromEnum termId) termIdMap+            == Just (TermIdInfo term docidset)+          | (term, (TermInfo termId docidset)) <- Map.assocs termMap ]+  &&  and [ case Map.lookup term termMap of+              Just (TermInfo termId' docidset') -> toEnum termId == termId'+                                                   && docidset == docidset'+              Nothing                           -> False+          | (termId, (TermIdInfo term docidset)) <- IntMap.assocs termIdMap ]+  &&  and [ case IntMap.lookup (fromEnum docId) docIdMap of+              Just (DocInfo docKey' _ _) -> docKey == docKey'+              Nothing                  -> False+          | (docKey, docId) <- Map.assocs docKeyMap ]+  &&  and [ Map.lookup docKey docKeyMap == Just (toEnum docId)+          | (docId, DocInfo docKey _ _) <- IntMap.assocs docIdMap ]+  &&  and [ DocIdSet.invariant docIdSet+          | (_term, (TermInfo _ docIdSet)) <- Map.assocs termMap ]+  &&  and [ any (\field -> DocTermIds.fieldTermCount docterms field termId > 0) fields+          | (_term, (TermInfo termId docIdSet)) <- Map.assocs termMap+          , docId <- DocIdSet.toList docIdSet+          , let DocInfo _ docterms _ = docIdMap IntMap.! fromEnum docId ]+  &&  and [ IntMap.member (fromEnum termid) termIdMap+          | (_docId, DocInfo _ docTerms _) <- IntMap.assocs docIdMap+          , field <- fields+          , termid <- DocTermIds.fieldElems docTerms field ]+  where+    fields = Ix.range (minBound, maxBound)+++-------------------+-- Lookups+--++docCount :: SearchIndex key field feature -> Int+docCount SearchIndex{docIdMap} = IntMap.size docIdMap++lookupTerm :: SearchIndex key field feature -> Term -> Maybe (TermId, DocIdSet)+lookupTerm SearchIndex{termMap} term =+    case Map.lookup term termMap of+      Nothing                         -> Nothing+      Just (TermInfo termid docidset) -> Just (termid, docidset)++lookupTermsByPrefix :: SearchIndex key field feature ->+                       Term -> [(TermId, DocIdSet)]+lookupTermsByPrefix SearchIndex{termMap} term =+    [ (termid, docidset)+    | (TermInfo termid docidset) <- lookupPrefix term termMap ]++lookupTermId :: SearchIndex key field feature -> TermId -> DocIdSet+lookupTermId SearchIndex{termIdMap} termid =+    case IntMap.lookup (fromEnum termid) termIdMap of+      Nothing -> error $ "lookupTermId: not found " ++ show termid+      Just (TermIdInfo _ docidset) -> docidset++lookupDocId :: SearchIndex key field feature ->+               DocId -> (key, DocTermIds field, DocFeatVals feature)+lookupDocId SearchIndex{docIdMap} docid =+    case IntMap.lookup (fromEnum docid) docIdMap of+      Nothing                                   -> errNotFound+      Just (DocInfo key doctermids docfeatvals) -> (key, doctermids, docfeatvals)+  where+    errNotFound = error $ "lookupDocId: not found " ++ show docid++lookupDocKey :: Ord key => SearchIndex key field feature ->+                key -> Maybe (DocTermIds field)+lookupDocKey SearchIndex{docKeyMap, docIdMap} key = do+    case Map.lookup key docKeyMap of+      Nothing    -> Nothing+      Just docid ->+        case IntMap.lookup (fromEnum docid) docIdMap of+          Nothing                          -> error "lookupDocKey: internal error"+          Just (DocInfo _key doctermids _) -> Just doctermids++lookupDocKeyDocId :: Ord key => SearchIndex key field feature -> key -> Maybe DocId+lookupDocKeyDocId SearchIndex{docKeyMap} key = Map.lookup key docKeyMap+++getTerm :: SearchIndex key field feature -> TermId -> Term+getTerm SearchIndex{termIdMap} termId =+    case termIdMap IntMap.! fromEnum termId of TermIdInfo term _ -> term++getTermId :: SearchIndex key field feature -> Term -> TermId+getTermId SearchIndex{termMap} term =+    case termMap Map.! term of TermInfo termid _ -> termid++getDocKey :: SearchIndex key field feature -> DocId -> key+getDocKey SearchIndex{docIdMap} docid =+    case docIdMap IntMap.! fromEnum docid of+      DocInfo dockey _ _ -> dockey++getDocTermIds :: SearchIndex key field feature -> DocId -> DocTermIds field+getDocTermIds SearchIndex{docIdMap} docid =+    case docIdMap IntMap.! fromEnum docid of+      DocInfo _ doctermids _ -> doctermids++--------------------+-- Insert & delete+--++-- Procedure for adding a new doc...+-- (key, field -> [Term])+-- alloc docid for key+-- add term occurences for docid (include rev map for termid)+-- construct indexdoc now that we have all the term -> termid entries+-- insert indexdoc++-- Procedure for updating a doc...+-- (key, field -> [Term])+-- find docid for key+-- lookup old terms for docid (using termid rev map)+-- calc term occurrences to add, term occurrences to delete+-- add new term occurrences, delete old term occurrences+-- construct indexdoc now that we have all the term -> termid entries+-- insert indexdoc++-- Procedure for deleting a doc...+-- (key, field -> [Term])+-- find docid for key+-- lookup old terms for docid (using termid rev map)+-- delete old term occurrences+-- delete indexdoc++-- | This is the representation for documents to be added to the index.+-- Documents may +--+type DocTerms         field   = field   -> [Term]+type DocFeatureValues feature = feature -> Float++insertDoc :: (Ord key, Ix field, Bounded field, Ix feature, Bounded feature) =>+              key -> DocTerms field -> DocFeatureValues feature ->+              SearchIndex key field feature -> SearchIndex key field feature+insertDoc key userDocTerms userDocFeats si@SearchIndex{docKeyMap}+  | Just docid <- Map.lookup key docKeyMap+  = -- Some older version of the doc is already present in the index,+    -- So we keep its docid. Now have to update the doc itself+    -- and update the terms by removing old ones and adding new ones.+    let oldTermsIds   = getDocTermIds si docid+        userDocTerms' = memoiseDocTerms userDocTerms+        newTerms      = docTermSet userDocTerms'+        oldTerms      = docTermIdsTermSet si oldTermsIds+        -- We optimise for the typical case of significant overlap between+        -- the terms in the old and new versions of the document.+        delTerms      = oldTerms `Set.difference` newTerms+        addTerms      = newTerms `Set.difference` oldTerms++     -- Note: adding the doc relies on all the terms being in the termMap+     -- already, so we first add all the term occurences for the docid.+     in checkInvariant+      . insertDocIdToDocEntry docid key userDocTerms' userDocFeats+      . insertTermToDocIdEntries (Set.toList addTerms) docid+      . deleteTermToDocIdEntries (Set.toList delTerms) docid+      $ si++  | otherwise+  = -- We're dealing with a new doc, so allocate a docid for the key+    let (si', docid)  = allocFreshDocId si+        userDocTerms' = memoiseDocTerms userDocTerms+        addTerms      = docTermSet userDocTerms'++     -- Note: adding the doc relies on all the terms being in the termMap+     -- already, so we first add all the term occurences for the docid.+     in checkInvariant+      . insertDocIdToDocEntry docid key userDocTerms' userDocFeats+      . insertDocKeyToIdEntry key docid+      . insertTermToDocIdEntries (Set.toList addTerms) docid+      $ si'++deleteDoc :: (Ord key, Ix field, Bounded field) =>+             key ->+             SearchIndex key field feature -> SearchIndex key field feature+deleteDoc key si@SearchIndex{docKeyMap}+  | Just docid <- Map.lookup key docKeyMap+  = let oldTermsIds = getDocTermIds si docid+        oldTerms    = docTermIdsTermSet si oldTermsIds+     in checkInvariant+      . deleteDocEntry docid key+      . deleteTermToDocIdEntries (Set.toList oldTerms) docid+      $ si+  +  | otherwise = si+++----------------------------------+-- Insert & delete support utils+--+++memoiseDocTerms :: (Ix field, Bounded field) => DocTerms field -> DocTerms field+memoiseDocTerms docTermsFn =+    \field -> vecIndexIx vec field+  where+    vec = vecCreateIx docTermsFn++docTermSet :: (Bounded t, Ix t) => DocTerms t -> Set.Set Term+docTermSet docterms =+    Set.unions [ Set.fromList (docterms field)+               | field <- Ix.range (minBound, maxBound) ]++docTermIdsTermSet :: (Bounded field, Ix field) =>+                     SearchIndex key field feature ->+                     DocTermIds field -> Set.Set Term+docTermIdsTermSet si doctermids =+    Set.unions [ Set.fromList terms+               | field <- Ix.range (minBound, maxBound)+               , let termids = DocTermIds.fieldElems doctermids field+                     terms   = map (getTerm si) termids ]++--+-- The Term <-> DocId mapping+--++-- | Add an entry into the 'Term' to 'DocId' mapping.+insertTermToDocIdEntry :: Term -> DocId ->+                          SearchIndex key field feature ->+                          SearchIndex key field feature+insertTermToDocIdEntry term !docid si@SearchIndex{termMap, termIdMap, nextTermId} =+    case Map.lookup term termMap of+      Nothing ->+        let docIdSet'    = DocIdSet.singleton docid+            !termInfo'   = TermInfo nextTermId docIdSet'+            !termIdInfo' = TermIdInfo term     docIdSet'+         in si { termMap    = Map.insert term termInfo' termMap+               , termIdMap  = IntMap.insert (fromEnum nextTermId)+                                            termIdInfo' termIdMap+               , nextTermId = succ nextTermId }++      Just (TermInfo termId docIdSet) ->+        let docIdSet'    = DocIdSet.insert docid docIdSet+            !termInfo'   = TermInfo termId docIdSet'+            !termIdInfo' = TermIdInfo term docIdSet'+         in si { termMap   = Map.insert term termInfo' termMap+               , termIdMap = IntMap.insert (fromEnum termId)+                                           termIdInfo' termIdMap+               }++-- | Add multiple entries into the 'Term' to 'DocId' mapping: many terms that+-- map to the same document.+insertTermToDocIdEntries :: [Term] -> DocId ->+                            SearchIndex key field feature ->+                            SearchIndex key field feature+insertTermToDocIdEntries terms !docid si =+    List.foldl' (\si' term -> insertTermToDocIdEntry term docid si') si terms++-- | Delete an entry from the 'Term' to 'DocId' mapping.+deleteTermToDocIdEntry :: Term -> DocId ->+                          SearchIndex key field feature ->+                          SearchIndex key field feature+deleteTermToDocIdEntry term !docid si@SearchIndex{termMap, termIdMap} =+    case  Map.lookup term termMap of+      Nothing -> si+      Just (TermInfo termId docIdSet) ->+        let docIdSet'    = DocIdSet.delete docid docIdSet+            !termInfo'   = TermInfo termId docIdSet'+            !termIdInfo' = TermIdInfo term docIdSet'+        in if DocIdSet.null docIdSet'+            then si { termMap = Map.delete term termMap+                    , termIdMap = IntMap.delete (fromEnum termId) termIdMap }+            else si { termMap   = Map.insert term termInfo' termMap+                    , termIdMap = IntMap.insert (fromEnum termId)+                                                termIdInfo' termIdMap+                    }++-- | Delete multiple entries from the 'Term' to 'DocId' mapping: many terms+-- that map to the same document.+deleteTermToDocIdEntries :: [Term] -> DocId ->+                            SearchIndex key field feature ->+                            SearchIndex key field feature+deleteTermToDocIdEntries terms !docid si =+    List.foldl' (\si' term -> deleteTermToDocIdEntry term docid si') si terms++--+-- The DocId <-> Doc mapping+--++allocFreshDocId :: SearchIndex key field feature ->+                  (SearchIndex key field feature, DocId)+allocFreshDocId si@SearchIndex{nextDocId} =+    let !si' = si { nextDocId = succ nextDocId }+     in (si', nextDocId)++insertDocKeyToIdEntry :: Ord key => key -> DocId ->+                         SearchIndex key field feature ->+                         SearchIndex key field feature+insertDocKeyToIdEntry dockey !docid si@SearchIndex{docKeyMap} =+    si { docKeyMap = Map.insert dockey docid docKeyMap }++insertDocIdToDocEntry :: (Ix field, Bounded field,+                          Ix feature, Bounded feature) =>+                         DocId -> key ->+                         DocTerms field ->+                         DocFeatureValues feature ->+                         SearchIndex key field feature ->+                         SearchIndex key field feature+insertDocIdToDocEntry !docid dockey userdocterms userdocfeats+                       si@SearchIndex{docIdMap} =+    let doctermids = DocTermIds.create (map (getTermId si) . userdocterms)+        docfeatvals= DocFeatVals.create userdocfeats+        !docinfo   = DocInfo dockey doctermids docfeatvals+     in si { docIdMap  = IntMap.insert (fromEnum docid) docinfo docIdMap }++deleteDocEntry :: Ord key => DocId -> key ->+                  SearchIndex key field feature -> SearchIndex key field feature+deleteDocEntry docid key si@SearchIndex{docIdMap, docKeyMap} =+     si { docIdMap  = IntMap.delete (fromEnum docid) docIdMap+        , docKeyMap = Map.delete key docKeyMap }++--+-- Data.Map utils+--++-- Data.Map does not support prefix lookups directly (unlike a trie)+-- but we can implement it reasonably efficiently using split:++-- | Lookup values for a range of keys (inclusive lower bound and exclusive+-- upper bound)+--+lookupRange :: Ord k => (k, k) -> Map k v -> [v]+lookupRange (lb, ub) m =+  let (_, mv, gt)  = Map.splitLookup lb m+      (between, _) = Map.split       ub gt+   in case mv of+        Just v  -> v : Map.elems between+        Nothing ->     Map.elems between++lookupPrefix :: Text -> Map Text v -> [v]+lookupPrefix t _ | T.null t = []+lookupPrefix t m = lookupRange (t, prefixUpperBound t) m++prefixUpperBound :: Text -> Text+prefixUpperBound = succLast . T.dropWhileEnd (== maxBound)+  where+    succLast t = T.init t `T.snoc` succ (T.last t)+
+ src/Data/SearchEngine/TermBag.hs view
@@ -0,0 +1,268 @@+{-# LANGUAGE BangPatterns, GeneralizedNewtypeDeriving, MultiParamTypeClasses,+             TypeFamilies #-}+{-# LANGUAGE CPP #-}+#if __GLASGOW_HASKELL__ >= 908+{-# OPTIONS_GHC -Wno-x-partial #-}+#endif++module Data.SearchEngine.TermBag (+    TermId(TermId), TermCount,+    TermBag,+    size,+    fromList,+    toList,+    elems,+    termCount,+    denseTable,+    invariant+  ) where++import qualified Data.Vector.Unboxed         as Vec+import qualified Data.Vector.Unboxed.Mutable as MVec+import qualified Data.Vector.Generic         as GVec+import qualified Data.Vector.Generic.Mutable as GMVec+import Control.Monad.ST+import Control.Monad (liftM)+import qualified Data.Map as Map+import Data.Word (Word32, Word8)+import Data.Bits+import qualified Data.List as List+import Data.Function (on)++newtype TermId = TermId { unTermId :: Word32 }+  deriving (Eq, Ord, Show, Enum)++instance Bounded TermId where+  minBound = TermId 0+  maxBound = TermId 0x00FFFFFF++data TermBag = TermBag !Int !(Vec.Vector TermIdAndCount)+  deriving Show++-- We sneakily stuff both the TermId and the bag count into one 32bit word+type TermIdAndCount = Word32+type TermCount      = Word8++-- Bottom 24 bits is the TermId, top 8 bits is the bag count+termIdAndCount :: TermId -> Int -> TermIdAndCount+termIdAndCount (TermId termid) freq =+      (min (fromIntegral freq) 255 `shiftL` 24)+  .|. (termid .&. 0x00FFFFFF)++getTermId :: TermIdAndCount -> TermId+getTermId word = TermId (word .&. 0x00FFFFFF)++getTermCount :: TermIdAndCount -> TermCount+getTermCount word = fromIntegral (word `shiftR` 24)++invariant :: TermBag -> Bool+invariant (TermBag _ vec) =+    strictlyAscending (Vec.toList vec)+  where+    strictlyAscending (a:xs@(b:_)) = getTermId a < getTermId b+                                  && strictlyAscending xs+    strictlyAscending _  = True++size :: TermBag -> Int+size (TermBag sz _) = sz++elems :: TermBag -> [TermId]+elems (TermBag _ vec) = map getTermId (Vec.toList vec)++toList :: TermBag -> [(TermId, TermCount)]+toList (TermBag _ vec) = [ (getTermId x, getTermCount x)+                         | x <- Vec.toList vec ]++termCount :: TermBag -> TermId -> TermCount+termCount (TermBag _ vec) =+    binarySearch 0 (Vec.length vec - 1)+  where+    binarySearch :: Int -> Int -> TermId -> TermCount+    binarySearch !a !b !key+      | a > b     = 0+      | otherwise =+        let mid         = (a + b) `div` 2+            tidAndCount = vec Vec.! mid+         in case compare key (getTermId tidAndCount) of+              LT -> binarySearch a (mid-1) key+              EQ -> getTermCount tidAndCount+              GT -> binarySearch (mid+1) b key++fromList :: [TermId] -> TermBag+fromList termids =+    let bag = Map.fromListWith (+) [ (t, 1) | t <- termids ]+        sz  = Map.foldl' (+) 0 bag+        vec = Vec.fromListN (Map.size bag)+                            [ termIdAndCount termid freq+                            | (termid, freq) <- Map.toAscList bag ]+     in TermBag sz vec++-- | Given a bunch of term bags, merge them into a table for easier subsequent+-- processing. This is bascially a sparse to dense conversion. Missing entries+-- are filled in with 0. We represent the table as one vector for the+-- term ids and a 2d array for the counts.+--+-- Unfortunately vector does not directly support 2d arrays and array does+-- not make it easy to trim arrays.+--+denseTable :: [TermBag] -> (Vec.Vector TermId, Vec.Vector TermCount)+denseTable termbags =+    (tids, tcts)+  where+    -- First merge the TermIds into one array+    -- then make a linear pass to create the counts array+    -- filling in 0s or the counts as we find them+    !numBags   = length termbags+    !tids      = unionsTermId termbags+    !numTerms  = Vec.length tids+    !numCounts = numTerms * numBags+    !tcts      = Vec.create (do+                   out <- MVec.new numCounts+                   sequence_+                     [ writeMergedTermCounts tids bag out i+                     | (n, TermBag _ bag) <- zip [0..] termbags+                     , let i = n * numTerms ]+                   return out+                 )++writeMergedTermCounts :: Vec.Vector TermId -> Vec.Vector TermIdAndCount ->+                         MVec.MVector s TermCount -> Int -> ST s ()+writeMergedTermCounts xs0 ys0 !out i0 =+    -- assume xs & ys are sorted, and ys contains a subset of xs+    go xs0 ys0 i0+  where+    go !xs !ys !i+      | Vec.null ys = MVec.set (MVec.slice i (Vec.length xs) out) 0+      | Vec.null xs = return ()+      | otherwise   = let x   = Vec.head xs+                          ytc = Vec.head ys+                          y   = getTermId ytc+                          c   = getTermCount ytc+                      in case x == y of+                           True  -> do MVec.write out i c+                                       go (Vec.tail xs) (Vec.tail ys) (i+1)+                           False -> do MVec.write out i 0+                                       go (Vec.tail xs)           ys  (i+1)++-- | Given a set of term bags, form the set of TermIds+--+unionsTermId :: [TermBag] -> Vec.Vector TermId+unionsTermId tbs =+    case List.sortBy (compare `on` bagVecLength) tbs of+      []             -> Vec.empty+      [TermBag _ xs] -> (Vec.map getTermId xs)+      (x0:x1:xs)     -> List.foldl' union3 (union2 x0 x1) xs+  where+    bagVecLength (TermBag _ vec) = Vec.length vec++union2 :: TermBag -> TermBag -> Vec.Vector TermId+union2 (TermBag _ xs) (TermBag _ ys) =+    Vec.create (MVec.new sizeBound >>= writeMergedUnion2 xs ys)+  where+    sizeBound = Vec.length xs + Vec.length ys++writeMergedUnion2 :: Vec.Vector TermIdAndCount -> Vec.Vector TermIdAndCount ->+                     MVec.MVector s TermId -> ST s (MVec.MVector s TermId)+writeMergedUnion2 xs0 ys0 !out = do+    i <- go xs0 ys0 0+    return $! MVec.take i out+  where+    go !xs !ys !i+      | Vec.null xs = do Vec.copy (MVec.slice i (Vec.length ys) out)+                                  (Vec.map getTermId ys)+                         return (i + Vec.length ys)+      | Vec.null ys = do Vec.copy (MVec.slice i (Vec.length xs) out)+                                  (Vec.map getTermId xs)+                         return (i + Vec.length xs)+      | otherwise   = let x = getTermId (Vec.head xs)+                          y = getTermId (Vec.head ys)+                      in case compare x y of+                          GT -> do MVec.write out i y+                                   go           xs  (Vec.tail ys) (i+1)+                          EQ -> do MVec.write out i x+                                   go (Vec.tail xs) (Vec.tail ys) (i+1)+                          LT -> do MVec.write out i x+                                   go (Vec.tail xs)           ys  (i+1)++union3 :: Vec.Vector TermId -> TermBag -> Vec.Vector TermId+union3 xs (TermBag _ ys) =+    Vec.create (MVec.new sizeBound >>= writeMergedUnion3 xs ys)+  where+    sizeBound = Vec.length xs + Vec.length ys++writeMergedUnion3 :: Vec.Vector TermId -> Vec.Vector TermIdAndCount ->+                     MVec.MVector s TermId -> ST s (MVec.MVector s TermId)+writeMergedUnion3 xs0 ys0 !out = do+    i <- go xs0 ys0 0+    return $! MVec.take i out+  where+    go !xs !ys !i+      | Vec.null xs = do Vec.copy (MVec.slice i (Vec.length ys) out)+                                  (Vec.map getTermId ys)+                         return (i + Vec.length ys)+      | Vec.null ys = do Vec.copy (MVec.slice i (Vec.length xs) out) xs+                         return (i + Vec.length xs)+      | otherwise   = let x =            Vec.head xs+                          y = getTermId (Vec.head ys)+                      in case compare x y of+                          GT -> do MVec.write out i y+                                   go           xs  (Vec.tail ys) (i+1)+                          EQ -> do MVec.write out i x+                                   go (Vec.tail xs) (Vec.tail ys) (i+1)+                          LT -> do MVec.write out i x+                                   go (Vec.tail xs)           ys  (i+1)++------------------------------------------------------------------------------+-- verbose Unbox instances+--++instance MVec.Unbox TermId++newtype instance MVec.MVector s TermId = MV_TermId (MVec.MVector s Word32)++instance GMVec.MVector MVec.MVector TermId where+    basicLength          (MV_TermId v) = GMVec.basicLength v+    basicUnsafeSlice i l (MV_TermId v) = MV_TermId (GMVec.basicUnsafeSlice i l v)+    basicUnsafeNew     l               = MV_TermId `liftM` GMVec.basicUnsafeNew l+    basicInitialize      (MV_TermId v) = GMVec.basicInitialize v+    basicUnsafeReplicate l x           = MV_TermId `liftM` GMVec.basicUnsafeReplicate l (unTermId x)+    basicUnsafeRead  (MV_TermId v) i   = TermId `liftM`    GMVec.basicUnsafeRead v i+    basicUnsafeWrite (MV_TermId v) i x = GMVec.basicUnsafeWrite v i (unTermId x)+    basicClear       (MV_TermId v)     = GMVec.basicClear v+    basicSet         (MV_TermId v) x   = GMVec.basicSet v (unTermId x)+    basicUnsafeGrow  (MV_TermId v) l   = MV_TermId `liftM` GMVec.basicUnsafeGrow v l+    basicUnsafeCopy  (MV_TermId v) (MV_TermId v') = GMVec.basicUnsafeCopy v v'+    basicUnsafeMove  (MV_TermId v) (MV_TermId v') = GMVec.basicUnsafeMove v v'+    basicOverlaps    (MV_TermId v) (MV_TermId v') = GMVec.basicOverlaps   v v'+    {-# INLINE basicLength #-}+    {-# INLINE basicUnsafeSlice #-}+    {-# INLINE basicOverlaps #-}+    {-# INLINE basicUnsafeNew #-}+    {-# INLINE basicInitialize #-}+    {-# INLINE basicUnsafeReplicate #-}+    {-# INLINE basicUnsafeRead #-}+    {-# INLINE basicUnsafeWrite #-}+    {-# INLINE basicClear #-}+    {-# INLINE basicSet #-}+    {-# INLINE basicUnsafeCopy #-}+    {-# INLINE basicUnsafeMove #-}+    {-# INLINE basicUnsafeGrow #-}++newtype instance Vec.Vector TermId = V_TermId (Vec.Vector Word32)++instance GVec.Vector Vec.Vector TermId where+    basicUnsafeFreeze (MV_TermId mv)  = V_TermId  `liftM` GVec.basicUnsafeFreeze mv+    basicUnsafeThaw   (V_TermId  v)   = MV_TermId `liftM` GVec.basicUnsafeThaw v+    basicLength       (V_TermId  v)   = GVec.basicLength v+    basicUnsafeSlice i l (V_TermId v) = V_TermId (GVec.basicUnsafeSlice i l v)+    basicUnsafeIndexM (V_TermId  v) i = TermId `liftM` GVec.basicUnsafeIndexM v i+    basicUnsafeCopy   (MV_TermId mv)+                      (V_TermId  v)   = GVec.basicUnsafeCopy mv v+    elemseq           (V_TermId  v) x = GVec.elemseq v (unTermId x)+    {-# INLINE basicUnsafeFreeze #-}+    {-# INLINE basicUnsafeThaw #-}+    {-# INLINE basicLength #-}+    {-# INLINE basicUnsafeSlice #-}+    {-# INLINE basicUnsafeIndexM #-}+    {-# INLINE basicUnsafeCopy #-}+    {-# INLINE elemseq #-}
+ src/Data/SearchEngine/Types.hs view
@@ -0,0 +1,124 @@+{-# LANGUAGE NamedFieldPuns, RecordWildCards #-}++module Data.SearchEngine.Types (+    -- * Search engine types and helper functions+    SearchEngine(..),+    SearchConfig(..),+    SearchRankParameters(..),+    BM25F.FeatureFunction(..),+    initSearchEngine,+    cacheBM25Context,++    -- ** Helper type for non-term features+    NoFeatures,+    noFeatures,++    -- * Re-export SearchIndex and other types+    SearchIndex, Term, TermId,+    DocIdSet, DocId,+    DocTermIds, DocFeatVals,++    -- * Internal sanity check+    invariant,+  ) where++import Data.SearchEngine.SearchIndex (SearchIndex, Term, TermId)+import qualified Data.SearchEngine.SearchIndex as SI+import Data.SearchEngine.DocIdSet (DocIdSet, DocId)+import qualified Data.SearchEngine.DocIdSet as DocIdSet+import Data.SearchEngine.DocFeatVals (DocFeatVals)+import Data.SearchEngine.DocTermIds  (DocTermIds)+import qualified Data.SearchEngine.BM25F as BM25F++import Data.Ix+import Data.Array.Unboxed++++data SearchConfig doc key field feature = SearchConfig {+       documentKey          :: doc -> key,+       extractDocumentTerms :: doc -> field -> [Term],+       transformQueryTerm   :: Term -> field -> Term,+       documentFeatureValue :: doc -> feature -> Float+     }++data SearchRankParameters field feature = SearchRankParameters {+       paramK1                 :: !Float,+       paramB                  :: field -> Float,+       paramFieldWeights       :: field -> Float,+       paramFeatureWeights     :: feature -> Float,+       paramFeatureFunctions   :: feature -> BM25F.FeatureFunction,++       paramResultsetSoftLimit   :: !Int,+       paramResultsetHardLimit   :: !Int,+       paramAutosuggestPrefilterLimit  :: !Int,+       paramAutosuggestPostfilterLimit :: !Int+     }++data SearchEngine doc key field feature = SearchEngine {+       searchIndex      :: !(SearchIndex      key field feature),+       searchConfig     :: !(SearchConfig doc key field feature),+       searchRankParams :: !(SearchRankParameters field feature),++       -- cached info+       sumFieldLengths :: !(UArray field Int),+       bm25Context     :: BM25F.Context TermId field feature+     }++invariant :: (Ord key, Ix field, Bounded field) =>+             SearchEngine doc key field feature -> Bool+invariant SearchEngine{searchIndex} =+    SI.invariant searchIndex+-- && check caches++initSearchEngine :: (Ix field, Bounded field, Ix feature, Bounded feature) =>+                    SearchConfig doc key field feature ->+                    SearchRankParameters field feature ->+                    SearchEngine doc key field feature+initSearchEngine config params =+    cacheBM25Context+      SearchEngine {+        searchIndex      = SI.emptySearchIndex,+        searchConfig     = config,+        searchRankParams = params,+        sumFieldLengths  = listArray (minBound, maxBound) (repeat 0),+        bm25Context      = undefined+      }++cacheBM25Context :: Ix field =>+                    SearchEngine doc key field feature ->+                    SearchEngine doc key field feature+cacheBM25Context+    se@SearchEngine {+      searchRankParams = SearchRankParameters{..},+      searchIndex,+      sumFieldLengths+    }+  = se { bm25Context = bm25Context' }+  where+    bm25Context' = BM25F.Context {+      BM25F.numDocsTotal    = SI.docCount searchIndex,+      BM25F.avgFieldLength  = \f -> fromIntegral (sumFieldLengths ! f)+                                  / fromIntegral (SI.docCount searchIndex),+      BM25F.numDocsWithTerm = DocIdSet.size . SI.lookupTermId searchIndex,+      BM25F.paramK1         = paramK1,+      BM25F.paramB          = paramB,+      BM25F.fieldWeight     = paramFieldWeights,+      BM25F.featureWeight   = paramFeatureWeights,+      BM25F.featureFunction = paramFeatureFunctions+    }+++-----------------------------++data NoFeatures = NoFeatures+  deriving (Eq, Ord, Bounded, Show)++instance Ix NoFeatures where+  range   _   = []+  inRange _ _ = False+  index   _ _ = -1++noFeatures :: NoFeatures -> a+noFeatures _ = error "noFeatures"+
+ src/Data/SearchEngine/Update.hs view
@@ -0,0 +1,90 @@+{-# LANGUAGE BangPatterns, NamedFieldPuns, RecordWildCards #-}++module Data.SearchEngine.Update (++    -- * Managing documents to be searched+    insertDoc,+    insertDocs,+    deleteDoc,++  ) where++import Data.SearchEngine.Types+import qualified Data.SearchEngine.SearchIndex as SI+import qualified Data.SearchEngine.DocTermIds as DocTermIds++import qualified Data.List as List+import Data.Ix+import Data.Array.Unboxed+++insertDocs :: (Ord key, Ix field, Bounded field, Ix feature, Bounded feature) =>+              [doc] ->+              SearchEngine doc key field feature ->+              SearchEngine doc key field feature+insertDocs docs se = List.foldl' (\se' doc -> insertDoc doc se') se docs+++insertDoc :: (Ord key, Ix field, Bounded field, Ix feature, Bounded feature) =>+             doc ->+             SearchEngine doc key field feature ->+             SearchEngine doc key field feature+insertDoc doc se@SearchEngine{ searchConfig = SearchConfig {+                                 documentKey,+                                 extractDocumentTerms,+                                 documentFeatureValue+                               }+                             , searchIndex } =+    let key = documentKey doc+        searchIndex' = SI.insertDoc key (extractDocumentTerms doc)+                                        (documentFeatureValue doc)+                                        searchIndex+        oldDoc       = SI.lookupDocKey searchIndex  key+        newDoc       = SI.lookupDocKey searchIndex' key++     in cacheBM25Context $+        updateCachedFieldLengths oldDoc newDoc $+          se { searchIndex = searchIndex' }+++deleteDoc :: (Ord key, Ix field, Bounded field) =>+             key ->+             SearchEngine doc key field feature ->+             SearchEngine doc key field feature+deleteDoc key se@SearchEngine{searchIndex} =+    let searchIndex' = SI.deleteDoc key searchIndex+        oldDoc       = SI.lookupDocKey searchIndex key++     in cacheBM25Context $+        updateCachedFieldLengths oldDoc Nothing $+          se { searchIndex = searchIndex' }+++updateCachedFieldLengths :: (Ix field, Bounded field) =>+                            Maybe (DocTermIds field) -> Maybe (DocTermIds field) ->+                            SearchEngine doc key field feature ->+                            SearchEngine doc key field feature+updateCachedFieldLengths Nothing (Just newDoc) se@SearchEngine{sumFieldLengths} =+    se {+      sumFieldLengths =+        array (bounds sumFieldLengths)+              [ (i, n + DocTermIds.fieldLength newDoc i)+              | (i, n) <- assocs sumFieldLengths ]+    }+updateCachedFieldLengths (Just oldDoc) (Just newDoc) se@SearchEngine{sumFieldLengths} =+    se {+      sumFieldLengths =+        array (bounds sumFieldLengths)+              [ (i, n - DocTermIds.fieldLength oldDoc i+                      + DocTermIds.fieldLength newDoc i)+              | (i, n) <- assocs sumFieldLengths ]+    }+updateCachedFieldLengths (Just oldDoc) Nothing se@SearchEngine{sumFieldLengths} =+    se {+      sumFieldLengths =+        array (bounds sumFieldLengths)+              [ (i, n - DocTermIds.fieldLength oldDoc i)+              | (i, n) <- assocs sumFieldLengths ]+    }+updateCachedFieldLengths Nothing Nothing se = se+
tests/Test/Data/SearchEngine/TermBag.hs view
@@ -1,4 +1,9 @@-{-# OPTIONS_GHC -fno-warn-orphans #-}+{-# OPTIONS_GHC -Wno-orphans #-}+{-# LANGUAGE CPP #-}+#if __GLASGOW_HASKELL__ >= 908+{-# OPTIONS_GHC -Wno-x-partial #-}+#endif+ module Test.Data.SearchEngine.TermBag where  import Data.SearchEngine.TermBag@@ -34,8 +39,8 @@  prop_termCount :: [TermId] -> Bool prop_termCount tids =-    and [ termCount bag tid == count -        | let bag = fromList tids +    and [ termCount bag tid == count+        | let bag = fromList tids         , (tid, count) <- toList bag         ] @@ -52,4 +57,4 @@               numTerms        = Vec.length terms         , (b, bag) <- zip [0..] bags         ,  t       <- [0..Vec.length terms - 1]-        ] +        ]