MeiliSearch/milli/src/search/criteria/typo.rs

488 lines
20 KiB
Rust
Raw Normal View History

2021-02-18 14:40:59 +01:00
use std::{borrow::Cow, collections::HashMap, mem::take};
2021-02-19 15:20:07 +01:00
use anyhow::bail;
2021-02-24 15:37:37 +01:00
use log::debug;
use roaring::RoaringBitmap;
2021-02-19 15:20:07 +01:00
2021-02-23 17:33:20 +01:00
use crate::search::query_tree::{maximum_typo, Operation, Query, QueryKind};
use crate::search::{word_derivations, WordDerivationsCache};
2021-02-17 18:16:21 +01:00
use super::{Candidates, Criterion, CriterionResult, Context, query_docids, query_pair_proximity_docids};
2021-02-19 15:20:07 +01:00
pub struct Typo<'t> {
2021-02-17 15:27:35 +01:00
ctx: &'t dyn Context,
2021-02-23 17:33:20 +01:00
query_tree: Option<(usize, Operation)>,
2021-02-19 15:20:07 +01:00
number_typos: u8,
candidates: Candidates,
2021-02-25 16:14:38 +01:00
bucket_candidates: RoaringBitmap,
2021-02-19 15:32:14 +01:00
parent: Option<Box<dyn Criterion + 't>>,
candidates_cache: HashMap<(Operation, u8), RoaringBitmap>,
2021-02-19 15:20:07 +01:00
}
impl<'t> Typo<'t> {
pub fn initial(
2021-02-17 15:27:35 +01:00
ctx: &'t dyn Context,
2021-02-19 15:20:07 +01:00
query_tree: Option<Operation>,
candidates: Option<RoaringBitmap>,
) -> Self
2021-02-19 15:20:07 +01:00
{
Typo {
2021-02-17 15:27:35 +01:00
ctx,
2021-02-23 17:33:20 +01:00
query_tree: query_tree.map(|op| (maximum_typo(&op), op)),
2021-02-19 15:20:07 +01:00
number_typos: 0,
candidates: candidates.map_or_else(Candidates::default, Candidates::Allowed),
2021-02-25 16:14:38 +01:00
bucket_candidates: RoaringBitmap::new(),
2021-02-19 15:20:07 +01:00
parent: None,
candidates_cache: HashMap::new(),
}
2021-02-19 15:20:07 +01:00
}
pub fn new(ctx: &'t dyn Context, parent: Box<dyn Criterion + 't>) -> Self {
Typo {
2021-02-17 15:27:35 +01:00
ctx,
2021-02-19 15:20:07 +01:00
query_tree: None,
number_typos: 0,
candidates: Candidates::default(),
2021-02-25 16:14:38 +01:00
bucket_candidates: RoaringBitmap::new(),
2021-02-19 15:20:07 +01:00
parent: Some(parent),
candidates_cache: HashMap::new(),
}
2021-02-19 15:20:07 +01:00
}
}
impl<'t> Criterion for Typo<'t> {
#[logging_timer::time("Typo::{}")]
fn next(&mut self, wdcache: &mut WordDerivationsCache) -> anyhow::Result<Option<CriterionResult>> {
2021-02-19 15:20:07 +01:00
use Candidates::{Allowed, Forbidden};
2021-02-23 17:33:20 +01:00
loop {
2021-02-24 15:37:37 +01:00
debug!("Typo at iteration {} ({:?})", self.number_typos, self.candidates);
2021-02-19 15:20:07 +01:00
match (&mut self.query_tree, &mut self.candidates) {
(_, Allowed(candidates)) if candidates.is_empty() => {
return Ok(Some(CriterionResult {
query_tree: self.query_tree.take().map(|(_, qt)| qt),
2021-03-09 12:04:52 +01:00
candidates: Some(take(&mut self.candidates).into_inner()),
bucket_candidates: take(&mut self.bucket_candidates),
}));
2021-02-19 15:20:07 +01:00
},
2021-02-23 17:33:20 +01:00
(Some((max_typos, query_tree)), Allowed(candidates)) => {
if self.number_typos as usize > *max_typos {
self.query_tree = None;
self.candidates = Candidates::default();
} else {
2021-02-23 17:33:20 +01:00
let fst = self.ctx.words_fst();
let new_query_tree = if self.number_typos < 2 {
2021-03-09 15:18:30 +01:00
alterate_query_tree(&fst, query_tree.clone(), self.number_typos, wdcache)?
2021-02-23 17:33:20 +01:00
} else if self.number_typos == 2 {
2021-03-09 15:18:30 +01:00
*query_tree = alterate_query_tree(&fst, query_tree.clone(), self.number_typos, wdcache)?;
2021-02-23 17:33:20 +01:00
query_tree.clone()
} else {
query_tree.clone()
};
let mut new_candidates = resolve_candidates(
self.ctx,
&new_query_tree,
self.number_typos,
&mut self.candidates_cache,
wdcache,
)?;
2021-02-23 17:33:20 +01:00
new_candidates.intersect_with(&candidates);
candidates.difference_with(&new_candidates);
self.number_typos += 1;
2021-02-19 15:20:07 +01:00
2021-02-23 17:33:20 +01:00
let bucket_candidates = match self.parent {
2021-02-25 16:14:38 +01:00
Some(_) => take(&mut self.bucket_candidates),
None => new_candidates.clone(),
2021-02-23 17:33:20 +01:00
};
2021-02-23 17:33:20 +01:00
return Ok(Some(CriterionResult {
query_tree: Some(new_query_tree),
2021-03-09 12:04:52 +01:00
candidates: Some(new_candidates),
2021-02-23 17:33:20 +01:00
bucket_candidates,
}));
}
2021-02-19 15:20:07 +01:00
},
2021-02-23 17:33:20 +01:00
(Some((max_typos, query_tree)), Forbidden(candidates)) => {
if self.number_typos as usize > *max_typos {
self.query_tree = None;
self.candidates = Candidates::default();
} else {
2021-02-23 17:33:20 +01:00
let fst = self.ctx.words_fst();
let new_query_tree = if self.number_typos < 2 {
2021-03-09 15:18:30 +01:00
alterate_query_tree(&fst, query_tree.clone(), self.number_typos, wdcache)?
2021-02-23 17:33:20 +01:00
} else if self.number_typos == 2 {
2021-03-09 15:18:30 +01:00
*query_tree = alterate_query_tree(&fst, query_tree.clone(), self.number_typos, wdcache)?;
2021-02-23 17:33:20 +01:00
query_tree.clone()
} else {
query_tree.clone()
};
let mut new_candidates = resolve_candidates(
self.ctx,
&new_query_tree,
self.number_typos,
&mut self.candidates_cache,
wdcache,
)?;
2021-02-23 17:33:20 +01:00
new_candidates.difference_with(&candidates);
candidates.union_with(&new_candidates);
self.number_typos += 1;
2021-03-09 15:55:59 +01:00
self.bucket_candidates.union_with(&new_candidates);
2021-02-23 17:33:20 +01:00
return Ok(Some(CriterionResult {
query_tree: Some(new_query_tree),
2021-03-09 12:04:52 +01:00
candidates: Some(new_candidates),
2021-03-09 15:55:59 +01:00
bucket_candidates: take(&mut self.bucket_candidates),
2021-02-23 17:33:20 +01:00
}));
}
2021-02-19 15:20:07 +01:00
},
(None, Allowed(_)) => {
let candidates = take(&mut self.candidates).into_inner();
return Ok(Some(CriterionResult {
query_tree: None,
2021-03-09 12:04:52 +01:00
candidates: Some(candidates.clone()),
2021-02-25 16:14:38 +01:00
bucket_candidates: candidates,
}));
2021-02-19 15:20:07 +01:00
},
(None, Forbidden(_)) => {
match self.parent.as_mut() {
Some(parent) => {
match parent.next(wdcache)? {
Some(CriterionResult { query_tree, candidates, bucket_candidates }) => {
2021-02-23 17:33:20 +01:00
self.query_tree = query_tree.map(|op| (maximum_typo(&op), op));
self.number_typos = 0;
2021-03-09 12:04:52 +01:00
self.candidates = candidates.map_or_else(Candidates::default, Candidates::Allowed);
2021-02-25 17:28:20 +01:00
self.bucket_candidates.union_with(&bucket_candidates);
2021-02-19 15:20:07 +01:00
},
None => return Ok(None),
}
},
None => return Ok(None),
}
},
}
}
}
}
/// Modify the query tree by replacing every tolerant query by an Or operation
/// containing all of the corresponding exact words in the words FST. Each tolerant
/// query will only be replaced by exact query with up to `number_typos` maximum typos.
fn alterate_query_tree(
words_fst: &fst::Set<Cow<[u8]>>,
mut query_tree: Operation,
number_typos: u8,
wdcache: &mut WordDerivationsCache,
2021-02-19 15:20:07 +01:00
) -> anyhow::Result<Operation>
{
fn recurse(
words_fst: &fst::Set<Cow<[u8]>>,
operation: &mut Operation,
number_typos: u8,
wdcache: &mut WordDerivationsCache,
2021-02-19 15:20:07 +01:00
) -> anyhow::Result<()>
{
use Operation::{And, Consecutive, Or};
match operation {
And(ops) | Consecutive(ops) | Or(_, ops) => {
2021-03-09 15:18:30 +01:00
ops.iter_mut().try_for_each(|op| recurse(words_fst, op, number_typos, wdcache))
2021-02-19 15:20:07 +01:00
},
Operation::Query(q) => {
2021-02-17 15:27:35 +01:00
// TODO may be optimized when number_typos == 0
2021-02-19 15:20:07 +01:00
if let QueryKind::Tolerant { typo, word } = &q.kind {
2021-02-18 16:31:10 +01:00
// if no typo is allowed we don't call word_derivations function,
// and directly create an Exact query
if number_typos == 0 {
*operation = Operation::Query(Query {
prefix: q.prefix,
kind: QueryKind::Exact { original_typo: 0, word: word.clone() },
});
} else {
let typo = *typo.min(&number_typos);
2021-03-09 15:18:30 +01:00
let words = word_derivations(word, q.prefix, typo, words_fst, wdcache)?;
let queries = words.into_iter().map(|(word, typo)| {
Operation::Query(Query {
prefix: false,
2021-03-09 15:18:30 +01:00
kind: QueryKind::Exact { original_typo: *typo, word: word.to_string() },
})
}).collect();
2021-02-19 15:20:07 +01:00
*operation = Operation::or(false, queries);
}
2021-02-19 15:20:07 +01:00
}
Ok(())
},
}
}
2021-03-09 15:18:30 +01:00
recurse(words_fst, &mut query_tree, number_typos, wdcache)?;
2021-02-19 15:20:07 +01:00
Ok(query_tree)
}
2021-02-17 15:27:35 +01:00
fn resolve_candidates<'t>(
ctx: &'t dyn Context,
2021-02-19 15:20:07 +01:00
query_tree: &Operation,
number_typos: u8,
2021-02-18 14:40:59 +01:00
cache: &mut HashMap<(Operation, u8), RoaringBitmap>,
wdcache: &mut WordDerivationsCache,
2021-02-19 15:20:07 +01:00
) -> anyhow::Result<RoaringBitmap>
{
2021-02-17 15:27:35 +01:00
fn resolve_operation<'t>(
ctx: &'t dyn Context,
2021-02-19 15:20:07 +01:00
query_tree: &Operation,
number_typos: u8,
2021-02-18 14:40:59 +01:00
cache: &mut HashMap<(Operation, u8), RoaringBitmap>,
wdcache: &mut WordDerivationsCache,
2021-02-19 15:20:07 +01:00
) -> anyhow::Result<RoaringBitmap>
{
use Operation::{And, Consecutive, Or, Query};
match query_tree {
And(ops) => {
mdfs(ctx, ops, number_typos, cache, wdcache)
2021-02-19 15:20:07 +01:00
},
Consecutive(ops) => {
let mut candidates = RoaringBitmap::new();
let mut first_loop = true;
for slice in ops.windows(2) {
match (&slice[0], &slice[1]) {
(Operation::Query(left), Operation::Query(right)) => {
match query_pair_proximity_docids(ctx, left, right, 1, wdcache)? {
2021-02-17 16:21:21 +01:00
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);
2021-02-19 15:20:07 +01:00
},
}
},
_ => bail!("invalid consecutive query type"),
}
}
Ok(candidates)
},
Or(_, ops) => {
let mut candidates = RoaringBitmap::new();
for op in ops {
let docids = resolve_operation(ctx, op, number_typos, cache, wdcache)?;
2021-02-19 15:20:07 +01:00
candidates.union_with(&docids);
}
Ok(candidates)
},
Query(q) => if q.kind.typo() == number_typos {
Ok(query_docids(ctx, q, wdcache)?)
2021-02-19 15:20:07 +01:00
} else {
Ok(RoaringBitmap::new())
},
}
}
2021-02-17 15:27:35 +01:00
fn mdfs<'t>(
ctx: &'t dyn Context,
2021-02-19 15:20:07 +01:00
branches: &[Operation],
mana: u8,
2021-02-18 14:40:59 +01:00
cache: &mut HashMap<(Operation, u8), RoaringBitmap>,
wdcache: &mut WordDerivationsCache,
2021-02-19 15:20:07 +01:00
) -> anyhow::Result<RoaringBitmap>
{
match branches.split_first() {
2021-02-18 14:40:59 +01:00
Some((head, [])) => {
let cache_key = (head.clone(), mana);
if let Some(candidates) = cache.get(&cache_key) {
2021-02-18 14:40:59 +01:00
Ok(candidates.clone())
} else {
let candidates = resolve_operation(ctx, head, mana, cache, wdcache)?;
cache.insert(cache_key, candidates.clone());
2021-02-18 14:40:59 +01:00
Ok(candidates)
}
},
2021-02-19 15:20:07 +01:00
Some((head, tail)) => {
let mut candidates = RoaringBitmap::new();
for m in 0..=mana {
2021-02-18 14:40:59 +01:00
let mut head_candidates = {
let cache_key = (head.clone(), m);
if let Some(candidates) = cache.get(&cache_key) {
2021-02-18 14:40:59 +01:00
candidates.clone()
} else {
let candidates = resolve_operation(ctx, head, m, cache, wdcache)?;
cache.insert(cache_key, candidates.clone());
2021-02-18 14:40:59 +01:00
candidates
}
};
2021-02-19 15:20:07 +01:00
if !head_candidates.is_empty() {
let tail_candidates = mdfs(ctx, tail, mana - m, cache, wdcache)?;
2021-02-19 15:20:07 +01:00
head_candidates.intersect_with(&tail_candidates);
candidates.union_with(&head_candidates);
}
}
Ok(candidates)
},
None => Ok(RoaringBitmap::new()),
}
}
resolve_operation(ctx, query_tree, number_typos, cache, wdcache)
2021-02-19 15:20:07 +01:00
}
2021-02-24 10:25:22 +01:00
#[cfg(test)]
mod test {
use super::*;
use super::super::test::TestContext;
#[test]
fn initial_placeholder_no_facets() {
let context = TestContext::default();
let query_tree = None;
let facet_candidates = None;
let mut wdcache = WordDerivationsCache::new();
let mut criteria = Typo::initial(&context, query_tree, facet_candidates);
2021-02-24 10:25:22 +01:00
assert!(criteria.next(&mut wdcache).unwrap().is_none());
2021-02-24 10:25:22 +01:00
}
#[test]
fn initial_query_tree_no_facets() {
let context = TestContext::default();
let query_tree = Operation::Or(false, vec![
Operation::And(vec![
Operation::Query(Query { prefix: false, kind: QueryKind::exact("split".to_string()) }),
Operation::Query(Query { prefix: false, kind: QueryKind::exact("this".to_string()) }),
Operation::Query(Query { prefix: false, kind: QueryKind::tolerant(1, "world".to_string()) }),
])
]);
let facet_candidates = None;
let mut wdcache = WordDerivationsCache::new();
let mut criteria = Typo::initial(&context, Some(query_tree), facet_candidates);
2021-02-24 10:25:22 +01:00
let candidates_1 = context.word_docids("split").unwrap().unwrap()
& context.word_docids("this").unwrap().unwrap()
& context.word_docids("world").unwrap().unwrap();
let expected_1 = CriterionResult {
query_tree: Some(Operation::Or(false, vec![
Operation::And(vec![
Operation::Query(Query { prefix: false, kind: QueryKind::exact("split".to_string()) }),
Operation::Query(Query { prefix: false, kind: QueryKind::exact("this".to_string()) }),
Operation::Query(Query { prefix: false, kind: QueryKind::exact("world".to_string()) }),
]),
])),
2021-03-09 12:04:52 +01:00
candidates: Some(candidates_1.clone()),
2021-02-25 16:14:38 +01:00
bucket_candidates: candidates_1,
2021-02-24 10:25:22 +01:00
};
assert_eq!(criteria.next(&mut wdcache).unwrap(), Some(expected_1));
2021-02-24 10:25:22 +01:00
let candidates_2 = (
context.word_docids("split").unwrap().unwrap()
& context.word_docids("this").unwrap().unwrap()
& context.word_docids("word").unwrap().unwrap()
) - context.word_docids("world").unwrap().unwrap();
let expected_2 = CriterionResult {
query_tree: Some(Operation::Or(false, vec![
Operation::And(vec![
Operation::Query(Query { prefix: false, kind: QueryKind::exact("split".to_string()) }),
Operation::Query(Query { prefix: false, kind: QueryKind::exact("this".to_string()) }),
Operation::Or(false, vec![
Operation::Query(Query { prefix: false, kind: QueryKind::exact_with_typo(1, "word".to_string()) }),
Operation::Query(Query { prefix: false, kind: QueryKind::exact("world".to_string()) }),
]),
]),
])),
2021-03-09 12:04:52 +01:00
candidates: Some(candidates_2.clone()),
2021-02-25 16:14:38 +01:00
bucket_candidates: candidates_2,
2021-02-24 10:25:22 +01:00
};
assert_eq!(criteria.next(&mut wdcache).unwrap(), Some(expected_2));
2021-02-24 10:25:22 +01:00
}
#[test]
fn initial_placeholder_with_facets() {
let context = TestContext::default();
let query_tree = None;
2021-02-25 16:14:38 +01:00
let facet_candidates = context.word_docids("earth").unwrap().unwrap();
2021-02-24 10:25:22 +01:00
let mut wdcache = WordDerivationsCache::new();
let mut criteria = Typo::initial(&context, query_tree, Some(facet_candidates.clone()));
2021-02-24 10:25:22 +01:00
let expected = CriterionResult {
query_tree: None,
2021-03-09 12:04:52 +01:00
candidates: Some(facet_candidates.clone()),
2021-02-24 10:25:22 +01:00
bucket_candidates: facet_candidates,
};
// first iteration, returns the facet candidates
assert_eq!(criteria.next(&mut wdcache).unwrap(), Some(expected));
2021-02-24 10:25:22 +01:00
// second iteration, returns None because there is no more things to do
assert!(criteria.next(&mut wdcache).unwrap().is_none());
2021-02-24 10:25:22 +01:00
}
#[test]
fn initial_query_tree_with_facets() {
let context = TestContext::default();
let query_tree = Operation::Or(false, vec![
Operation::And(vec![
Operation::Query(Query { prefix: false, kind: QueryKind::exact("split".to_string()) }),
Operation::Query(Query { prefix: false, kind: QueryKind::exact("this".to_string()) }),
Operation::Query(Query { prefix: false, kind: QueryKind::tolerant(1, "world".to_string()) }),
])
]);
let facet_candidates = context.word_docids("earth").unwrap().unwrap();
let mut wdcache = WordDerivationsCache::new();
let mut criteria = Typo::initial(&context, Some(query_tree), Some(facet_candidates.clone()));
2021-02-24 10:25:22 +01:00
let candidates_1 = context.word_docids("split").unwrap().unwrap()
& context.word_docids("this").unwrap().unwrap()
& context.word_docids("world").unwrap().unwrap();
let expected_1 = CriterionResult {
query_tree: Some(Operation::Or(false, vec![
Operation::And(vec![
Operation::Query(Query { prefix: false, kind: QueryKind::exact("split".to_string()) }),
Operation::Query(Query { prefix: false, kind: QueryKind::exact("this".to_string()) }),
Operation::Query(Query { prefix: false, kind: QueryKind::exact("world".to_string()) }),
]),
])),
2021-03-09 12:04:52 +01:00
candidates: Some(&candidates_1 & &facet_candidates),
2021-02-25 16:14:38 +01:00
bucket_candidates: candidates_1 & &facet_candidates,
2021-02-24 10:25:22 +01:00
};
assert_eq!(criteria.next(&mut wdcache).unwrap(), Some(expected_1));
2021-02-24 10:25:22 +01:00
let candidates_2 = (
context.word_docids("split").unwrap().unwrap()
& context.word_docids("this").unwrap().unwrap()
& context.word_docids("word").unwrap().unwrap()
) - context.word_docids("world").unwrap().unwrap();
let expected_2 = CriterionResult {
query_tree: Some(Operation::Or(false, vec![
Operation::And(vec![
Operation::Query(Query { prefix: false, kind: QueryKind::exact("split".to_string()) }),
Operation::Query(Query { prefix: false, kind: QueryKind::exact("this".to_string()) }),
Operation::Or(false, vec![
Operation::Query(Query { prefix: false, kind: QueryKind::exact_with_typo(1, "word".to_string()) }),
Operation::Query(Query { prefix: false, kind: QueryKind::exact("world".to_string()) }),
]),
]),
])),
2021-03-09 12:04:52 +01:00
candidates: Some(&candidates_2 & &facet_candidates),
2021-02-25 16:14:38 +01:00
bucket_candidates: candidates_2 & &facet_candidates,
2021-02-24 10:25:22 +01:00
};
assert_eq!(criteria.next(&mut wdcache).unwrap(), Some(expected_2));
2021-02-24 10:25:22 +01:00
}
}