Implement distinct attribute

This commit is contained in:
Loïc Lecrenier
2023-03-09 15:20:29 +01:00
parent 0465ba4a05
commit 78b9304d52
4 changed files with 176 additions and 37 deletions

View File

@ -3,6 +3,7 @@ use roaring::RoaringBitmap;
use super::logger::SearchLogger;
use super::{QueryGraph, SearchContext};
// use crate::search::new::sort::Sort;
use crate::search::new::distinct::{apply_distinct_rule, DistinctOutput};
use crate::Result;
/// An internal trait implemented by only [`PlaceholderQuery`] and [`QueryGraph`]
@ -80,6 +81,12 @@ pub fn bucket_sort<'search, Q: RankingRuleQueryTrait>(
logger.ranking_rules(&ranking_rules);
let distinct_fid = if let Some(field) = ctx.index.distinct_field(ctx.txn)? {
ctx.index.fields_ids_map(ctx.txn)?.id(field)
} else {
None
};
if universe.len() < from as u64 {
return Ok(vec![]);
}
@ -88,8 +95,9 @@ pub fn bucket_sort<'search, Q: RankingRuleQueryTrait>(
logger.start_iteration_ranking_rule(0, ranking_rules[0], query_graph, universe);
ranking_rules[0].start_iteration(ctx, logger, universe, query_graph)?;
let mut candidates: Vec<RoaringBitmap> = vec![RoaringBitmap::default(); ranking_rules_len];
candidates[0] = universe.clone();
let mut ranking_rule_universes: Vec<RoaringBitmap> =
vec![RoaringBitmap::default(); ranking_rules_len];
ranking_rule_universes[0] = universe.clone();
let mut cur_ranking_rule_index = 0;
@ -98,13 +106,13 @@ pub fn bucket_sort<'search, Q: RankingRuleQueryTrait>(
/// Update the candidates accordingly and inform the logger.
macro_rules! back {
() => {
assert!(candidates[cur_ranking_rule_index].is_empty());
assert!(ranking_rule_universes[cur_ranking_rule_index].is_empty());
logger.end_iteration_ranking_rule(
cur_ranking_rule_index,
ranking_rules[cur_ranking_rule_index],
&candidates[cur_ranking_rule_index],
&ranking_rule_universes[cur_ranking_rule_index],
);
candidates[cur_ranking_rule_index].clear();
ranking_rule_universes[cur_ranking_rule_index].clear();
ranking_rules[cur_ranking_rule_index].end_iteration(ctx, logger);
if cur_ranking_rule_index == 0 {
break;
@ -117,22 +125,35 @@ pub fn bucket_sort<'search, Q: RankingRuleQueryTrait>(
let mut results = vec![];
let mut cur_offset = 0usize;
/// Add the candidates to the results. Take the `from`, `limit`, and `cur_offset`
/// Add the candidates to the results. Take `distinct`, `from`, `limit`, and `cur_offset`
/// into account and inform the logger.
macro_rules! maybe_add_to_results {
($candidates:expr) => {
let candidates = $candidates;
// First apply the distinct rule on the candidates, reducing the universes if necessary
let candidates = if let Some(distinct_fid) = distinct_fid {
let DistinctOutput { remaining, excluded } = apply_distinct_rule(ctx, distinct_fid, $candidates)?;
for universe in ranking_rule_universes.iter_mut() {
*universe -= &excluded;
}
remaining
} else {
$candidates.clone()
};
let len = candidates.len();
// if the candidates are empty, there is nothing to do;
if !candidates.is_empty() {
// if we still haven't reached the first document to return
if cur_offset < from {
// and if no document from this bucket can be returned
if cur_offset + (candidates.len() as usize) < from {
// then just skip the bucket
logger.skip_bucket_ranking_rule(
cur_ranking_rule_index,
ranking_rules[cur_ranking_rule_index],
&candidates,
);
} else {
// otherwise, skip some of the documents and add some of the rest, in order of ids
let all_candidates = candidates.iter().collect::<Vec<_>>();
let (skipped_candidates, candidates) =
all_candidates.split_at(from - cur_offset);
@ -150,6 +171,7 @@ pub fn bucket_sort<'search, Q: RankingRuleQueryTrait>(
results.extend(&candidates);
}
} else {
// if we have passed the offset already, add some of the documents (up to the limit)
let candidates =
candidates.iter().take(length - results.len()).collect::<Vec<u32>>();
logger.add_to_results(&candidates);
@ -162,14 +184,14 @@ pub fn bucket_sort<'search, Q: RankingRuleQueryTrait>(
while results.len() < length {
// The universe for this bucket is zero or one element, so we don't need to sort
// anything, just extend the results and go back to the parent ranking rule.
if candidates[cur_ranking_rule_index].len() <= 1 {
maybe_add_to_results!(&candidates[cur_ranking_rule_index]);
candidates[cur_ranking_rule_index].clear();
if ranking_rule_universes[cur_ranking_rule_index].len() <= 1 {
maybe_add_to_results!(&ranking_rule_universes[cur_ranking_rule_index]);
ranking_rule_universes[cur_ranking_rule_index].clear();
back!();
continue;
}
let Some(next_bucket) = ranking_rules[cur_ranking_rule_index].next_bucket(ctx, logger, &candidates[cur_ranking_rule_index])? else {
let Some(next_bucket) = ranking_rules[cur_ranking_rule_index].next_bucket(ctx, logger, &ranking_rule_universes[cur_ranking_rule_index])? else {
back!();
continue;
};
@ -177,12 +199,12 @@ pub fn bucket_sort<'search, Q: RankingRuleQueryTrait>(
logger.next_bucket_ranking_rule(
cur_ranking_rule_index,
ranking_rules[cur_ranking_rule_index],
&candidates[cur_ranking_rule_index],
&ranking_rule_universes[cur_ranking_rule_index],
&next_bucket.candidates,
);
assert!(candidates[cur_ranking_rule_index].is_superset(&next_bucket.candidates));
candidates[cur_ranking_rule_index] -= &next_bucket.candidates;
assert!(ranking_rule_universes[cur_ranking_rule_index].is_superset(&next_bucket.candidates));
ranking_rule_universes[cur_ranking_rule_index] -= &next_bucket.candidates;
if cur_ranking_rule_index == ranking_rules_len - 1
|| next_bucket.candidates.len() <= 1
@ -193,12 +215,12 @@ pub fn bucket_sort<'search, Q: RankingRuleQueryTrait>(
}
cur_ranking_rule_index += 1;
candidates[cur_ranking_rule_index] = next_bucket.candidates.clone();
ranking_rule_universes[cur_ranking_rule_index] = next_bucket.candidates.clone();
logger.start_iteration_ranking_rule(
cur_ranking_rule_index,
ranking_rules[cur_ranking_rule_index],
&next_bucket.query,
&candidates[cur_ranking_rule_index],
&ranking_rule_universes[cur_ranking_rule_index],
);
ranking_rules[cur_ranking_rule_index].start_iteration(
ctx,