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 {
fn query_docids(&self, query: &Query) -> anyhow::Result<Option<RoaringBitmap>>;
fn query_pair_proximity_docids(&self, left: &Query, right: &Query, distance: u8) ->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<RoaringBitmap>;
fn words_fst<'t>(&self) -> &'t fst::Set<Cow<[u8]>>;
}
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<Option<RoaringBitmap>> {
fn query_docids(&self, query: &Query) -> anyhow::Result<RoaringBitmap> {
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<Option<RoaringBitmap>> {
// 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<RoaringBitmap> {
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<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]) {
(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())
},