clean docids fetchers

This commit is contained in:
many 2021-02-17 16:21:21 +01:00 committed by Kerollmops
parent 774a255f2e
commit 907482c8ac
No known key found for this signature in database
GPG Key ID: 92ADA4E935E71FA4
2 changed files with 41 additions and 65 deletions

View File

@ -36,8 +36,8 @@ impl Default for Candidates {
} }
} }
pub trait Context { pub trait Context {
fn query_docids(&self, query: &Query) -> anyhow::Result<Option<RoaringBitmap>>; fn query_docids(&self, query: &Query) -> anyhow::Result<RoaringBitmap>;
fn query_pair_proximity_docids(&self, left: &Query, right: &Query, distance: u8) ->anyhow::Result<Option<RoaringBitmap>>; fn query_pair_proximity_docids(&self, left: &Query, right: &Query, distance: u8) ->anyhow::Result<RoaringBitmap>;
fn words_fst<'t>(&self) -> &'t fst::Set<Cow<[u8]>>; fn words_fst<'t>(&self) -> &'t fst::Set<Cow<[u8]>>;
} }
pub struct HeedContext<'t> { pub struct HeedContext<'t> {
@ -48,10 +48,10 @@ pub struct HeedContext<'t> {
} }
impl<'a> Context for HeedContext<'a> { impl<'a> Context for HeedContext<'a> {
fn query_docids(&self, query: &Query) -> anyhow::Result<Option<RoaringBitmap>> { fn query_docids(&self, query: &Query) -> anyhow::Result<RoaringBitmap> {
match (&query.kind, query.prefix) { match (&query.kind, query.prefix) {
(QueryKind::Exact { word, .. }, true) if self.in_prefix_cache(&word) => { (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) => { (QueryKind::Exact { word, .. }, true) => {
let words = word_typos(&word, true, 0, &self.words_fst)?; 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 { for (word, _typo) in words {
docids.union_with(&self.index.word_docids.get(self.rtxn, &word)?.unwrap_or_default()); docids.union_with(&self.index.word_docids.get(self.rtxn, &word)?.unwrap_or_default());
} }
Ok(Some(docids)) Ok(docids)
}, },
(QueryKind::Exact { word, .. }, false) => { (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) => { (QueryKind::Tolerant { typo, word }, prefix) => {
let words = word_typos(&word, prefix, *typo, &self.words_fst)?; 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 { for (word, _typo) in words {
docids.union_with(&self.index.word_docids.get(self.rtxn, &word)?.unwrap_or_default()); 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<Option<RoaringBitmap>> { fn query_pair_proximity_docids(&self, left: &Query, right: &Query, distance: u8) -> anyhow::Result<RoaringBitmap> {
// TODO add prefix cache for Tolerant-Exact-true and Exact-Exact-true
match (&left.kind, &right.kind, right.prefix) { match (&left.kind, &right.kind, right.prefix) {
(QueryKind::Exact { word: left, .. }, QueryKind::Exact { word: right, .. }, true) if self.in_prefix_cache(&right) => { (QueryKind::Exact { word: left, .. }, QueryKind::Exact { word: right, .. }, true) if self.in_prefix_cache(&right) => {
let key = (left.as_str(), right.as_str(), distance); 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) => { (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 l_words = word_typos(&left, false, *typo, &self.words_fst)?;
let mut docids = RoaringBitmap::new(); self.all_word_pair_proximity_docids(&l_words, &[(right, 0)], distance)
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))
}, },
(QueryKind::Exact { word: left, .. }, QueryKind::Exact { word: right, .. }, true) => { (QueryKind::Exact { word: left, .. }, QueryKind::Exact { word: right, .. }, true) => {
let words = word_typos(&right, true, 0, &self.words_fst)?; let r_words = word_typos(&right, true, 0, &self.words_fst)?;
let mut docids = RoaringBitmap::new(); self.all_word_pair_proximity_docids(&[(left, 0)], &r_words, distance)
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))
}, },
(QueryKind::Tolerant { typo, word: left }, QueryKind::Exact { word: right, .. }, true) => { (QueryKind::Tolerant { typo, word: left }, QueryKind::Exact { word: right, .. }, true) => {
let l_words = word_typos(&left, false, *typo, &self.words_fst)?; let l_words = word_typos(&left, false, *typo, &self.words_fst)?;
let r_words = word_typos(&right, true, 0, &self.words_fst)?; let r_words = word_typos(&right, true, 0, &self.words_fst)?;
let mut docids = RoaringBitmap::new(); self.all_word_pair_proximity_docids(&l_words, &r_words, distance)
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))
}, },
(QueryKind::Tolerant { typo, word: left }, QueryKind::Exact { word: right, .. }, false) => { (QueryKind::Tolerant { typo, word: left }, QueryKind::Exact { word: right, .. }, false) => {
let words = word_typos(&left, false, *typo, &self.words_fst)?; let l_words = word_typos(&left, false, *typo, &self.words_fst)?;
let mut docids = RoaringBitmap::new(); self.all_word_pair_proximity_docids(&l_words, &[(right, 0)], distance)
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))
}, },
(QueryKind::Exact { word: left, .. }, QueryKind::Tolerant { typo, word: right }, prefix) => { (QueryKind::Exact { word: left, .. }, QueryKind::Tolerant { typo, word: right }, prefix) => {
let words = word_typos(&right, prefix, *typo, &self.words_fst)?; let r_words = word_typos(&right, prefix, *typo, &self.words_fst)?;
let mut docids = RoaringBitmap::new(); self.all_word_pair_proximity_docids(&[(left, 0)], &r_words, distance)
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))
}, },
(QueryKind::Tolerant { typo: l_typo, word: left }, QueryKind::Tolerant { typo: r_typo, word: right }, prefix) => { (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 l_words = word_typos(&left, false, *l_typo, &self.words_fst)?;
let r_words = word_typos(&right, prefix, *r_typo, &self.words_fst)?; let r_words = word_typos(&right, prefix, *r_typo, &self.words_fst)?;
let mut docids = RoaringBitmap::new(); self.all_word_pair_proximity_docids(&l_words, &r_words, distance)
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))
}, },
(QueryKind::Exact { word: left, .. }, QueryKind::Exact { word: right, .. }, false) => { (QueryKind::Exact { word: left, .. }, QueryKind::Exact { word: right, .. }, false) => {
let key = (left.as_str(), right.as_str(), distance); 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 { fn in_prefix_cache(&self, word: &str) -> bool {
self.words_prefixes_fst.contains(word) self.words_prefixes_fst.contains(word)
} }
fn all_word_pair_proximity_docids<T: AsRef<str>, U: AsRef<str>>(&self, left_words: &[(T, u8)], right_words: &[(U, u8)], distance: u8) -> anyhow::Result<RoaringBitmap> {
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)
}
} }

View File

@ -180,17 +180,17 @@ fn resolve_candidates<'t>(
match (&slice[0], &slice[1]) { match (&slice[0], &slice[1]) {
(Operation::Query(left), Operation::Query(right)) => { (Operation::Query(left), Operation::Query(right)) => {
match ctx.query_pair_proximity_docids(left, right, 1)? { match ctx.query_pair_proximity_docids(left, right, 1)? {
Some(pair_docids) => { pair_docids if pair_docids.is_empty() => {
if first_loop { return Ok(RoaringBitmap::new())
},
pair_docids if first_loop => {
candidates = pair_docids; candidates = pair_docids;
first_loop = false; first_loop = false;
} else {
candidates.intersect_with(&pair_docids)
}
}, },
None => return Ok(RoaringBitmap::new()), pair_docids => {
candidates.intersect_with(&pair_docids);
},
} }
}, },
_ => bail!("invalid consecutive query type"), _ => bail!("invalid consecutive query type"),
} }
@ -206,7 +206,7 @@ fn resolve_candidates<'t>(
Ok(candidates) Ok(candidates)
}, },
Query(q) => if q.kind.typo() == number_typos { Query(q) => if q.kind.typo() == number_typos {
Ok(ctx.query_docids(q)?.unwrap_or_default()) Ok(ctx.query_docids(q)?)
} else { } else {
Ok(RoaringBitmap::new()) Ok(RoaringBitmap::new())
}, },