From 999758f7a1cfd2c233a9eed38564216bdb517693 Mon Sep 17 00:00:00 2001
From: mpostma <postma.marin@protonmail.com>
Date: Wed, 17 Feb 2021 14:52:30 +0100
Subject: [PATCH] facets distribution

---
 src/data/search.rs   | 24 ++++++++++++++++++++----
 src/routes/search.rs |  8 ++++----
 2 files changed, 24 insertions(+), 8 deletions(-)

diff --git a/src/data/search.rs b/src/data/search.rs
index 391cc1960..f26730fcf 100644
--- a/src/data/search.rs
+++ b/src/data/search.rs
@@ -1,4 +1,4 @@
-use std::collections::HashSet;
+use std::collections::{HashSet, BTreeMap};
 use std::mem;
 use std::time::Instant;
 
@@ -6,7 +6,7 @@ use anyhow::{bail, Context};
 use either::Either;
 use heed::RoTxn;
 use meilisearch_tokenizer::{Analyzer, AnalyzerConfig};
-use milli::{obkv_to_json, FacetCondition, Index};
+use milli::{obkv_to_json, FacetCondition, Index, facet::FacetValue};
 use serde::{Deserialize, Serialize};
 use serde_json::{Map, Value};
 
@@ -34,7 +34,7 @@ pub struct SearchQuery {
     pub filters: Option<String>,
     pub matches: Option<bool>,
     pub facet_filters: Option<Value>,
-    pub facets_distribution: Option<Vec<String>>,
+    pub facet_distributions: Option<Vec<String>>,
 }
 
 impl SearchQuery {
@@ -96,13 +96,27 @@ impl SearchQuery {
             documents.push(object);
         }
 
+        let nb_hits = candidates.len();
+
+        let facet_distributions = match self.facet_distributions {
+            Some(ref fields) => {
+                let mut facet_distribution = index.facets_distribution(&rtxn);
+                if fields.iter().all(|f| f != "*") {
+                    facet_distribution.facets(fields);
+                }
+                Some(facet_distribution.candidates(candidates).execute()?)
+            }
+            None => None,
+        };
+
         Ok(SearchResult {
             hits: documents,
-            nb_hits: candidates.len(),
+            nb_hits,
             query: self.q.clone().unwrap_or_default(),
             limit: self.limit,
             offset: self.offset.unwrap_or_default(),
             processing_time_ms: before_search.elapsed().as_millis(),
+            facet_distributions,
         })
     }
 }
@@ -116,6 +130,8 @@ pub struct SearchResult {
     limit: usize,
     offset: usize,
     processing_time_ms: u128,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    facet_distributions: Option<BTreeMap<String, BTreeMap<FacetValue, u64>>>,
 }
 
 struct Highlighter<'a, A> {
diff --git a/src/routes/search.rs b/src/routes/search.rs
index 794c8ac74..7919c5412 100644
--- a/src/routes/search.rs
+++ b/src/routes/search.rs
@@ -27,7 +27,7 @@ pub struct SearchQueryGet {
     filters: Option<String>,
     matches: Option<bool>,
     facet_filters: Option<String>,
-    facets_distribution: Option<String>,
+    facet_distributions: Option<String>,
 }
 
 impl TryFrom<SearchQueryGet> for SearchQuery {
@@ -46,8 +46,8 @@ impl TryFrom<SearchQueryGet> for SearchQuery {
             .attributes_to_highlight
             .map(|attrs| attrs.split(",").map(String::from).collect::<HashSet<_>>());
 
-        let facets_distribution = other
-            .facets_distribution
+        let facet_distributions = other
+            .facet_distributions
             .map(|attrs| attrs.split(",").map(String::from).collect::<Vec<_>>());
 
         let facet_filters = match other.facet_filters {
@@ -66,7 +66,7 @@ impl TryFrom<SearchQueryGet> for SearchQuery {
             filters: other.filters,
             matches: other.matches,
             facet_filters,
-            facets_distribution,
+            facet_distributions,
         })
     }
 }