Merge #742

742: Add a "Criterion implementation strategy" parameter to Search r=irevoire a=loiclec

Add a parameter to search requests which determines the implementation strategy of the criteria. This can be either `set-based`, `iterative`, or `dynamic` (ie choosing between set-based or iterative at search time). See https://github.com/meilisearch/milli/issues/755 for more context about this change.


Co-authored-by: Loïc Lecrenier <[email protected]>

Author: bors[bot]

cli/src/main.rs

@@ -1,4 +1,5 @@
 use std::collections::BTreeMap;
+use std::fmt::Display;
 use std::fs::File;
 use std::io::{stdin, BufRead, BufReader, Cursor, Read, Write};
 use std::path::PathBuf;
@@ -13,7 +14,7 @@ use milli::update::UpdateIndexingStep::{
     ComputeIdsAndMergeDocuments, IndexDocuments, MergeDataIntoFinalDatabase, RemapDocumentAddition,
 };
 use milli::update::{self, IndexDocumentsConfig, IndexDocumentsMethod, IndexerConfig};
-use milli::{heed, Index, Object};
+use milli::{heed, CriterionImplementationStrategy, Index, Object};
 use structopt::StructOpt;
 
 #[global_allocator]
@@ -349,6 +350,29 @@ fn documents_from_csv(reader: impl Read) -> Result<Vec<u8>> {
     documents.into_inner().map_err(Into::into)
 }
 
+#[derive(Debug, Clone, Copy)]
+struct SearchStrategyOption(CriterionImplementationStrategy);
+impl FromStr for SearchStrategyOption {
+    type Err = String;
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        match s.to_lowercase().as_str() {
+            "dynamic" => Ok(SearchStrategyOption(CriterionImplementationStrategy::Dynamic)),
+            "set" => Ok(SearchStrategyOption(CriterionImplementationStrategy::OnlySetBased)),
+            "iterative" => Ok(SearchStrategyOption(CriterionImplementationStrategy::OnlyIterative)),
+            _ => Err("could not parse {s} as a criterion implementation strategy, available options are `dynamic`, `set`, and `iterative`".to_owned()),
+        }
+    }
+}
+impl Display for SearchStrategyOption {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self.0 {
+            CriterionImplementationStrategy::OnlyIterative => Display::fmt("iterative", f),
+            CriterionImplementationStrategy::OnlySetBased => Display::fmt("set", f),
+            CriterionImplementationStrategy::Dynamic => Display::fmt("dynamic", f),
+        }
+    }
+}
+
 #[derive(Debug, StructOpt)]
 struct Search {
     query: Option<String>,
@@ -360,6 +384,8 @@ struct Search {
     limit: Option<usize>,
     #[structopt(short, long, conflicts_with = "query")]
     interactive: bool,
+    #[structopt(short, long)]
+    strategy: Option<SearchStrategyOption>,
 }
 
 impl Performer for Search {
@@ -379,13 +405,15 @@ impl Performer for Search {
                             &self.filter,
                             &self.offset,
                             &self.limit,
+                            &self.strategy,
                         )?;
 
                         let time = now.elapsed();
 
                         let hits = serde_json::to_string_pretty(&jsons)?;
 
                         println!("{}", hits);
+
                         eprintln!("found {} results in {:.02?}", jsons.len(), time);
                     }
                     _ => break,
@@ -399,6 +427,7 @@ impl Performer for Search {
                 &self.filter,
                 &self.offset,
                 &self.limit,
+                &self.strategy,
             )?;
 
             let time = now.elapsed();
@@ -420,6 +449,7 @@ impl Search {
         filter: &Option<String>,
         offset: &Option<usize>,
         limit: &Option<usize>,
+        strategy: &Option<SearchStrategyOption>,
     ) -> Result<Vec<Object>> {
         let txn = index.read_txn()?;
         let mut search = index.search(&txn);
@@ -441,6 +471,9 @@ impl Search {
         if let Some(limit) = limit {
             search.limit(*limit);
         }
+        if let Some(strategy) = strategy {
+            search.criterion_implementation_strategy(strategy.0);
+        }
 
         let result = search.execute()?;
 

milli/src/lib.rs

@@ -42,8 +42,9 @@ pub use self::heed_codec::{
 };
 pub use self::index::Index;
 pub use self::search::{
-    FacetDistribution, Filter, FormatOptions, MatchBounds, MatcherBuilder, MatchingWord,
-    MatchingWords, Search, SearchResult, TermsMatchingStrategy, DEFAULT_VALUES_PER_FACET,
+    CriterionImplementationStrategy, FacetDistribution, Filter, FormatOptions, MatchBounds,
+    MatcherBuilder, MatchingWord, MatchingWords, Search, SearchResult, TermsMatchingStrategy,
+    DEFAULT_VALUES_PER_FACET,
 };
 
 pub type Result<T> = std::result::Result<T, error::Error>;

milli/src/search/criteria/asc_desc.rs

@@ -12,6 +12,7 @@ use crate::heed_codec::ByteSliceRefCodec;
 use crate::search::criteria::{resolve_query_tree, CriteriaBuilder, InitialCandidates};
 use crate::search::facet::{ascending_facet_sort, descending_facet_sort};
 use crate::search::query_tree::Operation;
+use crate::search::CriterionImplementationStrategy;
 use crate::{FieldId, Index, Result};
 
 /// Threshold on the number of candidates that will make
@@ -29,6 +30,7 @@ pub struct AscDesc<'t> {
     allowed_candidates: RoaringBitmap,
     initial_candidates: InitialCandidates,
     faceted_candidates: RoaringBitmap,
+    implementation_strategy: CriterionImplementationStrategy,
     parent: Box<dyn Criterion + 't>,
 }
 
@@ -38,17 +40,19 @@ impl<'t> AscDesc<'t> {
         rtxn: &'t heed::RoTxn,
         parent: Box<dyn Criterion + 't>,
         field_name: String,
+        implementation_strategy: CriterionImplementationStrategy,
     ) -> Result<Self> {
-        Self::new(index, rtxn, parent, field_name, true)
+        Self::new(index, rtxn, parent, field_name, true, implementation_strategy)
     }
 
     pub fn desc(
         index: &'t Index,
         rtxn: &'t heed::RoTxn,
         parent: Box<dyn Criterion + 't>,
         field_name: String,
+        implementation_strategy: CriterionImplementationStrategy,
     ) -> Result<Self> {
-        Self::new(index, rtxn, parent, field_name, false)
+        Self::new(index, rtxn, parent, field_name, false, implementation_strategy)
     }
 
     fn new(
@@ -57,6 +61,7 @@ impl<'t> AscDesc<'t> {
         parent: Box<dyn Criterion + 't>,
         field_name: String,
         is_ascending: bool,
+        implementation_strategy: CriterionImplementationStrategy,
     ) -> Result<Self> {
         let fields_ids_map = index.fields_ids_map(rtxn)?;
         let field_id = fields_ids_map.id(&field_name);
@@ -82,6 +87,7 @@ impl<'t> AscDesc<'t> {
             allowed_candidates: RoaringBitmap::new(),
             faceted_candidates,
             initial_candidates: InitialCandidates::Estimated(RoaringBitmap::new()),
+            implementation_strategy,
             parent,
         })
     }
@@ -149,6 +155,7 @@ impl<'t> Criterion for AscDesc<'t> {
                                 field_id,
                                 self.is_ascending,
                                 candidates & &self.faceted_candidates,
+                                self.implementation_strategy,
                             )?,
                             None => Box::new(std::iter::empty()),
                         };
@@ -170,6 +177,51 @@ impl<'t> Criterion for AscDesc<'t> {
     }
 }
 
+fn facet_ordered_iterative<'t>(
+    index: &'t Index,
+    rtxn: &'t heed::RoTxn,
+    field_id: FieldId,
+    is_ascending: bool,
+    candidates: RoaringBitmap,
+) -> Result<Box<dyn Iterator<Item = heed::Result<RoaringBitmap>> + 't>> {
+    let number_iter = iterative_facet_number_ordered_iter(
+        index,
+        rtxn,
+        field_id,
+        is_ascending,
+        candidates.clone(),
+    )?;
+    let string_iter =
+        iterative_facet_string_ordered_iter(index, rtxn, field_id, is_ascending, candidates)?;
+    Ok(Box::new(number_iter.chain(string_iter).map(Ok)) as Box<dyn Iterator<Item = _>>)
+}
+
+fn facet_ordered_set_based<'t>(
+    index: &'t Index,
+    rtxn: &'t heed::RoTxn,
+    field_id: FieldId,
+    is_ascending: bool,
+    candidates: RoaringBitmap,
+) -> Result<Box<dyn Iterator<Item = heed::Result<RoaringBitmap>> + 't>> {
+    let make_iter = if is_ascending { ascending_facet_sort } else { descending_facet_sort };
+
+    let number_iter = make_iter(
+        rtxn,
+        index.facet_id_f64_docids.remap_key_type::<FacetGroupKeyCodec<ByteSliceRefCodec>>(),
+        field_id,
+        candidates.clone(),
+    )?;
+
+    let string_iter = make_iter(
+        rtxn,
+        index.facet_id_string_docids.remap_key_type::<FacetGroupKeyCodec<ByteSliceRefCodec>>(),
+        field_id,
+        candidates,
+    )?;
+
+    Ok(Box::new(number_iter.chain(string_iter)))
+}
+
 /// Returns an iterator over groups of the given candidates in ascending or descending order.
 ///
 /// It will either use an iterative or a recursive method on the whole facet database depending
@@ -180,36 +232,22 @@ fn facet_ordered<'t>(
     field_id: FieldId,
     is_ascending: bool,
     candidates: RoaringBitmap,
+    implementation_strategy: CriterionImplementationStrategy,
 ) -> Result<Box<dyn Iterator<Item = heed::Result<RoaringBitmap>> + 't>> {
-    if candidates.len() <= CANDIDATES_THRESHOLD {
-        let number_iter = iterative_facet_number_ordered_iter(
-            index,
-            rtxn,
-            field_id,
-            is_ascending,
-            candidates.clone(),
-        )?;
-        let string_iter =
-            iterative_facet_string_ordered_iter(index, rtxn, field_id, is_ascending, candidates)?;
-        Ok(Box::new(number_iter.chain(string_iter).map(Ok)) as Box<dyn Iterator<Item = _>>)
-    } else {
-        let make_iter = if is_ascending { ascending_facet_sort } else { descending_facet_sort };
-
-        let number_iter = make_iter(
-            rtxn,
-            index.facet_id_f64_docids.remap_key_type::<FacetGroupKeyCodec<ByteSliceRefCodec>>(),
-            field_id,
-            candidates.clone(),
-        )?;
-
-        let string_iter = make_iter(
-            rtxn,
-            index.facet_id_string_docids.remap_key_type::<FacetGroupKeyCodec<ByteSliceRefCodec>>(),
-            field_id,
-            candidates,
-        )?;
-
-        Ok(Box::new(number_iter.chain(string_iter)))
+    match implementation_strategy {
+        CriterionImplementationStrategy::OnlyIterative => {
+            facet_ordered_iterative(index, rtxn, field_id, is_ascending, candidates)
+        }
+        CriterionImplementationStrategy::OnlySetBased => {
+            facet_ordered_set_based(index, rtxn, field_id, is_ascending, candidates)
+        }
+        CriterionImplementationStrategy::Dynamic => {
+            if candidates.len() <= CANDIDATES_THRESHOLD {
+                facet_ordered_iterative(index, rtxn, field_id, is_ascending, candidates)
+            } else {
+                facet_ordered_set_based(index, rtxn, field_id, is_ascending, candidates)
+            }
+        }
     }
 }
 

milli/src/search/criteria/attribute.rs

@@ -9,7 +9,9 @@ use roaring::RoaringBitmap;
 use super::{resolve_query_tree, Context, Criterion, CriterionParameters, CriterionResult};
 use crate::search::criteria::{InitialCandidates, Query};
 use crate::search::query_tree::{Operation, QueryKind};
-use crate::search::{build_dfa, word_derivations, WordDerivationsCache};
+use crate::search::{
+    build_dfa, word_derivations, CriterionImplementationStrategy, WordDerivationsCache,
+};
 use crate::Result;
 
 /// To be able to divide integers by the number of words in the query
@@ -30,17 +32,23 @@ pub struct Attribute<'t> {
     parent: Box<dyn Criterion + 't>,
     linear_buckets: Option<btree_map::IntoIter<u64, RoaringBitmap>>,
     set_buckets: Option<BinaryHeap<Branch<'t>>>,
+    implementation_strategy: CriterionImplementationStrategy,
 }
 
 impl<'t> Attribute<'t> {
-    pub fn new(ctx: &'t dyn Context<'t>, parent: Box<dyn Criterion + 't>) -> Self {
+    pub fn new(
+        ctx: &'t dyn Context<'t>,
+        parent: Box<dyn Criterion + 't>,
+        implementation_strategy: CriterionImplementationStrategy,
+    ) -> Self {
         Attribute {
             ctx,
             state: None,
             initial_candidates: InitialCandidates::Estimated(RoaringBitmap::new()),
             parent,
             linear_buckets: None,
             set_buckets: None,
+            implementation_strategy,
         }
     }
 }
@@ -64,7 +72,15 @@ impl<'t> Criterion for Attribute<'t> {
                     }));
                 }
                 Some((query_tree, flattened_query_tree, mut allowed_candidates)) => {
-                    let found_candidates = if allowed_candidates.len() < CANDIDATES_THRESHOLD {
+                    let found_candidates = if matches!(
+                        self.implementation_strategy,
+                        CriterionImplementationStrategy::OnlyIterative
+                    ) || (matches!(
+                        self.implementation_strategy,
+                        CriterionImplementationStrategy::Dynamic
+                    ) && allowed_candidates.len()
+                        < CANDIDATES_THRESHOLD)
+                    {
                         let linear_buckets = match self.linear_buckets.as_mut() {
                             Some(linear_buckets) => linear_buckets,
                             None => {