diff --git a/milli/src/search/criteria/mod.rs b/milli/src/search/criteria/mod.rs index 5c6561b62..9cb6547c9 100644 --- a/milli/src/search/criteria/mod.rs +++ b/milli/src/search/criteria/mod.rs @@ -36,8 +36,8 @@ impl Default for Candidates { } } pub trait Context { - fn query_docids(&self, query: &Query) -> anyhow::Result>; - fn query_pair_proximity_docids(&self, left: &Query, right: &Query, distance: u8) ->anyhow::Result>; + fn query_docids(&self, query: &Query) -> anyhow::Result; + fn query_pair_proximity_docids(&self, left: &Query, right: &Query, distance: u8) ->anyhow::Result; fn words_fst<'t>(&self) -> &'t fst::Set>; } pub struct HeedContext<'t> { @@ -48,10 +48,10 @@ pub struct HeedContext<'t> { } impl<'a> Context for HeedContext<'a> { - fn query_docids(&self, query: &Query) -> anyhow::Result> { + fn query_docids(&self, query: &Query) -> anyhow::Result { match (&query.kind, query.prefix) { (QueryKind::Exact { word, .. }, true) if self.in_prefix_cache(&word) => { - Ok(self.index.word_prefix_docids.get(self.rtxn, &word)?) + Ok(self.index.word_prefix_docids.get(self.rtxn, &word)?.unwrap_or_default()) }, (QueryKind::Exact { word, .. }, true) => { let words = word_typos(&word, true, 0, &self.words_fst)?; @@ -59,10 +59,10 @@ impl<'a> Context for HeedContext<'a> { for (word, _typo) in words { docids.union_with(&self.index.word_docids.get(self.rtxn, &word)?.unwrap_or_default()); } - Ok(Some(docids)) + Ok(docids) }, (QueryKind::Exact { word, .. }, false) => { - Ok(self.index.word_docids.get(self.rtxn, &word)?) + Ok(self.index.word_docids.get(self.rtxn, &word)?.unwrap_or_default()) }, (QueryKind::Tolerant { typo, word }, prefix) => { let words = word_typos(&word, prefix, *typo, &self.words_fst)?; @@ -70,81 +70,46 @@ impl<'a> Context for HeedContext<'a> { for (word, _typo) in words { docids.union_with(&self.index.word_docids.get(self.rtxn, &word)?.unwrap_or_default()); } - Ok(Some(docids)) + Ok(docids) }, } } - fn query_pair_proximity_docids(&self, left: &Query, right: &Query, distance: u8) -> anyhow::Result> { - // TODO add prefix cache for Tolerant-Exact-true and Exact-Exact-true + fn query_pair_proximity_docids(&self, left: &Query, right: &Query, distance: u8) -> anyhow::Result { match (&left.kind, &right.kind, right.prefix) { (QueryKind::Exact { word: left, .. }, QueryKind::Exact { word: right, .. }, true) if self.in_prefix_cache(&right) => { let key = (left.as_str(), right.as_str(), distance); - Ok(self.index.word_prefix_pair_proximity_docids.get(self.rtxn, &key)?) + Ok(self.index.word_prefix_pair_proximity_docids.get(self.rtxn, &key)?.unwrap_or_default()) }, (QueryKind::Tolerant { typo, word: left }, QueryKind::Exact { word: right, .. }, true) if self.in_prefix_cache(&right) => { - let words = word_typos(&left, false, *typo, &self.words_fst)?; - let mut docids = RoaringBitmap::new(); - for (word, _typo) in words { - let key = (word.as_str(), right.as_str(), distance); - docids.union_with(&self.index.word_prefix_pair_proximity_docids.get(self.rtxn, &key)?.unwrap_or_default()); - } - Ok(Some(docids)) + let l_words = word_typos(&left, false, *typo, &self.words_fst)?; + self.all_word_pair_proximity_docids(&l_words, &[(right, 0)], distance) }, (QueryKind::Exact { word: left, .. }, QueryKind::Exact { word: right, .. }, true) => { - let words = word_typos(&right, true, 0, &self.words_fst)?; - let mut docids = RoaringBitmap::new(); - for (word, _typo) in words { - let key = (left.as_str(), word.as_str(), distance); - docids.union_with(&self.index.word_pair_proximity_docids.get(self.rtxn, &key)?.unwrap_or_default()); - } - Ok(Some(docids)) + let r_words = word_typos(&right, true, 0, &self.words_fst)?; + self.all_word_pair_proximity_docids(&[(left, 0)], &r_words, distance) }, (QueryKind::Tolerant { typo, word: left }, QueryKind::Exact { word: right, .. }, true) => { let l_words = word_typos(&left, false, *typo, &self.words_fst)?; let r_words = word_typos(&right, true, 0, &self.words_fst)?; - let mut docids = RoaringBitmap::new(); - for (left, _typo) in l_words { - for (right, _typo) in r_words.iter() { - let key = (left.as_str(), right.as_str(), distance); - docids.union_with(&self.index.word_pair_proximity_docids.get(self.rtxn, &key)?.unwrap_or_default()); - } - } - Ok(Some(docids)) + self.all_word_pair_proximity_docids(&l_words, &r_words, distance) }, (QueryKind::Tolerant { typo, word: left }, QueryKind::Exact { word: right, .. }, false) => { - let words = word_typos(&left, false, *typo, &self.words_fst)?; - let mut docids = RoaringBitmap::new(); - for (word, _typo) in words { - let key = (word.as_str(), right.as_str(), distance); - docids.union_with(&self.index.word_pair_proximity_docids.get(self.rtxn, &key)?.unwrap_or_default()); - } - Ok(Some(docids)) + let l_words = word_typos(&left, false, *typo, &self.words_fst)?; + self.all_word_pair_proximity_docids(&l_words, &[(right, 0)], distance) }, (QueryKind::Exact { word: left, .. }, QueryKind::Tolerant { typo, word: right }, prefix) => { - let words = word_typos(&right, prefix, *typo, &self.words_fst)?; - let mut docids = RoaringBitmap::new(); - for (word, _typo) in words { - let key = (left.as_str(), word.as_str(), distance); - docids.union_with(&self.index.word_pair_proximity_docids.get(self.rtxn, &key)?.unwrap_or_default()); - } - Ok(Some(docids)) + let r_words = word_typos(&right, prefix, *typo, &self.words_fst)?; + self.all_word_pair_proximity_docids(&[(left, 0)], &r_words, distance) }, (QueryKind::Tolerant { typo: l_typo, word: left }, QueryKind::Tolerant { typo: r_typo, word: right }, prefix) => { let l_words = word_typos(&left, false, *l_typo, &self.words_fst)?; let r_words = word_typos(&right, prefix, *r_typo, &self.words_fst)?; - let mut docids = RoaringBitmap::new(); - for (left, _typo) in l_words { - for (right, _typo) in r_words.iter() { - let key = (left.as_str(), right.as_str(), distance); - docids.union_with(&self.index.word_pair_proximity_docids.get(self.rtxn, &key)?.unwrap_or_default()); - } - } - Ok(Some(docids)) + self.all_word_pair_proximity_docids(&l_words, &r_words, distance) }, (QueryKind::Exact { word: left, .. }, QueryKind::Exact { word: right, .. }, false) => { let key = (left.as_str(), right.as_str(), distance); - Ok(self.index.word_pair_proximity_docids.get(self.rtxn, &key)?) + Ok(self.index.word_pair_proximity_docids.get(self.rtxn, &key)?.unwrap_or_default()) }, } } @@ -170,4 +135,15 @@ impl<'t> HeedContext<'t> { fn in_prefix_cache(&self, word: &str) -> bool { self.words_prefixes_fst.contains(word) } + + fn all_word_pair_proximity_docids, U: AsRef>(&self, left_words: &[(T, u8)], right_words: &[(U, u8)], distance: u8) -> anyhow::Result { + let mut docids = RoaringBitmap::new(); + for (left, _l_typo) in left_words { + for (right, _r_typo) in right_words { + let key = (left.as_ref(), right.as_ref(), distance); + docids.union_with(&self.index.word_pair_proximity_docids.get(self.rtxn, &key)?.unwrap_or_default()); + } + } + Ok(docids) + } } diff --git a/milli/src/search/criteria/typo.rs b/milli/src/search/criteria/typo.rs index b3608c397..9834bbc21 100644 --- a/milli/src/search/criteria/typo.rs +++ b/milli/src/search/criteria/typo.rs @@ -180,17 +180,17 @@ fn resolve_candidates<'t>( match (&slice[0], &slice[1]) { (Operation::Query(left), Operation::Query(right)) => { match ctx.query_pair_proximity_docids(left, right, 1)? { - Some(pair_docids) => { - if first_loop { - candidates = pair_docids; - first_loop = false; - } else { - candidates.intersect_with(&pair_docids) - } + pair_docids if pair_docids.is_empty() => { + return Ok(RoaringBitmap::new()) + }, + pair_docids if first_loop => { + candidates = pair_docids; + first_loop = false; + }, + pair_docids => { + candidates.intersect_with(&pair_docids); }, - None => return Ok(RoaringBitmap::new()), } - }, _ => bail!("invalid consecutive query type"), } @@ -206,7 +206,7 @@ fn resolve_candidates<'t>( Ok(candidates) }, Query(q) => if q.kind.typo() == number_typos { - Ok(ctx.query_docids(q)?.unwrap_or_default()) + Ok(ctx.query_docids(q)?) } else { Ok(RoaringBitmap::new()) },