Skip to main content

zeppelin_embed/fts/
prefix.rs

1//! Prefix queries as a sorted-dictionary range scan.
2//!
3//! # The boundary rule, pinned (task 15 D2)
4//!
5//! Prefix matching is **by codepoint after analysis**, not by grapheme
6//! cluster. A prefix that ends in the middle of a grapheme — between a base
7//! letter and its combining mark, or inside a family emoji's zero-width
8//! joiner sequence — still matches, because the dictionary is a sorted byte
9//! array and the prefix is compared as bytes of the analyzed form.
10//!
11//! Two consequences worth stating rather than discovering:
12//!
13//! - The prefix is analyzed with the same pipeline as the documents, so
14//!   `Caf` and `caf` are the same prefix, and `café` folds to `cafe` before
15//!   the comparison. A user typing accented text gets the results they
16//!   expect.
17//! - Because folding can expand one character into several (`ß` becomes
18//!   `ss`), a prefix can match a term whose *surface* form it is not a
19//!   prefix of. That is correct: the index stores analyzed terms.
20//!
21//! # An empty prefix is an error, not a full scan
22//!
23//! Returning every term for an empty prefix is a denial-of-service wearing a
24//! query's clothes. It is a typed error.
25
26use super::dict::{DictError, TermDictionary};
27
28/// One prefix match.
29#[derive(Clone, Debug, Eq, PartialEq)]
30pub struct PrefixMatch {
31    /// Index of the term in the dictionary's flat order.
32    pub term_index: usize,
33    /// The matched term.
34    pub term: Vec<u8>,
35}
36
37/// Returns every dictionary term starting with `prefix`.
38///
39/// # Errors
40///
41/// Returns [`DictError::EmptyTerm`] for an empty prefix.
42pub fn search(dictionary: &TermDictionary, prefix: &[u8]) -> Result<Vec<PrefixMatch>, DictError> {
43    let range = dictionary.prefix_range(prefix)?;
44    Ok(range
45        .filter_map(|term_index| {
46            dictionary.term_at(term_index).map(|term| PrefixMatch {
47                term_index,
48                term: term.to_vec(),
49            })
50        })
51        .collect())
52}
53
54/// Returns the number of terms a prefix would match, without collecting.
55///
56/// # Errors
57///
58/// Returns [`DictError::EmptyTerm`] for an empty prefix.
59pub fn count(dictionary: &TermDictionary, prefix: &[u8]) -> Result<usize, DictError> {
60    Ok(dictionary.prefix_range(prefix)?.len())
61}
62
63#[cfg(test)]
64#[allow(
65    clippy::expect_used,
66    clippy::indexing_slicing,
67    clippy::panic,
68    clippy::unwrap_used
69)]
70mod tests {
71    use super::*;
72    use crate::fts::dict::TermInfo;
73    use crate::fts::tokenizer::{Analyzer, TokenizerConfig};
74
75    fn dictionary(terms: &[&str]) -> TermDictionary {
76        let mut sorted: Vec<&str> = terms.to_vec();
77        sorted.sort_unstable();
78        sorted.dedup();
79        let mut dictionary = TermDictionary::default();
80        for term in sorted {
81            dictionary
82                .push(term.as_bytes(), TermInfo::default())
83                .expect("sorted");
84        }
85        dictionary
86    }
87
88    fn matched(dictionary: &TermDictionary, prefix: &str) -> Vec<String> {
89        search(dictionary, prefix.as_bytes())
90            .expect("valid prefix")
91            .into_iter()
92            .filter_map(|entry| String::from_utf8(entry.term).ok())
93            .collect()
94    }
95
96    const VOCABULARY: [&str; 10] = [
97        "engine",
98        "engineer",
99        "engineering",
100        "engines",
101        "england",
102        "lexical",
103        "search",
104        "searching",
105        "seat",
106        "zebra",
107    ];
108
109    #[test]
110    fn prefix_results_equal_brute_force_over_the_dictionary() {
111        let dictionary = dictionary(&VOCABULARY);
112        for prefix in [
113            "e", "en", "eng", "engine", "s", "sea", "search", "z", "q", "engines",
114        ] {
115            let mut expected: Vec<String> = VOCABULARY
116                .iter()
117                .filter(|term| term.starts_with(prefix))
118                .map(|term| (*term).to_owned())
119                .collect();
120            expected.sort();
121            let mut actual = matched(&dictionary, prefix);
122            actual.sort();
123            assert_eq!(actual, expected, "prefix {prefix:?}");
124            assert_eq!(
125                count(&dictionary, prefix.as_bytes()).expect("valid"),
126                expected.len()
127            );
128        }
129    }
130
131    #[test]
132    fn empty_prefix_is_rejected_with_a_typed_error() {
133        let dictionary = dictionary(&VOCABULARY);
134        assert_eq!(search(&dictionary, b""), Err(DictError::EmptyTerm));
135        assert_eq!(count(&dictionary, b""), Err(DictError::EmptyTerm));
136    }
137
138    #[test]
139    fn a_prefix_matching_nothing_returns_an_empty_range_not_an_error() {
140        let dictionary = dictionary(&VOCABULARY);
141        assert!(matched(&dictionary, "qqq").is_empty());
142        assert!(matched(&dictionary, "zzzz").is_empty());
143    }
144
145    #[test]
146    fn a_term_is_a_prefix_of_itself() {
147        let dictionary = dictionary(&VOCABULARY);
148        assert!(matched(&dictionary, "zebra").contains(&String::from("zebra")));
149    }
150
151    #[test]
152    fn prefix_ending_mid_grapheme_matches_by_codepoint_rule() {
153        // "cafe\u{0301}" folds to "cafe"; the prefix "caf" is a codepoint
154        // prefix of the ANALYZED term, which is what is compared.
155        let analyzer = Analyzer::new(TokenizerConfig::text_default()).expect("valid config");
156        let folded: Vec<String> = analyzer
157            .analyze("cafe\u{0301}")
158            .into_iter()
159            .map(|token| token.term)
160            .collect();
161        let refs: Vec<&str> = folded.iter().map(String::as_str).collect();
162        let dictionary = dictionary(&refs);
163        // A prefix that stops before the accented position still matches.
164        assert!(!matched(&dictionary, "caf").is_empty());
165        // And the combining mark never appears in the dictionary at all,
166        // because folding removed it before indexing.
167        assert!(
168            folded.iter().all(|term| !term.contains('\u{0301}')),
169            "a combining mark survived into the dictionary"
170        );
171    }
172
173    #[test]
174    fn a_prefix_can_match_a_term_whose_surface_form_it_does_not_prefix() {
175        // "stra" is a prefix of the ANALYZED "strasse", though the surface
176        // form is "straße" and "stra" is a prefix of that too. The point is
177        // that the comparison is against the analyzed term.
178        let analyzer = Analyzer::new(TokenizerConfig::text_default()).expect("valid config");
179        let terms: Vec<String> = analyzer
180            .analyze("stra\u{00DF}e")
181            .into_iter()
182            .map(|token| token.term)
183            .collect();
184        let refs: Vec<&str> = terms.iter().map(String::as_str).collect();
185        let dictionary = dictionary(&refs);
186        assert!(!matched(&dictionary, "strass").is_empty());
187    }
188
189    #[test]
190    fn an_empty_dictionary_answers_every_prefix_with_nothing() {
191        let dictionary = TermDictionary::default();
192        assert!(search(&dictionary, b"anything").expect("valid").is_empty());
193    }
194
195    #[test]
196    fn matches_carry_the_index_the_dictionary_agrees_with() {
197        let dictionary = dictionary(&VOCABULARY);
198        for entry in search(&dictionary, b"engine").expect("valid") {
199            assert_eq!(
200                dictionary.term_at(entry.term_index),
201                Some(entry.term.as_slice())
202            );
203        }
204    }
205}