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

873 lines
34 KiB
Rust
Raw Normal View History

2021-06-16 18:33:33 +02:00
use std::borrow::Cow;
use std::cmp::{self, Ordering};
use std::collections::binary_heap::PeekMut;
2021-06-16 18:33:33 +02:00
use std::collections::{btree_map, BTreeMap, BinaryHeap, HashMap};
use std::mem::take;
2021-03-11 11:48:55 +01:00
use roaring::RoaringBitmap;
2021-06-16 18:33:33 +02:00
use super::{resolve_query_tree, Context, Criterion, CriterionParameters, CriterionResult};
2021-03-11 11:48:55 +01:00
use crate::search::criteria::Query;
use crate::search::query_tree::{Operation, QueryKind};
2021-06-16 18:33:33 +02:00
use crate::search::{build_dfa, word_derivations, WordDerivationsCache};
use crate::{Result, TreeLevel};
2021-03-11 11:48:55 +01:00
/// To be able to divide integers by the number of words in the query
/// we want to find a multiplier that allow us to divide by any number between 1 and 10.
/// We chose the LCM of all numbers between 1 and 10 as the multiplier (https://en.wikipedia.org/wiki/Least_common_multiple).
const LCM_10_FIRST_NUMBERS: u32 = 2520;
2021-04-15 10:44:27 +02:00
/// To compute the interval size of a level,
/// we use 4 as the exponentiation base and the level as the exponent.
const LEVEL_EXPONENTIATION_BASE: u32 = 4;
/// Threshold on the number of candidates that will make
/// the system to choose between one algorithm or another.
const CANDIDATES_THRESHOLD: u64 = 1000;
2021-05-05 20:46:56 +02:00
type FlattenedQueryTree = Vec<Vec<Vec<Query>>>;
2021-05-10 12:33:37 +02:00
2021-03-11 11:48:55 +01:00
pub struct Attribute<'t> {
ctx: &'t dyn Context<'t>,
2021-05-05 20:46:56 +02:00
state: Option<(Operation, FlattenedQueryTree, RoaringBitmap)>,
2021-03-11 11:48:55 +01:00
bucket_candidates: RoaringBitmap,
2021-03-23 15:25:46 +01:00
parent: Box<dyn Criterion + 't>,
current_buckets: Option<btree_map::IntoIter<u64, RoaringBitmap>>,
2021-03-11 11:48:55 +01:00
}
impl<'t> Attribute<'t> {
pub fn new(ctx: &'t dyn Context<'t>, parent: Box<dyn Criterion + 't>) -> Self {
2021-03-11 11:48:55 +01:00
Attribute {
ctx,
2021-05-05 20:46:56 +02:00
state: None,
2021-03-11 11:48:55 +01:00
bucket_candidates: RoaringBitmap::new(),
2021-03-23 15:25:46 +01:00
parent,
current_buckets: None,
2021-03-11 11:48:55 +01:00
}
}
}
impl<'t> Criterion for Attribute<'t> {
#[logging_timer::time("Attribute::{}")]
fn next(&mut self, params: &mut CriterionParameters) -> Result<Option<CriterionResult>> {
// remove excluded candidates when next is called, instead of doing it in the loop.
2021-05-05 20:46:56 +02:00
if let Some((_, _, allowed_candidates)) = self.state.as_mut() {
*allowed_candidates -= params.excluded_candidates;
}
loop {
2021-05-05 20:46:56 +02:00
match self.state.take() {
Some((query_tree, _, allowed_candidates)) if allowed_candidates.is_empty() => {
return Ok(Some(CriterionResult {
2021-05-05 20:46:56 +02:00
query_tree: Some(query_tree),
candidates: Some(RoaringBitmap::new()),
2021-05-10 12:33:37 +02:00
filtered_candidates: None,
2021-05-05 20:46:56 +02:00
bucket_candidates: Some(take(&mut self.bucket_candidates)),
}));
2021-06-16 18:33:33 +02:00
}
2021-05-05 20:46:56 +02:00
Some((query_tree, flattened_query_tree, mut allowed_candidates)) => {
let found_candidates = if allowed_candidates.len() < CANDIDATES_THRESHOLD {
let current_buckets = match self.current_buckets.as_mut() {
Some(current_buckets) => current_buckets,
None => {
2021-06-16 18:33:33 +02:00
let new_buckets = linear_compute_candidates(
self.ctx,
&flattened_query_tree,
&allowed_candidates,
)?;
self.current_buckets.get_or_insert(new_buckets.into_iter())
2021-06-16 18:33:33 +02:00
}
};
match current_buckets.next() {
Some((_score, candidates)) => candidates,
None => {
return Ok(Some(CriterionResult {
2021-05-05 20:46:56 +02:00
query_tree: Some(query_tree),
candidates: Some(RoaringBitmap::new()),
2021-05-10 12:33:37 +02:00
filtered_candidates: None,
2021-05-05 20:46:56 +02:00
bucket_candidates: Some(take(&mut self.bucket_candidates)),
}));
2021-06-16 18:33:33 +02:00
}
}
} else {
2021-06-16 18:33:33 +02:00
match set_compute_candidates(
self.ctx,
&flattened_query_tree,
&allowed_candidates,
params.wdcache,
)? {
Some(candidates) => candidates,
None => {
return Ok(Some(CriterionResult {
2021-05-05 20:46:56 +02:00
query_tree: Some(query_tree),
candidates: Some(RoaringBitmap::new()),
2021-05-10 12:33:37 +02:00
filtered_candidates: None,
2021-05-05 20:46:56 +02:00
bucket_candidates: Some(take(&mut self.bucket_candidates)),
}));
2021-06-16 18:33:33 +02:00
}
}
};
2021-05-05 20:46:56 +02:00
allowed_candidates -= &found_candidates;
2021-06-16 18:33:33 +02:00
self.state =
Some((query_tree.clone(), flattened_query_tree, allowed_candidates));
2021-03-23 15:25:46 +01:00
return Ok(Some(CriterionResult {
2021-05-05 20:46:56 +02:00
query_tree: Some(query_tree),
candidates: Some(found_candidates),
2021-05-10 12:33:37 +02:00
filtered_candidates: None,
2021-05-05 20:46:56 +02:00
bucket_candidates: Some(take(&mut self.bucket_candidates)),
}));
2021-06-16 18:33:33 +02:00
}
None => match self.parent.next(params)? {
Some(CriterionResult {
query_tree: Some(query_tree),
candidates,
filtered_candidates,
bucket_candidates,
}) => {
let mut candidates = match candidates {
Some(candidates) => candidates,
None => {
resolve_query_tree(self.ctx, &query_tree, params.wdcache)?
- params.excluded_candidates
2021-05-10 12:33:37 +02:00
}
2021-06-16 18:33:33 +02:00
};
2021-05-10 12:33:37 +02:00
2021-06-16 18:33:33 +02:00
if let Some(filtered_candidates) = filtered_candidates {
candidates &= filtered_candidates;
}
2021-05-05 20:46:56 +02:00
2021-06-16 18:33:33 +02:00
let flattened_query_tree = flatten_query_tree(&query_tree);
2021-05-05 20:46:56 +02:00
2021-06-16 18:33:33 +02:00
match bucket_candidates {
Some(bucket_candidates) => self.bucket_candidates |= bucket_candidates,
None => self.bucket_candidates |= &candidates,
}
self.state = Some((query_tree, flattened_query_tree, candidates));
self.current_buckets = None;
}
Some(CriterionResult {
query_tree: None,
candidates,
filtered_candidates,
bucket_candidates,
}) => {
return Ok(Some(CriterionResult {
query_tree: None,
candidates,
filtered_candidates,
bucket_candidates,
}));
}
2021-06-16 18:33:33 +02:00
None => return Ok(None),
},
}
}
2021-03-11 11:48:55 +01:00
}
}
2021-04-15 10:44:27 +02:00
/// WordLevelIterator is an pseudo-Iterator over intervals of word-position for one word,
/// it will begin at the first non-empty interval and will return every interval without
/// jumping over empty intervals.
struct WordLevelIterator<'t, 'q> {
2021-06-16 18:33:33 +02:00
inner: Box<
dyn Iterator<Item = heed::Result<((&'t str, TreeLevel, u32, u32), RoaringBitmap)>> + 't,
>,
level: TreeLevel,
interval_size: u32,
word: Cow<'q, str>,
in_prefix_cache: bool,
inner_next: Option<(u32, u32, RoaringBitmap)>,
current_interval: Option<(u32, u32)>,
}
impl<'t, 'q> WordLevelIterator<'t, 'q> {
2021-06-16 18:33:33 +02:00
fn new(
ctx: &'t dyn Context<'t>,
word: Cow<'q, str>,
in_prefix_cache: bool,
) -> heed::Result<Option<Self>> {
match ctx.word_position_last_level(&word, in_prefix_cache)? {
2021-09-29 14:41:56 +02:00
Some(_) => {
// HOTFIX Meilisearch#1707: it is better to only iterate over level 0 for performances reasons.
2021-09-29 14:41:56 +02:00
let level = TreeLevel::min_value();
2021-05-10 10:27:18 +02:00
let interval_size = LEVEL_EXPONENTIATION_BASE.pow(Into::<u8>::into(level) as u32);
2021-06-16 18:33:33 +02:00
let inner =
ctx.word_position_iterator(&word, level, in_prefix_cache, None, None)?;
Ok(Some(Self {
inner,
level,
interval_size,
word,
in_prefix_cache,
inner_next: None,
current_interval: None,
}))
}
None => Ok(None),
}
}
2021-06-16 18:33:33 +02:00
fn dig(
&self,
ctx: &'t dyn Context<'t>,
level: &TreeLevel,
left_interval: Option<u32>,
) -> heed::Result<Self> {
2021-05-10 10:27:18 +02:00
let level = *level.min(&self.level);
let interval_size = LEVEL_EXPONENTIATION_BASE.pow(Into::<u8>::into(level) as u32);
let word = self.word.clone();
let in_prefix_cache = self.in_prefix_cache;
2021-06-16 18:33:33 +02:00
let inner =
ctx.word_position_iterator(&word, level, in_prefix_cache, left_interval, None)?;
Ok(Self {
inner,
level,
interval_size,
word,
in_prefix_cache,
inner_next: None,
current_interval: None,
})
}
fn next(&mut self) -> heed::Result<Option<(u32, u32, RoaringBitmap)>> {
2021-06-16 18:33:33 +02:00
fn is_next_interval(last_right: u32, next_left: u32) -> bool {
last_right + 1 == next_left
}
let inner_next = match self.inner_next.take() {
Some(inner_next) => Some(inner_next),
2021-06-16 18:33:33 +02:00
None => self
.inner
.next()
.transpose()?
.map(|((_, _, left, right), docids)| (left, right, docids)),
};
match inner_next {
2021-06-16 18:33:33 +02:00
Some((left, right, docids)) => match self.current_interval {
Some((last_left, last_right)) if !is_next_interval(last_right, left) => {
let blank_left = last_left + self.interval_size;
let blank_right = last_right + self.interval_size;
self.current_interval = Some((blank_left, blank_right));
self.inner_next = Some((left, right, docids));
Ok(Some((blank_left, blank_right, RoaringBitmap::new())))
}
_ => {
self.current_interval = Some((left, right));
Ok(Some((left, right, docids)))
}
},
None => Ok(None),
}
}
}
2021-04-15 10:44:27 +02:00
/// QueryLevelIterator is an pseudo-Iterator for a Query,
/// It contains WordLevelIterators and is chainned with other QueryLevelIterator.
struct QueryLevelIterator<'t, 'q> {
2021-04-15 10:44:27 +02:00
parent: Option<Box<QueryLevelIterator<'t, 'q>>>,
inner: Vec<WordLevelIterator<'t, 'q>>,
level: TreeLevel,
accumulator: Vec<Option<(u32, u32, RoaringBitmap)>>,
2021-04-15 10:44:27 +02:00
parent_accumulator: Vec<Option<(u32, u32, RoaringBitmap)>>,
interval_to_skip: usize,
}
impl<'t, 'q> QueryLevelIterator<'t, 'q> {
fn new(
ctx: &'t dyn Context<'t>,
queries: &'q [Query],
wdcache: &mut WordDerivationsCache,
2021-06-16 18:33:33 +02:00
) -> Result<Option<Self>> {
let mut inner = Vec::with_capacity(queries.len());
for query in queries {
match &query.kind {
QueryKind::Exact { word, .. } => {
if !query.prefix || ctx.in_prefix_cache(&word) {
let word = Cow::Borrowed(query.kind.word());
2021-06-16 18:33:33 +02:00
if let Some(word_level_iterator) =
WordLevelIterator::new(ctx, word, query.prefix)?
{
inner.push(word_level_iterator);
}
} else {
2021-06-16 18:33:33 +02:00
for (word, _) in word_derivations(&word, true, 0, ctx.words_fst(), wdcache)?
{
let word = Cow::Owned(word.to_owned());
2021-06-16 18:33:33 +02:00
if let Some(word_level_iterator) =
WordLevelIterator::new(ctx, word, false)?
{
inner.push(word_level_iterator);
}
}
}
2021-06-16 18:33:33 +02:00
}
QueryKind::Tolerant { typo, word } => {
2021-06-16 18:33:33 +02:00
for (word, _) in
word_derivations(&word, query.prefix, *typo, ctx.words_fst(), wdcache)?
{
let word = Cow::Owned(word.to_owned());
2021-06-16 18:33:33 +02:00
if let Some(word_level_iterator) = WordLevelIterator::new(ctx, word, false)?
{
inner.push(word_level_iterator);
}
}
}
}
}
2021-05-10 10:27:18 +02:00
let highest = inner.iter().max_by_key(|wli| wli.level).map(|wli| wli.level);
match highest {
Some(level) => Ok(Some(Self {
2021-04-15 10:44:27 +02:00
parent: None,
inner,
level,
accumulator: vec![],
2021-04-15 10:44:27 +02:00
parent_accumulator: vec![],
interval_to_skip: 0,
})),
None => Ok(None),
}
}
2021-04-15 10:44:27 +02:00
fn parent(&mut self, parent: QueryLevelIterator<'t, 'q>) -> &Self {
self.parent = Some(Box::new(parent));
self
}
2021-04-15 10:44:27 +02:00
/// create a new QueryLevelIterator with a lower level than the current one.
fn dig(&self, ctx: &'t dyn Context<'t>) -> heed::Result<Self> {
2021-04-15 10:44:27 +02:00
let (level, parent) = match &self.parent {
Some(parent) => {
let parent = parent.dig(ctx)?;
(parent.level.min(self.level), Some(Box::new(parent)))
2021-06-16 18:33:33 +02:00
}
None => (self.level.saturating_sub(1), None),
};
2021-06-16 18:33:33 +02:00
let left_interval = self
.accumulator
.get(self.interval_to_skip)
.map(|opt| opt.as_ref().map(|(left, _, _)| *left))
.flatten();
let mut inner = Vec::with_capacity(self.inner.len());
for word_level_iterator in self.inner.iter() {
inner.push(word_level_iterator.dig(ctx, &level, left_interval)?);
}
2021-06-16 18:33:33 +02:00
Ok(Self {
parent,
inner,
level,
accumulator: vec![],
parent_accumulator: vec![],
interval_to_skip: 0,
})
}
fn inner_next(&mut self, level: TreeLevel) -> heed::Result<Option<(u32, u32, RoaringBitmap)>> {
let mut accumulated: Option<(u32, u32, RoaringBitmap)> = None;
let u8_level = Into::<u8>::into(level);
let interval_size = LEVEL_EXPONENTIATION_BASE.pow(u8_level as u32);
for wli in self.inner.iter_mut() {
2021-05-10 10:27:18 +02:00
let wli_u8_level = Into::<u8>::into(wli.level);
let accumulated_count = LEVEL_EXPONENTIATION_BASE.pow((u8_level - wli_u8_level) as u32);
for _ in 0..accumulated_count {
2021-06-16 18:33:33 +02:00
if let Some((next_left, _, next_docids)) = wli.next()? {
accumulated = match accumulated.take() {
Some((acc_left, acc_right, mut acc_docids)) => {
acc_docids |= next_docids;
Some((acc_left, acc_right, acc_docids))
2021-06-16 18:33:33 +02:00
}
None => Some((next_left, next_left + interval_size, next_docids)),
};
}
}
}
Ok(accumulated)
}
2021-04-15 10:44:27 +02:00
/// return the next meta-interval created from inner WordLevelIterators,
/// and from eventual chainned QueryLevelIterator.
2021-06-16 18:33:33 +02:00
fn next(
&mut self,
allowed_candidates: &RoaringBitmap,
tree_level: TreeLevel,
) -> heed::Result<Option<(u32, u32, RoaringBitmap)>> {
2021-04-15 10:44:27 +02:00
let parent_result = match self.parent.as_mut() {
Some(parent) => Some(parent.next(allowed_candidates, tree_level)?),
None => None,
};
2021-04-15 10:44:27 +02:00
match parent_result {
Some(parent_next) => {
let inner_next = self.inner_next(tree_level)?;
self.interval_to_skip += interval_to_skip(
&self.parent_accumulator,
&self.accumulator,
self.interval_to_skip,
2021-06-16 18:33:33 +02:00
allowed_candidates,
);
self.accumulator.push(inner_next);
2021-04-15 10:44:27 +02:00
self.parent_accumulator.push(parent_next);
let mut merged_interval: Option<(u32, u32, RoaringBitmap)> = None;
2021-06-16 18:33:33 +02:00
for current in self
.accumulator
.iter()
.rev()
.zip(self.parent_accumulator.iter())
.skip(self.interval_to_skip)
{
if let (Some((left_a, right_a, a)), Some((left_b, right_b, b))) = current {
match merged_interval.as_mut() {
Some((_, _, merged_docids)) => *merged_docids |= a & b,
2021-06-16 18:33:33 +02:00
None => {
merged_interval = Some((left_a + left_b, right_a + right_b, a & b))
}
}
}
}
Ok(merged_interval)
2021-06-16 18:33:33 +02:00
}
None => {
let level = self.level;
match self.inner_next(level)? {
Some((left, right, mut candidates)) => {
self.accumulator = vec![Some((left, right, RoaringBitmap::new()))];
candidates &= allowed_candidates;
Ok(Some((left, right, candidates)))
2021-06-16 18:33:33 +02:00
}
None => {
self.accumulator = vec![None];
Ok(None)
2021-06-16 18:33:33 +02:00
}
}
}
}
}
}
/// Count the number of interval that can be skiped when we make the cross-intersections
/// in order to compute the next meta-interval.
/// A pair of intervals is skiped when both intervals doesn't contain any allowed docids.
fn interval_to_skip(
parent_accumulator: &[Option<(u32, u32, RoaringBitmap)>],
current_accumulator: &[Option<(u32, u32, RoaringBitmap)>],
already_skiped: usize,
allowed_candidates: &RoaringBitmap,
) -> usize {
2021-06-16 18:33:33 +02:00
parent_accumulator
.iter()
2021-05-10 10:27:18 +02:00
.zip(current_accumulator.iter())
.skip(already_skiped)
.take_while(|(parent, current)| {
let skip_parent = parent.as_ref().map_or(true, |(_, _, docids)| docids.is_empty());
2021-06-16 18:33:33 +02:00
let skip_current = current
.as_ref()
.map_or(true, |(_, _, docids)| docids.is_disjoint(allowed_candidates));
skip_parent && skip_current
})
.count()
}
/// A Branch is represent a possible alternative of the original query and is build with the Query Tree,
/// This branch allows us to iterate over meta-interval of position and to dig in it if it contains interesting candidates.
struct Branch<'t, 'q> {
query_level_iterator: QueryLevelIterator<'t, 'q>,
2021-04-12 11:19:25 +02:00
last_result: (u32, u32, RoaringBitmap),
tree_level: TreeLevel,
branch_size: u32,
}
impl<'t, 'q> Branch<'t, 'q> {
/// return the next meta-interval of the branch,
/// and update inner interval in order to be ranked by the BinaryHeap.
fn next(&mut self, allowed_candidates: &RoaringBitmap) -> heed::Result<bool> {
let tree_level = self.query_level_iterator.level;
match self.query_level_iterator.next(allowed_candidates, tree_level)? {
Some(last_result) => {
self.last_result = last_result;
self.tree_level = tree_level;
Ok(true)
2021-06-16 18:33:33 +02:00
}
None => Ok(false),
}
}
/// make the current Branch iterate over smaller intervals.
fn dig(&mut self, ctx: &'t dyn Context<'t>) -> heed::Result<()> {
self.query_level_iterator = self.query_level_iterator.dig(ctx)?;
Ok(())
}
/// because next() method could be time consuming,
/// update inner interval in order to be ranked by the binary_heap without computing it,
/// the next() method should be called when the real interval is needed.
fn lazy_next(&mut self) {
2021-05-10 10:27:18 +02:00
let u8_level = Into::<u8>::into(self.tree_level);
let interval_size = LEVEL_EXPONENTIATION_BASE.pow(u8_level as u32);
let (left, right, _) = self.last_result;
2021-06-16 18:33:33 +02:00
self.last_result = (left + interval_size, right + interval_size, RoaringBitmap::new());
}
/// return the score of the current inner interval.
fn compute_rank(&self) -> u32 {
2021-04-12 11:19:25 +02:00
// we compute a rank from the left interval.
let (left, _, _) = self.last_result;
left.saturating_sub((0..self.branch_size).sum()) * LCM_10_FIRST_NUMBERS / self.branch_size
}
fn cmp(&self, other: &Self) -> Ordering {
let self_rank = self.compute_rank();
let other_rank = other.compute_rank();
2021-09-29 14:41:56 +02:00
let left_cmp = self_rank.cmp(&other_rank);
// on level: lower is better,
// we want to dig faster into levels on interesting branches.
2021-09-29 14:41:56 +02:00
let level_cmp = self.tree_level.cmp(&other.tree_level);
2021-04-12 11:19:25 +02:00
left_cmp.then(level_cmp).then(self.last_result.2.len().cmp(&other.last_result.2.len()))
}
}
impl<'t, 'q> Ord for Branch<'t, 'q> {
fn cmp(&self, other: &Self) -> Ordering {
self.cmp(other)
}
}
impl<'t, 'q> PartialOrd for Branch<'t, 'q> {
fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
Some(self.cmp(other))
}
}
impl<'t, 'q> PartialEq for Branch<'t, 'q> {
fn eq(&self, other: &Self) -> bool {
self.cmp(other) == Ordering::Equal
}
}
impl<'t, 'q> Eq for Branch<'t, 'q> {}
fn initialize_query_level_iterators<'t, 'q>(
ctx: &'t dyn Context<'t>,
2021-05-05 20:46:56 +02:00
branches: &'q FlattenedQueryTree,
allowed_candidates: &RoaringBitmap,
wdcache: &mut WordDerivationsCache,
) -> Result<BinaryHeap<Branch<'t, 'q>>> {
let mut positions = BinaryHeap::with_capacity(branches.len());
for branch in branches {
let mut branch_positions = Vec::with_capacity(branch.len());
2021-06-16 18:33:33 +02:00
for queries in branch {
match QueryLevelIterator::new(ctx, queries, wdcache)? {
Some(qli) => branch_positions.push(qli),
None => {
// the branch seems to be invalid, so we skip it.
branch_positions.clear();
break;
2021-06-16 18:33:33 +02:00
}
}
}
// QueryLevelIterator need to be sorted by level and folded in descending order.
branch_positions.sort_unstable_by_key(|qli| qli.level);
2021-06-16 18:33:33 +02:00
let folded_query_level_iterators =
branch_positions.into_iter().fold(None, |fold: Option<QueryLevelIterator>, mut qli| {
match fold {
Some(fold) => {
qli.parent(fold);
Some(qli)
}
None => Some(qli),
}
});
if let Some(mut folded_query_level_iterators) = folded_query_level_iterators {
let tree_level = folded_query_level_iterators.level;
let last_result = folded_query_level_iterators.next(allowed_candidates, tree_level)?;
2021-04-12 11:19:25 +02:00
if let Some(last_result) = last_result {
let branch = Branch {
last_result,
tree_level,
query_level_iterator: folded_query_level_iterators,
branch_size: branch.len() as u32,
};
positions.push(branch);
}
}
}
Ok(positions)
}
fn set_compute_candidates<'t>(
ctx: &'t dyn Context<'t>,
2021-05-05 20:46:56 +02:00
branches: &FlattenedQueryTree,
allowed_candidates: &RoaringBitmap,
wdcache: &mut WordDerivationsCache,
2021-06-16 18:33:33 +02:00
) -> Result<Option<RoaringBitmap>> {
let mut branches_heap =
initialize_query_level_iterators(ctx, branches, allowed_candidates, wdcache)?;
let lowest_level = TreeLevel::min_value();
let mut final_candidates: Option<(u32, RoaringBitmap)> = None;
let mut allowed_candidates = allowed_candidates.clone();
while let Some(mut branch) = branches_heap.peek_mut() {
let is_lowest_level = branch.tree_level == lowest_level;
let branch_rank = branch.compute_rank();
// if current is worst than best we break to return
// candidates that correspond to the best rank
if let Some((best_rank, _)) = final_candidates {
2021-06-16 18:33:33 +02:00
if branch_rank > best_rank {
break;
}
}
let _left = branch.last_result.0;
let candidates = take(&mut branch.last_result.2);
2021-04-12 11:19:25 +02:00
if candidates.is_empty() {
// we don't have candidates, get next interval.
2021-06-16 18:33:33 +02:00
if !branch.next(&allowed_candidates)? {
PeekMut::pop(branch);
}
} else if is_lowest_level {
// we have candidates, but we can't dig deeper.
allowed_candidates -= &candidates;
final_candidates = match final_candidates.take() {
// we add current candidates to best candidates
Some((best_rank, mut best_candidates)) => {
best_candidates |= candidates;
branch.lazy_next();
Some((best_rank, best_candidates))
2021-06-16 18:33:33 +02:00
}
// we take current candidates as best candidates
None => {
branch.lazy_next();
Some((branch_rank, candidates))
2021-06-16 18:33:33 +02:00
}
};
2021-04-12 11:19:25 +02:00
} else {
// we have candidates, lets dig deeper in levels.
branch.dig(ctx)?;
2021-06-16 18:33:33 +02:00
if !branch.next(&allowed_candidates)? {
PeekMut::pop(branch);
}
}
}
Ok(final_candidates.map(|(_rank, candidates)| candidates))
}
fn linear_compute_candidates(
ctx: &dyn Context,
2021-05-05 20:46:56 +02:00
branches: &FlattenedQueryTree,
allowed_candidates: &RoaringBitmap,
2021-06-16 18:33:33 +02:00
) -> Result<BTreeMap<u64, RoaringBitmap>> {
fn compute_candidate_rank(
branches: &FlattenedQueryTree,
words_positions: HashMap<String, RoaringBitmap>,
) -> u64 {
let mut min_rank = u64::max_value();
for branch in branches {
2021-03-24 18:20:13 +01:00
let branch_len = branch.len();
let mut branch_rank = Vec::with_capacity(branch_len);
for derivates in branch {
let mut position = None;
for Query { prefix, kind } in derivates {
// find the best position of the current word in the document.
let current_position = match kind {
QueryKind::Exact { word, .. } => {
if *prefix {
word_derivations(word, true, 0, &words_positions)
2021-06-16 18:33:33 +02:00
.flat_map(|positions| positions.iter().next())
.min()
} else {
2021-06-16 18:33:33 +02:00
words_positions
.get(word)
.map(|positions| positions.iter().next())
.flatten()
}
2021-06-16 18:33:33 +02:00
}
QueryKind::Tolerant { typo, word } => {
word_derivations(word, *prefix, *typo, &words_positions)
2021-06-16 18:33:33 +02:00
.flat_map(|positions| positions.iter().next())
.min()
}
};
match (position, current_position) {
(Some(p), Some(cp)) => position = Some(cmp::min(p, cp)),
(None, Some(cp)) => position = Some(cp),
_ => (),
}
}
// if a position is found, we add it to the branch score,
// otherwise the branch is considered as unfindable in this document and we break.
if let Some(position) = position {
2021-03-24 18:20:13 +01:00
branch_rank.push(position as u64);
} else {
2021-03-24 18:20:13 +01:00
branch_rank.clear();
break;
}
}
2021-03-24 18:20:13 +01:00
if !branch_rank.is_empty() {
branch_rank.sort_unstable();
// because several words in same query can't match all a the position 0,
// we substract the word index to the position.
2021-06-16 18:33:33 +02:00
let branch_rank: u64 =
branch_rank.into_iter().enumerate().map(|(i, r)| r - i as u64).sum();
2021-03-24 18:20:13 +01:00
// here we do the means of the words of the branch
2021-06-16 18:33:33 +02:00
min_rank =
min_rank.min(branch_rank * LCM_10_FIRST_NUMBERS as u64 / branch_len as u64);
2021-03-24 18:20:13 +01:00
}
}
min_rank
}
fn word_derivations<'a>(
word: &str,
is_prefix: bool,
max_typo: u8,
words_positions: &'a HashMap<String, RoaringBitmap>,
2021-06-16 18:33:33 +02:00
) -> impl Iterator<Item = &'a RoaringBitmap> {
let dfa = build_dfa(word, max_typo, is_prefix);
words_positions.iter().filter_map(move |(document_word, positions)| {
use levenshtein_automata::Distance;
match dfa.eval(document_word) {
Distance::Exact(_) => Some(positions),
Distance::AtLeast(_) => None,
}
})
}
let mut candidates = BTreeMap::new();
for docid in allowed_candidates {
let words_positions = ctx.docid_words_positions(docid)?;
let rank = compute_candidate_rank(branches, words_positions);
candidates.entry(rank).or_insert_with(RoaringBitmap::new).insert(docid);
}
Ok(candidates)
}
2021-03-11 11:48:55 +01:00
// TODO can we keep refs of Query
2021-05-05 20:46:56 +02:00
fn flatten_query_tree(query_tree: &Operation) -> FlattenedQueryTree {
use crate::search::criteria::Operation::{And, Or, Phrase};
2021-03-11 11:48:55 +01:00
2021-05-05 20:46:56 +02:00
fn and_recurse(head: &Operation, tail: &[Operation]) -> FlattenedQueryTree {
2021-03-11 11:48:55 +01:00
match tail.split_first() {
Some((thead, tail)) => {
let tail = and_recurse(thead, tail);
let mut out = Vec::new();
for array in recurse(head) {
for tail_array in &tail {
let mut array = array.clone();
array.extend(tail_array.iter().cloned());
out.push(array);
}
}
out
2021-06-16 18:33:33 +02:00
}
2021-03-11 11:48:55 +01:00
None => recurse(head),
}
}
2021-05-05 20:46:56 +02:00
fn recurse(op: &Operation) -> FlattenedQueryTree {
2021-03-11 11:48:55 +01:00
match op {
2021-06-16 18:33:33 +02:00
And(ops) => ops.split_first().map_or_else(Vec::new, |(h, t)| and_recurse(h, t)),
Or(_, ops) => {
if ops.iter().all(|op| op.query().is_some()) {
vec![vec![ops.iter().flat_map(|op| op.query()).cloned().collect()]]
} else {
ops.iter().map(recurse).flatten().collect()
}
}
Phrase(words) => {
2021-06-16 18:33:33 +02:00
let queries = words
.iter()
.map(|word| vec![Query { prefix: false, kind: QueryKind::exact(word.clone()) }])
.collect();
vec![queries]
}
Operation::Query(query) => vec![vec![vec![query.clone()]]],
2021-03-11 11:48:55 +01:00
}
}
recurse(query_tree)
}
#[cfg(test)]
mod tests {
use big_s::S;
use super::*;
2021-06-16 18:33:33 +02:00
use crate::search::criteria::QueryKind;
2021-03-11 11:48:55 +01:00
#[test]
fn simple_flatten_query_tree() {
2021-06-16 18:33:33 +02:00
let query_tree = Operation::Or(
false,
vec![
Operation::Query(Query { prefix: false, kind: QueryKind::exact(S("manythefish")) }),
Operation::And(vec![
Operation::Query(Query { prefix: false, kind: QueryKind::exact(S("manythe")) }),
Operation::Query(Query { prefix: false, kind: QueryKind::exact(S("fish")) }),
2021-03-11 11:48:55 +01:00
]),
2021-06-16 18:33:33 +02:00
Operation::And(vec![
Operation::Query(Query { prefix: false, kind: QueryKind::exact(S("many")) }),
Operation::Or(
false,
vec![
Operation::Query(Query {
prefix: false,
kind: QueryKind::exact(S("thefish")),
}),
Operation::And(vec![
Operation::Query(Query {
prefix: false,
kind: QueryKind::exact(S("the")),
}),
Operation::Query(Query {
prefix: false,
kind: QueryKind::exact(S("fish")),
}),
]),
],
),
]),
],
);
2021-03-11 11:48:55 +01:00
let expected = vec![
vec![vec![Query { prefix: false, kind: QueryKind::exact(S("manythefish")) }]],
2021-03-11 11:48:55 +01:00
vec![
vec![Query { prefix: false, kind: QueryKind::exact(S("manythe")) }],
vec![Query { prefix: false, kind: QueryKind::exact(S("fish")) }],
2021-03-11 11:48:55 +01:00
],
vec![
vec![Query { prefix: false, kind: QueryKind::exact(S("many")) }],
vec![Query { prefix: false, kind: QueryKind::exact(S("thefish")) }],
2021-03-11 11:48:55 +01:00
],
vec![
vec![Query { prefix: false, kind: QueryKind::exact(S("many")) }],
vec![Query { prefix: false, kind: QueryKind::exact(S("the")) }],
vec![Query { prefix: false, kind: QueryKind::exact(S("fish")) }],
2021-03-11 11:48:55 +01:00
],
];
let result = flatten_query_tree(&query_tree);
2021-03-11 11:48:55 +01:00
assert_eq!(expected, result);
}
}