Move the indexers into their own modules

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.

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 {

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<impl ParallelIterator<Item = Result<Option<DocumentChange>>> + '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(),
+                ))))
+            })
+        }))
+    }
+}

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<Payload>,
+    pub(crate) index_documents_method: IndexDocumentsMethod,
+}
+
+pub enum Payload {
+    Addition(File),
+    Deletion(Vec<String>),
+}
+
+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<Mmap>, // 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<PayloadStats> {
+        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<String>) {
+        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<impl ParallelIterator<Item = Result<Option<DocumentChange>>> + '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::<String, _>::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<Option<DocumentChange>> {
+    let mut document = BTreeMap::<_, Cow<_>>::new();
+    let current = index.documents.remap_data_type::<Bytes>().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<Option<DocumentChange>> {
+    let current = index.documents.remap_data_type::<Bytes>().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),
+    }
+}

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<impl ParallelIterator<Item = Result<Option<DocumentChange>>> + '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<PI>(
+    wtxn: &mut RwTxn,
+    index: &Index,
+    pool: &ThreadPool,
+    document_changes: PI,
+) -> Result<()>
+where
+    PI: IntoParallelIterator<Item = Result<DocumentChange>> + 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(())
+    })
+}

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<I> {
+    pub iter: I,
+}
+
+impl<I> PartialDumpIndexer<I> {
+    pub fn new_from_jsonlines(iter: I) -> Self {
+        PartialDumpIndexer { iter }
+    }
+}
+
+impl<'p, I> Indexer<'p> for PartialDumpIndexer<I>
+where
+    I: IntoIterator<Item = Object>,
+    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<impl ParallelIterator<Item = Result<Option<DocumentChange>>> + '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)))
+        }))
+    }
+}

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<impl ParallelIterator<Item = Result<Option<DocumentChange>>> + 'p> {
+        Ok(vec![].into_par_iter())
+    }
+}

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::<Bytes, Bytes>();
+
+                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<Operation> {
+    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<u8>) -> &'b [u8] {
+    buffer.clear();
+    CboRoaringBitmapCodec::serialize_into(bitmap, buffer);
+    buffer.as_slice()
+}

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<T, E> = std::result::Result<T, E>;
-
-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<FieldId>;
-    pub type KvReaderDelAdd = obkv::KvReader<DelAdd>;
-    pub type KvWriterFieldId<W> = obkv::KvWriter<W, FieldId>;
-    pub type KvWriterDelAdd<W> = obkv::KvWriter<W, DelAdd>;
-
-    pub struct DocumentOperationIndexer {
-        operations: Vec<Payload>,
-        index_documents_method: IndexDocumentsMethod,
-    }
-
-    enum Payload {
-        Addition(File),
-        Deletion(Vec<String>),
-    }
-
-    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<Mmap>, // 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<PayloadStats> {
-            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<String>) {
-            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<impl ParallelIterator<Item = Result<Option<DocumentChange>>> + 'a> {
-            let documents_ids = index.documents_ids(rtxn)?;
-            let mut available_docids = AvailableIds::new(&documents_ids);
-            let mut docids_version_offsets = HashMap::<String, _>::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<impl ParallelIterator<Item = Result<DocumentChange>> + '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<I> {
-        iter: I,
-    }
-
-    impl<I> PartialDumpIndexer<I>
-    where
-        I: IntoIterator<Item = Object>,
-        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<Item = Result<Option<DocumentChange>>> + '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<PI>(
-        wtxn: &mut RwTxn,
-        index: &Index,
-        pool: &ThreadPool,
-        document_changes: PI,
-    ) -> Result<()>
-    where
-        PI: IntoParallelIterator<Item = Result<DocumentChange>> + 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<Operation> {
-        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<u8>) -> &'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::<Bytes, Bytes>();
-
-                    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<Option<DocumentChange>> {
-        let mut document = BTreeMap::<_, Cow<_>>::new();
-        let current = index.documents.remap_data_type::<Bytes>().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<Option<DocumentChange>> {
-        let current = index.documents.remap_data_type::<Bytes>().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<DelAdd>;
+pub type KvReaderFieldId = obkv::KvReader<FieldId>;
+pub type KvWriterDelAdd<W> = obkv::KvWriter<W, DelAdd>;
+pub type KvWriterFieldId<W> = obkv::KvWriter<W, FieldId>;