From 26a9974667fc10cec2e27557f52e2e00c953642b Mon Sep 17 00:00:00 2001
From: many <maxime@meilisearch.com>
Date: Wed, 2 Jun 2021 16:30:56 +0200
Subject: [PATCH] Make asc/desc criterion return resting documents

Fix #161.2
---
 milli/src/search/criteria/asc_desc.rs | 24 +++++++++++++++++++-----
 milli/src/search/criteria/mod.rs      |  4 ++--
 2 files changed, 21 insertions(+), 7 deletions(-)

diff --git a/milli/src/search/criteria/asc_desc.rs b/milli/src/search/criteria/asc_desc.rs
index c80bb38f1..f90f3e421 100644
--- a/milli/src/search/criteria/asc_desc.rs
+++ b/milli/src/search/criteria/asc_desc.rs
@@ -24,6 +24,7 @@ pub struct AscDesc<'t> {
     ascending: bool,
     query_tree: Option<Operation>,
     candidates: Box<dyn Iterator<Item = heed::Result<RoaringBitmap>> + 't>,
+    allowed_candidates: RoaringBitmap,
     bucket_candidates: RoaringBitmap,
     faceted_candidates: RoaringBitmap,
     parent: Box<dyn Criterion + 't>,
@@ -68,6 +69,7 @@ impl<'t> AscDesc<'t> {
             ascending,
             query_tree: None,
             candidates: Box::new(std::iter::empty()),
+            allowed_candidates: RoaringBitmap::new(),
             faceted_candidates: index.number_faceted_documents_ids(rtxn, field_id)?,
             bucket_candidates: RoaringBitmap::new(),
             parent,
@@ -78,6 +80,9 @@ impl<'t> AscDesc<'t> {
 impl<'t> Criterion for AscDesc<'t> {
     #[logging_timer::time("AscDesc::{}")]
     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 {
             debug!(
                 "Facet {}({}) iteration",
@@ -86,18 +91,25 @@ impl<'t> Criterion for AscDesc<'t> {
             );
 
             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 => {
                     match self.parent.next(params)? {
                         Some(CriterionResult { query_tree, candidates, filtered_candidates, bucket_candidates }) => {
                             self.query_tree = query_tree;
                             let mut candidates = match (&self.query_tree, candidates) {
-                                (_, Some(candidates)) => candidates & &self.faceted_candidates,
+                                (_, Some(candidates)) => candidates,
                                 (Some(qt), None) => {
                                     let context = CriteriaBuilder::new(&self.rtxn, &self.index)?;
-                                    let candidates = resolve_query_tree(&context, qt, params.wdcache)?;
-                                    candidates & &self.faceted_candidates
+                                    resolve_query_tree(&context, qt, params.wdcache)?
                                 },
-                                (None, None) => take(&mut self.faceted_candidates),
+                                (None, None) => self.index.documents_ids(self.rtxn)?,
                             };
 
                             if let Some(filtered_candidates) = filtered_candidates {
@@ -113,12 +125,13 @@ impl<'t> Criterion for AscDesc<'t> {
                                 continue;
                             }
 
+                            self.allowed_candidates = &candidates - params.excluded_candidates;
                             self.candidates = facet_ordered(
                                 self.index,
                                 self.rtxn,
                                 self.field_id,
                                 self.ascending,
-                                candidates,
+                                candidates & &self.faceted_candidates,
                             )?;
                         },
                         None => return Ok(None),
@@ -126,6 +139,7 @@ impl<'t> Criterion for AscDesc<'t> {
                 },
                 Some(mut candidates) => {
                     candidates -= params.excluded_candidates;
+                    self.allowed_candidates -= &candidates;
                     return Ok(Some(CriterionResult {
                         query_tree: self.query_tree.clone(),
                         candidates: Some(candidates),
diff --git a/milli/src/search/criteria/mod.rs b/milli/src/search/criteria/mod.rs
index 456d16e1a..e4ca66b2c 100644
--- a/milli/src/search/criteria/mod.rs
+++ b/milli/src/search/criteria/mod.rs
@@ -203,14 +203,14 @@ impl<'t> CriteriaBuilder<'t> {
         &'t self,
         query_tree: Option<Operation>,
         primitive_query: Option<Vec<PrimitiveQueryPart>>,
-        facet_candidates: Option<RoaringBitmap>,
+        filtered_candidates: Option<RoaringBitmap>,
     ) -> anyhow::Result<Final<'t>>
     {
         use crate::criterion::Criterion as Name;
 
         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)? {
             criterion = match name {
                 Name::Typo => Box::new(Typo::new(self, criterion)),