Merge the two <facetttype>_faceted_documents_ids methods into one

milli/src/index.rs

@@ -12,6 +12,7 @@ use rstar::RTree;
 use time::OffsetDateTime;
 
 use crate::error::{InternalError, UserError};
+use crate::facet::FacetType;
 use crate::fields_ids_map::FieldsIdsMap;
 use crate::heed_codec::facet::new::ordered_f64_codec::OrderedF64Codec;
 use crate::heed_codec::facet::new::str_ref::StrRefCodec;
@@ -780,68 +781,38 @@ impl Index {
 
     /* faceted documents ids */
 
-    /// Writes the documents ids that are faceted with numbers under this field id.
+    /// Writes the documents ids that are faceted under this field id for the given facet type.
-    pub(crate) fn put_number_faceted_documents_ids(
+    pub fn put_faceted_documents_ids(
         &self,
         wtxn: &mut RwTxn,
         field_id: FieldId,
+        facet_type: FacetType,
         docids: &RoaringBitmap,
     ) -> heed::Result<()> {
-        let mut buffer =
+        let key = match facet_type {
-            [0u8; main_key::NUMBER_FACETED_DOCUMENTS_IDS_PREFIX.len() + size_of::<FieldId>()];
+            FacetType::String => main_key::STRING_FACETED_DOCUMENTS_IDS_PREFIX,
-        buffer[..main_key::NUMBER_FACETED_DOCUMENTS_IDS_PREFIX.len()]
+            FacetType::Number => main_key::NUMBER_FACETED_DOCUMENTS_IDS_PREFIX,
-            .copy_from_slice(main_key::NUMBER_FACETED_DOCUMENTS_IDS_PREFIX.as_bytes());
+        };
-        buffer[main_key::NUMBER_FACETED_DOCUMENTS_IDS_PREFIX.len()..]
+        let mut buffer = vec![0u8; key.len() + size_of::<FieldId>()];
-            .copy_from_slice(&field_id.to_be_bytes());
+        buffer[..key.len()].copy_from_slice(key.as_bytes());
+        buffer[key.len()..].copy_from_slice(&field_id.to_be_bytes());
         self.main.put::<_, ByteSlice, RoaringBitmapCodec>(wtxn, &buffer, docids)
     }
 
-    /// Retrieve all the documents ids that faceted with numbers under this field id.
+    /// Retrieve all the documents ids that are faceted under this field id for the given facet type.
-    pub fn number_faceted_documents_ids(
+    pub fn faceted_documents_ids(
         &self,
         rtxn: &RoTxn,
         field_id: FieldId,
+        facet_type: FacetType,
     ) -> heed::Result<RoaringBitmap> {
-        let mut buffer =
+        let key = match facet_type {
-            [0u8; main_key::NUMBER_FACETED_DOCUMENTS_IDS_PREFIX.len() + size_of::<FieldId>()];
+            FacetType::String => main_key::STRING_FACETED_DOCUMENTS_IDS_PREFIX,
-        buffer[..main_key::NUMBER_FACETED_DOCUMENTS_IDS_PREFIX.len()]
+            FacetType::Number => main_key::NUMBER_FACETED_DOCUMENTS_IDS_PREFIX,
-            .copy_from_slice(main_key::NUMBER_FACETED_DOCUMENTS_IDS_PREFIX.as_bytes());
+        };
-        buffer[main_key::NUMBER_FACETED_DOCUMENTS_IDS_PREFIX.len()..]
+        let mut buffer = vec![0u8; key.len() + size_of::<FieldId>()];
-            .copy_from_slice(&field_id.to_be_bytes());
+        buffer[..key.len()].copy_from_slice(key.as_bytes());
-        match self.main.get::<_, ByteSlice, RoaringBitmapCodec>(rtxn, &buffer)? {
+        buffer[key.len()..].copy_from_slice(&field_id.to_be_bytes());
-            Some(docids) => Ok(docids),
-            None => Ok(RoaringBitmap::new()),
-        }
-    }
-
-    /// Writes the documents ids that are faceted with strings under this field id.
-    pub(crate) fn put_string_faceted_documents_ids(
-        &self,
-        wtxn: &mut RwTxn,
-        field_id: FieldId,
-        docids: &RoaringBitmap,
-    ) -> heed::Result<()> {
-        let mut buffer =
-            [0u8; main_key::STRING_FACETED_DOCUMENTS_IDS_PREFIX.len() + size_of::<FieldId>()];
-        buffer[..main_key::STRING_FACETED_DOCUMENTS_IDS_PREFIX.len()]
-            .copy_from_slice(main_key::STRING_FACETED_DOCUMENTS_IDS_PREFIX.as_bytes());
-        buffer[main_key::STRING_FACETED_DOCUMENTS_IDS_PREFIX.len()..]
-            .copy_from_slice(&field_id.to_be_bytes());
-        self.main.put::<_, ByteSlice, RoaringBitmapCodec>(wtxn, &buffer, docids)
-    }
-
-    /// Retrieve all the documents ids that faceted with strings under this field id.
-    pub fn string_faceted_documents_ids(
-        &self,
-        rtxn: &RoTxn,
-        field_id: FieldId,
-    ) -> heed::Result<RoaringBitmap> {
-        let mut buffer =
-            [0u8; main_key::STRING_FACETED_DOCUMENTS_IDS_PREFIX.len() + size_of::<FieldId>()];
-        buffer[..main_key::STRING_FACETED_DOCUMENTS_IDS_PREFIX.len()]
-            .copy_from_slice(main_key::STRING_FACETED_DOCUMENTS_IDS_PREFIX.as_bytes());
-        buffer[main_key::STRING_FACETED_DOCUMENTS_IDS_PREFIX.len()..]
-            .copy_from_slice(&field_id.to_be_bytes());
         match self.main.get::<_, ByteSlice, RoaringBitmapCodec>(rtxn, &buffer)? {
             Some(docids) => Ok(docids),
             None => Ok(RoaringBitmap::new()),

milli/src/search/criteria/asc_desc.rs

@@ -6,6 +6,7 @@ use ordered_float::OrderedFloat;
 use roaring::RoaringBitmap;
 
 use super::{Criterion, CriterionParameters, CriterionResult};
+use crate::facet::FacetType;
 use crate::heed_codec::facet::new::{FacetKeyCodec, MyByteSlice};
 use crate::search::criteria::{resolve_query_tree, CriteriaBuilder};
 use crate::search::facet::facet_sort_ascending::ascending_facet_sort;
@@ -62,8 +63,10 @@ impl<'t> AscDesc<'t> {
         let field_id = fields_ids_map.id(&field_name);
         let faceted_candidates = match field_id {
             Some(field_id) => {
-                let number_faceted = index.number_faceted_documents_ids(rtxn, field_id)?;
+                let number_faceted =
-                let string_faceted = index.string_faceted_documents_ids(rtxn, field_id)?;
+                    index.faceted_documents_ids(rtxn, field_id, FacetType::Number)?;
+                let string_faceted =
+                    index.faceted_documents_ids(rtxn, field_id, FacetType::String)?;
                 number_faceted | string_faceted
             }
             None => RoaringBitmap::default(),

milli/src/snapshot_tests.rs

@@ -4,6 +4,7 @@ use std::path::Path;
 
 use roaring::RoaringBitmap;
 
+use crate::facet::FacetType;
 use crate::heed_codec::facet::new::{FacetGroupValue, FacetKey};
 use crate::{make_db_snap_from_iter, ExternalDocumentsIds, Index};
 
@@ -370,7 +371,7 @@ pub fn snap_number_faceted_documents_ids(index: &Index) -> String {
     let mut snap = String::new();
     for field_id in fields_ids_map.ids() {
         let number_faceted_documents_ids =
-            index.number_faceted_documents_ids(&rtxn, field_id).unwrap();
+            index.faceted_documents_ids(&rtxn, field_id, FacetType::Number).unwrap();
         writeln!(&mut snap, "{field_id:<3} {}", display_bitmap(&number_faceted_documents_ids))
             .unwrap();
     }
@@ -383,7 +384,7 @@ pub fn snap_string_faceted_documents_ids(index: &Index) -> String {
     let mut snap = String::new();
     for field_id in fields_ids_map.ids() {
         let string_faceted_documents_ids =
-            index.string_faceted_documents_ids(&rtxn, field_id).unwrap();
+            index.faceted_documents_ids(&rtxn, field_id, FacetType::String).unwrap();
         writeln!(&mut snap, "{field_id:<3} {}", display_bitmap(&string_faceted_documents_ids))
             .unwrap();
     }

milli/src/update/clear_documents.rs

@@ -1,7 +1,7 @@
 use roaring::RoaringBitmap;
 use time::OffsetDateTime;
 
-use crate::{ExternalDocumentsIds, FieldDistribution, Index, Result};
+use crate::{facet::FacetType, ExternalDocumentsIds, FieldDistribution, Index, Result};
 
 pub struct ClearDocuments<'t, 'u, 'i> {
     wtxn: &'t mut heed::RwTxn<'i, 'u>,
@@ -55,8 +55,18 @@ impl<'t, 'u, 'i> ClearDocuments<'t, 'u, 'i> {
 
         // We clean all the faceted documents ids.
         for field_id in faceted_fields {
-            self.index.put_number_faceted_documents_ids(self.wtxn, field_id, &empty_roaring)?;
+            self.index.put_faceted_documents_ids(
-            self.index.put_string_faceted_documents_ids(self.wtxn, field_id, &empty_roaring)?;
+                self.wtxn,
+                field_id,
+                FacetType::Number,
+                &empty_roaring,
+            )?;
+            self.index.put_faceted_documents_ids(
+                self.wtxn,
+                field_id,
+                FacetType::String,
+                &empty_roaring,
+            )?;
         }
 
         // Clear the other databases.

milli/src/update/delete_documents.rs

@@ -461,9 +461,15 @@ impl<'t, 'u, 'i> DeleteDocuments<'t, 'u, 'i> {
         // Remove the documents ids from the faceted documents ids.
         for field_id in self.index.faceted_fields_ids(self.wtxn)? {
             // Remove docids from the number faceted documents ids
-            let mut docids = self.index.number_faceted_documents_ids(self.wtxn, field_id)?;
+            let mut docids =
+                self.index.faceted_documents_ids(self.wtxn, field_id, FacetType::Number)?;
             docids -= &self.to_delete_docids;
-            self.index.put_number_faceted_documents_ids(self.wtxn, field_id, &docids)?;
+            self.index.put_faceted_documents_ids(
+                self.wtxn,
+                field_id,
+                FacetType::Number,
+                &docids,
+            )?;
 
             remove_docids_from_field_id_docid_facet_value(
                 self.wtxn,
@@ -474,9 +480,15 @@ impl<'t, 'u, 'i> DeleteDocuments<'t, 'u, 'i> {
             )?;
 
             // Remove docids from the string faceted documents ids
-            let mut docids = self.index.string_faceted_documents_ids(self.wtxn, field_id)?;
+            let mut docids =
+                self.index.faceted_documents_ids(self.wtxn, field_id, FacetType::String)?;
             docids -= &self.to_delete_docids;
-            self.index.put_string_faceted_documents_ids(self.wtxn, field_id, &docids)?;
+            self.index.put_faceted_documents_ids(
+                self.wtxn,
+                field_id,
+                FacetType::String,
+                &docids,
+            )?;
 
             remove_docids_from_field_id_docid_facet_value(
                 self.wtxn,
@@ -648,7 +660,7 @@ fn remove_docids_from_facet_id_docids<'a>(
     if !modified {
         return Ok(());
     }
-    let builder = FacetsUpdateBulk::new(index, facet_type);
+    let builder = FacetsUpdateBulk::new_not_updating_level_0(index, facet_type);
     builder.execute(wtxn)?;
 
     Ok(())