Adapt Document change

crates/milli/src/update/new/document.rs

@@ -9,6 +9,7 @@ use super::vector_document::VectorDocument;
 use super::{KvReaderFieldId, KvWriterFieldId};
 use crate::constants::{RESERVED_GEO_FIELD_NAME, RESERVED_VECTORS_FIELD_NAME};
 use crate::documents::FieldIdMapper;
+use crate::vector::settings::EmbedderAction;
 use crate::{DocumentId, GlobalFieldsIdsMap, Index, InternalError, Result, UserError};
 
 /// A view into a document that can represent either the current version from the DB,
@@ -309,6 +310,7 @@ where
 pub fn write_to_obkv<'s, 'a, 'map, 'buffer>(
     document: &'s impl Document<'s>,
     vector_document: Option<&'s impl VectorDocument<'s>>,
+    embedder_actions: &'a BTreeMap<String, EmbedderAction>,
     fields_ids_map: &'a mut GlobalFieldsIdsMap<'map>,
     mut document_buffer: &'a mut bumpalo::collections::Vec<'buffer, u8>,
 ) -> Result<&'a KvReaderFieldId>
@@ -338,7 +340,21 @@ where
         for res in vector_document.iter_vectors() {
             let (name, entry) = res?;
             if entry.has_configured_embedder {
-                continue; // we don't write vectors with configured embedder in documents
+                if let Some(action) = embedder_actions.get(name) {
+                    if action.write_back().is_some() && !entry.regenerate {
+                        vectors.insert(
+                            name,
+                            serde_json::json!({
+                                "regenerate": entry.regenerate,
+                                // TODO: consider optimizing the shape of embedders here to store an array of f32 rather than a JSON object
+                                "embeddings": entry.embeddings,
+                            }),
+                        );
+                    }
+                }
+            } else {
+                if embedder_actions.contains_key(name) {
+                    continue;
                 }
                 vectors.insert(
                     name,
@@ -353,6 +369,7 @@ where
                     },
                 );
             }
+        }
 
         if vectors.is_empty() {
             break 'inject_vectors;

crates/milli/src/update/new/extract/documents.rs

@@ -45,6 +45,7 @@ impl<'extractor> Extractor<'extractor> for DocumentsExtractor<'_, '_> {
     ) -> Result<()> {
         let mut document_buffer = bumpalo::collections::Vec::new_in(&context.doc_alloc);
         let mut document_extractor_data = context.data.0.borrow_mut_or_yield();
+        let embedder_actions = &Default::default();
 
         for change in changes {
             let change = change?;
@@ -121,9 +122,11 @@ impl<'extractor> Extractor<'extractor> for DocumentsExtractor<'_, '_> {
                     let content = write_to_obkv(
                         &content,
                         vector_content.as_ref(),
+                        embedder_actions,
                         &mut new_fields_ids_map,
                         &mut document_buffer,
                     )?;
+
                     self.document_sender.uncompressed(docid, external_docid, content).unwrap();
                 }
                 DocumentChange::Insertion(insertion) => {
@@ -146,6 +149,7 @@ impl<'extractor> Extractor<'extractor> for DocumentsExtractor<'_, '_> {
                     let content = write_to_obkv(
                         &content,
                         inserted_vectors.as_ref(),
+                        embedder_actions,
                         &mut new_fields_ids_map,
                         &mut document_buffer,
                     )?;