mirror of
https://github.com/meilisearch/MeiliSearch
synced 2024-12-22 12:40:04 +01:00
Add linear facet databases
This commit is contained in:
parent
b17896d899
commit
51b6293738
@ -27,17 +27,34 @@ 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 })?;
|
||||
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 {
|
||||
|
@ -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<'_> {
|
||||
|
@ -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,
|
||||
|
@ -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(),
|
||||
|
Loading…
x
Reference in New Issue
Block a user