Split the FacetDistribution facet_values method into three

src/search/facet/facet_distribution.rs

@@ -57,18 +57,15 @@ impl<'a> FacetDistribution<'a> {
         self
     }
 
-    fn facet_values(
+    /// There is a small amount of candidates OR we ask for facet string values so we
+    /// decide to iterate over the facet values of each one of them, one by one.
+    fn facet_values_from_documents(
         &self,
         field_id: FieldId,
         facet_type: FacetType,
+        candidates: &RoaringBitmap,
     ) -> heed::Result<BTreeMap<FacetValue, u64>>
     {
-        if let Some(candidates) = self.candidates.as_ref() {
-            // Classic search, candidates were specified, we must return
-            // facet values only related to those candidates.
-            if candidates.len() <= CANDIDATES_THRESHOLD || facet_type == FacetType::String {
-                // There is a small amount of candidates OR we ask for facet string values so we
-                // decide to iterate over the facet values of each one of them, one by one.
         let mut key_buffer = vec![field_id];
         match facet_type {
             FacetType::String => {
@@ -120,9 +117,17 @@ impl<'a> FacetDistribution<'a> {
                 Ok(facet_values)
             },
         }
-            } else {
+    }
-                // There is too much documents, we use the facet levels to move throught
+
-                // the facet values, to find the candidates and values associated.
+    /// There is too much documents, we use the facet levels to move throught
+    /// the facet values, to find the candidates and values associated.
+    fn facet_values_from_facet_levels(
+        &self,
+        field_id: FieldId,
+        facet_type: FacetType,
+        candidates: &RoaringBitmap,
+    ) -> heed::Result<BTreeMap<FacetValue, u64>>
+    {
         let iter = match facet_type {
             FacetType::String => unreachable!(),
             FacetType::Float => {
@@ -154,10 +159,17 @@ impl<'a> FacetDistribution<'a> {
 
         Ok(facet_values)
     }
-        } else {
+
-            // Placeholder search, a.k.a. no candidates were specified. We iterate throught the
+    /// Placeholder search, a.k.a. no candidates were specified. We iterate throught the
-            // facet values one by one and iterate on the facet level 0 for numbers.
+    /// facet values one by one and iterate on the facet level 0 for numbers.
+    fn facet_values_from_raw_facet_database(
+        &self,
+        field_id: FieldId,
+        facet_type: FacetType,
+    ) -> heed::Result<BTreeMap<FacetValue, u64>>
+    {
         let db = self.index.facet_field_id_value_docids;
+        let level = 0;
         let iter = match facet_type {
             FacetType::String => {
                 let iter = db
@@ -169,14 +181,14 @@ impl<'a> FacetDistribution<'a> {
             FacetType::Float => {
                 let db = db.remap_key_type::<FacetLevelValueF64Codec>();
                 let range = FacetRange::<f64, _>::new(
-                        self.rtxn, db, field_id, 0, Unbounded, Unbounded,
+                    self.rtxn, db, field_id, level, Unbounded, Unbounded,
                 )?;
                 Box::new(range.map(|r| r.map(|((_, _, v, _), docids)| (FacetValue::from(v), docids))))
             },
             FacetType::Integer => {
                 let db = db.remap_key_type::<FacetLevelValueI64Codec>();
                 let range = FacetRange::<i64, _>::new(
-                        self.rtxn, db, field_id, 0, Unbounded, Unbounded,
+                    self.rtxn, db, field_id, level, Unbounded, Unbounded,
                 )?;
                 Box::new(range.map(|r| r.map(|((_, _, v, _), docids)| (FacetValue::from(v), docids))))
             },
@@ -193,6 +205,24 @@ impl<'a> FacetDistribution<'a> {
 
         Ok(facet_values)
     }
+
+    fn facet_values(
+        &self,
+        field_id: FieldId,
+        facet_type: FacetType,
+    ) -> heed::Result<BTreeMap<FacetValue, u64>>
+    {
+        if let Some(candidates) = self.candidates.as_ref() {
+            // Classic search, candidates were specified, we must return facet values only related
+            // to those candidates. We also enter here for facet strings for performance reasons.
+            if candidates.len() <= CANDIDATES_THRESHOLD || facet_type == FacetType::String {
+                self.facet_values_from_documents(field_id, facet_type, candidates)
+            } else {
+                self.facet_values_from_facet_levels(field_id, facet_type, candidates)
+            }
+        } else {
+            self.facet_values_from_raw_facet_database(field_id, facet_type)
+        }
     }
 
     pub fn execute(&self) -> anyhow::Result<BTreeMap<String, BTreeMap<FacetValue, u64>>> {