From 02714ef5edb87f9fb371efeb2208526c2dbdb284 Mon Sep 17 00:00:00 2001
From: Louis Dureuil <louis@meilisearch.com>
Date: Tue, 14 May 2024 11:43:16 +0200
Subject: [PATCH] Add vectors from vector DB in dump

---
 index-scheduler/src/batch.rs | 54 ++++++++++++++++++++++++++++++++++--
 1 file changed, 52 insertions(+), 2 deletions(-)

diff --git a/index-scheduler/src/batch.rs b/index-scheduler/src/batch.rs
index 582497c15..40398dc37 100644
--- a/index-scheduler/src/batch.rs
+++ b/index-scheduler/src/batch.rs
@@ -31,6 +31,7 @@ use meilisearch_types::milli::heed::CompactionOption;
 use meilisearch_types::milli::update::{
     IndexDocumentsConfig, IndexDocumentsMethod, IndexerConfig, Settings as MilliSettings,
 };
+use meilisearch_types::milli::vector::parsed_vectors::RESERVED_VECTORS_FIELD_NAME;
 use meilisearch_types::milli::{self, Filter};
 use meilisearch_types::settings::{apply_settings_to_builder, Settings, Unchecked};
 use meilisearch_types::tasks::{Details, IndexSwap, Kind, KindWithContent, Status, Task};
@@ -916,8 +917,57 @@ impl IndexScheduler {
                         if self.must_stop_processing.get() {
                             return Err(Error::AbortedTask);
                         }
-                        let (_id, doc) = ret?;
-                        let document = milli::obkv_to_json(&all_fields, &fields_ids_map, doc)?;
+
+                        let (id, doc) = ret?;
+
+                        let mut document = milli::obkv_to_json(&all_fields, &fields_ids_map, doc)?;
+
+                        'inject_vectors: {
+                            let embeddings = index.embeddings(&rtxn, id)?;
+
+                            if embeddings.is_empty() {
+                                break 'inject_vectors;
+                            }
+
+                            let vectors = document
+                                .entry(RESERVED_VECTORS_FIELD_NAME.to_owned())
+                                .or_insert(serde_json::Value::Object(Default::default()));
+
+                            let serde_json::Value::Object(vectors) = vectors else {
+                                return Err(milli::Error::UserError(
+                                    milli::UserError::InvalidVectorsMapType {
+                                        document_id: {
+                                            if let Ok(Some(Ok(index))) = index
+                                                .external_id_of(&rtxn, std::iter::once(id))
+                                                .map(|it| it.into_iter().next())
+                                            {
+                                                index
+                                            } else {
+                                                format!("internal docid={id}")
+                                            }
+                                        },
+                                        value: vectors.clone(),
+                                    },
+                                )
+                                .into());
+                            };
+
+                            /// some tests to consider:
+                            ///
+                            /// - dump, then import, then change a document with autogenerated vectors
+                            for (embedder_name, embeddings) in embeddings {
+                                // don't change the entry if it already exists, because it was user-provided
+                                vectors.entry(embedder_name).or_insert_with(|| {
+
+                                        let embeddings = milli::vector::parsed_vectors::ExplicitVectors {
+                                            embeddings: milli::vector::parsed_vectors::VectorOrArrayOfVectors::from_array_of_vectors(embeddings),
+                                            user_provided: false,
+                                        };
+                                        serde_json::to_value(embeddings).unwrap()
+                                });
+                            }
+                        }
+
                         index_dumper.push_document(&document)?;
                     }