Make asc/desc criterion return resting documents

Fix #161.2
This commit is contained in:
many 2021-06-02 16:30:56 +02:00
parent 3c304c89d4
commit 26a9974667
No known key found for this signature in database
GPG Key ID: 2CEF23B75189EACA
2 changed files with 21 additions and 7 deletions

View File

@ -24,6 +24,7 @@ pub struct AscDesc<'t> {
ascending: bool, ascending: bool,
query_tree: Option<Operation>, query_tree: Option<Operation>,
candidates: Box<dyn Iterator<Item = heed::Result<RoaringBitmap>> + 't>, candidates: Box<dyn Iterator<Item = heed::Result<RoaringBitmap>> + 't>,
allowed_candidates: RoaringBitmap,
bucket_candidates: RoaringBitmap, bucket_candidates: RoaringBitmap,
faceted_candidates: RoaringBitmap, faceted_candidates: RoaringBitmap,
parent: Box<dyn Criterion + 't>, parent: Box<dyn Criterion + 't>,
@ -68,6 +69,7 @@ impl<'t> AscDesc<'t> {
ascending, ascending,
query_tree: None, query_tree: None,
candidates: Box::new(std::iter::empty()), candidates: Box::new(std::iter::empty()),
allowed_candidates: RoaringBitmap::new(),
faceted_candidates: index.number_faceted_documents_ids(rtxn, field_id)?, faceted_candidates: index.number_faceted_documents_ids(rtxn, field_id)?,
bucket_candidates: RoaringBitmap::new(), bucket_candidates: RoaringBitmap::new(),
parent, parent,
@ -78,6 +80,9 @@ impl<'t> AscDesc<'t> {
impl<'t> Criterion for AscDesc<'t> { impl<'t> Criterion for AscDesc<'t> {
#[logging_timer::time("AscDesc::{}")] #[logging_timer::time("AscDesc::{}")]
fn next(&mut self, params: &mut CriterionParameters) -> anyhow::Result<Option<CriterionResult>> { fn next(&mut self, params: &mut CriterionParameters) -> anyhow::Result<Option<CriterionResult>> {
// remove excluded candidates when next is called, instead of doing it in the loop.
self.allowed_candidates -= params.excluded_candidates;
loop { loop {
debug!( debug!(
"Facet {}({}) iteration", "Facet {}({}) iteration",
@ -86,18 +91,25 @@ impl<'t> Criterion for AscDesc<'t> {
); );
match self.candidates.next().transpose()? { match self.candidates.next().transpose()? {
None if !self.allowed_candidates.is_empty() => {
return Ok(Some(CriterionResult {
query_tree: self.query_tree.clone(),
candidates: Some(take(&mut self.allowed_candidates)),
filtered_candidates: None,
bucket_candidates: Some(take(&mut self.bucket_candidates)),
}));
},
None => { None => {
match self.parent.next(params)? { match self.parent.next(params)? {
Some(CriterionResult { query_tree, candidates, filtered_candidates, bucket_candidates }) => { Some(CriterionResult { query_tree, candidates, filtered_candidates, bucket_candidates }) => {
self.query_tree = query_tree; self.query_tree = query_tree;
let mut candidates = match (&self.query_tree, candidates) { let mut candidates = match (&self.query_tree, candidates) {
(_, Some(candidates)) => candidates & &self.faceted_candidates, (_, Some(candidates)) => candidates,
(Some(qt), None) => { (Some(qt), None) => {
let context = CriteriaBuilder::new(&self.rtxn, &self.index)?; let context = CriteriaBuilder::new(&self.rtxn, &self.index)?;
let candidates = resolve_query_tree(&context, qt, params.wdcache)?; resolve_query_tree(&context, qt, params.wdcache)?
candidates & &self.faceted_candidates
}, },
(None, None) => take(&mut self.faceted_candidates), (None, None) => self.index.documents_ids(self.rtxn)?,
}; };
if let Some(filtered_candidates) = filtered_candidates { if let Some(filtered_candidates) = filtered_candidates {
@ -113,12 +125,13 @@ impl<'t> Criterion for AscDesc<'t> {
continue; continue;
} }
self.allowed_candidates = &candidates - params.excluded_candidates;
self.candidates = facet_ordered( self.candidates = facet_ordered(
self.index, self.index,
self.rtxn, self.rtxn,
self.field_id, self.field_id,
self.ascending, self.ascending,
candidates, candidates & &self.faceted_candidates,
)?; )?;
}, },
None => return Ok(None), None => return Ok(None),
@ -126,6 +139,7 @@ impl<'t> Criterion for AscDesc<'t> {
}, },
Some(mut candidates) => { Some(mut candidates) => {
candidates -= params.excluded_candidates; candidates -= params.excluded_candidates;
self.allowed_candidates -= &candidates;
return Ok(Some(CriterionResult { return Ok(Some(CriterionResult {
query_tree: self.query_tree.clone(), query_tree: self.query_tree.clone(),
candidates: Some(candidates), candidates: Some(candidates),

View File

@ -203,14 +203,14 @@ impl<'t> CriteriaBuilder<'t> {
&'t self, &'t self,
query_tree: Option<Operation>, query_tree: Option<Operation>,
primitive_query: Option<Vec<PrimitiveQueryPart>>, primitive_query: Option<Vec<PrimitiveQueryPart>>,
facet_candidates: Option<RoaringBitmap>, filtered_candidates: Option<RoaringBitmap>,
) -> anyhow::Result<Final<'t>> ) -> anyhow::Result<Final<'t>>
{ {
use crate::criterion::Criterion as Name; use crate::criterion::Criterion as Name;
let primitive_query = primitive_query.unwrap_or_default(); let primitive_query = primitive_query.unwrap_or_default();
let mut criterion = Box::new(Initial::new(query_tree, facet_candidates)) as Box<dyn Criterion>; let mut criterion = Box::new(Initial::new(query_tree, filtered_candidates)) as Box<dyn Criterion>;
for name in self.index.criteria(&self.rtxn)? { for name in self.index.criteria(&self.rtxn)? {
criterion = match name { criterion = match name {
Name::Typo => Box::new(Typo::new(self, criterion)), Name::Typo => Box::new(Typo::new(self, criterion)),