Function preprocess_embedding

src/vector_store.rs:629–712

Converts raw embeddings into `IndexableEmbedding` i.e. ready to be indexed - with quantization performed and property values and metadata fields converted into appropriate types. If metadata filtering is supported for the collection, then one input raw embedding may result in multiple `IndexableEmbedding` instances.

(
    collection: &Collection,
    hnsw_index: &HNSWIndex,
    quantization_metric: &RwLock<QuantizationMetric>,
    raw_emb: &RawDenseVectorEmbedding,
)

/// input raw embedding may result in multiple `IndexableEmbedding`
/// instances.
fn preprocess_embedding(
    collection: &Collection,
    hnsw_index: &HNSWIndex,
    quantization_metric: &RwLock<QuantizationMetric>,
    raw_emb: &RawDenseVectorEmbedding,
) -> Result<Vec<IndexableEmbedding>, WaCustomError> {
    let quantization = quantization_metric.read().unwrap();
    let quantized_vec = Arc::new(quantization.quantize(
        &raw_emb.raw_vec,
        *hnsw_index.storage_type.read().unwrap(),
        *hnsw_index.values_range.read().unwrap(),
    )?);

    let base_id = raw_emb.hash_vec;

    let metadata_schema = collection.meta.metadata_schema.as_ref();
    let prop_file = &hnsw_index.cache.prop_file;

    let embeddings = if raw_emb.is_pseudo {
        let replicas = pseudo_metadata_replicas(metadata_schema.unwrap(), prop_file, &base_id)?;
        let num_levels = hnsw_index.levels_prob.len() - 1;
        let plp = pseudo_level_probs(num_levels as u8, replicas.len() as u16);

        // Find the pseudo root node so that we can reuse it's
        // prop_value in rest of the pseudo nodes
        let pseudo_root = hnsw_index.pseudo_root_vec.unwrap();
        let pseudo_root_lazy = unsafe { &*pseudo_root }.latest;
        let pseudo_root_node = unsafe { &*pseudo_root_lazy }.try_get_data(&hnsw_index.cache)?;

        let mut embeddings: Vec<IndexableEmbedding> = vec![];
        for prop_metadata in replicas.into_iter() {
            let emb = IndexableEmbedding {
                prop_value: pseudo_root_node.prop_value.clone(),
                prop_metadata: Some(Arc::new(prop_metadata)),
                overridden_level_probs: Some(plp.clone()),
            };
            embeddings.push(emb);
        }
        embeddings
    } else {
        // Write props to the prop file
        let mut prop_file_guard = hnsw_index.cache.prop_file.write().unwrap();
        let location = write_prop_value_to_file(&base_id, &quantized_vec, &mut prop_file_guard)
            .expect("failed to write prop");
        drop(prop_file_guard);

        let prop_value = Arc::new(NodePropValue {
            id: base_id,
            vec: quantized_vec.clone(),
            location,
        });

        let metadata_replicas = prop_metadata_replicas(
            collection.meta.metadata_schema.as_ref(),
            raw_emb.raw_metadata.as_ref(),
            &hnsw_index.cache.prop_file,
            &base_id,
        )?;