mirror of
https://github.com/meilisearch/meilisearch.git
synced 2025-07-27 00:31:02 +00:00
Implement attribute criterion
* Implement WordLevelIterator * Implement QueryLevelIterator * Implement set algorithm based on iterators Not tested + Some TODO to fix
This commit is contained in:
@ -1,17 +1,17 @@
|
||||
use std::cmp;
|
||||
use std::{cmp::{self, Ordering}, collections::BinaryHeap};
|
||||
use std::collections::{BTreeMap, HashMap, btree_map};
|
||||
use std::mem::take;
|
||||
|
||||
use roaring::RoaringBitmap;
|
||||
|
||||
use crate::{search::build_dfa};
|
||||
use crate::{TreeLevel, search::build_dfa};
|
||||
use crate::search::criteria::Query;
|
||||
use crate::search::query_tree::{Operation, QueryKind};
|
||||
use crate::search::WordDerivationsCache;
|
||||
use super::{Criterion, CriterionResult, Context, resolve_query_tree};
|
||||
|
||||
pub struct Attribute<'t> {
|
||||
ctx: &'t dyn Context,
|
||||
ctx: &'t dyn Context<'t>,
|
||||
query_tree: Option<Operation>,
|
||||
candidates: Option<RoaringBitmap>,
|
||||
bucket_candidates: RoaringBitmap,
|
||||
@ -21,7 +21,7 @@ pub struct Attribute<'t> {
|
||||
}
|
||||
|
||||
impl<'t> Attribute<'t> {
|
||||
pub fn new(ctx: &'t dyn Context, parent: Box<dyn Criterion + 't>) -> Self {
|
||||
pub fn new(ctx: &'t dyn Context<'t>, parent: Box<dyn Criterion + 't>) -> Self {
|
||||
Attribute {
|
||||
ctx,
|
||||
query_tree: None,
|
||||
@ -51,23 +51,27 @@ impl<'t> Criterion for Attribute<'t> {
|
||||
flatten_query_tree(&qt)
|
||||
});
|
||||
|
||||
let current_buckets = match self.current_buckets.as_mut() {
|
||||
Some(current_buckets) => current_buckets,
|
||||
None => {
|
||||
let new_buckets = linear_compute_candidates(self.ctx, flattened_query_tree, candidates)?;
|
||||
self.current_buckets.get_or_insert(new_buckets.into_iter())
|
||||
},
|
||||
};
|
||||
let found_candidates = if candidates.len() < 1000 {
|
||||
let current_buckets = match self.current_buckets.as_mut() {
|
||||
Some(current_buckets) => current_buckets,
|
||||
None => {
|
||||
let new_buckets = linear_compute_candidates(self.ctx, flattened_query_tree, candidates)?;
|
||||
self.current_buckets.get_or_insert(new_buckets.into_iter())
|
||||
},
|
||||
};
|
||||
|
||||
let found_candidates = match current_buckets.next() {
|
||||
Some((_score, candidates)) => candidates,
|
||||
None => {
|
||||
return Ok(Some(CriterionResult {
|
||||
query_tree: self.query_tree.take(),
|
||||
candidates: self.candidates.take(),
|
||||
bucket_candidates: take(&mut self.bucket_candidates),
|
||||
}));
|
||||
},
|
||||
match current_buckets.next() {
|
||||
Some((_score, candidates)) => candidates,
|
||||
None => {
|
||||
return Ok(Some(CriterionResult {
|
||||
query_tree: self.query_tree.take(),
|
||||
candidates: self.candidates.take(),
|
||||
bucket_candidates: take(&mut self.bucket_candidates),
|
||||
}));
|
||||
},
|
||||
}
|
||||
} else {
|
||||
set_compute_candidates(self.ctx, flattened_query_tree, candidates)?
|
||||
};
|
||||
|
||||
candidates.difference_with(&found_candidates);
|
||||
@ -114,6 +118,316 @@ impl<'t> Criterion for Attribute<'t> {
|
||||
}
|
||||
}
|
||||
|
||||
struct WordLevelIterator<'t, 'q> {
|
||||
inner: Box<dyn Iterator<Item =heed::Result<((&'t str, TreeLevel, u32, u32), RoaringBitmap)>> + 't>,
|
||||
level: TreeLevel,
|
||||
interval_size: u32,
|
||||
word: &'q str,
|
||||
in_prefix_cache: bool,
|
||||
inner_next: Option<(u32, u32, RoaringBitmap)>,
|
||||
current_interval: Option<(u32, u32)>,
|
||||
}
|
||||
|
||||
impl<'t, 'q> WordLevelIterator<'t, 'q> {
|
||||
fn new(ctx: &'t dyn Context<'t>, query: &'q Query) -> heed::Result<Option<Self>> {
|
||||
// TODO make it typo/prefix tolerant
|
||||
let word = query.kind.word();
|
||||
let in_prefix_cache = query.prefix && ctx.in_prefix_cache(word);
|
||||
match ctx.word_position_last_level(word, in_prefix_cache)? {
|
||||
Some(level) => {
|
||||
let interval_size = 4u32.pow(Into::<u8>::into(level.clone()) as u32);
|
||||
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),
|
||||
}
|
||||
}
|
||||
|
||||
fn dig(&self, ctx: &'t dyn Context<'t>, level: &TreeLevel) -> heed::Result<Self> {
|
||||
let level = level.min(&self.level).clone();
|
||||
let interval_size = 4u32.pow(Into::<u8>::into(level.clone()) as u32);
|
||||
let word = self.word;
|
||||
let in_prefix_cache = self.in_prefix_cache;
|
||||
// TODO try to dig starting from the current interval
|
||||
// let left = self.current_interval.map(|(left, _)| left);
|
||||
let inner = ctx.word_position_iterator(word, level, in_prefix_cache, None, 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)>> {
|
||||
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),
|
||||
None => self.inner.next().transpose()?.map(|((_, _, left, right), docids)| (left, right, docids)),
|
||||
};
|
||||
|
||||
match inner_next {
|
||||
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),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
struct QueryLevelIterator<'t, 'q> {
|
||||
previous: Option<Box<QueryLevelIterator<'t, 'q>>>,
|
||||
inner: Vec<WordLevelIterator<'t, 'q>>,
|
||||
level: TreeLevel,
|
||||
accumulator: Vec<Option<(u32, u32, RoaringBitmap)>>,
|
||||
previous_accumulator: Vec<Option<(u32, u32, RoaringBitmap)>>,
|
||||
}
|
||||
|
||||
impl<'t, 'q> QueryLevelIterator<'t, 'q> {
|
||||
fn new(ctx: &'t dyn Context<'t>, queries: &'q Vec<Query>) -> heed::Result<Option<Self>> {
|
||||
let mut inner = Vec::with_capacity(queries.len());
|
||||
for query in queries {
|
||||
if let Some(word_level_iterator) = WordLevelIterator::new(ctx, query)? {
|
||||
inner.push(word_level_iterator);
|
||||
}
|
||||
}
|
||||
|
||||
let highest = inner.iter().max_by_key(|wli| wli.level).map(|wli| wli.level.clone());
|
||||
match highest {
|
||||
Some(level) => Ok(Some(Self {
|
||||
previous: None,
|
||||
inner,
|
||||
level,
|
||||
accumulator: vec![],
|
||||
previous_accumulator: vec![],
|
||||
})),
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
fn previous(&mut self, previous: QueryLevelIterator<'t, 'q>) -> &Self {
|
||||
self.previous = Some(Box::new(previous));
|
||||
self
|
||||
}
|
||||
|
||||
fn dig(&self, ctx: &'t dyn Context<'t>) -> heed::Result<Self> {
|
||||
let (level, previous) = match &self.previous {
|
||||
Some(previous) => {
|
||||
let previous = previous.dig(ctx)?;
|
||||
(previous.level.min(self.level), Some(Box::new(previous)))
|
||||
},
|
||||
None => (self.level.saturating_sub(1), None),
|
||||
};
|
||||
|
||||
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)?);
|
||||
}
|
||||
|
||||
Ok(Self {previous, inner, level, accumulator: vec![], previous_accumulator: vec![]})
|
||||
}
|
||||
|
||||
|
||||
|
||||
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 = 4u32.pow(u8_level as u32);
|
||||
for wli in self.inner.iter_mut() {
|
||||
let wli_u8_level = Into::<u8>::into(wli.level.clone());
|
||||
let accumulated_count = 4u32.pow((u8_level - wli_u8_level) as u32);
|
||||
for _ in 0..accumulated_count {
|
||||
if let Some((next_left, _, next_docids)) = wli.next()? {
|
||||
accumulated = accumulated.take().map(
|
||||
|(acc_left, acc_right, mut acc_docids)| {
|
||||
acc_docids.union_with(&next_docids);
|
||||
(acc_left, acc_right, acc_docids)
|
||||
}
|
||||
).or_else(|| Some((next_left, next_left + interval_size, next_docids)));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(accumulated)
|
||||
}
|
||||
|
||||
fn next(&mut self) -> heed::Result<(TreeLevel, Option<(u32, u32, RoaringBitmap)>)> {
|
||||
let previous_result = match self.previous.as_mut() {
|
||||
Some(previous) => {
|
||||
Some(previous.next()?)
|
||||
},
|
||||
None => None,
|
||||
};
|
||||
|
||||
match previous_result {
|
||||
Some((previous_level, previous_next)) => {
|
||||
let inner_next = self.inner_next(previous_level)?;
|
||||
self.accumulator.push(inner_next);
|
||||
self.previous_accumulator.push(previous_next);
|
||||
// TODO @many clean firsts intervals of both accumulators when both RoaringBitmap are empty,
|
||||
// WARNING the cleaned intervals count needs to be kept to skip at the end
|
||||
let mut merged_interval = None;
|
||||
for current in self.accumulator.iter().rev().zip(self.previous_accumulator.iter()) {
|
||||
if let (Some((left_a, right_a, a)), Some((left_b, right_b, b))) = current {
|
||||
let (_, _, merged_docids) = merged_interval.get_or_insert_with(|| (left_a + left_b, right_a + right_b, RoaringBitmap::new()));
|
||||
merged_docids.union_with(&(a & b));
|
||||
}
|
||||
}
|
||||
Ok((previous_level, merged_interval))
|
||||
},
|
||||
None => {
|
||||
let level = self.level.clone();
|
||||
let next_interval = self.inner_next(level.clone())?;
|
||||
self.accumulator = vec![next_interval.clone()];
|
||||
Ok((level, next_interval))
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
struct Branch<'t, 'q> {
|
||||
query_level_iterator: QueryLevelIterator<'t, 'q>,
|
||||
last_result: Option<(u32, u32, RoaringBitmap)>,
|
||||
tree_level: TreeLevel,
|
||||
branch_size: u32,
|
||||
}
|
||||
|
||||
impl<'t, 'q> Branch<'t, 'q> {
|
||||
fn cmp(&self, other: &Self) -> Ordering {
|
||||
fn compute_rank(left: u32, branch_size: u32) -> u32 { left.saturating_sub((1..branch_size).sum()) / branch_size }
|
||||
match (&self.last_result, &other.last_result) {
|
||||
(Some((s_left, _, _)), Some((o_left, _, _))) => {
|
||||
// we compute a rank form the left interval.
|
||||
let self_rank = compute_rank(*s_left, self.branch_size);
|
||||
let other_rank = compute_rank(*o_left, other.branch_size);
|
||||
let left_cmp = self_rank.cmp(&other_rank).reverse();
|
||||
// on level: higher is better,
|
||||
// we want to reduce highest levels first.
|
||||
let level_cmp = self.tree_level.cmp(&other.tree_level);
|
||||
|
||||
left_cmp.then(level_cmp)
|
||||
},
|
||||
(Some(_), None) => Ordering::Greater,
|
||||
(None, Some(_)) => Ordering::Less,
|
||||
(None, None) => Ordering::Equal,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
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>,
|
||||
branches: &'q Vec<Vec<Vec<Query>>>,
|
||||
) -> heed::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());
|
||||
for query in branch {
|
||||
match QueryLevelIterator::new(ctx, query)? {
|
||||
Some(qli) => branch_positions.push(qli),
|
||||
None => {
|
||||
// the branch seems to be invalid, so we skip it.
|
||||
branch_positions.clear();
|
||||
break;
|
||||
},
|
||||
}
|
||||
}
|
||||
// QueryLevelIterator need to be sorted by level and folded in descending order.
|
||||
branch_positions.sort_unstable_by_key(|qli| qli.level);
|
||||
let folded_query_level_iterators = branch_positions
|
||||
.into_iter()
|
||||
.rev()
|
||||
.fold(None, |fold: Option<QueryLevelIterator>, qli| match fold {
|
||||
Some(mut fold) => {
|
||||
fold.previous(qli);
|
||||
Some(fold)
|
||||
},
|
||||
None => Some(qli),
|
||||
});
|
||||
|
||||
if let Some(mut folded_query_level_iterators) = folded_query_level_iterators {
|
||||
let (tree_level, last_result) = folded_query_level_iterators.next()?;
|
||||
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>,
|
||||
branches: &Vec<Vec<Vec<Query>>>,
|
||||
allowed_candidates: &RoaringBitmap,
|
||||
) -> anyhow::Result<RoaringBitmap>
|
||||
{
|
||||
let mut branches_heap = initialize_query_level_iterators(ctx, branches)?;
|
||||
let lowest_level = TreeLevel::min_value();
|
||||
|
||||
while let Some(mut branch) = branches_heap.peek_mut() {
|
||||
let is_lowest_level = branch.tree_level == lowest_level;
|
||||
match branch.last_result.as_mut() {
|
||||
Some((_, _, candidates)) => {
|
||||
candidates.intersect_with(&allowed_candidates);
|
||||
if candidates.len() > 0 && is_lowest_level {
|
||||
// we have candidates, but we can't dig deeper, return candidates.
|
||||
return Ok(std::mem::take(candidates));
|
||||
} else if candidates.len() > 0 {
|
||||
// we have candidates, lets dig deeper in levels.
|
||||
let mut query_level_iterator = branch.query_level_iterator.dig(ctx)?;
|
||||
let (tree_level, last_result) = query_level_iterator.next()?;
|
||||
branch.query_level_iterator = query_level_iterator;
|
||||
branch.tree_level = tree_level;
|
||||
branch.last_result = last_result;
|
||||
} else {
|
||||
// we don't have candidates, get next interval.
|
||||
let (_, last_result) = branch.query_level_iterator.next()?;
|
||||
branch.last_result = last_result;
|
||||
}
|
||||
},
|
||||
// None = no candidates to find.
|
||||
None => return Ok(RoaringBitmap::new()),
|
||||
}
|
||||
}
|
||||
|
||||
// we made all iterations without finding anything.
|
||||
Ok(RoaringBitmap::new())
|
||||
}
|
||||
|
||||
fn linear_compute_candidates(
|
||||
ctx: &dyn Context,
|
||||
branches: &Vec<Vec<Vec<Query>>>,
|
||||
|
Reference in New Issue
Block a user