pub use self::facet::{FacetDistribution, Filter, DEFAULT_VALUES_PER_FACET}; use self::fst_utils::{Complement, Intersection, StartsWith, Union}; pub use self::matches::{ FormatOptions, MatchBounds, Matcher, MatcherBuilder, MatchingWord, MatchingWords, }; use crate::{ execute_search, AscDesc, DefaultSearchLogger, DocumentId, Index, Result, SearchContext, }; use fst::automaton::Str; use fst::{Automaton, IntoStreamer, Streamer}; use levenshtein_automata::{LevenshteinAutomatonBuilder as LevBuilder, DFA}; use once_cell::sync::Lazy; use roaring::bitmap::RoaringBitmap; use std::borrow::Cow; use std::collections::hash_map::{Entry, HashMap}; use std::fmt; use std::result::Result as StdResult; use std::str::Utf8Error; // Building these factories is not free. static LEVDIST0: Lazy = Lazy::new(|| LevBuilder::new(0, true)); static LEVDIST1: Lazy = Lazy::new(|| LevBuilder::new(1, true)); static LEVDIST2: Lazy = Lazy::new(|| LevBuilder::new(2, true)); pub mod facet; mod fst_utils; mod matches; pub mod new; mod query_tree; pub struct Search<'a> { query: Option, // this should be linked to the String in the query filter: Option>, offset: usize, limit: usize, sort_criteria: Option>, terms_matching_strategy: TermsMatchingStrategy, authorize_typos: bool, words_limit: usize, exhaustive_number_hits: bool, criterion_implementation_strategy: CriterionImplementationStrategy, rtxn: &'a heed::RoTxn<'a>, index: &'a Index, } impl<'a> Search<'a> { pub fn new(rtxn: &'a heed::RoTxn, index: &'a Index) -> Search<'a> { Search { query: None, filter: None, offset: 0, limit: 20, sort_criteria: None, terms_matching_strategy: TermsMatchingStrategy::default(), authorize_typos: true, exhaustive_number_hits: false, words_limit: 10, criterion_implementation_strategy: CriterionImplementationStrategy::default(), rtxn, index, } } pub fn query(&mut self, query: impl Into) -> &mut Search<'a> { self.query = Some(query.into()); self } pub fn offset(&mut self, offset: usize) -> &mut Search<'a> { self.offset = offset; self } pub fn limit(&mut self, limit: usize) -> &mut Search<'a> { self.limit = limit; self } pub fn sort_criteria(&mut self, criteria: Vec) -> &mut Search<'a> { self.sort_criteria = Some(criteria); self } pub fn terms_matching_strategy(&mut self, value: TermsMatchingStrategy) -> &mut Search<'a> { self.terms_matching_strategy = value; self } pub fn authorize_typos(&mut self, value: bool) -> &mut Search<'a> { self.authorize_typos = value; self } pub fn words_limit(&mut self, value: usize) -> &mut Search<'a> { self.words_limit = value; self } pub fn filter(&mut self, condition: Filter<'a>) -> &mut Search<'a> { self.filter = Some(condition); self } /// Force the search to exhastivelly compute the number of candidates, /// this will increase the search time but allows finite pagination. pub fn exhaustive_number_hits(&mut self, exhaustive_number_hits: bool) -> &mut Search<'a> { self.exhaustive_number_hits = exhaustive_number_hits; self } pub fn criterion_implementation_strategy( &mut self, strategy: CriterionImplementationStrategy, ) -> &mut Search<'a> { self.criterion_implementation_strategy = strategy; self } fn is_typo_authorized(&self) -> Result { let index_authorizes_typos = self.index.authorize_typos(self.rtxn)?; // only authorize typos if both the index and the query allow it. Ok(self.authorize_typos && index_authorizes_typos) } pub fn execute(&self) -> Result { let mut ctx = SearchContext::new(self.index, self.rtxn); execute_search( &mut ctx, &self.query, self.terms_matching_strategy, &self.filter, self.offset, self.limit, Some(self.words_limit), &mut DefaultSearchLogger, &mut DefaultSearchLogger, ) } } impl fmt::Debug for Search<'_> { fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { let Search { query, filter, offset, limit, sort_criteria, terms_matching_strategy, authorize_typos, words_limit, exhaustive_number_hits, criterion_implementation_strategy, rtxn: _, index: _, } = self; f.debug_struct("Search") .field("query", query) .field("filter", filter) .field("offset", offset) .field("limit", limit) .field("sort_criteria", sort_criteria) .field("terms_matching_strategy", terms_matching_strategy) .field("authorize_typos", authorize_typos) .field("exhaustive_number_hits", exhaustive_number_hits) .field("criterion_implementation_strategy", criterion_implementation_strategy) .field("words_limit", words_limit) .finish() } } #[derive(Default, Debug)] pub struct SearchResult { pub matching_words: MatchingWords, pub candidates: RoaringBitmap, // TODO those documents ids should be associated with their criteria scores. pub documents_ids: Vec, } #[derive(Debug, Default, Clone, Copy)] pub enum CriterionImplementationStrategy { OnlyIterative, OnlySetBased, #[default] Dynamic, } #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum TermsMatchingStrategy { // remove last word first Last, // all words are mandatory All, } impl Default for TermsMatchingStrategy { fn default() -> Self { Self::Last } } pub type WordDerivationsCache = HashMap<(String, bool, u8), Vec<(String, u8)>>; pub fn word_derivations<'c>( word: &str, is_prefix: bool, max_typo: u8, fst: &fst::Set>, cache: &'c mut WordDerivationsCache, ) -> StdResult<&'c [(String, u8)], Utf8Error> { match cache.entry((word.to_string(), is_prefix, max_typo)) { Entry::Occupied(entry) => Ok(entry.into_mut()), Entry::Vacant(entry) => { // println!("word derivations {word} {is_prefix} {max_typo}"); let mut derived_words = Vec::new(); if max_typo == 0 { if is_prefix { let prefix = Str::new(word).starts_with(); let mut stream = fst.search(prefix).into_stream(); while let Some(word) = stream.next() { let word = std::str::from_utf8(word)?; derived_words.push((word.to_string(), 0)); } } else if fst.contains(word) { derived_words.push((word.to_string(), 0)); } } else if max_typo == 1 { let dfa = build_dfa(word, 1, is_prefix); let starts = StartsWith(Str::new(get_first(word))); let mut stream = fst.search_with_state(Intersection(starts, &dfa)).into_stream(); while let Some((word, state)) = stream.next() { let word = std::str::from_utf8(word)?; let d = dfa.distance(state.1); derived_words.push((word.to_string(), d.to_u8())); } } else { let starts = StartsWith(Str::new(get_first(word))); let first = Intersection(build_dfa(word, 1, is_prefix), Complement(&starts)); let second_dfa = build_dfa(word, 2, is_prefix); let second = Intersection(&second_dfa, &starts); let automaton = Union(first, &second); let mut stream = fst.search_with_state(automaton).into_stream(); while let Some((found_word, state)) = stream.next() { let found_word = std::str::from_utf8(found_word)?; // in the case the typo is on the first letter, we know the number of typo // is two if get_first(found_word) != get_first(word) { derived_words.push((found_word.to_string(), 2)); } else { // Else, we know that it is the second dfa that matched and compute the // correct distance let d = second_dfa.distance((state.1).0); derived_words.push((found_word.to_string(), d.to_u8())); } } } Ok(entry.insert(derived_words)) } } } fn get_first(s: &str) -> &str { match s.chars().next() { Some(c) => &s[..c.len_utf8()], None => panic!("unexpected empty query"), } } pub fn build_dfa(word: &str, typos: u8, is_prefix: bool) -> DFA { let lev = match typos { 0 => &LEVDIST0, 1 => &LEVDIST1, _ => &LEVDIST2, }; if is_prefix { lev.build_prefix_dfa(word) } else { lev.build_dfa(word) } } #[cfg(test)] mod test { use super::*; use crate::index::tests::TempIndex; #[cfg(feature = "default")] #[test] fn test_kanji_language_detection() { let index = TempIndex::new(); index .add_documents(documents!([ { "id": 0, "title": "The quick (\"brown\") fox can't jump 32.3 feet, right? Brr, it's 29.3°F!" }, { "id": 1, "title": "東京のお寿司。" }, { "id": 2, "title": "הַשּׁוּעָל הַמָּהִיר (״הַחוּם״) לֹא יָכוֹל לִקְפֹּץ 9.94 מֶטְרִים, נָכוֹן? ברר, 1.5°C- בַּחוּץ!" } ])) .unwrap(); let txn = index.write_txn().unwrap(); let mut search = Search::new(&txn, &index); search.query("東京"); let SearchResult { documents_ids, .. } = search.execute().unwrap(); assert_eq!(documents_ids, vec![1]); } #[test] fn test_is_authorized_typos() { let index = TempIndex::new(); let mut txn = index.write_txn().unwrap(); let mut search = Search::new(&txn, &index); // default is authorized assert!(search.is_typo_authorized().unwrap()); search.authorize_typos(false); assert!(!search.is_typo_authorized().unwrap()); index.put_authorize_typos(&mut txn, false).unwrap(); txn.commit().unwrap(); let txn = index.read_txn().unwrap(); let mut search = Search::new(&txn, &index); assert!(!search.is_typo_authorized().unwrap()); search.authorize_typos(true); assert!(!search.is_typo_authorized().unwrap()); } #[test] fn test_one_typos_tolerance() { let fst = fst::Set::from_iter(["zealand"].iter()).unwrap().map_data(Cow::Owned).unwrap(); let mut cache = HashMap::new(); let found = word_derivations("zealend", false, 1, &fst, &mut cache).unwrap(); assert_eq!(found, &[("zealand".to_string(), 1)]); } #[test] fn test_one_typos_first_letter() { let fst = fst::Set::from_iter(["zealand"].iter()).unwrap().map_data(Cow::Owned).unwrap(); let mut cache = HashMap::new(); let found = word_derivations("sealand", false, 1, &fst, &mut cache).unwrap(); assert_eq!(found, &[]); } #[test] fn test_two_typos_tolerance() { let fst = fst::Set::from_iter(["zealand"].iter()).unwrap().map_data(Cow::Owned).unwrap(); let mut cache = HashMap::new(); let found = word_derivations("zealemd", false, 2, &fst, &mut cache).unwrap(); assert_eq!(found, &[("zealand".to_string(), 2)]); } #[test] fn test_two_typos_first_letter() { let fst = fst::Set::from_iter(["zealand"].iter()).unwrap().map_data(Cow::Owned).unwrap(); let mut cache = HashMap::new(); let found = word_derivations("sealand", false, 2, &fst, &mut cache).unwrap(); assert_eq!(found, &[("zealand".to_string(), 2)]); } #[test] fn test_prefix() { let fst = fst::Set::from_iter(["zealand"].iter()).unwrap().map_data(Cow::Owned).unwrap(); let mut cache = HashMap::new(); let found = word_derivations("ze", true, 0, &fst, &mut cache).unwrap(); assert_eq!(found, &[("zealand".to_string(), 0)]); } #[test] fn test_bad_prefix() { let fst = fst::Set::from_iter(["zealand"].iter()).unwrap().map_data(Cow::Owned).unwrap(); let mut cache = HashMap::new(); let found = word_derivations("se", true, 0, &fst, &mut cache).unwrap(); assert_eq!(found, &[]); } #[test] fn test_prefix_with_typo() { let fst = fst::Set::from_iter(["zealand"].iter()).unwrap().map_data(Cow::Owned).unwrap(); let mut cache = HashMap::new(); let found = word_derivations("zae", true, 1, &fst, &mut cache).unwrap(); assert_eq!(found, &[("zealand".to_string(), 1)]); } }