Speed-up facet aggregation on low number of candidates

src/search/facet/facet_distribution.rs

@@ -2,11 +2,13 @@ use std::collections::{HashSet, HashMap};
 use std::{cmp, fmt};
 use std::ops::Bound::Unbounded;
 
+use ordered_float::OrderedFloat;
 use roaring::RoaringBitmap;
 use serde_json::Value;
 
 use crate::facet::FacetType;
 use crate::heed_codec::facet::{FacetValueStringCodec, FacetLevelValueF64Codec, FacetLevelValueI64Codec};
+use crate::heed_codec::facet::{FieldDocIdFacetStringCodec, FieldDocIdFacetF64Codec, FieldDocIdFacetI64Codec};
 use crate::search::facet::FacetRange;
 use crate::{Index, FieldId};
 
@@ -38,9 +40,59 @@ impl<'a> FacetDistribution<'a> {
         self
     }
 
-    fn facet_values(&self, field_id: FieldId, field_type: FacetType) -> heed::Result<Vec<Value>> {
+    fn facet_values(&self, field_id: FieldId, facet_type: FacetType) -> heed::Result<Vec<Value>> {
+        if let Some(candidates) = self.candidates.as_ref().filter(|c| c.len() <= 1000) {
+            let mut key_buffer = vec![field_id];
+            match facet_type {
+                FacetType::Float => {
+                    let mut facet_values = HashSet::new();
+                    for docid in candidates {
+                        key_buffer.truncate(1);
+                        key_buffer.extend_from_slice(&docid.to_be_bytes());
+                        let iter = self.index.field_id_docid_facet_values
+                            .prefix_iter(self.rtxn, &key_buffer)?
+                            .remap_key_type::<FieldDocIdFacetF64Codec>();
+                        for result in iter {
+                            let ((_, _, value), ()) = result?;
+                            facet_values.insert(OrderedFloat(value));
+                        }
+                    }
+                    Ok(facet_values.into_iter().map(|f| Value::from(*f)).collect())
+                },
+                FacetType::Integer => {
+                    let mut facet_values = HashSet::new();
+                    for docid in candidates {
+                        key_buffer.truncate(1);
+                        key_buffer.extend_from_slice(&docid.to_be_bytes());
+                        let iter = self.index.field_id_docid_facet_values
+                            .prefix_iter(self.rtxn, &key_buffer)?
+                            .remap_key_type::<FieldDocIdFacetI64Codec>();
+                        for result in iter {
+                            let ((_, _, value), ()) = result?;
+                            facet_values.insert(value);
+                        }
+                    }
+                    Ok(facet_values.into_iter().map(Value::from).collect())
+                },
+                FacetType::String => {
+                    let mut facet_values = HashSet::new();
+                    for docid in candidates {
+                        key_buffer.truncate(1);
+                        key_buffer.extend_from_slice(&docid.to_be_bytes());
+                        let iter = self.index.field_id_docid_facet_values
+                            .prefix_iter(self.rtxn, &key_buffer)?
+                            .remap_key_type::<FieldDocIdFacetStringCodec>();
+                        for result in iter {
+                            let ((_, _, value), ()) = result?;
+                            facet_values.insert(value);
+                        }
+                    }
+                    Ok(facet_values.into_iter().map(Value::from).collect())
+                },
+            }
+        } else {
             let db = self.index.facet_field_id_value_docids;
-        let iter = match field_type {
+            let iter = match facet_type {
                 FacetType::String => {
                     let iter = db
                         .prefix_iter(&self.rtxn, &[field_id])?
@@ -80,6 +132,7 @@ impl<'a> FacetDistribution<'a> {
 
             Ok(facet_values)
         }
+    }
 
     pub fn execute(&self) -> heed::Result<HashMap<String, Vec<Value>>> {
         let fields_ids_map = self.index.fields_ids_map(self.rtxn)?;