lb_tantivy/query/
fuzzy_query.rs

1use levenshtein_automata::{Distance, LevenshteinAutomatonBuilder, DFA};
2use once_cell::sync::OnceCell;
3use tantivy_fst::Automaton;
4
5use crate::query::{AutomatonWeight, EnableScoring, Query, Weight};
6use crate::schema::{Term, Type};
7use crate::TantivyError::InvalidArgument;
8
9pub(crate) struct DfaWrapper(pub DFA);
10
11impl Automaton for DfaWrapper {
12    type State = u32;
13
14    fn start(&self) -> Self::State {
15        self.0.initial_state()
16    }
17
18    fn is_match(&self, state: &Self::State) -> bool {
19        match self.0.distance(*state) {
20            Distance::Exact(_) => true,
21            Distance::AtLeast(_) => false,
22        }
23    }
24
25    fn can_match(&self, state: &u32) -> bool {
26        *state != levenshtein_automata::SINK_STATE
27    }
28
29    fn accept(&self, state: &Self::State, byte: u8) -> Self::State {
30        self.0.transition(*state, byte)
31    }
32}
33
34/// A Fuzzy Query matches all of the documents
35/// containing a specific term that is within
36/// Levenshtein distance
37/// ```rust
38/// use tantivy::collector::{Count, TopDocs};
39/// use tantivy::query::FuzzyTermQuery;
40/// use tantivy::schema::{Schema, TEXT};
41/// use tantivy::{doc, Index, IndexWriter, Term};
42///
43/// fn example() -> tantivy::Result<()> {
44///     let mut schema_builder = Schema::builder();
45///     let title = schema_builder.add_text_field("title", TEXT);
46///     let schema = schema_builder.build();
47///     let index = Index::create_in_ram(schema);
48///     {
49///         let mut index_writer: IndexWriter = index.writer(15_000_000)?;
50///         index_writer.add_document(doc!(
51///             title => "The Name of the Wind",
52///         ))?;
53///         index_writer.add_document(doc!(
54///             title => "The Diary of Muadib",
55///         ))?;
56///         index_writer.add_document(doc!(
57///             title => "A Dairy Cow",
58///         ))?;
59///         index_writer.add_document(doc!(
60///             title => "The Diary of a Young Girl",
61///         ))?;
62///         index_writer.commit()?;
63///     }
64///     let reader = index.reader()?;
65///     let searcher = reader.searcher();
66///
67///     {
68///         let term = Term::from_field_text(title, "Diary");
69///         let query = FuzzyTermQuery::new(term, 1, true);
70///         let (top_docs, count) = searcher.search(&query, &(TopDocs::with_limit(2), Count)).unwrap();
71///         assert_eq!(count, 2);
72///         assert_eq!(top_docs.len(), 2);
73///     }
74///
75///     Ok(())
76/// }
77/// # assert!(example().is_ok());
78/// ```
79#[derive(Debug, Clone)]
80pub struct FuzzyTermQuery {
81    /// What term are we searching
82    term: Term,
83    /// How many changes are we going to allow
84    distance: u8,
85    /// Should a transposition cost 1 or 2?
86    transposition_cost_one: bool,
87    /// is a starts with query
88    prefix: bool,
89}
90
91impl FuzzyTermQuery {
92    /// Creates a new Fuzzy Query
93    pub fn new(term: Term, distance: u8, transposition_cost_one: bool) -> FuzzyTermQuery {
94        FuzzyTermQuery {
95            term,
96            distance,
97            transposition_cost_one,
98            prefix: false,
99        }
100    }
101
102    /// Creates a new Fuzzy Query of the Term prefix
103    pub fn new_prefix(term: Term, distance: u8, transposition_cost_one: bool) -> FuzzyTermQuery {
104        FuzzyTermQuery {
105            term,
106            distance,
107            transposition_cost_one,
108            prefix: true,
109        }
110    }
111
112    fn specialized_weight(&self) -> crate::Result<AutomatonWeight<DfaWrapper>> {
113        static AUTOMATON_BUILDER: [[OnceCell<LevenshteinAutomatonBuilder>; 2]; 3] = [
114            [OnceCell::new(), OnceCell::new()],
115            [OnceCell::new(), OnceCell::new()],
116            [OnceCell::new(), OnceCell::new()],
117        ];
118
119        let automaton_builder = AUTOMATON_BUILDER
120            .get(self.distance as usize)
121            .ok_or_else(|| {
122                InvalidArgument(format!(
123                    "Levenshtein distance of {} is not allowed. Choose a value less than {}",
124                    self.distance,
125                    AUTOMATON_BUILDER.len()
126                ))
127            })?
128            .get(self.transposition_cost_one as usize)
129            .unwrap()
130            .get_or_init(|| {
131                LevenshteinAutomatonBuilder::new(self.distance, self.transposition_cost_one)
132            });
133
134        let term_value = self.term.value();
135
136        let term_text = if term_value.typ() == Type::Json {
137            if let Some(json_path_type) = term_value.json_path_type() {
138                if json_path_type != Type::Str {
139                    return Err(InvalidArgument(format!(
140                        "The fuzzy term query requires a string path type for a json term. Found \
141                         {json_path_type:?}"
142                    )));
143                }
144            }
145
146            std::str::from_utf8(self.term.serialized_value_bytes()).map_err(|_| {
147                InvalidArgument(
148                    "Failed to convert json term value bytes to utf8 string.".to_string(),
149                )
150            })?
151        } else {
152            term_value.as_str().ok_or_else(|| {
153                InvalidArgument("The fuzzy term query requires a string term.".to_string())
154            })?
155        };
156        let automaton = if self.prefix {
157            automaton_builder.build_prefix_dfa(term_text)
158        } else {
159            automaton_builder.build_dfa(term_text)
160        };
161
162        if let Some((json_path_bytes, _)) = term_value.as_json() {
163            Ok(AutomatonWeight::new_for_json_path(
164                self.term.field(),
165                DfaWrapper(automaton),
166                json_path_bytes,
167            ))
168        } else {
169            Ok(AutomatonWeight::new(
170                self.term.field(),
171                DfaWrapper(automaton),
172            ))
173        }
174    }
175}
176
177impl Query for FuzzyTermQuery {
178    fn weight(&self, _enable_scoring: EnableScoring<'_>) -> crate::Result<Box<dyn Weight>> {
179        Ok(Box::new(self.specialized_weight()?))
180    }
181}
182
183#[cfg(test)]
184mod test {
185    use super::FuzzyTermQuery;
186    use crate::collector::{Count, TopDocs};
187    use crate::indexer::NoMergePolicy;
188    use crate::query::QueryParser;
189    use crate::schema::{Schema, STORED, TEXT};
190    use crate::{assert_nearly_equals, Index, IndexWriter, TantivyDocument, Term};
191
192    #[test]
193    pub fn test_fuzzy_json_path() -> crate::Result<()> {
194        // # Defining the schema
195        let mut schema_builder = Schema::builder();
196        let attributes = schema_builder.add_json_field("attributes", TEXT | STORED);
197        let schema = schema_builder.build();
198
199        // # Indexing documents
200        let index = Index::create_in_ram(schema.clone());
201
202        let mut index_writer = index.writer_for_tests()?;
203        index_writer.set_merge_policy(Box::new(NoMergePolicy));
204        let doc = TantivyDocument::parse_json(
205            &schema,
206            r#"{
207            "attributes": {
208                "a": "japan"
209            }
210        }"#,
211        )?;
212        index_writer.add_document(doc)?;
213        let doc = TantivyDocument::parse_json(
214            &schema,
215            r#"{
216            "attributes": {
217                "aa": "japan"
218            }
219        }"#,
220        )?;
221        index_writer.add_document(doc)?;
222        index_writer.commit()?;
223
224        let reader = index.reader()?;
225        let searcher = reader.searcher();
226
227        // # Fuzzy search
228        let query_parser = QueryParser::for_index(&index, vec![attributes]);
229
230        let get_json_path_term = |query: &str| -> crate::Result<Term> {
231            let query = query_parser.parse_query(query)?;
232            let mut terms = Vec::new();
233            query.query_terms(&mut |term, _| {
234                terms.push(term.clone());
235            });
236
237            Ok(terms[0].clone())
238        };
239
240        // shall not match the first document due to json path mismatch
241        {
242            let term = get_json_path_term("attributes.aa:japan")?;
243            let fuzzy_query = FuzzyTermQuery::new(term, 2, true);
244            let top_docs = searcher.search(&fuzzy_query, &TopDocs::with_limit(2))?;
245            assert_eq!(top_docs.len(), 1, "Expected only 1 document");
246            assert_eq!(top_docs[0].1.doc_id, 1, "Expected the second document");
247        }
248
249        // shall match the first document because Levenshtein distance is 1 (substitute 'o' with
250        // 'a')
251        {
252            let term = get_json_path_term("attributes.a:japon")?;
253
254            let fuzzy_query = FuzzyTermQuery::new(term, 1, true);
255            let top_docs = searcher.search(&fuzzy_query, &TopDocs::with_limit(2))?;
256            assert_eq!(top_docs.len(), 1, "Expected only 1 document");
257            assert_eq!(top_docs[0].1.doc_id, 0, "Expected the first document");
258        }
259
260        // shall not match because non-prefix Levenshtein distance is more than 1 (add 'a' and 'n')
261        {
262            let term = get_json_path_term("attributes.a:jap")?;
263
264            let fuzzy_query = FuzzyTermQuery::new(term, 1, true);
265            let top_docs = searcher.search(&fuzzy_query, &TopDocs::with_limit(2))?;
266            assert_eq!(top_docs.len(), 0, "Expected no document");
267        }
268
269        Ok(())
270    }
271
272    #[test]
273    pub fn test_fuzzy_term() -> crate::Result<()> {
274        let mut schema_builder = Schema::builder();
275        let country_field = schema_builder.add_text_field("country", TEXT);
276        let schema = schema_builder.build();
277        let index = Index::create_in_ram(schema);
278        {
279            let mut index_writer: IndexWriter = index.writer_for_tests()?;
280            index_writer.add_document(doc!(
281                country_field => "japan",
282            ))?;
283            index_writer.add_document(doc!(
284                country_field => "korea",
285            ))?;
286            index_writer.commit()?;
287        }
288        let reader = index.reader()?;
289        let searcher = reader.searcher();
290
291        // passes because Levenshtein distance is 1 (substitute 'o' with 'a')
292        {
293            let term = Term::from_field_text(country_field, "japon");
294            let fuzzy_query = FuzzyTermQuery::new(term, 1, true);
295            let top_docs = searcher.search(&fuzzy_query, &TopDocs::with_limit(2))?;
296            assert_eq!(top_docs.len(), 1, "Expected only 1 document");
297            let (score, _) = top_docs[0];
298            assert_nearly_equals!(1.0, score);
299        }
300
301        // fails because non-prefix Levenshtein distance is more than 1 (add 'a' and 'n')
302        {
303            let term = Term::from_field_text(country_field, "jap");
304
305            let fuzzy_query = FuzzyTermQuery::new(term, 1, true);
306            let top_docs = searcher.search(&fuzzy_query, &TopDocs::with_limit(2))?;
307            assert_eq!(top_docs.len(), 0, "Expected no document");
308        }
309
310        // passes because prefix Levenshtein distance is 0
311        {
312            let term = Term::from_field_text(country_field, "jap");
313            let fuzzy_query = FuzzyTermQuery::new_prefix(term, 1, true);
314            let top_docs = searcher.search(&fuzzy_query, &TopDocs::with_limit(2))?;
315            assert_eq!(top_docs.len(), 1, "Expected only 1 document");
316            let (score, _) = top_docs[0];
317            assert_nearly_equals!(1.0, score);
318        }
319        Ok(())
320    }
321
322    #[test]
323    pub fn test_fuzzy_term_transposition_cost_one() -> crate::Result<()> {
324        let mut schema_builder = Schema::builder();
325        let country_field = schema_builder.add_text_field("country", TEXT);
326        let schema = schema_builder.build();
327        let index = Index::create_in_ram(schema);
328        let mut index_writer: IndexWriter = index.writer_for_tests()?;
329        index_writer.add_document(doc!(country_field => "japan"))?;
330        index_writer.commit()?;
331        let reader = index.reader()?;
332        let searcher = reader.searcher();
333        let term_jaapn = Term::from_field_text(country_field, "jaapn");
334        {
335            let fuzzy_query_transposition = FuzzyTermQuery::new(term_jaapn.clone(), 1, true);
336            let count = searcher.search(&fuzzy_query_transposition, &Count)?;
337            assert_eq!(count, 1);
338        }
339        {
340            let fuzzy_query_transposition = FuzzyTermQuery::new(term_jaapn, 1, false);
341            let count = searcher.search(&fuzzy_query_transposition, &Count)?;
342            assert_eq!(count, 0);
343        }
344        Ok(())
345    }
346}
lb_tantivy/query/fuzzy_query.rs

lb_tantivy/query/
fuzzy_query.rs
