From 51b6293738168e620f127daea0159311d30bea76 Mon Sep 17 00:00:00 2001 From: ManyTheFish Date: Wed, 13 Nov 2024 11:34:49 +0100 Subject: [PATCH] Add linear facet databases --- .../src/heed_codec/facet/ordered_f64_codec.rs | 29 +++- crates/milli/src/update/new/channel.rs | 40 +++++ .../new/extract/faceted/extract_facets.rs | 154 ++++++++++++++++-- crates/milli/src/update/new/indexer/mod.rs | 2 +- 4 files changed, 202 insertions(+), 23 deletions(-) diff --git a/crates/milli/src/heed_codec/facet/ordered_f64_codec.rs b/crates/milli/src/heed_codec/facet/ordered_f64_codec.rs index 4eccdb68b..19ba7a460 100644 --- a/crates/milli/src/heed_codec/facet/ordered_f64_codec.rs +++ b/crates/milli/src/heed_codec/facet/ordered_f64_codec.rs @@ -27,17 +27,34 @@ impl heed::BytesEncode<'_> for OrderedF64Codec { fn bytes_encode(f: &Self::EItem) -> Result, BoxedError> { let mut buffer = [0u8; 16]; - // write the globally ordered float - 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[..]); + 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(()) +} + #[derive(Error, Debug)] #[error("the float {float} cannot be converted to a globally ordered representation")] pub struct InvalidGloballyOrderedFloatError { diff --git a/crates/milli/src/update/new/channel.rs b/crates/milli/src/update/new/channel.rs index 2027b4db8..3287a1f7f 100644 --- a/crates/milli/src/update/new/channel.rs +++ b/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<'_> { diff --git a/crates/milli/src/update/new/extract/faceted/extract_facets.rs b/crates/milli/src/update/new/extract/faceted/extract_facets.rs index d0dc425ae..0e7dcc4b9 100644 --- a/crates/milli/src/update/new/extract/faceted/extract_facets.rs +++ b/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>, 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>, + 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, diff --git a/crates/milli/src/update/new/indexer/mod.rs b/crates/milli/src/update/new/indexer/mod.rs index e3b24642e..2cdeca76d 100644 --- a/crates/milli/src/update/new/indexer/mod.rs +++ b/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(),