diff --git a/Data/SearchEngine.hs b/Data/SearchEngine.hs deleted file mode 100644 --- a/Data/SearchEngine.hs +++ /dev/null @@ -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 - diff --git a/Data/SearchEngine/Autosuggest.hs b/Data/SearchEngine/Autosuggest.hs deleted file mode 100644 --- a/Data/SearchEngine/Autosuggest.hs +++ /dev/null @@ -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 - ] - diff --git a/Data/SearchEngine/BM25F.hs b/Data/SearchEngine/BM25F.hs deleted file mode 100644 --- a/Data/SearchEngine/BM25F.hs +++ /dev/null @@ -1,253 +0,0 @@ -{-# LANGUAGE RecordWildCards, BangPatterns, ScopedTypeVariables #-} - --- | An implementation of BM25F ranking. See: --- --- * A quick overview: --- --- * /The Probabilistic Relevance Framework: BM25 and Beyond/ --- --- --- * /An Introduction to Information Retrieval/ --- --- -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' diff --git a/Data/SearchEngine/DocFeatVals.hs b/Data/SearchEngine/DocFeatVals.hs deleted file mode 100644 --- a/Data/SearchEngine/DocFeatVals.hs +++ /dev/null @@ -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) - diff --git a/Data/SearchEngine/DocIdSet.hs b/Data/SearchEngine/DocIdSet.hs deleted file mode 100644 --- a/Data/SearchEngine/DocIdSet.hs +++ /dev/null @@ -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 #-} diff --git a/Data/SearchEngine/DocTermIds.hs b/Data/SearchEngine/DocTermIds.hs deleted file mode 100644 --- a/Data/SearchEngine/DocTermIds.hs +++ /dev/null @@ -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) - diff --git a/Data/SearchEngine/Query.hs b/Data/SearchEngine/Query.hs deleted file mode 100644 --- a/Data/SearchEngine/Query.hs +++ /dev/null @@ -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 - diff --git a/Data/SearchEngine/SearchIndex.hs b/Data/SearchEngine/SearchIndex.hs deleted file mode 100644 --- a/Data/SearchEngine/SearchIndex.hs +++ /dev/null @@ -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) - diff --git a/Data/SearchEngine/TermBag.hs b/Data/SearchEngine/TermBag.hs deleted file mode 100644 --- a/Data/SearchEngine/TermBag.hs +++ /dev/null @@ -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 #-} diff --git a/Data/SearchEngine/Types.hs b/Data/SearchEngine/Types.hs deleted file mode 100644 --- a/Data/SearchEngine/Types.hs +++ /dev/null @@ -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" - diff --git a/Data/SearchEngine/Update.hs b/Data/SearchEngine/Update.hs deleted file mode 100644 --- a/Data/SearchEngine/Update.hs +++ /dev/null @@ -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 - diff --git a/changelog b/changelog --- a/changelog +++ b/changelog @@ -1,3 +1,7 @@ +0.2.2.3 Ben Gamari January 2025 + * Introduce compatibility with GHCs up to 9.12 + * Drop compatibility with GHC versions pre-9.2 + 0.2.2.2 Adam Gundry March 2023 * Fix bug in 0.2.2.1 autosuggest patch diff --git a/demo/ExtractNameTerms.hs b/demo/ExtractNameTerms.hs --- a/demo/ExtractNameTerms.hs +++ b/demo/ExtractNameTerms.hs @@ -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) diff --git a/full-text-search.cabal b/full-text-search.cabal --- a/full-text-search.cabal +++ b/full-text-search.cabal @@ -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 - and - . - . - 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 - - 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 + and + . + + 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 + + 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 , @@ -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, diff --git a/src/Data/SearchEngine.hs b/src/Data/SearchEngine.hs new file mode 100644 --- /dev/null +++ b/src/Data/SearchEngine.hs @@ -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 + diff --git a/src/Data/SearchEngine/Autosuggest.hs b/src/Data/SearchEngine/Autosuggest.hs new file mode 100644 --- /dev/null +++ b/src/Data/SearchEngine/Autosuggest.hs @@ -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 + ] + diff --git a/src/Data/SearchEngine/BM25F.hs b/src/Data/SearchEngine/BM25F.hs new file mode 100644 --- /dev/null +++ b/src/Data/SearchEngine/BM25F.hs @@ -0,0 +1,253 @@ +{-# LANGUAGE RecordWildCards, BangPatterns, ScopedTypeVariables #-} + +-- | An implementation of BM25F ranking. See: +-- +-- * A quick overview: +-- +-- * /The Probabilistic Relevance Framework: BM25 and Beyond/ +-- +-- +-- * /An Introduction to Information Retrieval/ +-- +-- +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' diff --git a/src/Data/SearchEngine/DocFeatVals.hs b/src/Data/SearchEngine/DocFeatVals.hs new file mode 100644 --- /dev/null +++ b/src/Data/SearchEngine/DocFeatVals.hs @@ -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) + diff --git a/src/Data/SearchEngine/DocIdSet.hs b/src/Data/SearchEngine/DocIdSet.hs new file mode 100644 --- /dev/null +++ b/src/Data/SearchEngine/DocIdSet.hs @@ -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 #-} diff --git a/src/Data/SearchEngine/DocTermIds.hs b/src/Data/SearchEngine/DocTermIds.hs new file mode 100644 --- /dev/null +++ b/src/Data/SearchEngine/DocTermIds.hs @@ -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) + diff --git a/src/Data/SearchEngine/Query.hs b/src/Data/SearchEngine/Query.hs new file mode 100644 --- /dev/null +++ b/src/Data/SearchEngine/Query.hs @@ -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 + diff --git a/src/Data/SearchEngine/SearchIndex.hs b/src/Data/SearchEngine/SearchIndex.hs new file mode 100644 --- /dev/null +++ b/src/Data/SearchEngine/SearchIndex.hs @@ -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) + diff --git a/src/Data/SearchEngine/TermBag.hs b/src/Data/SearchEngine/TermBag.hs new file mode 100644 --- /dev/null +++ b/src/Data/SearchEngine/TermBag.hs @@ -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 #-} diff --git a/src/Data/SearchEngine/Types.hs b/src/Data/SearchEngine/Types.hs new file mode 100644 --- /dev/null +++ b/src/Data/SearchEngine/Types.hs @@ -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" + diff --git a/src/Data/SearchEngine/Update.hs b/src/Data/SearchEngine/Update.hs new file mode 100644 --- /dev/null +++ b/src/Data/SearchEngine/Update.hs @@ -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 + diff --git a/tests/Test/Data/SearchEngine/TermBag.hs b/tests/Test/Data/SearchEngine/TermBag.hs --- a/tests/Test/Data/SearchEngine/TermBag.hs +++ b/tests/Test/Data/SearchEngine/TermBag.hs @@ -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] - ] + ]