nodedb_lite/engine/fts/manager/
phrase.rs

1// SPDX-License-Identifier: Apache-2.0
2
3//! Exact phrase search for [`FtsCollectionManager`].
4
5use std::collections::HashMap;
6
7use nodedb_fts::IndexScope;
8use nodedb_fts::posting::QueryMode as FtsQueryMode;
9use nodedb_types::Surrogate;
10use nodedb_types::text_search::TextSearchParams;
11
12use super::registry::{FtsCollectionManager, FtsResult, read_err};
13use super::search::Resolved;
14use crate::error::LiteError;
15
16/// Positions of one analyzed term, per document.
17type TermPositions = HashMap<Surrogate, Vec<u32>>;
18
19/// Postings of `term` in the resolved index, as document → positions.
20///
21/// The resolved index key is the memtable collection, with the
22/// whole-document scope. Lite is single-database and single-tenant, so both
23/// ids are 0.
24fn term_positions(index: &Resolved<'_>, term: &str) -> TermPositions {
25    index
26        .idx
27        .memtable()
28        .get_postings(0, 0, IndexScope::document(&index.key), term)
29        .into_iter()
30        .map(|p| (p.doc_id, p.positions))
31        .collect()
32}
33
34/// The earliest position `p` at which term `i` of the phrase sits at
35/// `p + i` for every term, if any.
36fn phrase_anchor(doc: Surrogate, per_term: &[TermPositions]) -> Option<u32> {
37    let (first, rest) = per_term.split_first()?;
38    let anchors = first.get(&doc)?;
39    let rest: Vec<&Vec<u32>> = rest.iter().map(|t| t.get(&doc)).collect::<Option<_>>()?;
40    anchors.iter().copied().find(|&p| {
41        rest.iter().enumerate().all(|(i, positions)| {
42            let offset = i as u32 + 1;
43            p.checked_add(offset)
44                .is_some_and(|want| positions.binary_search(&want).is_ok())
45        })
46    })
47}
48
49impl FtsCollectionManager {
50    /// Search the `field` index for documents where `terms` appear as an
51    /// exact consecutive phrase. An empty `field` searches the
52    /// whole-document index.
53    ///
54    /// The phrase runs through the index's analyzer first, so it matches the
55    /// same stemmed, lower-cased tokens the documents were indexed as.
56    /// Candidates are the documents holding every token at consecutive
57    /// positions. Each one scores its BM25 score for the phrase with an
58    /// earlier-position bonus.
59    pub fn phrase_search(
60        &self,
61        collection: &str,
62        field: &str,
63        terms: &[String],
64        top_k: usize,
65        params: &TextSearchParams,
66    ) -> Result<Vec<FtsResult>, LiteError> {
67        if terms.is_empty() {
68            return Ok(Vec::new());
69        }
70        let Some(index) = self.resolve(collection, field)? else {
71            return Ok(Vec::new());
72        };
73        let phrase = terms.join(" ");
74        let tokens = index
75            .idx
76            .analyze_for_collection(0, 0, &index.key, &phrase)
77            .map_err(|e| read_err(collection, e))?;
78        if tokens.is_empty() {
79            return Ok(Vec::new());
80        }
81
82        let per_term: Vec<TermPositions> =
83            tokens.iter().map(|t| term_positions(&index, t)).collect();
84        let Some(first) = per_term.first() else {
85            return Ok(Vec::new());
86        };
87        let anchored: Vec<(Surrogate, u32)> = first
88            .keys()
89            .filter_map(|&doc| phrase_anchor(doc, &per_term).map(|p| (doc, p)))
90            .collect();
91        if anchored.is_empty() {
92            return Ok(Vec::new());
93        }
94
95        // Every phrase match holds all tokens, so an AND query sized to the
96        // first token's document count returns a score for each of them.
97        let scores: HashMap<Surrogate, (f32, bool)> = Self::bm25(
98            collection,
99            &index,
100            &phrase,
101            first.len(),
102            params.fuzzy,
103            FtsQueryMode::And,
104        )?
105        .into_iter()
106        .map(|r| (r.doc_id, (r.score, r.fuzzy)))
107        .collect();
108
109        let mut hits: Vec<FtsResult> = anchored
110            .into_iter()
111            .map(|(doc, earliest)| -> Result<FtsResult, LiteError> {
112                let (score, fuzzy) = scores.get(&doc).copied().unwrap_or((0.0, false));
113                let position_bonus = 1.0 / (1.0 + earliest as f32 * 0.01);
114                Ok(FtsResult {
115                    doc_id: self.doc_id_of(collection, doc)?.to_owned(),
116                    score: score * position_bonus,
117                    fuzzy,
118                })
119            })
120            .collect::<Result<_, _>>()?;
121
122        hits.sort_by(|a, b| b.score.total_cmp(&a.score));
123        hits.truncate(top_k);
124        Ok(hits)
125    }
126}
127
128#[cfg(test)]
129mod tests {
130    use nodedb_types::text_search::{QueryMode, TextSearchParams};
131
132    use super::super::registry::test_governor;
133    use super::FtsCollectionManager;
134
135    fn default_params() -> TextSearchParams {
136        TextSearchParams {
137            fuzzy: false,
138            mode: QueryMode::Or,
139        }
140    }
141
142    fn phrase(mgr: &FtsCollectionManager, field: &str, words: &[&str]) -> Vec<String> {
143        let terms: Vec<String> = words.iter().map(|w| (*w).to_owned()).collect();
144        mgr.phrase_search("col", field, &terms, 10, &default_params())
145            .expect("phrase search must succeed")
146            .into_iter()
147            .map(|r| r.doc_id)
148            .collect()
149    }
150
151    #[test]
152    fn phrase_search_finds_exact_phrase() {
153        let mut mgr = FtsCollectionManager::new(test_governor());
154        mgr.index_document("col", "doc1", "the quick brown fox jumps over")
155            .expect("index update must succeed");
156        mgr.index_document("col", "doc2", "the brown quick fox")
157            .expect("index update must succeed");
158
159        let ids = phrase(&mgr, "", &["quick", "brown"]);
160        assert!(ids.contains(&"doc1".to_owned()), "doc1 holds 'quick brown'");
161        assert!(
162            !ids.contains(&"doc2".to_owned()),
163            "doc2 holds 'brown quick'"
164        );
165    }
166
167    #[test]
168    fn phrase_search_no_results_for_nonexistent_phrase() {
169        let mut mgr = FtsCollectionManager::new(test_governor());
170        mgr.index_document("col", "doc1", "the quick brown fox")
171            .expect("index update must succeed");
172        assert!(phrase(&mgr, "", &["fox", "jumps"]).is_empty());
173    }
174
175    #[test]
176    fn phrase_terms_match_the_analyzed_tokens() {
177        let mut mgr = FtsCollectionManager::new(test_governor());
178        mgr.index_document("col", "doc1", "the quick brown fox")
179            .expect("index update must succeed");
180        assert_eq!(phrase(&mgr, "", &["Quick", "BROWN"]), vec!["doc1"]);
181    }
182
183    #[test]
184    fn phrase_search_is_scoped_to_the_named_field() {
185        let mut mgr = FtsCollectionManager::new(test_governor());
186        mgr.index_field("col", "title", "doc1", "quick brown fox")
187            .expect("index update must succeed");
188        mgr.index_field("col", "body", "doc2", "quick brown fox")
189            .expect("index update must succeed");
190        assert_eq!(phrase(&mgr, "title", &["quick", "brown"]), vec!["doc1"]);
191    }
192}
nodedb_lite/engine/fts/manager/phrase.rs

nodedb_lite/engine/fts/manager/
phrase.rs
