nodedb-lite 0.1.0

// SPDX-License-Identifier: Apache-2.0
//! The vector-primary insert family: `DirectInsert`, `DirectInsertIfAbsent`,
//! `DirectUpsert`.

use std::sync::Arc;

use nodedb_physical::physical_plan::{UpdateValue, VectorDirectWriteIntent};
use nodedb_types::collection_config::VectorPrimaryConfig;
use nodedb_types::result::QueryResult;
use nodedb_types::vector_dtype::VectorStorageDtype;
use nodedb_types::{Surrogate, VectorQuantization};

use crate::error::LiteError;
use crate::nodedb::LockExt;
use crate::query::engine::LiteQueryEngine;
use crate::query::text_index::reindex_documents;
use crate::storage::engine::StorageEngine;

use super::super::adapter::LitePhysicalFut;
use super::common::{
    apply_patch, decode_payload, doc_id_for, index_key, insert_node, read_row, write_row,
};

/// One row of the insert family, as the three ops carry it.
pub(in crate::query::physical_visitor) struct DirectWriteArgs {
    pub collection: String,
    pub field: String,
    pub surrogate: Surrogate,
    pub pk_bytes: Vec<u8>,
    pub vector: Vec<f32>,
    pub payload: Vec<u8>,
    pub quantization: VectorQuantization,
    pub storage_dtype: VectorStorageDtype,
    pub intent: VectorDirectWriteIntent,
    /// `ON CONFLICT (pk) DO UPDATE SET` assignments; only `Upsert` carries
    /// any. Empty means whole-row replace.
    pub on_conflict_updates: Vec<(String, UpdateValue)>,
}

/// Write one vector-primary row by `intent`.
///
/// `payload_indexes` have no Lite bitmap index: every payload predicate is
/// evaluated by brute force against the stored row, so the row itself is
/// the whole index.
pub(in crate::query::physical_visitor) fn vector_direct_write<'a, S>(
    engine: &'a LiteQueryEngine<S>,
    args: DirectWriteArgs,
) -> Result<LitePhysicalFut<'a>, LiteError>
where
    S: StorageEngine + 'a,
{
    vector_direct_write_coordinated(engine, None, args)
}

pub(in crate::query::physical_visitor) fn vector_direct_write_coordinated<'a, S>(
    engine: &'a LiteQueryEngine<S>,
    permit: Option<&'a crate::engine::fts::coordinator::TextMutationPermit>,
    args: DirectWriteArgs,
) -> Result<LitePhysicalFut<'a>, LiteError>
where
    S: StorageEngine + 'a,
{
    let DirectWriteArgs {
        collection,
        field,
        surrogate,
        pk_bytes,
        vector,
        payload,
        quantization,
        storage_dtype,
        intent,
        on_conflict_updates,
    } = args;
    let op_name = match intent {
        VectorDirectWriteIntent::Insert => "DirectInsert",
        VectorDirectWriteIntent::InsertIfAbsent => "DirectInsertIfAbsent",
        VectorDirectWriteIntent::Upsert => "DirectUpsert",
    };
    if vector.is_empty() {
        return Err(LiteError::BadRequest {
            detail: format!("{op_name}: the vector column is empty"),
        });
    }
    let doc_id = doc_id_for(&pk_bytes, surrogate)?;
    let incoming = decode_payload(&payload)?;
    let key = index_key(&collection, &field);
    let vector_state = Arc::clone(&engine.vector_state);
    let crdt = Arc::clone(&engine.crdt);

    Ok(Box::pin(async move {
        let guard = if permit.is_none() {
            Some(engine.fts_state.admit_mutation().await)
        } else {
            None
        };
        let result = async {
            let dim = vector.len();
            // First-insert config (quantization + storage dtype) when absent.
            vector_state
                .per_index_config
                .lock_or_recover()
                .entry(key.clone())
                .or_insert_with(|| VectorPrimaryConfig {
                    vector_field: field.clone(),
                    dim: dim as u32,
                    quantization,
                    storage_dtype,
                    ..VectorPrimaryConfig::default()
                });

            let existing = read_row(&crdt, &collection, &doc_id);
            let has_conflict_updates = !on_conflict_updates.is_empty();
            let (row, command): (_, &'static str) = match (intent, existing) {
                (VectorDirectWriteIntent::Insert, Some(_)) => {
                    return Err(LiteError::UniqueViolation {
                        collection: collection.clone(),
                        detail: format!("key '{doc_id}' already exists"),
                    });
                }
                (VectorDirectWriteIntent::InsertIfAbsent, Some(_)) => {
                    return Ok(QueryResult {
                        columns: vec![],
                        rows: vec![],
                        rows_affected: 0,
                        command: Some("INSERT".into()),
                    });
                }
                (VectorDirectWriteIntent::Upsert, Some(mut stored)) if has_conflict_updates => {
                    apply_patch(&mut stored, &on_conflict_updates, &incoming)?;
                    (stored, "UPDATE")
                }
                (VectorDirectWriteIntent::Upsert, Some(_)) => (incoming, "UPSERT"),
                (VectorDirectWriteIntent::Insert, None) => (incoming, "INSERT"),
                (VectorDirectWriteIntent::InsertIfAbsent, None) => (incoming, "INSERT"),
                (VectorDirectWriteIntent::Upsert, None) if has_conflict_updates => {
                    (incoming, "INSERT")
                }
                (VectorDirectWriteIntent::Upsert, None) => (incoming, "UPSERT"),
            };

            // A stored row keeps exactly one live node: binding the new vector
            // tombstones the old one under the same lock.
            insert_node(&vector_state, &key, &doc_id, &vector, op_name).await?;
            write_row(&crdt, &collection, &doc_id, dim, &row, op_name)?;
            reindex_documents(engine, &collection, [doc_id.as_str()])?;

            Ok(QueryResult {
                columns: vec![],
                rows: vec![],
                rows_affected: 1,
                command: Some(command.into()),
            })
        }
        .await;
        match guard {
            Some(guard) => guard.finish(result),
            None => result,
        }
    }))
}