Rework the Asc/Desc criteria to be facet iterator based

This commit is contained in:
Clément Renault 2021-03-04 11:00:18 +01:00 committed by Kerollmops
parent a58d2b6137
commit 3c76b3548d
No known key found for this signature in database
GPG Key ID: 92ADA4E935E71FA4
2 changed files with 92 additions and 81 deletions

View File

@ -2,8 +2,10 @@ use std::collections::HashMap;
use std::mem::take;
use anyhow::{bail, Context as _};
use heed::{BytesDecode, BytesEncode};
use itertools::Itertools;
use log::debug;
use num_traits::Bounded;
use ordered_float::OrderedFloat;
use roaring::RoaringBitmap;
@ -24,7 +26,7 @@ pub struct AscDesc<'t> {
facet_type: FacetType,
ascending: bool,
query_tree: Option<Operation>,
candidates: RoaringBitmap,
candidates: Box<dyn Iterator<Item = heed::Result<RoaringBitmap>> + 't>,
bucket_candidates: RoaringBitmap,
faceted_candidates: RoaringBitmap,
parent: Option<Box<dyn Criterion + 't>>,
@ -107,7 +109,7 @@ impl<'t> AscDesc<'t> {
facet_type,
ascending,
query_tree,
candidates,
candidates: facet_ordered(index, rtxn, field_id, facet_type, ascending, candidates)?,
faceted_candidates,
bucket_candidates: RoaringBitmap::new(),
parent: None,
@ -134,7 +136,7 @@ impl<'t> AscDesc<'t> {
facet_type,
ascending,
query_tree: None,
candidates: RoaringBitmap::new(),
candidates: Box::new(std::iter::empty()),
faceted_candidates: index.faceted_documents_ids(rtxn, field_id)?,
bucket_candidates: RoaringBitmap::new(),
parent: Some(parent),
@ -145,23 +147,28 @@ impl<'t> AscDesc<'t> {
impl<'t> Criterion for AscDesc<'t> {
fn next(&mut self) -> anyhow::Result<Option<CriterionResult>> {
loop {
debug!("Facet {}({}) iteration ({:?})",
if self.ascending { "Asc" } else { "Desc" }, self.field_name, self.candidates,
debug!("Facet {}({}) iteration",
if self.ascending { "Asc" } else { "Desc" }, self.field_name
);
match &mut self.candidates {
candidates if candidates.is_empty() => {
match self.candidates.next().transpose()? {
None => {
let query_tree = self.query_tree.take();
let candidates = take(&mut self.candidates);
let bucket_candidates = take(&mut self.bucket_candidates);
match self.parent.as_mut() {
Some(parent) => {
match parent.next()? {
Some(CriterionResult { query_tree, mut candidates, bucket_candidates }) => {
self.query_tree = query_tree;
candidates.intersect_with(&self.faceted_candidates);
self.candidates = candidates;
self.candidates = facet_ordered(
self.index,
self.rtxn,
self.field_id,
self.facet_type,
self.ascending,
candidates,
)?;
self.bucket_candidates = bucket_candidates;
},
None => return Ok(None),
@ -172,28 +179,21 @@ impl<'t> Criterion for AscDesc<'t> {
},
}
return Ok(Some(CriterionResult { query_tree, candidates, bucket_candidates }));
return Ok(Some(CriterionResult {
query_tree,
candidates: RoaringBitmap::new(),
bucket_candidates,
}));
},
candidates => {
Some(candidates) => {
let bucket_candidates = match self.parent {
Some(_) => take(&mut self.bucket_candidates),
None => candidates.clone(),
};
let found_candidates = facet_ordered(
self.index,
self.rtxn,
self.field_id,
self.facet_type,
self.ascending,
candidates.clone(),
)?;
candidates.difference_with(&found_candidates);
return Ok(Some(CriterionResult {
query_tree: self.query_tree.clone(),
candidates: found_candidates,
candidates,
bucket_candidates,
}));
},
@ -217,86 +217,98 @@ fn field_id_facet_type(
Ok((id, *facet_type))
}
fn facet_ordered(
index: &Index,
rtxn: &heed::RoTxn,
/// Returns an iterator over groups of the given candidates in ascending or descending order.
///
/// It will either use an iterative or a recusrsive method on the whole facet database depending
/// on the number of candidates to rank.
fn facet_ordered<'t>(
index: &'t Index,
rtxn: &'t heed::RoTxn,
field_id: FieldId,
facet_type: FacetType,
ascending: bool,
candidates: RoaringBitmap,
) -> anyhow::Result<RoaringBitmap>
) -> anyhow::Result<Box<dyn Iterator<Item = heed::Result<RoaringBitmap>> + 't>>
{
match facet_type {
FacetType::Float => {
if candidates.len() <= 1000 {
let db = index.field_id_docid_facet_values.remap_key_type::<FieldDocIdFacetF64Codec>();
let mut docids_values = Vec::with_capacity(candidates.len() as usize);
for docid in candidates.iter() {
let left = (field_id, docid, f64::MIN);
let right = (field_id, docid, f64::MAX);
let mut iter = db.range(rtxn, &(left..=right))?;
let entry = if ascending { iter.next() } else { iter.last() };
if let Some(((_, _, value), ())) = entry.transpose()? {
docids_values.push((docid, OrderedFloat(value)));
}
}
docids_values.sort_unstable_by_key(|(_, value)| *value);
let iter = docids_values.into_iter();
let iter = if ascending {
Box::new(iter) as Box<dyn Iterator<Item = _>>
} else {
Box::new(iter.rev())
};
match iter.group_by(|(_, v)| *v).into_iter().next() {
Some((_, ids)) => Ok(ids.map(|(id, _)| id).into_iter().collect()),
None => Ok(RoaringBitmap::new())
}
let iter = iterative_facet_ordered_iter::<FieldDocIdFacetF64Codec, f64, OrderedFloat<f64>>(
index, rtxn, field_id, ascending, candidates,
)?;
Ok(Box::new(iter.map(Ok)) as Box<dyn Iterator<Item = _>>)
} else {
let facet_fn = if ascending {
FacetIter::<f64, FacetLevelValueF64Codec>::new_reducing
} else {
FacetIter::<f64, FacetLevelValueF64Codec>::new_reverse_reducing
};
let mut iter = facet_fn(rtxn, index, field_id, candidates)?;
Ok(iter.next().transpose()?.map(|(_, docids)| docids).unwrap_or_default())
let iter = facet_fn(rtxn, index, field_id, candidates)?;
Ok(Box::new(iter.map(|res| res.map(|(_, docids)| docids))))
}
},
FacetType::Integer => {
if candidates.len() <= 1000 {
let db = index.field_id_docid_facet_values.remap_key_type::<FieldDocIdFacetI64Codec>();
let mut docids_values = Vec::with_capacity(candidates.len() as usize);
for docid in candidates.iter() {
let left = (field_id, docid, i64::MIN);
let right = (field_id, docid, i64::MAX);
let mut iter = db.range(rtxn, &(left..=right))?;
let entry = if ascending { iter.next() } else { iter.last() };
if let Some(((_, _, value), ())) = entry.transpose()? {
docids_values.push((docid, value));
}
}
docids_values.sort_unstable_by_key(|(_, value)| *value);
let iter = docids_values.into_iter();
let iter = if ascending {
Box::new(iter) as Box<dyn Iterator<Item = _>>
} else {
Box::new(iter.rev())
};
match iter.group_by(|(_, v)| *v).into_iter().next() {
Some((_, ids)) => Ok(ids.map(|(id, _)| id).into_iter().collect()),
None => Ok(RoaringBitmap::new())
}
let iter = iterative_facet_ordered_iter::<FieldDocIdFacetI64Codec, i64, i64>(
index, rtxn, field_id, ascending, candidates,
)?;
Ok(Box::new(iter.map(Ok)) as Box<dyn Iterator<Item = _>>)
} else {
let facet_fn = if ascending {
FacetIter::<i64, FacetLevelValueI64Codec>::new_reducing
} else {
FacetIter::<i64, FacetLevelValueI64Codec>::new_reverse_reducing
};
let mut iter = facet_fn(rtxn, index, field_id, candidates)?;
Ok(iter.next().transpose()?.map(|(_, docids)| docids).unwrap_or_default())
let iter = facet_fn(rtxn, index, field_id, candidates)?;
Ok(Box::new(iter.map(|res| res.map(|(_, docids)| docids))))
}
},
FacetType::String => bail!("criteria facet type must be a number"),
}
}
/// Fetch the whole list of candidates facet values one by one and order them by it.
///
/// This function is fast when the amount of candidates to rank is small.
fn iterative_facet_ordered_iter<'t, KC, T, U>(
index: &'t Index,
rtxn: &'t heed::RoTxn,
field_id: FieldId,
ascending: bool,
candidates: RoaringBitmap,
) -> anyhow::Result<impl Iterator<Item = RoaringBitmap> + 't>
where
KC: BytesDecode<'t, DItem = (FieldId, u32, T)>,
KC: for<'a> BytesEncode<'a, EItem = (FieldId, u32, T)>,
T: Bounded,
U: From<T> + Ord + Clone + 't,
{
let db = index.field_id_docid_facet_values.remap_key_type::<KC>();
let mut docids_values = Vec::with_capacity(candidates.len() as usize);
for docid in candidates.iter() {
let left = (field_id, docid, T::min_value());
let right = (field_id, docid, T::max_value());
let mut iter = db.range(rtxn, &(left..=right))?;
let entry = if ascending { iter.next() } else { iter.last() };
if let Some(((_, _, value), ())) = entry.transpose()? {
docids_values.push((docid, U::from(value)));
}
}
docids_values.sort_unstable_by_key(|(_, v)| v.clone());
let iter = docids_values.into_iter();
let iter = if ascending {
Box::new(iter) as Box<dyn Iterator<Item = _>>
} else {
Box::new(iter.rev())
};
// The itertools GroupBy iterator doesn't provide an owned version, we are therefore
// required to collect the result into an owned collection (a Vec).
// https://github.com/rust-itertools/itertools/issues/499
let vec: Vec<_> = iter.group_by(|(_, v)| v.clone())
.into_iter()
.map(|(_, ids)| ids.map(|(id, _)| id).collect())
.collect();
Ok(vec.into_iter())
}

View File

@ -1,12 +1,11 @@
use std::collections::HashMap;
use std::borrow::Cow;
use anyhow::{bail, Context as _};
use anyhow::bail;
use roaring::RoaringBitmap;
use crate::facet::FacetType;
use crate::search::word_derivations;
use crate::{Index, FieldId};
use crate::Index;
use super::query_tree::{Operation, Query, QueryKind};
use self::typo::Typo;