Introduce the new update by function

This commit is contained in:
Clément Renault
2024-10-14 15:41:10 +02:00
parent 96658ec775
commit 7e1dc8439b
7 changed files with 336 additions and 88 deletions

View File

@ -287,7 +287,7 @@ pub trait DocumentChanges<'pl // lifetime of the underlying payload
&'doc self,
context: &'doc DocumentChangeContext<T>,
item: Self::Item,
) -> Result<DocumentChange<'doc>> where 'pl: 'doc // the payload must survive the process calls
) -> Result<Option<DocumentChange<'doc>>> where 'pl: 'doc // the payload must survive the process calls
;
}
@ -352,8 +352,11 @@ where
// Clean up and reuse the document-specific allocator
context.doc_alloc.reset();
let change =
document_changes.item_to_document_change(context, item).map_err(Arc::new)?;
let Some(change) =
document_changes.item_to_document_change(context, item).map_err(Arc::new)?
else {
return Ok(());
};
let res = extractor.process(change, context).map_err(Arc::new);

View File

@ -55,7 +55,7 @@ impl<'pl> DocumentChanges<'pl> for DocumentDeletionChanges<'pl> {
&'doc self,
context: &'doc DocumentChangeContext<T>,
docid: Self::Item,
) -> Result<DocumentChange<'doc>>
) -> Result<Option<DocumentChange<'doc>>>
where
'pl: 'doc, // the payload must survive the process calls
{
@ -69,7 +69,7 @@ impl<'pl> DocumentChanges<'pl> for DocumentDeletionChanges<'pl> {
let external_document_id = external_document_id.to_bump(&context.doc_alloc);
Ok(DocumentChange::Deletion(Deletion::create(docid, external_document_id)))
Ok(Some(DocumentChange::Deletion(Deletion::create(docid, external_document_id))))
}
}

View File

@ -240,7 +240,7 @@ impl<'pl> DocumentChanges<'pl> for DocumentOperationChanges<'pl> {
&'doc self,
context: &'doc DocumentChangeContext<T>,
item: Self::Item,
) -> Result<DocumentChange<'doc>>
) -> Result<Option<DocumentChange<'doc>>>
where
'pl: 'doc,
{
@ -275,7 +275,7 @@ trait MergeChanges {
is_new: bool,
doc_alloc: &'doc Bump,
operations: &'doc [InnerDocOp],
) -> Result<DocumentChange<'doc>>;
) -> Result<Option<DocumentChange<'doc>>>;
}
struct MergeDocumentForReplacement;
@ -301,7 +301,7 @@ impl MergeChanges for MergeDocumentForReplacement {
is_new: bool,
doc_alloc: &'doc Bump,
operations: &'doc [InnerDocOp],
) -> Result<DocumentChange<'doc>> {
) -> Result<Option<DocumentChange<'doc>>> {
match operations.last() {
Some(InnerDocOp::Addition(DocumentOffset { content })) => {
let document = serde_json::from_slice(content).unwrap();
@ -312,18 +312,27 @@ impl MergeChanges for MergeDocumentForReplacement {
let document = DocumentFromVersions::new(Versions::Single(document));
if is_new {
Ok(DocumentChange::Insertion(Insertion::create(docid, external_doc, document)))
Ok(Some(DocumentChange::Insertion(Insertion::create(
docid,
external_doc,
document,
))))
} else {
Ok(DocumentChange::Update(Update::create(docid, external_doc, document, true)))
Ok(Some(DocumentChange::Update(Update::create(
docid,
external_doc,
document,
true,
))))
}
}
Some(InnerDocOp::Deletion) => {
let deletion = if is_new {
Deletion::create(docid, external_doc)
return if is_new {
let deletion = Deletion::create(docid, external_doc);
Ok(Some(DocumentChange::Deletion(deletion)))
} else {
todo!("Do that with Louis")
Ok(None)
};
Ok(DocumentChange::Deletion(deletion))
}
None => unreachable!("We must not have empty set of operations on a document"),
}
@ -354,7 +363,7 @@ impl MergeChanges for MergeDocumentForUpdates {
is_new: bool,
doc_alloc: &'doc Bump,
operations: &'doc [InnerDocOp],
) -> Result<DocumentChange<'doc>> {
) -> Result<Option<DocumentChange<'doc>>> {
if operations.is_empty() {
unreachable!("We must not have empty set of operations on a document");
}
@ -365,13 +374,12 @@ impl MergeChanges for MergeDocumentForUpdates {
let has_deletion = last_deletion.is_some();
if operations.is_empty() {
let deletion = if !is_new {
Deletion::create(docid, external_docid)
return if !is_new {
let deletion = Deletion::create(docid, external_docid);
Ok(Some(DocumentChange::Deletion(deletion)))
} else {
todo!("Do that with Louis")
Ok(None)
};
return Ok(DocumentChange::Deletion(deletion));
}
let mut versions = bumpalo::collections::Vec::with_capacity_in(operations.len(), doc_alloc);
@ -401,14 +409,14 @@ impl MergeChanges for MergeDocumentForUpdates {
let document = DocumentFromVersions::new(versions);
if is_new {
Ok(DocumentChange::Insertion(Insertion::create(docid, external_docid, document)))
Ok(Some(DocumentChange::Insertion(Insertion::create(docid, external_docid, document))))
} else {
Ok(DocumentChange::Update(Update::create(
Ok(Some(DocumentChange::Update(Update::create(
docid,
external_docid,
document,
has_deletion,
)))
))))
}
}
}

View File

@ -53,7 +53,7 @@ where
&'doc self,
context: &'doc DocumentChangeContext<T>,
document: Self::Item,
) -> Result<DocumentChange<'doc>>
) -> Result<Option<DocumentChange<'doc>>>
where
'index: 'doc,
{
@ -80,6 +80,6 @@ where
let document = DocumentFromVersions::new(Versions::Single(document));
let insertion = Insertion::create(docid, external_document_id, document);
Ok(DocumentChange::Insertion(insertion))
Ok(Some(DocumentChange::Insertion(insertion)))
}
}

View File

@ -1,33 +1,236 @@
use rayon::iter::IntoParallelIterator;
use std::collections::BTreeMap;
use super::document_changes::{DocumentChangeContext, DocumentChanges};
use crate::Result;
use raw_collections::RawMap;
use rayon::iter::{IndexedParallelIterator, IntoParallelRefIterator, ParallelIterator};
use rhai::{Dynamic, Engine, OptimizationLevel, Scope, AST};
use roaring::RoaringBitmap;
pub struct UpdateByFunction;
use super::document_changes::{DocumentChangeContext, MostlySend};
use super::DocumentChanges;
use crate::documents::Error::InvalidDocumentFormat;
use crate::documents::PrimaryKey;
use crate::error::{FieldIdMapMissingEntry, InternalError};
use crate::update::new::document::DocumentFromVersions;
use crate::update::new::document_change::Versions;
use crate::update::new::{Deletion, DocumentChange, KvReaderFieldId, KvWriterFieldId, Update};
use crate::{all_obkv_to_json, Error, FieldsIdsMap, GlobalFieldsIdsMap, Object, Result, UserError};
pub struct UpdateByFunction {
documents: RoaringBitmap,
context: Option<Object>,
code: String,
}
pub struct UpdateByFunctionChanges<'doc> {
primary_key: &'doc PrimaryKey<'doc>,
engine: Engine,
ast: AST,
context: Option<Dynamic>,
// It is sad that the RoaringBitmap doesn't
// implement IndexedParallelIterator
documents: Vec<u32>,
}
impl UpdateByFunction {
pub fn into_changes(self) -> UpdateByFunctionChanges {
UpdateByFunctionChanges
pub fn new(documents: RoaringBitmap, context: Option<Object>, code: String) -> Self {
UpdateByFunction { documents, context, code }
}
pub fn into_changes<'index>(
self,
primary_key: &'index PrimaryKey,
) -> Result<UpdateByFunctionChanges<'index>> {
let Self { documents, context, code } = self;
// Setup the security and limits of the Engine
let mut engine = Engine::new();
engine.set_optimization_level(OptimizationLevel::Full);
engine.set_max_call_levels(1000);
// It is an arbitrary value. We need to let users define this in the settings.
engine.set_max_operations(1_000_000);
engine.set_max_variables(1000);
engine.set_max_functions(30);
engine.set_max_expr_depths(100, 1000);
engine.set_max_string_size(1024 * 1024 * 1024); // 1 GiB
engine.set_max_array_size(10_000);
engine.set_max_map_size(10_000);
let ast = engine.compile(code).map_err(UserError::DocumentEditionCompilationError)?;
let context = match context {
Some(context) => {
Some(serde_json::from_value(context.into()).map_err(InternalError::SerdeJson)?)
}
None => None,
};
Ok(UpdateByFunctionChanges {
primary_key,
engine,
ast,
context,
documents: documents.into_iter().collect(),
})
}
}
pub struct UpdateByFunctionChanges;
impl<'index> DocumentChanges<'index> for UpdateByFunctionChanges {
impl<'index> DocumentChanges<'index> for UpdateByFunctionChanges<'index> {
type Item = u32;
fn iter(&self) -> impl rayon::prelude::IndexedParallelIterator<Item = Self::Item> {
(0..100).into_par_iter()
fn iter(&self) -> impl IndexedParallelIterator<Item = Self::Item> {
self.documents.par_iter().copied()
}
fn item_to_document_change<'doc, T: super::document_changes::MostlySend + 'doc>(
fn item_to_document_change<'doc, T: MostlySend + 'doc>(
&self,
_context: &'doc DocumentChangeContext<T>,
_item: Self::Item,
) -> Result<crate::update::new::DocumentChange<'doc>>
context: &'doc DocumentChangeContext<T>,
docid: Self::Item,
) -> Result<Option<DocumentChange<'doc>>>
where
'index: 'doc,
{
todo!()
let DocumentChangeContext {
index,
db_fields_ids_map,
txn,
new_fields_ids_map,
doc_alloc,
..
} = context;
// safety: Both documents *must* exists in the database as
// their IDs comes from the list of documents ids.
let document = index.document(txn, docid)?;
let rhai_document = obkv_to_rhaimap(document, db_fields_ids_map)?;
let json_document = all_obkv_to_json(document, db_fields_ids_map)?;
let document_id = self
.primary_key
.document_id(document, db_fields_ids_map)?
.map_err(|_| InvalidDocumentFormat)?;
let mut scope = Scope::new();
if let Some(context) = self.context.as_ref().cloned() {
scope.push_constant_dynamic("context", context.clone());
}
scope.push("doc", rhai_document);
// That's were the magic happens. We run the user script
// which edits "doc" scope variable reprensenting the document
// and ignore the output and even the type of it, i.e., Dynamic.
let _ = self
.engine
.eval_ast_with_scope::<Dynamic>(&mut scope, &self.ast)
.map_err(UserError::DocumentEditionRuntimeError)?;
match scope.remove::<Dynamic>("doc") {
// If the "doc" variable has set to (), we effectively delete the document.
Some(doc) if doc.is_unit() => Ok(Some(DocumentChange::Deletion(Deletion::create(
docid,
doc_alloc.alloc_str(&document_id),
)))),
None => unreachable!("missing doc variable from the Rhai scope"),
Some(new_document) => match new_document.try_cast() {
Some(new_rhai_document) => {
let mut buffer = bumpalo::collections::Vec::new_in(doc_alloc);
serde_json::to_writer(&mut buffer, &new_rhai_document)
.map_err(InternalError::SerdeJson)?;
let raw_new_doc = serde_json::from_slice(buffer.into_bump_slice())
.map_err(InternalError::SerdeJson)?;
// Note: This condition is not perfect. Sometimes it detect changes
// like with floating points numbers and consider updating
// the document even if nothing actually changed.
//
// Future: Use a custom function rhai function to track changes.
// <https://docs.rs/rhai/latest/rhai/struct.Engine.html#method.register_indexer_set>
if json_document != rhaimap_to_object(new_rhai_document) {
let mut global_fields_ids_map = new_fields_ids_map.borrow_mut();
let new_document_id = self
.primary_key
.extract_fields_and_docid(
raw_new_doc,
&mut *global_fields_ids_map,
doc_alloc,
)?
.to_de();
if document_id != new_document_id {
Err(Error::UserError(UserError::DocumentEditionCannotModifyPrimaryKey))
} else {
let raw_new_doc = RawMap::from_raw_value(raw_new_doc, doc_alloc)
.map_err(InternalError::SerdeJson)?;
let new_doc_version = DocumentFromVersions::new(Versions::Single(
raw_new_doc.into_bump_slice(),
));
Ok(Some(DocumentChange::Update(Update::create(
docid,
new_document_id,
new_doc_version,
true, // It is like document replacement
))))
}
} else {
Ok(None)
}
}
None => Err(Error::UserError(UserError::DocumentEditionDocumentMustBeObject)),
},
}
}
}
fn obkv_to_rhaimap(obkv: &KvReaderFieldId, fields_ids_map: &FieldsIdsMap) -> Result<rhai::Map> {
let all_keys = obkv.iter().map(|(k, _v)| k).collect::<Vec<_>>();
let map: Result<rhai::Map> = all_keys
.iter()
.copied()
.flat_map(|id| obkv.get(id).map(|value| (id, value)))
.map(|(id, value)| {
let name = fields_ids_map.name(id).ok_or(FieldIdMapMissingEntry::FieldId {
field_id: id,
process: "all_obkv_to_rhaimap",
})?;
let value = serde_json::from_slice(value).map_err(InternalError::SerdeJson)?;
Ok((name.into(), value))
})
.collect();
map
}
fn rhaimap_to_object(map: rhai::Map) -> Object {
let mut output = Object::new();
for (key, value) in map {
let value = serde_json::to_value(&value).unwrap();
output.insert(key.into(), value);
}
output
}
fn rhaimap_to_obkv(
map: rhai::Map,
global_fields_ids_map: &mut GlobalFieldsIdsMap,
buffer: &mut Vec<u8>,
) -> Result<Box<KvReaderFieldId>> {
let result: Result<BTreeMap<_, _>> = map
.keys()
.map(|key| {
global_fields_ids_map
.id_or_insert(key)
.ok_or(UserError::AttributeLimitReached)
.map_err(Error::from)
.map(|fid| (fid, key))
})
.collect();
let ordered_fields = result?;
let mut writer = KvWriterFieldId::memory();
for (fid, key) in ordered_fields {
let value = map.get(key).unwrap();
let value = serde_json::to_value(value).unwrap();
buffer.clear();
serde_json::to_writer(&mut *buffer, &value).unwrap();
writer.insert(fid, &buffer)?;
}
Ok(writer.into_boxed())
}

View File

@ -1,8 +1,8 @@
use std::collections::HashSet;
use hashbrown::HashMap;
use heed::{types::Bytes, RwTxn};
use heed::{BytesDecode, Database};
use heed::types::Bytes;
use heed::{BytesDecode, Database, RwTxn};
use roaring::RoaringBitmap;
use crate::heed_codec::StrBEU16Codec;
@ -110,7 +110,7 @@ impl WordPrefixIntegerDocids {
key_buffer.extend_from_slice(prefix);
key_buffer.push(0);
key_buffer.extend_from_slice(&pos.to_be_bytes());
self.prefix_database.put(wtxn, &key_buffer, &docids)?;
self.prefix_database.put(wtxn, &key_buffer, docids)?;
}
docids.clear();
}