haematite-cli 0.21.7

//! Whole-list supplied vector preparation without publication or field changes.

use std::collections::{BTreeMap, BTreeSet};

use haematite::db::ReadOnlyDatabase;

use super::embedded::Field;
use super::mutations::PreparedWrite;
use super::{Captured, Failure, reads, vectors};
use crate::serve::protocol::VectorImport;

pub(super) fn prepare(
    reader: &ReadOnlyDatabase,
    at: &Captured,
    import: &VectorImport,
) -> Result<Vec<PreparedWrite>, Failure> {
    if import.vectors.is_empty() {
        return Err(Failure {
            path: Some("/vectors".to_owned()),
            ..Failure::new(
                "import_empty",
                "a vector import must contain at least one vector",
            )
        });
    }
    let mut fields = BTreeMap::new();
    let mut identities = BTreeSet::new();
    let mut writes = Vec::with_capacity(import.vectors.len());
    for (index, vector) in import.vectors.iter().enumerate() {
        let in_row = |part: &str, failure: Failure| Failure {
            path: Some(format!("/vectors/{index}{part}")),
            ..failure
        };
        for (part, value) in [
            ("/collection", &vector.collection),
            ("/field", &vector.field),
            ("/document_id", &vector.document_id),
        ] {
            super::identifier(value).map_err(|failure| in_row(part, failure))?;
        }
        vectors::finite_nonempty(&vector.vector).map_err(|failure| in_row("/vector", failure))?;
        let key = vectors::vector_key(&vector.collection, &vector.field, &vector.document_id);
        if !identities.insert(key.clone()) {
            return Err(in_row(
                "",
                Failure::new(
                    "import_duplicate",
                    "repeated collection, document and field in vector import",
                ),
            ));
        }
        if !reads::document_exists_for_constraint(
            reader,
            at,
            &vector.collection,
            &vector.document_id,
        )
        .map_err(|failure| in_row("/document_id", failure))?
        {
            return Err(in_row(
                "/document_id",
                Failure::new("document_missing", &vector.document_id),
            ));
        }
        let field = (&vector.collection, &vector.field);
        let dimensions = if let Some(dimensions) = fields.get(&field) {
            *dimensions
        } else {
            let dimensions = supplied_dimensions(reader, at, field.0, field.1)
                .map_err(|failure| in_row("/field", failure))?;
            fields.insert(field, dimensions);
            dimensions
        };
        if dimensions != vector.vector.len() {
            return Err(in_row(
                "/vector",
                Failure::new(
                    "vector_dimension",
                    "vector dimension differs from the field",
                ),
            ));
        }
        writes.push((
            key,
            Some(
                vector
                    .vector
                    .iter()
                    .flat_map(|member| member.to_le_bytes())
                    .collect(),
            ),
        ));
    }
    Ok(writes)
}

fn supplied_dimensions(
    reader: &ReadOnlyDatabase,
    at: &Captured,
    collection: &str,
    field: &str,
) -> Result<usize, Failure> {
    match vectors::recorded_field(reader, at, collection, field)? {
        Some(Field::Supplied { dimensions }) => Ok(dimensions),
        Some(Field::Embedded(_)) => Err(Failure::new(
            "vector_field_embedded",
            "an embedded field does not accept supplied vectors",
        )),
        None => Err(Failure::new(
            "vector_dimension_missing",
            "a vector import requires an existing supplied field",
        )),
    }
}