implements a first version of the cutoff without settings

milli/src/search/new/bucket_sort.rs

@@ -5,12 +5,14 @@ use super::ranking_rules::{BoxRankingRule, RankingRuleQueryTrait};
 use super::SearchContext;
 use crate::score_details::{ScoreDetails, ScoringStrategy};
 use crate::search::new::distinct::{apply_distinct_rule, distinct_single_docid, DistinctOutput};
-use crate::Result;
+use crate::{Result, TimeBudget};
 
 pub struct BucketSortOutput {
     pub docids: Vec<u32>,
     pub scores: Vec<Vec<ScoreDetails>>,
     pub all_candidates: RoaringBitmap,
+
+    pub degraded: bool,
 }
 
 // TODO: would probably be good to regroup some of these inside of a struct?
@@ -25,6 +27,7 @@ pub fn bucket_sort<'ctx, Q: RankingRuleQueryTrait>(
     length: usize,
     scoring_strategy: ScoringStrategy,
     logger: &mut dyn SearchLogger<Q>,
+    time_budget: TimeBudget,
 ) -> Result<BucketSortOutput> {
     logger.initial_query(query);
     logger.ranking_rules(&ranking_rules);
@@ -41,6 +44,7 @@ pub fn bucket_sort<'ctx, Q: RankingRuleQueryTrait>(
             docids: vec![],
             scores: vec![],
             all_candidates: universe.clone(),
+            degraded: false,
         });
     }
     if ranking_rules.is_empty() {
@@ -74,6 +78,7 @@ pub fn bucket_sort<'ctx, Q: RankingRuleQueryTrait>(
                 scores: vec![Default::default(); results.len()],
                 docids: results,
                 all_candidates,
+                degraded: false,
             });
         } else {
             let docids: Vec<u32> = universe.iter().skip(from).take(length).collect();
@@ -81,6 +86,7 @@ pub fn bucket_sort<'ctx, Q: RankingRuleQueryTrait>(
                 scores: vec![Default::default(); docids.len()],
                 docids,
                 all_candidates: universe.clone(),
+                degraded: false,
             });
         };
     }
@@ -154,6 +160,18 @@ pub fn bucket_sort<'ctx, Q: RankingRuleQueryTrait>(
     }
 
     while valid_docids.len() < length {
+        if time_budget.exceeded() {
+            let bucket = std::mem::take(&mut ranking_rule_universes[cur_ranking_rule_index]);
+            maybe_add_to_results!(bucket);
+
+            return Ok(BucketSortOutput {
+                scores: vec![Default::default(); valid_docids.len()],
+                docids: valid_docids,
+                all_candidates,
+                degraded: true,
+            });
+        }
+
         // The universe for this bucket is zero, so we don't need to sort
         // anything, just go back to the parent ranking rule.
         if ranking_rule_universes[cur_ranking_rule_index].is_empty()
@@ -219,7 +237,12 @@ pub fn bucket_sort<'ctx, Q: RankingRuleQueryTrait>(
         )?;
     }
 
-    Ok(BucketSortOutput { docids: valid_docids, scores: valid_scores, all_candidates })
+    Ok(BucketSortOutput {
+        docids: valid_docids,
+        scores: valid_scores,
+        all_candidates,
+        degraded: false,
+    })
 }
 
 /// Add the candidates to the results. Take `distinct`, `from`, `length`, and `cur_offset`

milli/src/search/new/matches/mod.rs

@@ -502,7 +502,7 @@ mod tests {
 
     use super::*;
     use crate::index::tests::TempIndex;
-    use crate::{execute_search, filtered_universe, SearchContext};
+    use crate::{execute_search, filtered_universe, SearchContext, TimeBudget};
 
     impl<'a> MatcherBuilder<'a> {
         fn new_test(rtxn: &'a heed::RoTxn, index: &'a TempIndex, query: &str) -> Self {
@@ -522,6 +522,7 @@ mod tests {
                 Some(10),
                 &mut crate::DefaultSearchLogger,
                 &mut crate::DefaultSearchLogger,
+                TimeBudget::max(),
             )
             .unwrap();
 

milli/src/search/new/mod.rs

@@ -52,7 +52,8 @@ use crate::score_details::{ScoreDetails, ScoringStrategy};
 use crate::search::new::distinct::apply_distinct_rule;
 use crate::vector::DistributionShift;
 use crate::{
-    AscDesc, DocumentId, FieldId, Filter, Index, Member, Result, TermsMatchingStrategy, UserError,
+    AscDesc, DocumentId, FieldId, Filter, Index, Member, Result, TermsMatchingStrategy, TimeBudget,
+    UserError,
 };
 
 /// A structure used throughout the execution of a search query.
@@ -518,6 +519,7 @@ pub fn execute_vector_search(
     length: usize,
     distribution_shift: Option<DistributionShift>,
     embedder_name: &str,
+    time_budget: TimeBudget,
 ) -> Result<PartialSearchResult> {
     check_sort_criteria(ctx, sort_criteria.as_ref())?;
 
@@ -537,7 +539,7 @@ pub fn execute_vector_search(
     let placeholder_search_logger: &mut dyn SearchLogger<PlaceholderQuery> =
         &mut placeholder_search_logger;
 
-    let BucketSortOutput { docids, scores, all_candidates } = bucket_sort(
+    let BucketSortOutput { docids, scores, all_candidates, degraded } = bucket_sort(
         ctx,
         ranking_rules,
         &PlaceholderQuery,
@@ -546,6 +548,7 @@ pub fn execute_vector_search(
         length,
         scoring_strategy,
         placeholder_search_logger,
+        time_budget,
     )?;
 
     Ok(PartialSearchResult {
@@ -553,6 +556,7 @@ pub fn execute_vector_search(
         document_scores: scores,
         documents_ids: docids,
         located_query_terms: None,
+        degraded,
     })
 }
 
@@ -572,6 +576,7 @@ pub fn execute_search(
     words_limit: Option<usize>,
     placeholder_search_logger: &mut dyn SearchLogger<PlaceholderQuery>,
     query_graph_logger: &mut dyn SearchLogger<QueryGraph>,
+    time_budget: TimeBudget,
 ) -> Result<PartialSearchResult> {
     check_sort_criteria(ctx, sort_criteria.as_ref())?;
 
@@ -648,6 +653,7 @@ pub fn execute_search(
             length,
             scoring_strategy,
             query_graph_logger,
+            time_budget,
         )?
     } else {
         let ranking_rules =
@@ -661,10 +667,11 @@ pub fn execute_search(
             length,
             scoring_strategy,
             placeholder_search_logger,
+            time_budget,
         )?
     };
 
-    let BucketSortOutput { docids, scores, mut all_candidates } = bucket_sort_output;
+    let BucketSortOutput { docids, scores, mut all_candidates, degraded } = bucket_sort_output;
     let fields_ids_map = ctx.index.fields_ids_map(ctx.txn)?;
 
     // The candidates is the universe unless the exhaustive number of hits
@@ -682,6 +689,7 @@ pub fn execute_search(
         document_scores: scores,
         documents_ids: docids,
         located_query_terms,
+        degraded,
     })
 }
 
@@ -742,4 +750,6 @@ pub struct PartialSearchResult {
     pub candidates: RoaringBitmap,
     pub documents_ids: Vec<DocumentId>,
     pub document_scores: Vec<Vec<ScoreDetails>>,
+
+    pub degraded: bool,
 }