implement Context trait for criteria

This commit is contained in:
many
2021-02-17 15:27:35 +01:00
committed by Kerollmops
parent f091f370d0
commit 98e69e63d2
3 changed files with 152 additions and 40 deletions

View File

@ -1,8 +1,11 @@
use std::borrow::Cow;
use crate::Index;
use crate::search::word_typos;
use roaring::RoaringBitmap;
use super::query_tree::Operation;
use super::query_tree::{Operation, Query, QueryKind};
pub mod typo;
@ -32,3 +35,120 @@ impl Default for Candidates {
Self::Forbidden(RoaringBitmap::new())
}
}
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 words_fst<'t>(&self) -> &'t fst::Set<Cow<[u8]>>;
}
pub struct HeedContext<'t> {
rtxn: &'t heed::RoTxn<'t>,
index: &'t Index,
words_fst: fst::Set<Cow<'t, [u8]>>,
}
impl<'a> Context for HeedContext<'a> {
fn query_docids(&self, query: &Query) -> anyhow::Result<Option<RoaringBitmap>> {
match (&query.kind, query.prefix) {
// TODO de-comment when ~ready
// (QueryKind::Exact { word, .. }, true) if in_prefix_cache(&word) => {
// Ok(self.index.prefix_docids.get(self.rtxn, &word)?)
// },
(QueryKind::Exact { word, .. }, true) => {
let words = word_typos(&word, true, 0, &self.words_fst)?;
let mut docids = RoaringBitmap::new();
for (word, _typo) in words {
docids.union_with(&self.index.word_docids.get(self.rtxn, &word)?.unwrap_or_default());
}
Ok(Some(docids))
},
(QueryKind::Exact { word, .. }, false) => {
Ok(self.index.word_docids.get(self.rtxn, &word)?)
},
(QueryKind::Tolerant { typo, word }, prefix) => {
let words = word_typos(&word, prefix, *typo, &self.words_fst)?;
let mut docids = RoaringBitmap::new();
for (word, _typo) in words {
docids.union_with(&self.index.word_docids.get(self.rtxn, &word)?.unwrap_or_default());
}
Ok(Some(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
match (&left.kind, &right.kind, right.prefix) {
(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))
},
(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))
},
(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))
},
(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))
},
(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))
},
(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)?)
},
}
}
fn words_fst<'t>(&self) -> &'t fst::Set<Cow<[u8]>> {
&self.words_fst
}
}
impl<'t> HeedContext<'t> {
pub fn new(rtxn: &'t heed::RoTxn<'t>, index: &'t Index) -> anyhow::Result<Self> {
let words_fst = index.words_fst(rtxn)?;
Ok(Self {
rtxn,
index,
words_fst,
})
}
}