From e639ec79d1eb291803269c3091940d52ff0a5f43 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Cl=C3=A9ment=20Renault?= Date: Mon, 2 Sep 2024 10:42:19 +0200 Subject: [PATCH] Move the indexers into their own modules --- milli/src/update/new/channel.rs | 2 +- milli/src/update/new/document_change.rs | 2 +- .../update/new/indexer/document_deletion.rs | 53 ++ .../update/new/indexer/document_operation.rs | 325 +++++++++ milli/src/update/new/indexer/mod.rs | 82 +++ milli/src/update/new/indexer/partial_dump.rs | 73 ++ .../update/new/indexer/update_by_function.rs | 18 + milli/src/update/new/merger.rs | 97 +++ milli/src/update/new/mod.rs | 622 +----------------- 9 files changed, 668 insertions(+), 606 deletions(-) create mode 100644 milli/src/update/new/indexer/document_deletion.rs create mode 100644 milli/src/update/new/indexer/document_operation.rs create mode 100644 milli/src/update/new/indexer/mod.rs create mode 100644 milli/src/update/new/indexer/partial_dump.rs create mode 100644 milli/src/update/new/indexer/update_by_function.rs create mode 100644 milli/src/update/new/merger.rs diff --git a/milli/src/update/new/channel.rs b/milli/src/update/new/channel.rs index 0aeabcfa7..088303fb3 100644 --- a/milli/src/update/new/channel.rs +++ b/milli/src/update/new/channel.rs @@ -3,8 +3,8 @@ use std::fs::File; use crossbeam_channel::{IntoIter, Receiver, SendError, Sender}; use heed::types::Bytes; -use super::indexer::KvReaderFieldId; use super::StdResult; +use crate::update::new::KvReaderFieldId; use crate::{DocumentId, Index}; /// The capacity of the channel is currently in number of messages. diff --git a/milli/src/update/new/document_change.rs b/milli/src/update/new/document_change.rs index 1764b6ee7..6f9d767cb 100644 --- a/milli/src/update/new/document_change.rs +++ b/milli/src/update/new/document_change.rs @@ -1,7 +1,7 @@ use heed::RoTxn; use obkv::KvReader; -use super::indexer::KvReaderFieldId; +use crate::update::new::KvReaderFieldId; use crate::{DocumentId, FieldId}; pub enum DocumentChange { diff --git a/milli/src/update/new/indexer/document_deletion.rs b/milli/src/update/new/indexer/document_deletion.rs new file mode 100644 index 000000000..24ba0c671 --- /dev/null +++ b/milli/src/update/new/indexer/document_deletion.rs @@ -0,0 +1,53 @@ +use std::sync::Arc; + +use rayon::iter::{ParallelBridge, ParallelIterator}; +use roaring::RoaringBitmap; + +use super::Indexer; +use crate::documents::PrimaryKey; +use crate::update::new::{Deletion, DocumentChange, ItemsPool}; +use crate::{FieldsIdsMap, Index, InternalError, Result}; + +pub struct DocumentDeletionIndexer { + pub to_delete: RoaringBitmap, +} + +impl DocumentDeletionIndexer { + pub fn new() -> Self { + Self { to_delete: Default::default() } + } + + pub fn delete_documents_by_docids(&mut self, docids: RoaringBitmap) { + self.to_delete |= docids; + } +} + +impl<'p> Indexer<'p> for DocumentDeletionIndexer { + type Parameter = (&'p Index, &'p FieldsIdsMap, &'p PrimaryKey<'p>); + + fn document_changes( + self, + param: Self::Parameter, + ) -> Result>> + 'p> { + let (index, fields, primary_key) = param; + let items = Arc::new(ItemsPool::new(|| index.read_txn().map_err(crate::Error::from))); + Ok(self.to_delete.into_iter().par_bridge().map_with(items, |items, docid| { + items.with(|rtxn| { + let current = index.document(rtxn, docid)?; + let external_docid = match primary_key.document_id(current, fields)? { + Ok(document_id) => Ok(document_id) as Result<_>, + Err(_) => Err(InternalError::DocumentsError( + crate::documents::Error::InvalidDocumentFormat, + ) + .into()), + }?; + + Ok(Some(DocumentChange::Deletion(Deletion::create( + docid, + external_docid, + current.boxed(), + )))) + }) + })) + } +} diff --git a/milli/src/update/new/indexer/document_operation.rs b/milli/src/update/new/indexer/document_operation.rs new file mode 100644 index 000000000..ba93915ea --- /dev/null +++ b/milli/src/update/new/indexer/document_operation.rs @@ -0,0 +1,325 @@ +use std::borrow::Cow; +use std::collections::{BTreeMap, HashMap}; +use std::fs::File; +use std::io::Cursor; +use std::sync::Arc; + +use heed::types::Bytes; +use heed::RoTxn; +use memmap2::Mmap; +use rayon::iter::{IntoParallelIterator, ParallelIterator}; + +use super::super::document_change::DocumentChange; +use super::super::items_pool::ItemsPool; +use super::Indexer; +use crate::documents::{ + obkv_to_object, DocumentIdExtractionError, DocumentsBatchReader, PrimaryKey, +}; +use crate::update::new::{Deletion, Insertion, KvReaderFieldId, KvWriterFieldId, Update}; +use crate::update::{AvailableIds, IndexDocumentsMethod}; +use crate::{DocumentId, Error, FieldsIdsMap, Index, Result, UserError}; + +pub struct DocumentOperationIndexer { + pub(crate) operations: Vec, + pub(crate) index_documents_method: IndexDocumentsMethod, +} + +pub enum Payload { + Addition(File), + Deletion(Vec), +} + +pub struct PayloadStats { + pub document_count: usize, + pub bytes: u64, +} + +pub enum DocumentOperation { + Addition(DocumentOffset), + Deletion, +} + +/// Represents an offset where a document lives +/// in an mmapped grenad reader file. +pub struct DocumentOffset { + /// The mmapped grenad reader file. + pub content: Arc, // grenad::Reader + /// The offset of the document in the file. + pub offset: u32, +} + +impl DocumentOperationIndexer { + pub fn new(method: IndexDocumentsMethod) -> Self { + Self { operations: Default::default(), index_documents_method: method } + } + + /// TODO please give me a type + /// The payload is expected to be in the grenad format + pub fn add_documents(&mut self, payload: File) -> Result { + let reader = DocumentsBatchReader::from_reader(&payload)?; + let bytes = payload.metadata()?.len(); + let document_count = reader.documents_count() as usize; + + self.operations.push(Payload::Addition(payload)); + + Ok(PayloadStats { bytes, document_count }) + } + + pub fn delete_documents(&mut self, to_delete: Vec) { + self.operations.push(Payload::Deletion(to_delete)) + } +} + +impl<'p> Indexer<'p> for DocumentOperationIndexer { + type Parameter = (&'p Index, &'p RoTxn<'static>, &'p mut FieldsIdsMap, &'p PrimaryKey<'p>); + + fn document_changes( + self, + param: Self::Parameter, + ) -> Result>> + 'p> { + let (index, rtxn, fields_ids_map, primary_key) = param; + + let documents_ids = index.documents_ids(rtxn)?; + let mut available_docids = AvailableIds::new(&documents_ids); + let mut docids_version_offsets = HashMap::::new(); + + for operation in self.operations { + match operation { + Payload::Addition(payload) => { + let content = unsafe { Mmap::map(&payload).map(Arc::new)? }; + let cursor = Cursor::new(content.as_ref()); + let reader = DocumentsBatchReader::from_reader(cursor)?; + + let (mut batch_cursor, batch_index) = reader.into_cursor_and_fields_index(); + // TODO Fetch all document fields to fill the fields ids map + batch_index.iter().for_each(|(_, name)| { + fields_ids_map.insert(name); + }); + + let mut offset: u32 = 0; + while let Some(document) = batch_cursor.next_document()? { + let external_document_id = + match primary_key.document_id(document, &batch_index)? { + Ok(document_id) => Ok(document_id), + Err(DocumentIdExtractionError::InvalidDocumentId(user_error)) => { + Err(user_error) + } + Err(DocumentIdExtractionError::MissingDocumentId) => { + Err(UserError::MissingDocumentId { + primary_key: primary_key.name().to_string(), + document: obkv_to_object(document, &batch_index)?, + }) + } + Err(DocumentIdExtractionError::TooManyDocumentIds(_)) => { + Err(UserError::TooManyDocumentIds { + primary_key: primary_key.name().to_string(), + document: obkv_to_object(document, &batch_index)?, + }) + } + }?; + + let content = content.clone(); + let document_offset = DocumentOffset { content, offset }; + let document_operation = DocumentOperation::Addition(document_offset); + + match docids_version_offsets.get_mut(&external_document_id) { + None => { + let docid = match index + .external_documents_ids() + .get(rtxn, &external_document_id)? + { + Some(docid) => docid, + None => available_docids + .next() + .ok_or(Error::UserError(UserError::DocumentLimitReached))?, + }; + + docids_version_offsets.insert( + external_document_id, + (docid, vec![document_operation]), + ); + } + Some((_, offsets)) => offsets.push(document_operation), + } + offset += 1; + } + } + Payload::Deletion(to_delete) => { + for external_document_id in to_delete { + match docids_version_offsets.get_mut(&external_document_id) { + None => { + let docid = match index + .external_documents_ids() + .get(rtxn, &external_document_id)? + { + Some(docid) => docid, + None => available_docids + .next() + .ok_or(Error::UserError(UserError::DocumentLimitReached))?, + }; + + docids_version_offsets.insert( + external_document_id, + (docid, vec![DocumentOperation::Deletion]), + ); + } + Some((_, offsets)) => offsets.push(DocumentOperation::Deletion), + } + } + } + } + } + + Ok(docids_version_offsets.into_par_iter().map_with( + Arc::new(ItemsPool::new(|| index.read_txn().map_err(crate::Error::from))), + move |context_pool, (external_docid, (internal_docid, operations))| { + context_pool.with(|rtxn| { + use IndexDocumentsMethod as Idm; + let document_merge_function = match self.index_documents_method { + Idm::ReplaceDocuments => merge_document_for_replacements, + Idm::UpdateDocuments => merge_document_for_updates, + }; + + document_merge_function( + rtxn, + index, + fields_ids_map, + internal_docid, + external_docid, + &operations, + ) + }) + }, + )) + } +} + +/// Reads the previous version of a document from the database, the new versions +/// in the grenad update files and merges them to generate a new boxed obkv. +/// +/// This function is only meant to be used when doing an update and not a replacement. +fn merge_document_for_updates( + rtxn: &RoTxn, + index: &Index, + fields_ids_map: &FieldsIdsMap, + docid: DocumentId, + external_docid: String, + operations: &[DocumentOperation], +) -> Result> { + let mut document = BTreeMap::<_, Cow<_>>::new(); + let current = index.documents.remap_data_type::().get(rtxn, &docid)?; + let current: Option<&KvReaderFieldId> = current.map(Into::into); + + if let Some(current) = current { + current.into_iter().for_each(|(k, v)| { + document.insert(k, v.into()); + }); + } + + let last_deletion = + operations.iter().rposition(|operation| matches!(operation, DocumentOperation::Deletion)); + + let operations = &operations[last_deletion.map_or(0, |i| i + 1)..]; + + if operations.is_empty() { + match current { + Some(current) => { + return Ok(Some(DocumentChange::Deletion(Deletion::create( + docid, + external_docid, + current.boxed(), + )))); + } + None => return Ok(None), + } + } + + for operation in operations { + let DocumentOffset { content, offset } = match operation { + DocumentOperation::Addition(offset) => offset, + DocumentOperation::Deletion => { + unreachable!("Deletion in document operations") + } + }; + + let reader = DocumentsBatchReader::from_reader(Cursor::new(content.as_ref()))?; + let (mut cursor, batch_index) = reader.into_cursor_and_fields_index(); + let update = cursor.get(*offset)?.expect("must exists"); + + update.into_iter().for_each(|(k, v)| { + let field_name = batch_index.name(k).unwrap(); + let id = fields_ids_map.id(field_name).unwrap(); + document.insert(id, v.to_vec().into()); + }); + } + + let mut writer = KvWriterFieldId::memory(); + document.into_iter().for_each(|(id, value)| writer.insert(id, value).unwrap()); + let new = writer.into_boxed(); + + match current { + Some(current) => { + let update = Update::create(docid, external_docid, current.boxed(), new); + Ok(Some(DocumentChange::Update(update))) + } + None => { + let insertion = Insertion::create(docid, external_docid, new); + Ok(Some(DocumentChange::Insertion(insertion))) + } + } +} + +/// Returns only the most recent version of a document based on the updates from the payloads. +/// +/// This function is only meant to be used when doing a replacement and not an update. +fn merge_document_for_replacements( + rtxn: &RoTxn, + index: &Index, + fields_ids_map: &FieldsIdsMap, + docid: DocumentId, + external_docid: String, + operations: &[DocumentOperation], +) -> Result> { + let current = index.documents.remap_data_type::().get(rtxn, &docid)?; + let current: Option<&KvReaderFieldId> = current.map(Into::into); + + match operations.last() { + Some(DocumentOperation::Addition(DocumentOffset { content, offset })) => { + let reader = DocumentsBatchReader::from_reader(Cursor::new(content.as_ref()))?; + let (mut cursor, batch_index) = reader.into_cursor_and_fields_index(); + let update = cursor.get(*offset)?.expect("must exists"); + + let mut document_entries = Vec::new(); + update.into_iter().for_each(|(k, v)| { + let field_name = batch_index.name(k).unwrap(); + let id = fields_ids_map.id(field_name).unwrap(); + document_entries.push((id, v)); + }); + + document_entries.sort_unstable_by_key(|(id, _)| *id); + + let mut writer = KvWriterFieldId::memory(); + document_entries.into_iter().for_each(|(id, value)| writer.insert(id, value).unwrap()); + let new = writer.into_boxed(); + + match current { + Some(current) => { + let update = Update::create(docid, external_docid, current.boxed(), new); + Ok(Some(DocumentChange::Update(update))) + } + None => { + let insertion = Insertion::create(docid, external_docid, new); + Ok(Some(DocumentChange::Insertion(insertion))) + } + } + } + Some(DocumentOperation::Deletion) => match current { + Some(current) => { + let deletion = Deletion::create(docid, external_docid, current.boxed()); + Ok(Some(DocumentChange::Deletion(deletion))) + } + None => Ok(None), + }, + None => Ok(None), + } +} diff --git a/milli/src/update/new/indexer/mod.rs b/milli/src/update/new/indexer/mod.rs new file mode 100644 index 000000000..998793b49 --- /dev/null +++ b/milli/src/update/new/indexer/mod.rs @@ -0,0 +1,82 @@ +use std::thread; + +use big_s::S; +pub use document_deletion::DocumentDeletionIndexer; +pub use document_operation::DocumentOperationIndexer; +use heed::RwTxn; +pub use partial_dump::PartialDumpIndexer; +use rayon::iter::{IntoParallelIterator, ParallelIterator}; +use rayon::ThreadPool; +pub use update_by_function::UpdateByFunctionIndexer; + +use super::channel::{ + extractors_merger_channels, merger_writer_channels, EntryOperation, ExtractorsMergerChannels, + WriterOperation, +}; +use super::document_change::DocumentChange; +use super::merger::merge_grenad_entries; +use crate::{Index, Result}; + +mod document_deletion; +mod document_operation; +mod partial_dump; +mod update_by_function; + +pub trait Indexer<'p> { + type Parameter: 'p; + + fn document_changes( + self, + param: Self::Parameter, + ) -> Result>> + 'p>; +} + +/// This is the main function of this crate. +/// +/// Give it the output of the [`Indexer::document_changes`] method and it will execute it in the [`rayon::ThreadPool`]. +/// +/// TODO return stats +/// TODO take the rayon ThreadPool +pub fn index( + wtxn: &mut RwTxn, + index: &Index, + pool: &ThreadPool, + document_changes: PI, +) -> Result<()> +where + PI: IntoParallelIterator> + Send, + PI::Iter: Clone, +{ + let (merger_sender, writer_receiver) = merger_writer_channels(100); + let ExtractorsMergerChannels { merger_receiver, deladd_cbo_roaring_bitmap_sender } = + extractors_merger_channels(100); + + thread::scope(|s| { + // TODO manage the errors correctly + thread::Builder::new().name(S("indexer-extractors")).spawn_scoped(s, || { + pool.in_place_scope(|_s| { + document_changes.into_par_iter().for_each(|_dc| ()); + }) + })?; + + // TODO manage the errors correctly + thread::Builder::new().name(S("indexer-merger")).spawn_scoped(s, || { + let rtxn = index.read_txn().unwrap(); + merge_grenad_entries(merger_receiver, merger_sender, &rtxn, index).unwrap() + })?; + + // TODO Split this code into another function + for operation in writer_receiver { + let database = operation.database(index); + match operation { + WriterOperation::WordDocids(operation) => match operation { + EntryOperation::Delete(e) => database.delete(wtxn, e.entry()).map(drop)?, + EntryOperation::Write(e) => database.put(wtxn, e.key(), e.value())?, + }, + WriterOperation::Document(e) => database.put(wtxn, &e.key(), e.content())?, + } + } + + Ok(()) + }) +} diff --git a/milli/src/update/new/indexer/partial_dump.rs b/milli/src/update/new/indexer/partial_dump.rs new file mode 100644 index 000000000..24ba70bcb --- /dev/null +++ b/milli/src/update/new/indexer/partial_dump.rs @@ -0,0 +1,73 @@ +use rayon::iter::{ParallelBridge, ParallelIterator}; + +use super::Indexer; +use crate::documents::{DocumentIdExtractionError, PrimaryKey}; +use crate::update::concurrent_available_ids::ConcurrentAvailableIds; +use crate::update::new::{DocumentChange, Insertion, KvWriterFieldId}; +use crate::{all_obkv_to_json, Error, FieldsIdsMap, Object, Result, UserError}; + +pub struct PartialDumpIndexer { + pub iter: I, +} + +impl PartialDumpIndexer { + pub fn new_from_jsonlines(iter: I) -> Self { + PartialDumpIndexer { iter } + } +} + +impl<'p, I> Indexer<'p> for PartialDumpIndexer +where + I: IntoIterator, + I::IntoIter: Send + 'p, + I::Item: Send, +{ + type Parameter = (&'p FieldsIdsMap, &'p ConcurrentAvailableIds, &'p PrimaryKey<'p>); + + /// Note for future self: + /// - the field ids map must already be valid so you must have to generate it beforehand. + /// - We should probably expose another method that generates the fields ids map from an iterator of JSON objects. + /// - We recommend sending chunks of documents in this `PartialDumpIndexer` we therefore need to create a custom take_while_size method (that doesn't drop items). + fn document_changes( + self, + param: Self::Parameter, + ) -> Result>> + 'p> { + let (fields_ids_map, concurrent_available_ids, primary_key) = param; + + Ok(self.iter.into_iter().par_bridge().map(|object| { + let docid = match concurrent_available_ids.next() { + Some(id) => id, + None => return Err(Error::UserError(UserError::DocumentLimitReached)), + }; + + let mut writer = KvWriterFieldId::memory(); + object.iter().for_each(|(key, value)| { + let key = fields_ids_map.id(key).unwrap(); + /// TODO better error management + let value = serde_json::to_vec(&value).unwrap(); + writer.insert(key, value).unwrap(); + }); + + let document = writer.into_boxed(); + let external_docid = match primary_key.document_id(&document, fields_ids_map)? { + Ok(document_id) => Ok(document_id), + Err(DocumentIdExtractionError::InvalidDocumentId(user_error)) => Err(user_error), + Err(DocumentIdExtractionError::MissingDocumentId) => { + Err(UserError::MissingDocumentId { + primary_key: primary_key.name().to_string(), + document: all_obkv_to_json(&document, fields_ids_map)?, + }) + } + Err(DocumentIdExtractionError::TooManyDocumentIds(_)) => { + Err(UserError::TooManyDocumentIds { + primary_key: primary_key.name().to_string(), + document: all_obkv_to_json(&document, fields_ids_map)?, + }) + } + }?; + + let insertion = Insertion::create(docid, external_docid, document); + Ok(Some(DocumentChange::Insertion(insertion))) + })) + } +} diff --git a/milli/src/update/new/indexer/update_by_function.rs b/milli/src/update/new/indexer/update_by_function.rs new file mode 100644 index 000000000..c8c434b72 --- /dev/null +++ b/milli/src/update/new/indexer/update_by_function.rs @@ -0,0 +1,18 @@ +use rayon::iter::{IntoParallelIterator, ParallelIterator}; + +use super::Indexer; +use crate::update::new::DocumentChange; +use crate::Result; + +pub struct UpdateByFunctionIndexer; + +impl<'p> Indexer<'p> for UpdateByFunctionIndexer { + type Parameter = (); + + fn document_changes( + self, + _param: Self::Parameter, + ) -> Result>> + 'p> { + Ok(vec![].into_par_iter()) + } +} diff --git a/milli/src/update/new/merger.rs b/milli/src/update/new/merger.rs new file mode 100644 index 000000000..97f9e6ac6 --- /dev/null +++ b/milli/src/update/new/merger.rs @@ -0,0 +1,97 @@ +use heed::types::Bytes; +use heed::RoTxn; +use roaring::RoaringBitmap; + +use super::channel::{MergerReceiver, MergerSender}; +use super::KvReaderDelAdd; +use crate::update::del_add::DelAdd; +use crate::update::new::channel::MergerOperation; +use crate::update::MergeDeladdCboRoaringBitmaps; +use crate::{CboRoaringBitmapCodec, Index, Result}; + +/// TODO We must return some infos/stats +pub fn merge_grenad_entries( + receiver: MergerReceiver, + sender: MergerSender, + rtxn: &RoTxn, + index: &Index, +) -> Result<()> { + let mut buffer = Vec::new(); + + for merger_operation in receiver { + match merger_operation { + MergerOperation::WordDocidsCursors(cursors) => { + let sender = sender.word_docids(); + let database = index.word_docids.remap_types::(); + + let mut builder = grenad::MergerBuilder::new(MergeDeladdCboRoaringBitmaps); + builder.extend(cursors); + /// TODO manage the error correctly + let mut merger_iter = builder.build().into_stream_merger_iter().unwrap(); + + // TODO manage the error correctly + while let Some((key, deladd)) = merger_iter.next().unwrap() { + let current = database.get(rtxn, key)?; + let deladd: &KvReaderDelAdd = deladd.into(); + let del = deladd.get(DelAdd::Deletion); + let add = deladd.get(DelAdd::Addition); + + match merge_cbo_bitmaps(current, del, add)? { + Operation::Write(bitmap) => { + let value = cbo_serialize_into_vec(&bitmap, &mut buffer); + sender.write(key, value).unwrap(); + } + Operation::Delete => sender.delete(key).unwrap(), + Operation::Ignore => (), + } + } + } + } + } + + Ok(()) +} + +enum Operation { + Write(RoaringBitmap), + Delete, + Ignore, +} + +/// A function that merges the DelAdd CboRoaringBitmaps with the current bitmap. +fn merge_cbo_bitmaps( + current: Option<&[u8]>, + del: Option<&[u8]>, + add: Option<&[u8]>, +) -> Result { + let current = current.map(CboRoaringBitmapCodec::deserialize_from).transpose()?; + let del = del.map(CboRoaringBitmapCodec::deserialize_from).transpose()?; + let add = add.map(CboRoaringBitmapCodec::deserialize_from).transpose()?; + + match (current, del, add) { + (None, None, None) => Ok(Operation::Ignore), // but it's strange + (None, None, Some(add)) => Ok(Operation::Write(add)), + (None, Some(_del), None) => Ok(Operation::Ignore), // but it's strange + (None, Some(_del), Some(add)) => Ok(Operation::Write(add)), + (Some(_current), None, None) => Ok(Operation::Ignore), // but it's strange + (Some(current), None, Some(add)) => Ok(Operation::Write(current | add)), + (Some(current), Some(del), add) => { + let output = match add { + Some(add) => (current - del) | add, + None => current - del, + }; + if output.is_empty() { + Ok(Operation::Delete) + } else { + Ok(Operation::Write(output)) + } + } + } +} + +/// Return the slice directly from the serialize_into method +fn cbo_serialize_into_vec<'b>(bitmap: &RoaringBitmap, buffer: &'b mut Vec) -> &'b [u8] { + buffer.clear(); + CboRoaringBitmapCodec::serialize_into(bitmap, buffer); + buffer.as_slice() +} diff --git a/milli/src/update/new/mod.rs b/milli/src/update/new/mod.rs index a6ac38abe..830565368 100644 --- a/milli/src/update/new/mod.rs +++ b/milli/src/update/new/mod.rs @@ -1,610 +1,24 @@ +pub use document_change::{Deletion, DocumentChange, Insertion, Update}; +pub use indexer::{ + index, DocumentDeletionIndexer, DocumentOperationIndexer, PartialDumpIndexer, + UpdateByFunctionIndexer, +}; +pub use items_pool::ItemsPool; + +use super::del_add::DelAdd; +use crate::FieldId; + mod document_change; +mod merger; // mod extract; +// mod global_fields_ids_map; mod channel; +mod indexer; mod items_pool; -/// TODO remove this -// mod global_fields_ids_map; - +/// TODO move them elsewhere pub type StdResult = std::result::Result; - -mod indexer { - use std::borrow::Cow; - use std::collections::{BTreeMap, HashMap}; - use std::fs::File; - use std::io::Cursor; - use std::os::unix::fs::MetadataExt; - use std::sync::Arc; - use std::thread; - - use big_s::S; - use heed::types::Bytes; - use heed::{RoTxn, RwTxn}; - use memmap2::Mmap; - use obkv::KvWriter; - use rayon::iter::{IntoParallelIterator, ParallelBridge, ParallelIterator}; - use rayon::ThreadPool; - use roaring::RoaringBitmap; - use serde_json::Value; - - use super::channel::{ - extractors_merger_channels, merger_writer_channels, EntryOperation, - ExtractorsMergerChannels, MergerReceiver, MergerSender, WriterOperation, - }; - use super::document_change::{Deletion, DocumentChange, Insertion, Update}; - use super::items_pool::ItemsPool; - use crate::documents::{ - obkv_to_object, DocumentIdExtractionError, DocumentsBatchReader, PrimaryKey, - }; - use crate::update::concurrent_available_ids::ConcurrentAvailableIds; - use crate::update::del_add::DelAdd; - use crate::update::new::channel::MergerOperation; - use crate::update::{AvailableIds, IndexDocumentsMethod, MergeDeladdCboRoaringBitmaps}; - use crate::{ - all_obkv_to_json, obkv_to_json, CboRoaringBitmapCodec, DocumentId, Error, FieldId, - FieldsIdsMap, Index, InternalError, Object, Result, UserError, - }; - - pub type KvReaderFieldId = obkv::KvReader; - pub type KvReaderDelAdd = obkv::KvReader; - pub type KvWriterFieldId = obkv::KvWriter; - pub type KvWriterDelAdd = obkv::KvWriter; - - pub struct DocumentOperationIndexer { - operations: Vec, - index_documents_method: IndexDocumentsMethod, - } - - enum Payload { - Addition(File), - Deletion(Vec), - } - - pub struct PayloadStats { - pub document_count: usize, - pub bytes: u64, - } - - enum DocumentOperation { - Addition(DocumentOffset), - Deletion, - } - - /// Represents an offset where a document lives - /// in an mmapped grenad reader file. - struct DocumentOffset { - /// The mmapped grenad reader file. - pub content: Arc, // grenad::Reader - /// The offset of the document in the file. - pub offset: u32, - } - - impl DocumentOperationIndexer { - pub fn new(method: IndexDocumentsMethod) -> Self { - Self { operations: Default::default(), index_documents_method: method } - } - - /// TODO please give me a type - /// The payload is expected to be in the grenad format - pub fn add_documents(&mut self, payload: File) -> Result { - let reader = DocumentsBatchReader::from_reader(&payload)?; - let bytes = payload.metadata()?.size(); - let document_count = reader.documents_count() as usize; - - self.operations.push(Payload::Addition(payload)); - - Ok(PayloadStats { bytes, document_count }) - } - - pub fn delete_documents(&mut self, to_delete: Vec) { - self.operations.push(Payload::Deletion(to_delete)) - } - - pub fn document_changes<'a>( - self, - index: &'a Index, - rtxn: &'a RoTxn, - fields_ids_map: &'a mut FieldsIdsMap, - primary_key: &'a PrimaryKey<'a>, - ) -> Result>> + 'a> { - let documents_ids = index.documents_ids(rtxn)?; - let mut available_docids = AvailableIds::new(&documents_ids); - let mut docids_version_offsets = HashMap::::new(); - - for operation in self.operations { - match operation { - Payload::Addition(payload) => { - let content = unsafe { Mmap::map(&payload).map(Arc::new)? }; - let cursor = Cursor::new(content.as_ref()); - let reader = DocumentsBatchReader::from_reader(cursor)?; - - let (mut batch_cursor, batch_index) = reader.into_cursor_and_fields_index(); - // TODO Fetch all document fields to fill the fields ids map - batch_index.iter().for_each(|(_, name)| { - fields_ids_map.insert(name); - }); - - let mut offset: u32 = 0; - while let Some(document) = batch_cursor.next_document()? { - let external_document_id = - match primary_key.document_id(document, &batch_index)? { - Ok(document_id) => Ok(document_id), - Err(DocumentIdExtractionError::InvalidDocumentId( - user_error, - )) => Err(user_error), - Err(DocumentIdExtractionError::MissingDocumentId) => { - Err(UserError::MissingDocumentId { - primary_key: primary_key.name().to_string(), - document: obkv_to_object(document, &batch_index)?, - }) - } - Err(DocumentIdExtractionError::TooManyDocumentIds(_)) => { - Err(UserError::TooManyDocumentIds { - primary_key: primary_key.name().to_string(), - document: obkv_to_object(document, &batch_index)?, - }) - } - }?; - - let content = content.clone(); - let document_offset = DocumentOffset { content, offset }; - let document_operation = DocumentOperation::Addition(document_offset); - - match docids_version_offsets.get_mut(&external_document_id) { - None => { - let docid = match index - .external_documents_ids() - .get(rtxn, &external_document_id)? - { - Some(docid) => docid, - None => available_docids.next().ok_or(Error::UserError( - UserError::DocumentLimitReached, - ))?, - }; - - docids_version_offsets.insert( - external_document_id, - (docid, vec![document_operation]), - ); - } - Some((_, offsets)) => offsets.push(document_operation), - } - offset += 1; - } - } - Payload::Deletion(to_delete) => { - for external_document_id in to_delete { - match docids_version_offsets.get_mut(&external_document_id) { - None => { - let docid = match index - .external_documents_ids() - .get(rtxn, &external_document_id)? - { - Some(docid) => docid, - None => available_docids.next().ok_or(Error::UserError( - UserError::DocumentLimitReached, - ))?, - }; - - docids_version_offsets.insert( - external_document_id, - (docid, vec![DocumentOperation::Deletion]), - ); - } - Some((_, offsets)) => offsets.push(DocumentOperation::Deletion), - } - } - } - } - } - - Ok(docids_version_offsets.into_par_iter().map_with( - Arc::new(ItemsPool::new(|| index.read_txn().map_err(crate::Error::from))), - move |context_pool, (external_docid, (internal_docid, operations))| { - context_pool.with(|rtxn| { - use IndexDocumentsMethod as Idm; - let document_merge_function = match self.index_documents_method { - Idm::ReplaceDocuments => merge_document_for_replacements, - Idm::UpdateDocuments => merge_document_for_updates, - }; - - document_merge_function( - rtxn, - index, - fields_ids_map, - internal_docid, - external_docid, - &operations, - ) - }) - }, - )) - } - } - - pub struct DeleteDocumentIndexer { - to_delete: RoaringBitmap, - } - - impl DeleteDocumentIndexer { - pub fn new() -> Self { - Self { to_delete: Default::default() } - } - - pub fn delete_documents_by_docids(&mut self, docids: RoaringBitmap) { - self.to_delete |= docids; - } - - // let fields = index.fields_ids_map(rtxn)?; - // let primary_key = - // index.primary_key(rtxn)?.ok_or(InternalError::DatabaseMissingEntry { - // db_name: db_name::MAIN, - // key: Some(main_key::PRIMARY_KEY_KEY), - // })?; - // let primary_key = PrimaryKey::new(primary_key, &fields).ok_or_else(|| { - // InternalError::FieldIdMapMissingEntry(crate::FieldIdMapMissingEntry::FieldName { - // field_name: primary_key.to_owned(), - // process: "external_id_of", - // }) - // })?; - pub fn document_changes<'a>( - self, - index: &'a Index, - fields: &'a FieldsIdsMap, - primary_key: &'a PrimaryKey<'a>, - ) -> Result> + 'a> { - let items = Arc::new(ItemsPool::new(|| index.read_txn().map_err(crate::Error::from))); - Ok(self.to_delete.into_iter().par_bridge().map_with(items, |items, docid| { - items.with(|rtxn| { - let current = index.document(rtxn, docid)?; - let external_docid = match primary_key.document_id(current, fields)? { - Ok(document_id) => Ok(document_id) as Result<_>, - Err(_) => Err(InternalError::DocumentsError( - crate::documents::Error::InvalidDocumentFormat, - ) - .into()), - }?; - - Ok(DocumentChange::Deletion(Deletion::create( - docid, - external_docid, - current.boxed(), - ))) - }) - })) - } - } - - pub struct PartialDumpIndexer { - iter: I, - } - - impl PartialDumpIndexer - where - I: IntoIterator, - I::IntoIter: Send, - I::Item: Send, - { - pub fn new_from_jsonlines(iter: I) -> Self { - PartialDumpIndexer { iter } - } - - /// Note for future self: - /// - the field ids map must already be valid so you must have to generate it beforehand. - /// - We should probably expose another method that generates the fields ids map from an iterator of JSON objects. - /// - We recommend sending chunks of documents in this `PartialDumpIndexer` we therefore need to create a custom take_while_size method (that doesn't drop items). - pub fn document_changes<'a>( - self, - fields_ids_map: &'a FieldsIdsMap, - concurrent_available_ids: &'a ConcurrentAvailableIds, - primary_key: &'a PrimaryKey<'a>, - ) -> impl ParallelIterator>> + 'a - where - // I don't like this, it will not fit in the future trait easily - I::IntoIter: 'a, - { - self.iter.into_iter().par_bridge().map(|object| { - let docid = match concurrent_available_ids.next() { - Some(id) => id, - None => return Err(Error::UserError(UserError::DocumentLimitReached)), - }; - - let mut writer = KvWriterFieldId::memory(); - object.iter().for_each(|(key, value)| { - let key = fields_ids_map.id(key).unwrap(); - /// TODO better error management - let value = serde_json::to_vec(&value).unwrap(); - writer.insert(key, value).unwrap(); - }); - - let document = writer.into_boxed(); - let external_docid = match primary_key.document_id(&document, fields_ids_map)? { - Ok(document_id) => Ok(document_id), - Err(DocumentIdExtractionError::InvalidDocumentId(user_error)) => { - Err(user_error) - } - Err(DocumentIdExtractionError::MissingDocumentId) => { - Err(UserError::MissingDocumentId { - primary_key: primary_key.name().to_string(), - document: all_obkv_to_json(&document, fields_ids_map)?, - }) - } - Err(DocumentIdExtractionError::TooManyDocumentIds(_)) => { - Err(UserError::TooManyDocumentIds { - primary_key: primary_key.name().to_string(), - document: all_obkv_to_json(&document, fields_ids_map)?, - }) - } - }?; - - let insertion = Insertion::create(docid, external_docid, document); - Ok(Some(DocumentChange::Insertion(insertion))) - }) - } - } - - pub struct UpdateByFunctionIndexer; - - /// TODO return stats - /// TODO take the rayon ThreadPool - pub fn index( - wtxn: &mut RwTxn, - index: &Index, - pool: &ThreadPool, - document_changes: PI, - ) -> Result<()> - where - PI: IntoParallelIterator> + Send, - PI::Iter: Clone, - { - let (merger_sender, writer_receiver) = merger_writer_channels(100); - let ExtractorsMergerChannels { merger_receiver, deladd_cbo_roaring_bitmap_sender } = - extractors_merger_channels(100); - - thread::scope(|s| { - thread::Builder::new().name(S("indexer-extractors")).spawn_scoped(s, || { - pool.in_place_scope(|_s| { - document_changes.into_par_iter().for_each(|_dc| ()); - }) - })?; - - // TODO manage the errors correctly - thread::Builder::new().name(S("indexer-merger")).spawn_scoped(s, || { - let rtxn = index.read_txn().unwrap(); - merge_grenad_entries(merger_receiver, merger_sender, &rtxn, index).unwrap() - })?; - - // TODO Split this code into another function - for operation in writer_receiver { - let database = operation.database(index); - match operation { - WriterOperation::WordDocids(operation) => match operation { - EntryOperation::Delete(e) => database.delete(wtxn, e.entry()).map(drop)?, - EntryOperation::Write(e) => database.put(wtxn, e.key(), e.value())?, - }, - WriterOperation::Document(e) => database.put(wtxn, &e.key(), e.content())?, - } - } - - Ok(()) - }) - } - - enum Operation { - Write(RoaringBitmap), - Delete, - Ignore, - } - - /// A function that merges the DelAdd CboRoaringBitmaps with the current bitmap. - fn merge_cbo_bitmaps( - current: Option<&[u8]>, - del: Option<&[u8]>, - add: Option<&[u8]>, - ) -> Result { - let current = current.map(CboRoaringBitmapCodec::deserialize_from).transpose()?; - let del = del.map(CboRoaringBitmapCodec::deserialize_from).transpose()?; - let add = add.map(CboRoaringBitmapCodec::deserialize_from).transpose()?; - - match (current, del, add) { - (None, None, None) => Ok(Operation::Ignore), // but it's strange - (None, None, Some(add)) => Ok(Operation::Write(add)), - (None, Some(_del), None) => Ok(Operation::Ignore), // but it's strange - (None, Some(_del), Some(add)) => Ok(Operation::Write(add)), - (Some(_current), None, None) => Ok(Operation::Ignore), // but it's strange - (Some(current), None, Some(add)) => Ok(Operation::Write(current | add)), - (Some(current), Some(del), add) => { - let output = match add { - Some(add) => (current - del) | add, - None => current - del, - }; - if output.is_empty() { - Ok(Operation::Delete) - } else { - Ok(Operation::Write(output)) - } - } - } - } - - /// Return the slice directly from the serialize_into method - fn cbo_serialize_into_vec<'b>(bitmap: &RoaringBitmap, buffer: &'b mut Vec) -> &'b [u8] { - buffer.clear(); - CboRoaringBitmapCodec::serialize_into(bitmap, buffer); - buffer.as_slice() - } - - /// TODO We must return some infos/stats - fn merge_grenad_entries( - receiver: MergerReceiver, - sender: MergerSender, - rtxn: &RoTxn, - index: &Index, - ) -> Result<()> { - let mut buffer = Vec::new(); - - for merger_operation in receiver { - match merger_operation { - MergerOperation::WordDocidsCursors(cursors) => { - let sender = sender.word_docids(); - let database = index.word_docids.remap_types::(); - - let mut builder = grenad::MergerBuilder::new(MergeDeladdCboRoaringBitmaps); - builder.extend(cursors); - /// TODO manage the error correctly - let mut merger_iter = builder.build().into_stream_merger_iter().unwrap(); - - // TODO manage the error correctly - while let Some((key, deladd)) = merger_iter.next().unwrap() { - let current = database.get(rtxn, key)?; - let deladd: &KvReaderDelAdd = deladd.into(); - let del = deladd.get(DelAdd::Deletion); - let add = deladd.get(DelAdd::Addition); - - match merge_cbo_bitmaps(current, del, add)? { - Operation::Write(bitmap) => { - let value = cbo_serialize_into_vec(&bitmap, &mut buffer); - sender.write(key, value).unwrap(); - } - Operation::Delete => sender.delete(key).unwrap(), - Operation::Ignore => (), - } - } - } - } - } - - Ok(()) - } - - /// Reads the previous version of a document from the database, the new versions - /// in the grenad update files and merges them to generate a new boxed obkv. - /// - /// This function is only meant to be used when doing an update and not a replacement. - fn merge_document_for_updates( - rtxn: &RoTxn, - index: &Index, - fields_ids_map: &FieldsIdsMap, - docid: DocumentId, - external_docid: String, - operations: &[DocumentOperation], - ) -> Result> { - let mut document = BTreeMap::<_, Cow<_>>::new(); - let current = index.documents.remap_data_type::().get(rtxn, &docid)?; - let current: Option<&KvReaderFieldId> = current.map(Into::into); - - if let Some(current) = current { - current.into_iter().for_each(|(k, v)| { - document.insert(k, v.into()); - }); - } - - let last_deletion = operations - .iter() - .rposition(|operation| matches!(operation, DocumentOperation::Deletion)); - - let operations = &operations[last_deletion.map_or(0, |i| i + 1)..]; - - if operations.is_empty() { - match current { - Some(current) => { - return Ok(Some(DocumentChange::Deletion(Deletion::create( - docid, - external_docid, - current.boxed(), - )))); - } - None => return Ok(None), - } - } - - for operation in operations { - let DocumentOffset { content, offset } = match operation { - DocumentOperation::Addition(offset) => offset, - DocumentOperation::Deletion => unreachable!("Deletion in document operations"), - }; - - let reader = DocumentsBatchReader::from_reader(Cursor::new(content.as_ref()))?; - let (mut cursor, batch_index) = reader.into_cursor_and_fields_index(); - let update = cursor.get(*offset)?.expect("must exists"); - - update.into_iter().for_each(|(k, v)| { - let field_name = batch_index.name(k).unwrap(); - let id = fields_ids_map.id(field_name).unwrap(); - document.insert(id, v.to_vec().into()); - }); - } - - let mut writer = KvWriterFieldId::memory(); - document.into_iter().for_each(|(id, value)| writer.insert(id, value).unwrap()); - let new = writer.into_boxed(); - - match current { - Some(current) => { - let update = Update::create(docid, external_docid, current.boxed(), new); - Ok(Some(DocumentChange::Update(update))) - } - None => { - let insertion = Insertion::create(docid, external_docid, new); - Ok(Some(DocumentChange::Insertion(insertion))) - } - } - } - - /// Returns only the most recent version of a document based on the updates from the payloads. - /// - /// This function is only meant to be used when doing a replacement and not an update. - fn merge_document_for_replacements( - rtxn: &RoTxn, - index: &Index, - fields_ids_map: &FieldsIdsMap, - docid: DocumentId, - external_docid: String, - operations: &[DocumentOperation], - ) -> Result> { - let current = index.documents.remap_data_type::().get(rtxn, &docid)?; - let current: Option<&KvReaderFieldId> = current.map(Into::into); - - match operations.last() { - Some(DocumentOperation::Addition(DocumentOffset { content, offset })) => { - let reader = DocumentsBatchReader::from_reader(Cursor::new(content.as_ref()))?; - let (mut cursor, batch_index) = reader.into_cursor_and_fields_index(); - let update = cursor.get(*offset)?.expect("must exists"); - - let mut document_entries = Vec::new(); - update.into_iter().for_each(|(k, v)| { - let field_name = batch_index.name(k).unwrap(); - let id = fields_ids_map.id(field_name).unwrap(); - document_entries.push((id, v)); - }); - - document_entries.sort_unstable_by_key(|(id, _)| *id); - - let mut writer = KvWriterFieldId::memory(); - document_entries - .into_iter() - .for_each(|(id, value)| writer.insert(id, value).unwrap()); - let new = writer.into_boxed(); - - match current { - Some(current) => { - let update = Update::create(docid, external_docid, current.boxed(), new); - Ok(Some(DocumentChange::Update(update))) - } - None => { - let insertion = Insertion::create(docid, external_docid, new); - Ok(Some(DocumentChange::Insertion(insertion))) - } - } - } - Some(DocumentOperation::Deletion) => match current { - Some(current) => { - let deletion = Deletion::create(docid, external_docid, current.boxed()); - Ok(Some(DocumentChange::Deletion(deletion))) - } - None => Ok(None), - }, - None => Ok(None), - } - } -} +pub type KvReaderDelAdd = obkv::KvReader; +pub type KvReaderFieldId = obkv::KvReader; +pub type KvWriterDelAdd = obkv::KvWriter; +pub type KvWriterFieldId = obkv::KvWriter;