Merge branch 'main' into change-proximity-precision-settings

milli/src/update/index_documents/extract/extract_vector_points.rs

@@ -1,9 +1,10 @@
 use std::cmp::Ordering;
-use std::convert::TryFrom;
+use std::convert::{TryFrom, TryInto};
 use std::fs::File;
 use std::io::{self, BufReader, BufWriter};
 use std::mem::size_of;
 use std::str::from_utf8;
+use std::sync::Arc;
 
 use bytemuck::cast_slice;
 use grenad::Writer;
@@ -13,13 +14,56 @@ use serde_json::{from_slice, Value};
 
 use super::helpers::{create_writer, writer_into_reader, GrenadParameters};
 use crate::error::UserError;
+use crate::prompt::Prompt;
 use crate::update::del_add::{DelAdd, KvReaderDelAdd, KvWriterDelAdd};
 use crate::update::index_documents::helpers::try_split_at;
-use crate::{DocumentId, FieldId, InternalError, Result, VectorOrArrayOfVectors};
+use crate::vector::Embedder;
+use crate::{DocumentId, FieldsIdsMap, InternalError, Result, VectorOrArrayOfVectors};
 
 /// The length of the elements that are always in the buffer when inserting new values.
 const TRUNCATE_SIZE: usize = size_of::<DocumentId>();
 
+pub struct ExtractedVectorPoints {
+    // docid, _index -> KvWriterDelAdd -> Vector
+    pub manual_vectors: grenad::Reader<BufReader<File>>,
+    // docid -> ()
+    pub remove_vectors: grenad::Reader<BufReader<File>>,
+    // docid -> prompt
+    pub prompts: grenad::Reader<BufReader<File>>,
+}
+
+enum VectorStateDelta {
+    NoChange,
+    // Remove all vectors, generated or manual, from this document
+    NowRemoved,
+
+    // Add the manually specified vectors, passed in the other grenad
+    // Remove any previously generated vectors
+    // Note: changing the value of the manually specified vector **should not record** this delta
+    WasGeneratedNowManual(Vec<Vec<f32>>),
+
+    ManualDelta(Vec<Vec<f32>>, Vec<Vec<f32>>),
+
+    // Add the vector computed from the specified prompt
+    // Remove any previous vector
+    // Note: changing the value of the prompt **does require** recording this delta
+    NowGenerated(String),
+}
+
+impl VectorStateDelta {
+    fn into_values(self) -> (bool, String, (Vec<Vec<f32>>, Vec<Vec<f32>>)) {
+        match self {
+            VectorStateDelta::NoChange => Default::default(),
+            VectorStateDelta::NowRemoved => (true, Default::default(), Default::default()),
+            VectorStateDelta::WasGeneratedNowManual(add) => {
+                (true, Default::default(), (Default::default(), add))
+            }
+            VectorStateDelta::ManualDelta(del, add) => (false, Default::default(), (del, add)),
+            VectorStateDelta::NowGenerated(prompt) => (true, prompt, Default::default()),
+        }
+    }
+}
+
 /// Extracts the embedding vector contained in each document under the `_vectors` field.
 ///
 /// Returns the generated grenad reader containing the docid as key associated to the Vec<f32>
@@ -27,16 +71,35 @@ const TRUNCATE_SIZE: usize = size_of::<DocumentId>();
 pub fn extract_vector_points<R: io::Read + io::Seek>(
     obkv_documents: grenad::Reader<R>,
     indexer: GrenadParameters,
-    vectors_fid: FieldId,
-) -> Result<grenad::Reader<BufReader<File>>> {
+    field_id_map: &FieldsIdsMap,
+    prompt: &Prompt,
+    embedder_name: &str,
+) -> Result<ExtractedVectorPoints> {
     puffin::profile_function!();
 
-    let mut writer = create_writer(
+    // (docid, _index) -> KvWriterDelAdd -> Vector
+    let mut manual_vectors_writer = create_writer(
         indexer.chunk_compression_type,
         indexer.chunk_compression_level,
         tempfile::tempfile()?,
     );
 
+    // (docid) -> (prompt)
+    let mut prompts_writer = create_writer(
+        indexer.chunk_compression_type,
+        indexer.chunk_compression_level,
+        tempfile::tempfile()?,
+    );
+
+    // (docid) -> ()
+    let mut remove_vectors_writer = create_writer(
+        indexer.chunk_compression_type,
+        indexer.chunk_compression_level,
+        tempfile::tempfile()?,
+    );
+
+    let vectors_fid = field_id_map.id("_vectors");
+
     let mut key_buffer = Vec::new();
     let mut cursor = obkv_documents.into_cursor()?;
     while let Some((key, value)) = cursor.move_on_next()? {
@@ -53,43 +116,157 @@ pub fn extract_vector_points<R: io::Read + io::Seek>(
         // lazily get it when needed
         let document_id = || -> Value { from_utf8(external_id_bytes).unwrap().into() };
 
-        // first we retrieve the _vectors field
-        if let Some(value) = obkv.get(vectors_fid) {
-            let vectors_obkv = KvReaderDelAdd::new(value);
+        let vectors_field = vectors_fid
+            .and_then(|vectors_fid| obkv.get(vectors_fid))
+            .map(KvReaderDelAdd::new)
+            .map(|obkv| to_vector_maps(obkv, document_id))
+            .transpose()?;
 
-            // then we extract the values
-            let del_vectors = vectors_obkv
-                .get(DelAdd::Deletion)
-                .map(|vectors| extract_vectors(vectors, document_id))
-                .transpose()?
-                .flatten();
-            let add_vectors = vectors_obkv
-                .get(DelAdd::Addition)
-                .map(|vectors| extract_vectors(vectors, document_id))
-                .transpose()?
-                .flatten();
+        let (del_map, add_map) = vectors_field.unzip();
+        let del_map = del_map.flatten();
+        let add_map = add_map.flatten();
 
-            // and we finally push the unique vectors into the writer
-            push_vectors_diff(
-                &mut writer,
-                &mut key_buffer,
-                del_vectors.unwrap_or_default(),
-                add_vectors.unwrap_or_default(),
-            )?;
-        }
+        let del_value = del_map.and_then(|mut map| map.remove(embedder_name));
+        let add_value = add_map.and_then(|mut map| map.remove(embedder_name));
+
+        let delta = match (del_value, add_value) {
+            (Some(old), Some(new)) => {
+                // no autogeneration
+                let del_vectors = extract_vectors(old, document_id, embedder_name)?;
+                let add_vectors = extract_vectors(new, document_id, embedder_name)?;
+
+                if add_vectors.len() > u8::MAX.into() {
+                    return Err(crate::Error::UserError(crate::UserError::TooManyVectors(
+                        document_id().to_string(),
+                        add_vectors.len(),
+                    )));
+                }
+
+                VectorStateDelta::ManualDelta(del_vectors, add_vectors)
+            }
+            (Some(_old), None) => {
+                // Do we keep this document?
+                let document_is_kept = obkv
+                    .iter()
+                    .map(|(_, deladd)| KvReaderDelAdd::new(deladd))
+                    .any(|deladd| deladd.get(DelAdd::Addition).is_some());
+                if document_is_kept {
+                    // becomes autogenerated
+                    VectorStateDelta::NowGenerated(prompt.render(
+                        obkv,
+                        DelAdd::Addition,
+                        field_id_map,
+                    )?)
+                } else {
+                    VectorStateDelta::NowRemoved
+                }
+            }
+            (None, Some(new)) => {
+                // was possibly autogenerated, remove all vectors for that document
+                let add_vectors = extract_vectors(new, document_id, embedder_name)?;
+                if add_vectors.len() > u8::MAX.into() {
+                    return Err(crate::Error::UserError(crate::UserError::TooManyVectors(
+                        document_id().to_string(),
+                        add_vectors.len(),
+                    )));
+                }
+
+                VectorStateDelta::WasGeneratedNowManual(add_vectors)
+            }
+            (None, None) => {
+                // Do we keep this document?
+                let document_is_kept = obkv
+                    .iter()
+                    .map(|(_, deladd)| KvReaderDelAdd::new(deladd))
+                    .any(|deladd| deladd.get(DelAdd::Addition).is_some());
+
+                if document_is_kept {
+                    // Don't give up if the old prompt was failing
+                    let old_prompt =
+                        prompt.render(obkv, DelAdd::Deletion, field_id_map).unwrap_or_default();
+                    let new_prompt = prompt.render(obkv, DelAdd::Addition, field_id_map)?;
+                    if old_prompt != new_prompt {
+                        log::trace!(
+                            "🚀 Changing prompt from\n{old_prompt}\n===to===\n{new_prompt}"
+                        );
+                        VectorStateDelta::NowGenerated(new_prompt)
+                    } else {
+                        log::trace!("⏭️ Prompt unmodified, skipping");
+                        VectorStateDelta::NoChange
+                    }
+                } else {
+                    VectorStateDelta::NowRemoved
+                }
+            }
+        };
+
+        // and we finally push the unique vectors into the writer
+        push_vectors_diff(
+            &mut remove_vectors_writer,
+            &mut prompts_writer,
+            &mut manual_vectors_writer,
+            &mut key_buffer,
+            delta,
+        )?;
     }
 
-    writer_into_reader(writer)
+    Ok(ExtractedVectorPoints {
+        // docid, _index -> KvWriterDelAdd -> Vector
+        manual_vectors: writer_into_reader(manual_vectors_writer)?,
+        // docid -> ()
+        remove_vectors: writer_into_reader(remove_vectors_writer)?,
+        // docid -> prompt
+        prompts: writer_into_reader(prompts_writer)?,
+    })
+}
+
+fn to_vector_maps(
+    obkv: KvReaderDelAdd,
+    document_id: impl Fn() -> Value,
+) -> Result<(Option<serde_json::Map<String, Value>>, Option<serde_json::Map<String, Value>>)> {
+    let del = to_vector_map(obkv, DelAdd::Deletion, &document_id)?;
+    let add = to_vector_map(obkv, DelAdd::Addition, &document_id)?;
+    Ok((del, add))
+}
+
+fn to_vector_map(
+    obkv: KvReaderDelAdd,
+    side: DelAdd,
+    document_id: &impl Fn() -> Value,
+) -> Result<Option<serde_json::Map<String, Value>>> {
+    Ok(if let Some(value) = obkv.get(side) {
+        let Ok(value) = from_slice(value) else {
+            let value = from_slice(value).map_err(InternalError::SerdeJson)?;
+            return Err(crate::Error::UserError(UserError::InvalidVectorsMapType {
+                document_id: document_id(),
+                value,
+            }));
+        };
+        Some(value)
+    } else {
+        None
+    })
 }
 
 /// Computes the diff between both Del and Add numbers and
 /// only inserts the parts that differ in the sorter.
 fn push_vectors_diff(
-    writer: &mut Writer<BufWriter<File>>,
+    remove_vectors_writer: &mut Writer<BufWriter<File>>,
+    prompts_writer: &mut Writer<BufWriter<File>>,
+    manual_vectors_writer: &mut Writer<BufWriter<File>>,
     key_buffer: &mut Vec<u8>,
-    mut del_vectors: Vec<Vec<f32>>,
-    mut add_vectors: Vec<Vec<f32>>,
+    delta: VectorStateDelta,
 ) -> Result<()> {
+    let (must_remove, prompt, (mut del_vectors, mut add_vectors)) = delta.into_values();
+    if must_remove {
+        key_buffer.truncate(TRUNCATE_SIZE);
+        remove_vectors_writer.insert(&key_buffer, [])?;
+    }
+    if !prompt.is_empty() {
+        key_buffer.truncate(TRUNCATE_SIZE);
+        prompts_writer.insert(&key_buffer, prompt.as_bytes())?;
+    }
+
     // We sort and dedup the vectors
     del_vectors.sort_unstable_by(|a, b| compare_vectors(a, b));
     add_vectors.sort_unstable_by(|a, b| compare_vectors(a, b));
@@ -114,7 +291,7 @@ fn push_vectors_diff(
                 let mut obkv = KvWriterDelAdd::memory();
                 obkv.insert(DelAdd::Deletion, cast_slice(&vector))?;
                 let bytes = obkv.into_inner()?;
-                writer.insert(&key_buffer, bytes)?;
+                manual_vectors_writer.insert(&key_buffer, bytes)?;
             }
             EitherOrBoth::Right(vector) => {
                 // We insert only the Add part of the Obkv to inform
@@ -122,7 +299,7 @@ fn push_vectors_diff(
                 let mut obkv = KvWriterDelAdd::memory();
                 obkv.insert(DelAdd::Addition, cast_slice(&vector))?;
                 let bytes = obkv.into_inner()?;
-                writer.insert(&key_buffer, bytes)?;
+                manual_vectors_writer.insert(&key_buffer, bytes)?;
             }
         }
     }
@@ -136,13 +313,112 @@ fn compare_vectors(a: &[f32], b: &[f32]) -> Ordering {
 }
 
 /// Extracts the vectors from a JSON value.
-fn extract_vectors(value: &[u8], document_id: impl Fn() -> Value) -> Result<Option<Vec<Vec<f32>>>> {
-    match from_slice(value) {
-        Ok(vectors) => Ok(VectorOrArrayOfVectors::into_array_of_vectors(vectors)),
+fn extract_vectors(
+    value: Value,
+    document_id: impl Fn() -> Value,
+    name: &str,
+) -> Result<Vec<Vec<f32>>> {
+    // FIXME: ugly clone of the vectors here
+    match serde_json::from_value(value.clone()) {
+        Ok(vectors) => {
+            Ok(VectorOrArrayOfVectors::into_array_of_vectors(vectors).unwrap_or_default())
+        }
         Err(_) => Err(UserError::InvalidVectorsType {
             document_id: document_id(),
-            value: from_slice(value).map_err(InternalError::SerdeJson)?,
+            value,
+            subfield: name.to_owned(),
         }
         .into()),
     }
 }
+
+#[logging_timer::time]
+pub fn extract_embeddings<R: io::Read + io::Seek>(
+    // docid, prompt
+    prompt_reader: grenad::Reader<R>,
+    indexer: GrenadParameters,
+    embedder: Arc<Embedder>,
+) -> Result<grenad::Reader<BufReader<File>>> {
+    let rt = tokio::runtime::Builder::new_current_thread().enable_io().enable_time().build()?;
+
+    let n_chunks = embedder.chunk_count_hint(); // chunk level parellelism
+    let n_vectors_per_chunk = embedder.prompt_count_in_chunk_hint(); // number of vectors in a single chunk
+
+    // docid, state with embedding
+    let mut state_writer = create_writer(
+        indexer.chunk_compression_type,
+        indexer.chunk_compression_level,
+        tempfile::tempfile()?,
+    );
+
+    let mut chunks = Vec::with_capacity(n_chunks);
+    let mut current_chunk = Vec::with_capacity(n_vectors_per_chunk);
+    let mut current_chunk_ids = Vec::with_capacity(n_vectors_per_chunk);
+    let mut chunks_ids = Vec::with_capacity(n_chunks);
+    let mut cursor = prompt_reader.into_cursor()?;
+
+    while let Some((key, value)) = cursor.move_on_next()? {
+        let docid = key.try_into().map(DocumentId::from_be_bytes).unwrap();
+        // SAFETY: precondition, the grenad value was saved from a string
+        let prompt = unsafe { std::str::from_utf8_unchecked(value) };
+        if current_chunk.len() == current_chunk.capacity() {
+            chunks.push(std::mem::replace(
+                &mut current_chunk,
+                Vec::with_capacity(n_vectors_per_chunk),
+            ));
+            chunks_ids.push(std::mem::replace(
+                &mut current_chunk_ids,
+                Vec::with_capacity(n_vectors_per_chunk),
+            ));
+        };
+        current_chunk.push(prompt.to_owned());
+        current_chunk_ids.push(docid);
+
+        if chunks.len() == chunks.capacity() {
+            let chunked_embeds = rt
+                .block_on(
+                    embedder
+                        .embed_chunks(std::mem::replace(&mut chunks, Vec::with_capacity(n_chunks))),
+                )
+                .map_err(crate::vector::Error::from)
+                .map_err(crate::Error::from)?;
+
+            for (docid, embeddings) in chunks_ids
+                .iter()
+                .flat_map(|docids| docids.iter())
+                .zip(chunked_embeds.iter().flat_map(|embeds| embeds.iter()))
+            {
+                state_writer.insert(docid.to_be_bytes(), cast_slice(embeddings.as_inner()))?;
+            }
+            chunks_ids.clear();
+        }
+    }
+
+    // send last chunk
+    if !chunks.is_empty() {
+        let chunked_embeds = rt
+            .block_on(embedder.embed_chunks(std::mem::take(&mut chunks)))
+            .map_err(crate::vector::Error::from)
+            .map_err(crate::Error::from)?;
+        for (docid, embeddings) in chunks_ids
+            .iter()
+            .flat_map(|docids| docids.iter())
+            .zip(chunked_embeds.iter().flat_map(|embeds| embeds.iter()))
+        {
+            state_writer.insert(docid.to_be_bytes(), cast_slice(embeddings.as_inner()))?;
+        }
+    }
+
+    if !current_chunk.is_empty() {
+        let embeds = rt
+            .block_on(embedder.embed(std::mem::take(&mut current_chunk)))
+            .map_err(crate::vector::Error::from)
+            .map_err(crate::Error::from)?;
+
+        for (docid, embeddings) in current_chunk_ids.iter().zip(embeds.iter()) {
+            state_writer.insert(docid.to_be_bytes(), cast_slice(embeddings.as_inner()))?;
+        }
+    }
+
+    writer_into_reader(state_writer)
+}

milli/src/update/index_documents/extract/mod.rs

@@ -23,7 +23,9 @@ use self::extract_facet_string_docids::extract_facet_string_docids;
 use self::extract_fid_docid_facet_values::{extract_fid_docid_facet_values, ExtractedFacetValues};
 use self::extract_fid_word_count_docids::extract_fid_word_count_docids;
 use self::extract_geo_points::extract_geo_points;
-use self::extract_vector_points::extract_vector_points;
+use self::extract_vector_points::{
+    extract_embeddings, extract_vector_points, ExtractedVectorPoints,
+};
 use self::extract_word_docids::extract_word_docids;
 use self::extract_word_pair_proximity_docids::extract_word_pair_proximity_docids;
 use self::extract_word_position_docids::extract_word_position_docids;
@@ -33,7 +35,8 @@ use super::helpers::{
 };
 use super::{helpers, TypedChunk};
 use crate::proximity::ProximityPrecision;
-use crate::{FieldId, Result};
+use crate::vector::EmbeddingConfigs;
+use crate::{FieldId, FieldsIdsMap, Result};
 
 /// Extract data for each databases from obkv documents in parallel.
 /// Send data in grenad file over provided Sender.
@@ -47,13 +50,14 @@ pub(crate) fn data_from_obkv_documents(
     faceted_fields: HashSet<FieldId>,
     primary_key_id: FieldId,
     geo_fields_ids: Option<(FieldId, FieldId)>,
-    vectors_field_id: Option<FieldId>,
+    field_id_map: FieldsIdsMap,
     stop_words: Option<fst::Set<&[u8]>>,
     allowed_separators: Option<&[&str]>,
     dictionary: Option<&[&str]>,
     max_positions_per_attributes: Option<u32>,
     exact_attributes: HashSet<FieldId>,
     proximity_precision: ProximityPrecision,
+    embedders: EmbeddingConfigs,
 ) -> Result<()> {
     puffin::profile_function!();
 
@@ -64,7 +68,8 @@ pub(crate) fn data_from_obkv_documents(
                 original_documents_chunk,
                 indexer,
                 lmdb_writer_sx.clone(),
-                vectors_field_id,
+                field_id_map.clone(),
+                embedders.clone(),
             )
         })
         .collect::<Result<()>>()?;
@@ -276,24 +281,53 @@ fn send_original_documents_data(
     original_documents_chunk: Result<grenad::Reader<BufReader<File>>>,
     indexer: GrenadParameters,
     lmdb_writer_sx: Sender<Result<TypedChunk>>,
-    vectors_field_id: Option<FieldId>,
+    field_id_map: FieldsIdsMap,
+    embedders: EmbeddingConfigs,
 ) -> Result<()> {
     let original_documents_chunk =
         original_documents_chunk.and_then(|c| unsafe { as_cloneable_grenad(&c) })?;
 
-    if let Some(vectors_field_id) = vectors_field_id {
-        let documents_chunk_cloned = original_documents_chunk.clone();
-        let lmdb_writer_sx_cloned = lmdb_writer_sx.clone();
-        rayon::spawn(move || {
-            let result = extract_vector_points(documents_chunk_cloned, indexer, vectors_field_id);
-            let _ = match result {
-                Ok(vector_points) => {
-                    lmdb_writer_sx_cloned.send(Ok(TypedChunk::VectorPoints(vector_points)))
+    let documents_chunk_cloned = original_documents_chunk.clone();
+    let lmdb_writer_sx_cloned = lmdb_writer_sx.clone();
+    rayon::spawn(move || {
+        for (name, (embedder, prompt)) in embedders {
+            let result = extract_vector_points(
+                documents_chunk_cloned.clone(),
+                indexer,
+                &field_id_map,
+                &prompt,
+                &name,
+            );
+            match result {
+                Ok(ExtractedVectorPoints { manual_vectors, remove_vectors, prompts }) => {
+                    let embeddings = match extract_embeddings(prompts, indexer, embedder.clone()) {
+                        Ok(results) => Some(results),
+                        Err(error) => {
+                            let _ = lmdb_writer_sx_cloned.send(Err(error));
+                            None
+                        }
+                    };
+
+                    if !(remove_vectors.is_empty()
+                        && manual_vectors.is_empty()
+                        && embeddings.as_ref().map_or(true, |e| e.is_empty()))
+                    {
+                        let _ = lmdb_writer_sx_cloned.send(Ok(TypedChunk::VectorPoints {
+                            remove_vectors,
+                            embeddings,
+                            expected_dimension: embedder.dimensions(),
+                            manual_vectors,
+                            embedder_name: name,
+                        }));
+                    }
                 }
-                Err(error) => lmdb_writer_sx_cloned.send(Err(error)),
-            };
-        });
-    }
+
+                Err(error) => {
+                    let _ = lmdb_writer_sx_cloned.send(Err(error));
+                }
+            }
+        }
+    });
 
     // TODO: create a custom internal error
     lmdb_writer_sx.send(Ok(TypedChunk::Documents(original_documents_chunk))).unwrap();