nodedb_fts/search/
doc_scorer.rs

1// SPDX-License-Identifier: Apache-2.0
2
3//! Per-document `bm25_score(field, query)` values.
4//!
5//! A [`DocScorer`] resolves its query once, then scores any document by
6//! point lookups: the document's postings of the query terms, and its
7//! recorded length for index membership. It reads no corpus-wide score map,
8//! so scoring the rows a query emits costs memory in the query terms'
9//! postings, never in the collection.
10//!
11//! The score is the one a search of the same query returns for the
12//! document, under the same match-mode decision over the same admitted rows.
13
14use nodedb_types::{Surrogate, SurrogateBitmap};
15
16use crate::backend::FtsBackend;
17use crate::index::FtsIndex;
18use crate::index::error::FtsIndexError;
19use crate::scope::IndexScope;
20use crate::search::doc_score::term_score;
21use crate::search::match_mode::{MatchMode, staged_candidates};
22use crate::search::query_terms::{ResolvedQuery, TextQuery};
23use crate::search::staged::{StagedView, staged_contributions};
24
25/// One document's score against one index.
26#[derive(Debug, Clone, Copy, PartialEq)]
27pub enum DocScore {
28    /// The document matches the query with this score.
29    Match(f32),
30    /// The index holds the document, and the query does not match it.
31    Miss,
32    /// The index does not hold the document.
33    Absent,
34}
35
36/// Scores documents against one resolved query.
37pub struct DocScorer<'a, B: FtsBackend> {
38    fts: &'a FtsIndex<B>,
39    database_id: u64,
40    tid: u64,
41    index: IndexScope<'a>,
42    staged: Option<StagedView>,
43    /// `None` when the query has no positive term: it matches nothing.
44    resolved: Option<(ResolvedQuery, MatchMode)>,
45}
46
47impl<B: FtsBackend> FtsIndex<B> {
48    /// A scorer of `query` against `index`. `eligible` is the set of rows the
49    /// reading query admits: the AND-mode fallback is decided over it, the
50    /// same way a search with that prefilter decides it. `staged` is the
51    /// open transaction's view of the index.
52    pub fn doc_scorer<'a>(
53        &'a self,
54        database_id: u64,
55        tid: u64,
56        index: impl Into<IndexScope<'a>>,
57        query: TextQuery<'_>,
58        eligible: Option<&SurrogateBitmap>,
59        staged: Option<StagedView>,
60    ) -> Result<DocScorer<'a, B>, FtsIndexError<B::Error>> {
61        let index = index.into();
62        let resolved = match self.resolve_query(database_id, tid, index, query, staged.as_ref())? {
63            Some(resolved) => {
64                let mut deny = resolved.negated.clone();
65                if let Some(view) = staged.as_ref() {
66                    deny.union_in_place(view.hidden());
67                }
68                let index_visible = staged.as_ref().is_none_or(|view| !view.hides_all());
69                let candidates =
70                    staged_candidates(&resolved, staged.as_ref(), eligible, &self.bm25_params);
71                let mode = resolved.match_mode(index_visible, eligible, &deny, &candidates);
72                Some((resolved, mode))
73            }
74            None => None,
75        };
76        Ok(DocScorer {
77            fts: self,
78            database_id,
79            tid,
80            index,
81            staged,
82            resolved,
83        })
84    }
85}
86
87impl<B: FtsBackend> DocScorer<'_, B> {
88    /// The score of each of `docs`, parallel to `docs`. Index membership of
89    /// the documents that do not match is read in one batch.
90    pub fn score(&self, docs: &[Surrogate]) -> Result<Vec<DocScore>, FtsIndexError<B::Error>> {
91        let mut scores = vec![DocScore::Absent; docs.len()];
92        let mut membership: Vec<(usize, Surrogate)> = Vec::new();
93        for (slot, doc_id) in docs.iter().enumerate() {
94            if let Some(view) = self.staged.as_ref() {
95                if let Some(doc) = view.doc(*doc_id) {
96                    scores[slot] = self
97                        .staged_match(doc)
98                        .map_or(DocScore::Miss, DocScore::Match);
99                    continue;
100                }
101                if view.hides(*doc_id) {
102                    continue;
103                }
104            }
105            match self.indexed_match(*doc_id) {
106                Some(score) => scores[slot] = DocScore::Match(score),
107                None => membership.push((slot, *doc_id)),
108            }
109        }
110        if membership.is_empty() {
111            return Ok(scores);
112        }
113        let ids: Vec<Surrogate> = membership.iter().map(|(_, doc_id)| *doc_id).collect();
114        let lengths = self
115            .fts
116            .backend
117            .read_doc_lengths(self.database_id, self.tid, self.index, &ids)
118            .map_err(FtsIndexError::backend)?;
119        for ((slot, _), length) in membership.into_iter().zip(lengths) {
120            if length.is_some() {
121                scores[slot] = DocScore::Miss;
122            }
123        }
124        Ok(scores)
125    }
126
127    fn staged_match(&self, doc: &crate::search::staged::StagedDoc) -> Option<f32> {
128        let (query, mode) = self.resolved.as_ref()?;
129        let contributions = staged_contributions(query, doc, &self.fts.bm25_params)?;
130        mode.score(query, &contributions)
131    }
132
133    fn indexed_match(&self, doc_id: Surrogate) -> Option<f32> {
134        let (query, mode) = self.resolved.as_ref()?;
135        if query.negated.contains(doc_id) {
136            return None;
137        }
138        let params = &self.fts.bm25_params;
139        let contributions: Vec<Option<f32>> = query
140            .blocks
141            .iter()
142            .map(|blocks| {
143                blocks.lookup(doc_id).map(|(tf, fieldnorm)| {
144                    term_score(
145                        crate::bm25::idf(blocks.df, query.total_docs),
146                        tf,
147                        fieldnorm,
148                        query.avg_doc_len,
149                        params,
150                    )
151                })
152            })
153            .collect();
154        mode.score(query, &contributions)
155    }
156}
157
158#[cfg(test)]
159mod tests {
160    use nodedb_types::{Surrogate, SurrogateBitmap};
161
162    use super::DocScore;
163    use crate::backend::memory::MemoryBackend;
164    use crate::index::FtsIndex;
165    use crate::posting::QueryMode;
166    use crate::search::bm25_search::FtsSearchParams;
167    use crate::search::query_terms::TextQuery;
168    use crate::search::staged::{StagedDoc, StagedView};
169    use crate::test_support::test_governor;
170
171    const DB: u64 = 0;
172    const T: u64 = 1;
173
174    fn query(text: &str) -> TextQuery<'_> {
175        TextQuery {
176            query: text,
177            fuzzy_enabled: false,
178            mode: QueryMode::And,
179        }
180    }
181
182    /// A document's point score equals the score a search returns for it.
183    #[test]
184    fn point_scores_equal_search_scores() {
185        let idx = FtsIndex::new(MemoryBackend::new(), test_governor());
186        idx.index_document(DB, T, "docs", Surrogate(1), "rust systems language")
187            .unwrap();
188        idx.index_document(DB, T, "docs", Surrogate(2), "rust rust web")
189            .unwrap();
190        idx.index_document(DB, T, "docs", Surrogate(3), "python scripting")
191            .unwrap();
192
193        let hits = idx
194            .search(
195                DB,
196                T,
197                "docs",
198                FtsSearchParams {
199                    query: "rust",
200                    top_k: 10,
201                    fuzzy_enabled: false,
202                    mode: QueryMode::And,
203                    prefilter: None,
204                },
205            )
206            .unwrap();
207        let scorer = idx
208            .doc_scorer(DB, T, "docs", query("rust"), None, None)
209            .unwrap();
210        let docs: Vec<Surrogate> = hits.iter().map(|h| h.doc_id).collect();
211        let scores = scorer.score(&docs).unwrap();
212        for (hit, score) in hits.iter().zip(scores) {
213            assert_eq!(score, DocScore::Match(hit.score));
214        }
215        assert_eq!(
216            scorer.score(&[Surrogate(3), Surrogate(99)]).unwrap(),
217            vec![DocScore::Miss, DocScore::Absent]
218        );
219    }
220
221    /// The fallback decision reads the eligible rows: with the only AND
222    /// match outside them, a one-word match scores under OR fallback.
223    #[test]
224    fn fallback_follows_the_eligible_rows() {
225        let idx = FtsIndex::new(MemoryBackend::new(), test_governor());
226        idx.index_document(DB, T, "docs", Surrogate(1), "alpha bravo")
227            .unwrap();
228        idx.index_document(DB, T, "docs", Surrogate(2), "alpha charlie")
229            .unwrap();
230
231        let all = idx
232            .doc_scorer(DB, T, "docs", query("alpha bravo"), None, None)
233            .unwrap();
234        assert_eq!(all.score(&[Surrogate(2)]).unwrap(), vec![DocScore::Miss]);
235
236        let mut eligible = SurrogateBitmap::new();
237        eligible.insert(Surrogate(2));
238        let scoped = idx
239            .doc_scorer(DB, T, "docs", query("alpha bravo"), Some(&eligible), None)
240            .unwrap();
241        assert!(matches!(
242            scoped.score(&[Surrogate(2)]).unwrap()[0],
243            DocScore::Match(_)
244        ));
245    }
246
247    /// A staged document scores from its tokens; a hidden indexed document
248    /// the transaction removed is absent.
249    #[test]
250    fn staged_documents_score_and_hidden_documents_are_absent() {
251        let idx = FtsIndex::new(MemoryBackend::new(), test_governor());
252        idx.index_document(DB, T, "docs", Surrogate(1), "rust language")
253            .unwrap();
254        let mut hidden = SurrogateBitmap::new();
255        hidden.insert(Surrogate(1));
256        let view = StagedView::new(
257            hidden,
258            false,
259            vec![StagedDoc {
260                doc_id: Surrogate(2),
261                tokens: vec!["rust".into()],
262            }],
263        );
264        let scorer = idx
265            .doc_scorer(DB, T, "docs", query("rust"), None, Some(view))
266            .unwrap();
267        let scores = scorer.score(&[Surrogate(1), Surrogate(2)]).unwrap();
268        assert_eq!(scores[0], DocScore::Absent);
269        assert!(matches!(scores[1], DocScore::Match(_)));
270    }
271}
nodedb_fts/search/doc_scorer.rs

nodedb_fts/search/
doc_scorer.rs
