mirror of
https://github.com/meilisearch/MeiliSearch
synced 2024-11-12 16:08:55 +01:00
Move the indexers into their own modules
This commit is contained in:
parent
bb885a5810
commit
e639ec79d1
@ -3,8 +3,8 @@ use std::fs::File;
|
|||||||
use crossbeam_channel::{IntoIter, Receiver, SendError, Sender};
|
use crossbeam_channel::{IntoIter, Receiver, SendError, Sender};
|
||||||
use heed::types::Bytes;
|
use heed::types::Bytes;
|
||||||
|
|
||||||
use super::indexer::KvReaderFieldId;
|
|
||||||
use super::StdResult;
|
use super::StdResult;
|
||||||
|
use crate::update::new::KvReaderFieldId;
|
||||||
use crate::{DocumentId, Index};
|
use crate::{DocumentId, Index};
|
||||||
|
|
||||||
/// The capacity of the channel is currently in number of messages.
|
/// The capacity of the channel is currently in number of messages.
|
||||||
|
@ -1,7 +1,7 @@
|
|||||||
use heed::RoTxn;
|
use heed::RoTxn;
|
||||||
use obkv::KvReader;
|
use obkv::KvReader;
|
||||||
|
|
||||||
use super::indexer::KvReaderFieldId;
|
use crate::update::new::KvReaderFieldId;
|
||||||
use crate::{DocumentId, FieldId};
|
use crate::{DocumentId, FieldId};
|
||||||
|
|
||||||
pub enum DocumentChange {
|
pub enum DocumentChange {
|
||||||
|
53
milli/src/update/new/indexer/document_deletion.rs
Normal file
53
milli/src/update/new/indexer/document_deletion.rs
Normal file
@ -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(),
|
||||||
|
))))
|
||||||
|
})
|
||||||
|
}))
|
||||||
|
}
|
||||||
|
}
|
325
milli/src/update/new/indexer/document_operation.rs
Normal file
325
milli/src/update/new/indexer/document_operation.rs
Normal file
@ -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),
|
||||||
|
}
|
||||||
|
}
|
82
milli/src/update/new/indexer/mod.rs
Normal file
82
milli/src/update/new/indexer/mod.rs
Normal file
@ -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(())
|
||||||
|
})
|
||||||
|
}
|
73
milli/src/update/new/indexer/partial_dump.rs
Normal file
73
milli/src/update/new/indexer/partial_dump.rs
Normal file
@ -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)))
|
||||||
|
}))
|
||||||
|
}
|
||||||
|
}
|
18
milli/src/update/new/indexer/update_by_function.rs
Normal file
18
milli/src/update/new/indexer/update_by_function.rs
Normal file
@ -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())
|
||||||
|
}
|
||||||
|
}
|
97
milli/src/update/new/merger.rs
Normal file
97
milli/src/update/new/merger.rs
Normal file
@ -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()
|
||||||
|
}
|
@ -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 document_change;
|
||||||
|
mod merger;
|
||||||
// mod extract;
|
// mod extract;
|
||||||
|
// mod global_fields_ids_map;
|
||||||
mod channel;
|
mod channel;
|
||||||
|
mod indexer;
|
||||||
mod items_pool;
|
mod items_pool;
|
||||||
|
|
||||||
/// TODO remove this
|
/// TODO move them elsewhere
|
||||||
// mod global_fields_ids_map;
|
|
||||||
|
|
||||||
pub type StdResult<T, E> = std::result::Result<T, E>;
|
pub type StdResult<T, E> = std::result::Result<T, E>;
|
||||||
|
pub type KvReaderDelAdd = obkv::KvReader<DelAdd>;
|
||||||
mod indexer {
|
pub type KvReaderFieldId = obkv::KvReader<FieldId>;
|
||||||
use std::borrow::Cow;
|
pub type KvWriterDelAdd<W> = obkv::KvWriter<W, DelAdd>;
|
||||||
use std::collections::{BTreeMap, HashMap};
|
pub type KvWriterFieldId<W> = obkv::KvWriter<W, FieldId>;
|
||||||
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),
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
Loading…
Reference in New Issue
Block a user