From 081278dfd6b99f720f409db92bdef8193bdadf16 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Cl=C3=A9ment=20Renault?= <clement@meilisearch.com>
Date: Sun, 4 Jul 2021 18:11:26 +0200
Subject: [PATCH] Use the facet string levels when computing the facet
 distribution

---
 milli/src/search/facet/facet_distribution.rs | 28 +++++++++++++++++---
 milli/src/search/facet/facet_string.rs       |  2 +-
 milli/src/search/facet/mod.rs                |  1 +
 3 files changed, 27 insertions(+), 4 deletions(-)

diff --git a/milli/src/search/facet/facet_distribution.rs b/milli/src/search/facet/facet_distribution.rs
index ceefe785b..6382e15e1 100644
--- a/milli/src/search/facet/facet_distribution.rs
+++ b/milli/src/search/facet/facet_distribution.rs
@@ -9,7 +9,7 @@ use roaring::RoaringBitmap;
 use crate::error::{FieldIdMapMissingEntry, UserError};
 use crate::facet::FacetType;
 use crate::heed_codec::facet::FacetStringLevelZeroCodec;
-use crate::search::facet::{FacetNumberIter, FacetNumberRange};
+use crate::search::facet::{FacetNumberIter, FacetNumberRange, FacetStringIter};
 use crate::{DocumentId, FieldId, Index, Result};
 
 /// The default number of values by facets that will
@@ -134,6 +134,29 @@ impl<'a> FacetDistribution<'a> {
         Ok(())
     }
 
+    fn facet_strings_distribution_from_facet_levels(
+        &self,
+        field_id: FieldId,
+        candidates: &RoaringBitmap,
+        distribution: &mut BTreeMap<String, u64>,
+    ) -> heed::Result<()> {
+        let iter =
+            FacetStringIter::new_non_reducing(self.rtxn, self.index, field_id, candidates.clone())?;
+
+        for result in iter {
+            let (value, mut docids) = result?;
+            docids &= candidates;
+            if !docids.is_empty() {
+                distribution.insert(value.to_string(), docids.len());
+            }
+            if distribution.len() == self.max_values_by_facet {
+                break;
+            }
+        }
+
+        Ok(())
+    }
+
     /// 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.
     fn facet_values_from_raw_facet_database(
@@ -198,9 +221,8 @@ impl<'a> FacetDistribution<'a> {
                         candidates,
                         &mut distribution,
                     )?;
-                    self.facet_distribution_from_documents(
+                    self.facet_strings_distribution_from_facet_levels(
                         field_id,
-                        String,
                         candidates,
                         &mut distribution,
                     )?;
diff --git a/milli/src/search/facet/facet_string.rs b/milli/src/search/facet/facet_string.rs
index 509bb4f0c..e1fe6ab74 100644
--- a/milli/src/search/facet/facet_string.rs
+++ b/milli/src/search/facet/facet_string.rs
@@ -192,7 +192,7 @@ impl<'t> Iterator for FacetStringGroupRange<'t> {
                 };
                 if must_be_returned {
                     match docids.decode() {
-                        Ok(docids) => Some(Ok(((level, left, right), docids))),
+                        Ok((bounds, docids)) => Some(Ok(((level, left, right), (bounds, docids)))),
                         Err(e) => Some(Err(e)),
                     }
                 } else {
diff --git a/milli/src/search/facet/mod.rs b/milli/src/search/facet/mod.rs
index d92a8e4bd..ddf710e32 100644
--- a/milli/src/search/facet/mod.rs
+++ b/milli/src/search/facet/mod.rs
@@ -1,5 +1,6 @@
 pub use self::facet_distribution::FacetDistribution;
 pub use self::facet_number::{FacetNumberIter, FacetNumberRange, FacetNumberRevRange};
+pub use self::facet_string::FacetStringIter;
 pub use self::filter_condition::{FilterCondition, Operator};
 pub(crate) use self::parser::Rule as ParserRule;