nodedb 0.2.0

// SPDX-License-Identifier: BUSL-1.1

//! Vector write handlers: VectorInsert, VectorBatchInsert, VectorDelete,
//! SetVectorParams.

use nodedb_types::Surrogate;
use tracing::{debug, warn};

use crate::bridge::envelope::{ErrorCode, Response};
use crate::data::executor::core_loop::CoreLoop;
use crate::data::executor::handlers::vector_upsert::decode_payload_lowercased;
use crate::data::executor::task::ExecutionTask;
use crate::engine::vector::collection::VectorCollection;
use crate::engine::vector::distance::DistanceMetric;
use crate::engine::vector::hnsw::HnswParams;
use crate::types::TenantId;

/// Parameters for configuring vector index settings.
pub(in crate::data::executor) struct SetVectorParamsInput<'a> {
    pub task: &'a ExecutionTask,
    pub tid: u64,
    pub collection: &'a str,
    pub m: usize,
    pub ef_construction: usize,
    pub metric: &'a str,
    pub index_type: &'a str,
    pub pq_m: usize,
    pub ivf_cells: usize,
    pub ivf_nprobe: usize,
}

/// Parameters for a vector insert operation.
pub(in crate::data::executor) struct VectorInsertParams<'a> {
    pub task: &'a ExecutionTask,
    pub tid: u64,
    pub collection: &'a str,
    pub vector: &'a [f32],
    pub dim: usize,
    pub field_name: &'a str,
    pub surrogate: Surrogate,
}

impl CoreLoop {
    /// Get or create a vector collection, validating dimension compatibility.
    pub(in crate::data::executor) fn get_or_create_vector_index(
        &mut self,
        tid: u64,
        collection: &str,
        dim: usize,
        field_name: &str,
    ) -> Result<&mut VectorCollection, ErrorCode> {
        let index_key = CoreLoop::vector_index_key(tid, collection, field_name);
        if let Some(existing) = self.vector_collections.get(&index_key)
            && existing.dim() != dim
        {
            return Err(ErrorCode::RejectedConstraint {
                detail: String::new(),
                constraint: format!(
                    "dimension mismatch: index has {}, got {dim}",
                    existing.dim()
                ),
            });
        }
        let core_id = self.core_id;
        let params = self
            .vector_params
            .get(&index_key)
            .cloned()
            .unwrap_or_default();
        Ok(self.vector_collections.entry(index_key).or_insert_with(|| {
            debug!(core = core_id, dim, m = params.m, ef = params.ef_construction, ?params.metric, "creating vector collection");
            VectorCollection::new(dim, params)
        }))
    }

    pub(in crate::data::executor) fn execute_vector_insert(
        &mut self,
        params: VectorInsertParams<'_>,
    ) -> Response {
        let VectorInsertParams {
            task,
            tid,
            collection,
            vector,
            dim,
            field_name,
            surrogate,
        } = params;
        debug!(core = self.core_id, %collection, dim, "vector insert");
        if vector.len() != dim {
            return self.response_error(
                task,
                ErrorCode::RejectedConstraint {
                    detail: String::new(),
                    constraint: format!(
                        "vector dimension mismatch: expected {dim}, got {}",
                        vector.len()
                    ),
                },
            );
        }
        let index_key = CoreLoop::vector_index_key(tid, collection, field_name);

        // Check if this collection uses IVF-PQ index.
        if let Some(cfg) = self.index_configs.get(&index_key)
            && cfg.index_type == crate::engine::vector::index_config::IndexType::IvfPq
        {
            let key = index_key.clone();
            return self.ivf_insert(task, &key, vector, dim, surrogate);
        }

        // Default: HNSW (with or without PQ).
        match self.get_or_create_vector_index(tid, collection, dim, field_name) {
            Ok(collection_ref) => {
                collection_ref.insert_with_surrogate(vector.to_vec(), surrogate);
                let seal_key = CoreLoop::vector_checkpoint_filename(&index_key);
                if collection_ref.needs_seal()
                    && let Some(req) = collection_ref.seal(&seal_key)
                    && let Some(tx) = &self.build_tx
                    && let Err(e) = tx.send(req)
                {
                    warn!(core = self.core_id, error = %e, "failed to send HNSW build request");
                }
                self.checkpoint_coordinator.mark_dirty("vector", 1);
                self.response_ok(task)
            }
            Err(err) => self.response_error(task, err),
        }
    }

    /// Insert into an IVF-PQ index, returning the assigned vector ID.
    fn ivf_insert(
        &mut self,
        task: &ExecutionTask,
        index_key: &(TenantId, String),
        vector: &[f32],
        dim: usize,
        surrogate: Surrogate,
    ) -> Response {
        let ivf = self
            .ivf_indexes
            .entry(index_key.clone())
            .or_insert_with(|| {
                let cfg = self
                    .index_configs
                    .get(index_key)
                    .cloned()
                    .unwrap_or_default();
                let params = cfg.to_ivf_params();
                debug!(
                    core = self.core_id,
                    key = %index_key.1,
                    "creating IVF-PQ index"
                );
                crate::engine::vector::ivf::IvfPqIndex::new(dim, params)
            });

        // IVF-PQ requires training before the first insert.
        if ivf.n_cells() == 0 {
            let refs: Vec<&[f32]> = vec![vector];
            ivf.train(&refs);
        }

        let vector_id = ivf.add(vector);

        // Register surrogate mapping using the actual IVF-assigned vector ID.
        if surrogate != Surrogate::ZERO {
            let coll = self
                .vector_collections
                .entry(index_key.clone())
                .or_insert_with(|| VectorCollection::new(dim, Default::default()));
            coll.surrogate_map.insert(vector_id, surrogate);
            coll.surrogate_to_local.insert(surrogate, vector_id);
        }

        self.checkpoint_coordinator.mark_dirty("vector", 1);
        self.response_ok(task)
    }

    /// Execute batch vector insert (always to the default/unnamed field).
    pub(in crate::data::executor) fn execute_vector_batch_insert(
        &mut self,
        task: &ExecutionTask,
        tid: u64,
        collection: &str,
        vectors: &[Vec<f32>],
        dim: usize,
        surrogates: &[Surrogate],
    ) -> Response {
        debug!(core = self.core_id, %collection, dim, count = vectors.len(), "vector batch insert");
        let index_key = CoreLoop::vector_index_key(tid, collection, "");
        match self.get_or_create_vector_index(tid, collection, dim, "") {
            Ok(collection_ref) => {
                for (i, vector) in vectors.iter().enumerate() {
                    if vector.len() != dim {
                        return self.response_error(
                            task,
                            ErrorCode::RejectedConstraint {
                                detail: String::new(),
                                constraint: format!(
                                    "dimension mismatch in batch: expected {dim}, got {}",
                                    vector.len()
                                ),
                            },
                        );
                    }
                    let s = surrogates.get(i).copied().unwrap_or(Surrogate::ZERO);
                    collection_ref.insert_with_surrogate(vector.clone(), s);
                }
                let seal_key = CoreLoop::vector_checkpoint_filename(&index_key);
                if collection_ref.needs_seal()
                    && let Some(req) = collection_ref.seal(&seal_key)
                    && let Some(tx) = &self.build_tx
                    && let Err(e) = tx.send(req)
                {
                    warn!(core = self.core_id, error = %e, "failed to send HNSW build request");
                }
                self.checkpoint_coordinator
                    .mark_dirty("vector", vectors.len());
                match super::super::response_codec::encode_count("inserted", vectors.len()) {
                    Ok(bytes) => self.response_with_payload(task, bytes),
                    Err(e) => self.response_error(
                        task,
                        ErrorCode::Internal {
                            detail: e.to_string(),
                        },
                    ),
                }
            }
            Err(err) => self.response_error(task, err),
        }
    }

    pub(in crate::data::executor) fn execute_vector_delete(
        &mut self,
        task: &ExecutionTask,
        tid: u64,
        collection: &str,
        vector_id: u32,
    ) -> Response {
        debug!(core = self.core_id, %collection, vector_id, "vector delete");
        // Resolve the actual index key. Legacy `CREATE VECTOR INDEX` uses
        // an empty field segment; vector-primary collections use
        // `"{collection}:{field}"`. Try the legacy key first, then scan
        // for any field-suffixed key under the same (tenant, collection).
        let tenant = TenantId::new(tid);
        let plain_key = (tenant, collection.to_string());
        let prefix = format!("{collection}:");
        let resolved_key = if self.vector_collections.contains_key(&plain_key) {
            Some(plain_key)
        } else {
            self.vector_collections
                .keys()
                .find(|(t, c)| *t == tenant && c.starts_with(&prefix))
                .cloned()
        };
        let Some(index_key) = resolved_key else {
            return self.response_error(task, ErrorCode::NotFound);
        };

        // Capture the surrogate before deletion so we can fetch the
        // payload row from the sparse store and update the bitmap. The
        // bitmap stores node-id -> field-value membership; without the
        // original field values we cannot remove the entries cleanly.
        //
        // Asymmetric with the insert path (`vector_upsert`): insert
        // atomically rolls back the HNSW node if the sparse write fails,
        // because a phantom node would be returned by future searches.
        // Delete is best-effort cleanup — if the sparse read or decode
        // fails we still drop the HNSW node and skip bitmap cleanup.
        // Phantom bitmap entries are safe (the bitmap is filtered against
        // live node ids on read), so leaving them is preferable to
        // aborting the delete and leaking the vector.
        let surrogate_opt = self
            .vector_collections
            .get(&index_key)
            .and_then(|c| c.get_surrogate(vector_id));

        if let Some(surrogate) = surrogate_opt {
            let row_key = format!("{:08x}", surrogate.as_u32());
            let fields = match self.sparse.get(tid, collection, &row_key) {
                Ok(Some(bytes)) => decode_payload_lowercased(&bytes).ok(),
                _ => None,
            };
            if let Some(fields) = fields
                && let Some(coll) = self.vector_collections.get_mut(&index_key)
            {
                coll.payload.delete_row(vector_id, &fields);
            }
        }

        let Some(collection_ref) = self.vector_collections.get_mut(&index_key) else {
            return self.response_error(task, ErrorCode::NotFound);
        };
        if collection_ref.delete(vector_id) {
            self.checkpoint_coordinator.mark_dirty("vector", 1);
            self.response_ok(task)
        } else {
            self.response_error(task, ErrorCode::NotFound)
        }
    }

    pub(in crate::data::executor) fn execute_set_vector_params(
        &mut self,
        params: SetVectorParamsInput<'_>,
    ) -> Response {
        let SetVectorParamsInput {
            task,
            tid,
            collection,
            m,
            ef_construction,
            metric,
            index_type,
            pq_m,
            ivf_cells,
            ivf_nprobe,
        } = params;
        debug!(core = self.core_id, %collection, m, ef_construction, %metric, %index_type, "set vector params");
        let index_key = CoreLoop::vector_index_key(tid, collection, "");

        if self.vector_collections.contains_key(&index_key) {
            return self.response_error(
                task,
                ErrorCode::RejectedConstraint {
                detail: String::new(),
                    constraint: "cannot change index params after creation; drop and recreate the collection".into(),
                },
            );
        }

        // Zero / empty inputs mean "preserve existing value if present, else default".
        // This keeps ALTER SET (index_type = ...) from clobbering m / ef_construction
        // that were set at CREATE time but not re-specified in the ALTER clause.
        let existing = self.index_configs.get(&index_key).cloned();

        let resolved_metric_str: String = if metric.is_empty() {
            existing
                .as_ref()
                .map(|c| {
                    match c.hnsw.metric {
                        DistanceMetric::L2 => "l2",
                        DistanceMetric::Cosine => "cosine",
                        DistanceMetric::InnerProduct => "inner_product",
                        DistanceMetric::Manhattan => "manhattan",
                        DistanceMetric::Chebyshev => "chebyshev",
                        DistanceMetric::Hamming => "hamming",
                        DistanceMetric::Jaccard => "jaccard",
                        DistanceMetric::Pearson => "pearson",
                        _ => "cosine",
                    }
                    .to_string()
                })
                .unwrap_or_else(|| "cosine".into())
        } else {
            metric.to_string()
        };

        let metric_enum = match resolved_metric_str.as_str() {
            "l2" | "euclidean" => DistanceMetric::L2,
            "cosine" => DistanceMetric::Cosine,
            "inner_product" | "ip" | "dot" => DistanceMetric::InnerProduct,
            "manhattan" | "l1" => DistanceMetric::Manhattan,
            "chebyshev" | "linf" => DistanceMetric::Chebyshev,
            "hamming" => DistanceMetric::Hamming,
            "jaccard" => DistanceMetric::Jaccard,
            "pearson" => DistanceMetric::Pearson,
            _ => {
                return self.response_error(
                    task,
                    ErrorCode::RejectedConstraint {
                detail: String::new(),
                        constraint: format!(
                            "unknown metric '{resolved_metric_str}'; supported: l2, cosine, inner_product, manhattan, chebyshev, hamming, jaccard, pearson"
                        ),
                    },
                );
            }
        };

        let idx_type = if index_type.is_empty() {
            existing
                .as_ref()
                .map(|c| c.index_type.clone())
                .unwrap_or_default()
        } else {
            match crate::engine::vector::index_config::IndexType::parse(index_type) {
                Some(t) => t,
                None => {
                    return self.response_error(
                        task,
                        ErrorCode::RejectedConstraint {
                detail: String::new(),
                            constraint: format!(
                                "unknown index_type '{index_type}'; supported: hnsw, hnsw_pq, ivf_pq"
                            ),
                        },
                    );
                }
            }
        };

        let resolved_m = if m > 0 {
            m
        } else {
            existing.as_ref().map(|c| c.hnsw.m).unwrap_or(16)
        };
        let resolved_ef = if ef_construction > 0 {
            ef_construction
        } else {
            existing
                .as_ref()
                .map(|c| c.hnsw.ef_construction)
                .unwrap_or(200)
        };
        let resolved_pq_m = if pq_m > 0 {
            pq_m
        } else {
            existing.as_ref().map(|c| c.pq_m).unwrap_or(8)
        };
        let resolved_ivf_cells = if ivf_cells > 0 {
            ivf_cells
        } else {
            existing.as_ref().map(|c| c.ivf_cells).unwrap_or(256)
        };
        let resolved_ivf_nprobe = if ivf_nprobe > 0 {
            ivf_nprobe
        } else {
            existing.as_ref().map(|c| c.ivf_nprobe).unwrap_or(16)
        };

        let params = HnswParams {
            m: resolved_m,
            m0: resolved_m * 2,
            ef_construction: resolved_ef,
            metric: metric_enum,
        };

        let config = crate::engine::vector::index_config::IndexConfig {
            hnsw: params.clone(),
            index_type: idx_type,
            pq_m: resolved_pq_m,
            ivf_cells: resolved_ivf_cells,
            ivf_nprobe: resolved_ivf_nprobe,
        };

        self.vector_params.insert(index_key.clone(), params);
        self.index_configs.insert(index_key, config);
        self.response_ok(task)
    }
}