Add linear facet databases

crates/milli/src/heed_codec/facet/ordered_f64_codec.rs

@@ -27,15 +27,32 @@ impl heed::BytesEncode<'_> for OrderedF64Codec {
     fn bytes_encode(f: &Self::EItem) -> Result<Cow<'_, [u8]>, BoxedError> {
         let mut buffer = [0u8; 16];
 
-        // write the globally ordered float
-        let bytes = f64_into_bytes(*f).ok_or(InvalidGloballyOrderedFloatError { float: *f })?;
+        encode_f64_into_ordered_bytes(*f, &mut buffer)?;
+
+        Ok(Cow::Owned(buffer.to_vec()))
+    }
+}
+
+impl OrderedF64Codec {
+    pub fn serialize_into(
+        f: f64,
+        buffer: &mut [u8; 16],
+    ) -> Result<(), InvalidGloballyOrderedFloatError> {
+        encode_f64_into_ordered_bytes(f, buffer)
+    }
+}
+
+fn encode_f64_into_ordered_bytes(
+    f: f64,
+    buffer: &mut [u8; 16],
+) -> Result<(), InvalidGloballyOrderedFloatError> {
+    let bytes = f64_into_bytes(f).ok_or(InvalidGloballyOrderedFloatError { float: f })?;
     buffer[..8].copy_from_slice(&bytes[..]);
     // Then the f64 value just to be able to read it back
     let bytes = f.to_be_bytes();
     buffer[8..16].copy_from_slice(&bytes[..]);
 
-        Ok(Cow::Owned(buffer.to_vec()))
-    }
+    Ok(())
 }
 
 #[derive(Error, Debug)]

crates/milli/src/update/new/channel.rs

@@ -3,11 +3,13 @@ use std::sync::atomic::{AtomicUsize, Ordering};
 
 use crossbeam_channel::{IntoIter, Receiver, SendError, Sender};
 use heed::types::Bytes;
+use heed::BytesDecode;
 use memmap2::Mmap;
 use roaring::RoaringBitmap;
 
 use super::extract::FacetKind;
 use super::StdResult;
+use crate::heed_codec::facet::{FieldDocIdFacetF64Codec, FieldDocIdFacetStringCodec};
 use crate::index::main_key::{GEO_FACETED_DOCUMENTS_IDS_KEY, GEO_RTREE_KEY};
 use crate::index::IndexEmbeddingConfig;
 use crate::update::new::KvReaderFieldId;
@@ -125,6 +127,8 @@ pub enum Database {
     FacetIdExistsDocids,
     FacetIdF64NumberDocids,
     FacetIdStringDocids,
+    FieldIdDocidFacetStrings,
+    FieldIdDocidFacetF64s,
 }
 
 impl Database {
@@ -144,6 +148,8 @@ impl Database {
             Database::FacetIdExistsDocids => index.facet_id_exists_docids.remap_types(),
             Database::FacetIdF64NumberDocids => index.facet_id_f64_docids.remap_types(),
             Database::FacetIdStringDocids => index.facet_id_string_docids.remap_types(),
+            Database::FieldIdDocidFacetStrings => index.field_id_docid_facet_strings.remap_types(),
+            Database::FieldIdDocidFacetF64s => index.field_id_docid_facet_f64s.remap_types(),
         }
     }
 }
@@ -215,6 +221,10 @@ impl ExtractorSender {
         FacetDocidsSender { sender: self }
     }
 
+    pub fn field_id_docid_facet_sender(&self) -> FieldIdDocidFacetSender<'_> {
+        FieldIdDocidFacetSender(self)
+    }
+
     pub fn documents(&self) -> DocumentsSender<'_> {
         DocumentsSender(self)
     }
@@ -351,6 +361,36 @@ impl DocidsSender for FacetDocidsSender<'_> {
     }
 }
 
+pub struct FieldIdDocidFacetSender<'a>(&'a ExtractorSender);
+
+impl FieldIdDocidFacetSender<'_> {
+    pub fn write_facet_string(&self, key: &[u8]) -> StdResult<(), SendError<()>> {
+        debug_assert!(FieldDocIdFacetStringCodec::bytes_decode(key).is_ok());
+        let entry = EntryOperation::Write(KeyValueEntry::from_small_key_value(&key, &[]));
+        self.0
+            .send_db_operation(DbOperation { database: Database::FieldIdDocidFacetStrings, entry })
+    }
+
+    pub fn write_facet_f64(&self, key: &[u8]) -> StdResult<(), SendError<()>> {
+        debug_assert!(FieldDocIdFacetF64Codec::bytes_decode(key).is_ok());
+        let entry = EntryOperation::Write(KeyValueEntry::from_small_key_value(&key, &[]));
+        self.0.send_db_operation(DbOperation { database: Database::FieldIdDocidFacetF64s, entry })
+    }
+
+    pub fn delete_facet_string(&self, key: &[u8]) -> StdResult<(), SendError<()>> {
+        debug_assert!(FieldDocIdFacetStringCodec::bytes_decode(key).is_ok());
+        let entry = EntryOperation::Delete(KeyEntry::from_key(key));
+        self.0
+            .send_db_operation(DbOperation { database: Database::FieldIdDocidFacetStrings, entry })
+    }
+
+    pub fn delete_facet_f64(&self, key: &[u8]) -> StdResult<(), SendError<()>> {
+        debug_assert!(FieldDocIdFacetF64Codec::bytes_decode(key).is_ok());
+        let entry = EntryOperation::Delete(KeyEntry::from_key(key));
+        self.0.send_db_operation(DbOperation { database: Database::FieldIdDocidFacetF64s, entry })
+    }
+}
+
 pub struct DocumentsSender<'a>(&'a ExtractorSender);
 
 impl DocumentsSender<'_> {

crates/milli/src/update/new/extract/faceted/extract_facets.rs

@@ -1,16 +1,21 @@
 use std::cell::RefCell;
 use std::collections::HashSet;
+use std::mem::size_of;
 use std::ops::DerefMut as _;
 
+use bumpalo::collections::Vec as BVec;
 use bumpalo::Bump;
-use heed::RoTxn;
+use hashbrown::HashMap;
+use heed::{BytesDecode, RoTxn};
 use serde_json::Value;
 
 use super::super::cache::BalancedCaches;
 use super::facet_document::extract_document_facets;
 use super::FacetKind;
 use crate::facet::value_encoding::f64_into_bytes;
-use crate::update::new::extract::DocidsExtractor;
+use crate::heed_codec::facet::OrderedF64Codec;
+use crate::update::del_add::DelAdd;
+use crate::update::new::channel::FieldIdDocidFacetSender;
 use crate::update::new::indexer::document_changes::{
     extract, DocumentChangeContext, DocumentChanges, Extractor, FullySend, IndexingContext,
     Progress, ThreadLocal,
@@ -22,6 +27,7 @@ use crate::{DocumentId, FieldId, Index, Result, MAX_FACET_VALUE_LENGTH};
 
 pub struct FacetedExtractorData<'a> {
     attributes_to_extract: &'a [&'a str],
+    sender: &'a FieldIdDocidFacetSender<'a>,
     grenad_parameters: GrenadParameters,
     buckets: usize,
 }
@@ -48,6 +54,7 @@ impl<'a, 'extractor> Extractor<'extractor> for FacetedExtractorData<'a> {
                 context,
                 self.attributes_to_extract,
                 change,
+                self.sender,
             )?
         }
         Ok(())
@@ -61,12 +68,15 @@ impl FacetedDocidsExtractor {
         context: &DocumentChangeContext<RefCell<BalancedCaches>>,
         attributes_to_extract: &[&str],
         document_change: DocumentChange,
+        sender: &FieldIdDocidFacetSender,
     ) -> Result<()> {
         let index = &context.index;
         let rtxn = &context.rtxn;
         let mut new_fields_ids_map = context.new_fields_ids_map.borrow_mut_or_yield();
         let mut cached_sorter = context.data.borrow_mut_or_yield();
-        match document_change {
+        let mut del_add_facet_value = DelAddFacetValue::new(&context.doc_alloc);
+        let docid = document_change.docid();
+        let res = match document_change {
             DocumentChange::Deletion(inner) => extract_document_facets(
                 attributes_to_extract,
                 inner.current(rtxn, index, context.db_fields_ids_map)?,
@@ -76,7 +86,9 @@ impl FacetedDocidsExtractor {
                         &context.doc_alloc,
                         cached_sorter.deref_mut(),
                         BalancedCaches::insert_del_u32,
-                        inner.docid(),
+                        &mut del_add_facet_value,
+                        DelAddFacetValue::insert_del,
+                        docid,
                         fid,
                         value,
                     )
@@ -92,7 +104,9 @@ impl FacetedDocidsExtractor {
                             &context.doc_alloc,
                             cached_sorter.deref_mut(),
                             BalancedCaches::insert_del_u32,
-                            inner.docid(),
+                            &mut del_add_facet_value,
+                            DelAddFacetValue::insert_del,
+                            docid,
                             fid,
                             value,
                         )
@@ -108,7 +122,9 @@ impl FacetedDocidsExtractor {
                             &context.doc_alloc,
                             cached_sorter.deref_mut(),
                             BalancedCaches::insert_add_u32,
-                            inner.docid(),
+                            &mut del_add_facet_value,
+                            DelAddFacetValue::insert_add,
+                            docid,
                             fid,
                             value,
                         )
@@ -124,24 +140,31 @@ impl FacetedDocidsExtractor {
                         &context.doc_alloc,
                         cached_sorter.deref_mut(),
                         BalancedCaches::insert_add_u32,
-                        inner.docid(),
+                        &mut del_add_facet_value,
+                        DelAddFacetValue::insert_add,
+                        docid,
                         fid,
                         value,
                     )
                 },
             ),
-        }
+        };
+
+        del_add_facet_value.send_data(docid, sender, &context.doc_alloc).unwrap();
+        res
     }
 
-    fn facet_fn_with_options<'extractor>(
-        doc_alloc: &Bump,
+    fn facet_fn_with_options<'extractor, 'doc>(
+        doc_alloc: &'doc Bump,
         cached_sorter: &mut BalancedCaches<'extractor>,
         cache_fn: impl Fn(&mut BalancedCaches<'extractor>, &[u8], u32) -> Result<()>,
+        del_add_facet_value: &mut DelAddFacetValue<'doc>,
+        facet_fn: impl Fn(&mut DelAddFacetValue<'doc>, FieldId, BVec<'doc, u8>, FacetKind),
         docid: DocumentId,
         fid: FieldId,
         value: &Value,
     ) -> Result<()> {
-        let mut buffer = bumpalo::collections::Vec::new_in(doc_alloc);
+        let mut buffer = BVec::new_in(doc_alloc);
         // Exists
         // key: fid
         buffer.push(FacetKind::Exists as u8);
@@ -152,15 +175,21 @@ impl FacetedDocidsExtractor {
             // Number
             // key: fid - level - orderedf64 - orignalf64
             Value::Number(number) => {
-                if let Some((n, ordered)) =
-                    number.as_f64().and_then(|n| f64_into_bytes(n).map(|ordered| (n, ordered)))
+                let mut ordered = [0u8; 16];
+                if number
+                    .as_f64()
+                    .and_then(|n| OrderedF64Codec::serialize_into(n, &mut ordered).ok())
+                    .is_some()
                 {
+                    let mut number = BVec::with_capacity_in(16, doc_alloc);
+                    number.extend_from_slice(&ordered);
+                    facet_fn(del_add_facet_value, fid, number, FacetKind::Number);
+
                     buffer.clear();
                     buffer.push(FacetKind::Number as u8);
                     buffer.extend_from_slice(&fid.to_be_bytes());
                     buffer.push(0); // level 0
                     buffer.extend_from_slice(&ordered);
-                    buffer.extend_from_slice(&n.to_be_bytes());
                     cache_fn(cached_sorter, &buffer, docid)
                 } else {
                     Ok(())
@@ -169,6 +198,10 @@ impl FacetedDocidsExtractor {
             // String
             // key: fid - level - truncated_string
             Value::String(s) => {
+                let mut string = BVec::new_in(doc_alloc);
+                string.extend_from_slice(s.as_bytes());
+                facet_fn(del_add_facet_value, fid, string, FacetKind::String);
+
                 let normalized = crate::normalize_facet(s);
                 let truncated = truncate_str(&normalized);
                 buffer.clear();
@@ -211,6 +244,84 @@ impl FacetedDocidsExtractor {
     }
 }
 
+struct DelAddFacetValue<'doc> {
+    strings: HashMap<(FieldId, BVec<'doc, u8>), DelAdd, hashbrown::DefaultHashBuilder, &'doc Bump>,
+    f64s: HashMap<(FieldId, BVec<'doc, u8>), DelAdd, hashbrown::DefaultHashBuilder, &'doc Bump>,
+}
+
+impl<'doc> DelAddFacetValue<'doc> {
+    fn new(doc_alloc: &'doc Bump) -> Self {
+        Self { strings: HashMap::new_in(doc_alloc), f64s: HashMap::new_in(doc_alloc) }
+    }
+
+    fn insert_add(&mut self, fid: FieldId, value: BVec<'doc, u8>, kind: FacetKind) {
+        let cache = match kind {
+            FacetKind::String => &mut self.strings,
+            FacetKind::Number => &mut self.f64s,
+            _ => return,
+        };
+
+        let key = (fid, value);
+        if let Some(DelAdd::Deletion) = cache.get(&key) {
+            cache.remove(&key);
+        } else {
+            cache.insert(key, DelAdd::Addition);
+        }
+    }
+
+    fn insert_del(&mut self, fid: FieldId, value: BVec<'doc, u8>, kind: FacetKind) {
+        let cache = match kind {
+            FacetKind::String => &mut self.strings,
+            FacetKind::Number => &mut self.f64s,
+            _ => return,
+        };
+
+        let key = (fid, value);
+        if let Some(DelAdd::Addition) = cache.get(&key) {
+            cache.remove(&key);
+        } else {
+            cache.insert(key, DelAdd::Deletion);
+        }
+    }
+
+    fn send_data(
+        self,
+        docid: DocumentId,
+        sender: &FieldIdDocidFacetSender,
+        doc_alloc: &Bump,
+    ) -> std::result::Result<(), crossbeam_channel::SendError<()>> {
+        println!("sending FieldIdDocidFacet data");
+        let mut count = 0;
+        let mut buffer = bumpalo::collections::Vec::new_in(doc_alloc);
+        for ((fid, value), deladd) in self.strings {
+            buffer.clear();
+            buffer.extend_from_slice(&fid.to_be_bytes());
+            buffer.extend_from_slice(&docid.to_be_bytes());
+            buffer.extend_from_slice(&value);
+            match deladd {
+                DelAdd::Deletion => sender.delete_facet_string(&buffer)?,
+                DelAdd::Addition => sender.write_facet_string(&buffer)?,
+            }
+            count += 1;
+        }
+
+        count = 0;
+        for ((fid, value), deladd) in self.f64s {
+            buffer.clear();
+            buffer.extend_from_slice(&fid.to_be_bytes());
+            buffer.extend_from_slice(&docid.to_be_bytes());
+            buffer.extend_from_slice(&value);
+            match deladd {
+                DelAdd::Deletion => sender.delete_facet_f64(&buffer)?,
+                DelAdd::Addition => sender.write_facet_f64(&buffer)?,
+            }
+            count += 1;
+        }
+
+        Ok(())
+    }
+}
+
 /// Truncates a string to the biggest valid LMDB key size.
 fn truncate_str(s: &str) -> &str {
     let index = s
@@ -223,13 +334,23 @@ fn truncate_str(s: &str) -> &str {
     &s[..index.unwrap_or(0)]
 }
 
-impl DocidsExtractor for FacetedDocidsExtractor {
+impl FacetedDocidsExtractor {
     #[tracing::instrument(level = "trace", skip_all, target = "indexing::extract::faceted")]
-    fn run_extraction<'pl, 'fid, 'indexer, 'index, 'extractor, DC: DocumentChanges<'pl>, MSP, SP>(
+    pub fn run_extraction<
+        'pl,
+        'fid,
+        'indexer,
+        'index,
+        'extractor,
+        DC: DocumentChanges<'pl>,
+        MSP,
+        SP,
+    >(
         grenad_parameters: GrenadParameters,
         document_changes: &DC,
         indexing_context: IndexingContext<'fid, 'indexer, 'index, MSP, SP>,
         extractor_allocs: &'extractor mut ThreadLocal<FullySend<Bump>>,
+        sender: &FieldIdDocidFacetSender,
         finished_steps: u16,
         total_steps: u16,
         step_name: &'static str,
@@ -254,6 +375,7 @@ impl DocidsExtractor for FacetedDocidsExtractor {
                 attributes_to_extract: &attributes_to_extract,
                 grenad_parameters,
                 buckets: rayon::current_num_threads(),
+                sender,
             };
             extract(
                 document_changes,

crates/milli/src/update/new/indexer/mod.rs

@@ -223,7 +223,7 @@ where
                     let (finished_steps, step_name) = steps::extract_facets();
 
                     facet_field_ids_delta = merge_and_send_facet_docids(
-                        FacetedDocidsExtractor::run_extraction(grenad_parameters, document_changes, indexing_context, &mut extractor_allocs, finished_steps, total_steps, step_name)?,
+                        FacetedDocidsExtractor::run_extraction(grenad_parameters, document_changes, indexing_context, &mut extractor_allocs, &extractor_sender.field_id_docid_facet_sender(), finished_steps, total_steps, step_name)?,
                         FacetDatabases::new(index),
                         index,
                         extractor_sender.facet_docids(),