velesdb-core 7.0.0

//! `TRAIN QUANTIZER` statement execution (PQ, OPQ, `RaBitQ`, SQ8).

use crate::index::hnsw::native::TraversalCodec;
use crate::{Error, Result, SearchResult, StorageMode};

use super::Database;

impl Database {
    /// Executes a `TRAIN QUANTIZER` statement.
    ///
    /// Trains a PQ/OPQ/RaBitQ codebook on the collection's vectors, stores the
    /// resulting quantizer, updates storage mode, and persists the codebook to disk.
    ///
    /// # Lock Ordering
    ///
    /// Vectors are extracted under `vector_storage` read lock, which is released
    /// before acquiring `pq_quantizer` write lock (respects canonical lock order).
    ///
    /// # Errors
    ///
    /// - `Error::CollectionNotFound` if the target collection does not exist.
    /// - `Error::InvalidQuantizerConfig` for invalid params (m=0, dim%m!=0, already trained).
    /// - `Error::TrainingFailed` if the underlying training algorithm errors.
    #[allow(clippy::cast_possible_truncation, clippy::cast_sign_loss)]
    pub(super) fn execute_train(
        &self,
        stmt: &crate::velesql::TrainStatement,
    ) -> Result<Vec<SearchResult>> {
        let collection = self.resolve_train_collection(stmt)?;
        let train_params = TrainParams::from_statement(stmt);

        train_params.validate_basic()?;

        let config = collection.config();
        let dim = config.dimension;
        train_params.validate_dimension(dim)?;

        Self::check_already_trained(&collection, train_params.force)?;

        let vectors = Self::extract_training_vectors(&collection, train_params.sample_limit)?;

        match train_params.train_type.as_str() {
            "pq" => Self::train_pq(&collection, &vectors, &train_params),
            "opq" => Self::train_opq(&collection, &vectors, &train_params),
            "rabitq" => Self::train_rabitq(&collection, vectors, dim),
            "sq8" => Self::train_sq8(&collection, vectors, dim),
            other => Err(Error::InvalidQuantizerConfig(format!(
                "unknown quantizer type: '{other}'. Supported: pq, opq, rabitq, sq8"
            ))),
        }
    }

    /// Normalizes training vectors in place for cosine collections.
    ///
    /// The traversal backends encode the prepared (stored) vector form,
    /// which for cosine engines is unit-normalized — so the quantizer must
    /// be trained on the same form, or its centroid/ranges would describe a
    /// space no stored code lives in. A no-op for every other metric.
    fn normalize_training_for_cosine(
        collection: &crate::collection::Collection,
        vectors: &mut [Vec<f32>],
    ) {
        if collection.config().metric == crate::DistanceMetric::Cosine {
            for vector in vectors.iter_mut() {
                crate::simd_native::normalize_inplace_native(vector);
            }
        }
    }

    /// Resolves the collection referenced by a TRAIN statement.
    ///
    /// Uses `resolve_writable_collection` for call-site clarity (training is
    /// a write-path operation); it resolves the same collections as
    /// `resolve_collection`. The actual "no vectors" guard for metadata-only
    /// collections lives in `extract_training_vectors`, which filters out
    /// points with an empty vector before training.
    fn resolve_train_collection(
        &self,
        stmt: &crate::velesql::TrainStatement,
    ) -> Result<crate::collection::Collection> {
        self.resolve_writable_collection(&stmt.collection)
    }

    /// Checks if a quantizer is already trained, returning an error unless `force` is set.
    fn check_already_trained(
        collection: &crate::collection::Collection,
        force: bool,
    ) -> Result<()> {
        let already_trained = collection.pq_quantizer_read().is_some();
        if already_trained && !force {
            return Err(Error::InvalidQuantizerConfig(
                "Quantizer already trained. Use force=true to retrain.".into(),
            ));
        }
        Ok(())
    }

    /// Extracts vectors from the collection for training, with optional sampling.
    fn extract_training_vectors(
        collection: &crate::collection::Collection,
        sample_limit: Option<usize>,
    ) -> Result<Vec<Vec<f32>>> {
        // `all_point_ids`, not `all_ids`: the latter enumerates the payload
        // store and documents that "points inserted with `None` payload may
        // not appear". A collection used for pure vector search carries no
        // payloads at all, so training saw an empty corpus and every
        // `TRAIN QUANTIZER` on it — pq, opq, rabitq, sq8 — failed with "no
        // vectors available for training" while the collection held its full
        // count. `all_point_ids` unions vector and payload storage and is the
        // authoritative set; the `is_empty` filter below already drops the
        // metadata-only points it also brings in.
        let all_ids = collection.all_point_ids();
        // get_raw: PQ training only consumes vectors; TTL-expired points are
        // still valid training samples and must not shrink the corpus.
        let points = collection.get_raw(&all_ids);
        let mut vectors: Vec<Vec<f32>> = points
            .into_iter()
            .flatten()
            .filter(|p| !p.vector.is_empty())
            .map(|p| p.vector)
            .collect();

        if let Some(limit) = sample_limit {
            if vectors.len() > limit {
                vectors.truncate(limit);
            }
        }

        if vectors.is_empty() {
            return Err(Error::TrainingFailed(
                "no vectors available for training".into(),
            ));
        }
        Ok(vectors)
    }

    /// Trains a standard Product Quantizer and persists it.
    fn train_pq(
        collection: &crate::collection::Collection,
        vectors: &[Vec<f32>],
        params: &TrainParams,
    ) -> Result<Vec<SearchResult>> {
        let pq = crate::quantization::ProductQuantizer::train(vectors, params.m, params.k)
            .map_err(|e| Error::TrainingFailed(e.to_string()))?;

        let codebook_size = pq.codebook.num_subspaces * pq.codebook.num_centroids;

        pq.save_codebook(collection.data_path())
            .map_err(|e| Error::TrainingFailed(e.to_string()))?;

        *collection.pq_quantizer_write() = Some(pq);

        Self::finalize_pq_config(
            collection,
            StorageMode::ProductQuantization,
            params.oversampling,
        )?;

        Ok(train_result_response(serde_json::json!({
            "status": "trained",
            "type": "pq",
            "m": params.m,
            "k": params.k,
            "codebook_size": codebook_size,
            "training_vectors": vectors.len()
        })))
    }

    /// Trains an Optimized Product Quantizer (with rotation) and persists it.
    ///
    /// OPQ keeps its codes in a rotated basis and rescoring rotates the query
    /// to match, so training is rejected on metrics an orthogonal transform
    /// does not preserve — see [`crate::DistanceMetric::is_rotation_invariant`].
    /// Without this, `TRAIN QUANTIZER ... TYPE opq` on a Hamming or Jaccard
    /// collection succeeded and every later rescore measured the metric in a
    /// space where it means nothing.
    fn train_opq(
        collection: &crate::collection::Collection,
        vectors: &[Vec<f32>],
        params: &TrainParams,
    ) -> Result<Vec<SearchResult>> {
        let metric = collection.config().metric;
        if !metric.is_rotation_invariant() {
            return Err(Error::InvalidQuantizerConfig(format!(
                "OPQ rotates vectors into another basis, which does not preserve {metric:?}; \
                 train plain PQ instead, or use a cosine/euclidean/dot-product collection"
            )));
        }

        let pq = crate::quantization::train_opq(vectors, params.m, params.k, true, 10)
            .map_err(|e| Error::TrainingFailed(e.to_string()))?;

        let codebook_size = pq.codebook.num_subspaces * pq.codebook.num_centroids;

        pq.save_codebook(collection.data_path())
            .map_err(|e| Error::TrainingFailed(e.to_string()))?;
        pq.save_rotation(collection.data_path())
            .map_err(|e| Error::TrainingFailed(e.to_string()))?;

        *collection.pq_quantizer_write() = Some(pq);

        Self::finalize_pq_config(
            collection,
            StorageMode::ProductQuantization,
            params.oversampling,
        )?;

        Ok(train_result_response(serde_json::json!({
            "status": "trained",
            "type": "opq",
            "m": params.m,
            "k": params.k,
            "codebook_size": codebook_size,
            "training_vectors": vectors.len()
        })))
    }

    /// Trains a `RaBitQ` quantizer, persists it, and installs it into the
    /// live index when the collection's HNSW backend is `RaBitQ`.
    ///
    /// For collections created with another storage mode the backend is
    /// Standard: training still persists `rabitq.idx` and flips the config
    /// below, and the wiring takes effect at the next open (the open path
    /// rebuilds the backend from the storage mode and reloads `rabitq.idx`).
    fn train_rabitq(
        collection: &crate::collection::Collection,
        mut vectors: Vec<Vec<f32>>,
        dim: usize,
    ) -> Result<Vec<SearchResult>> {
        Self::normalize_training_for_cosine(collection, &mut vectors);
        let rbq = crate::quantization::RaBitQIndex::train(&vectors, 42)
            .map_err(|e| Error::TrainingFailed(e.to_string()))?;

        rbq.save(collection.data_path())
            .map_err(|e| Error::TrainingFailed(e.to_string()))?;

        // Live install: re-encodes the collection's vectors (O(n·d)) so
        // RaBitQ traversal activates without a restart.
        let installed = collection
            .install_rabitq_quantizer(std::sync::Arc::new(rbq))
            .map_err(|e| Error::TrainingFailed(e.to_string()))?;
        if !installed {
            tracing::debug!(
                "RaBitQ trained on a Standard-backend collection; takes effect at next open"
            );
        }

        // RaBitQ uses default oversampling of 4.
        Self::finalize_pq_config(collection, StorageMode::RaBitQ, 4)?;

        Ok(train_result_response(serde_json::json!({
            "status": "trained",
            "type": "rabitq",
            "dimension": dim,
            "training_vectors": vectors.len()
        })))
    }

    /// Trains an SQ8 scalar quantizer, persists it to `sq8.idx`, and installs
    /// it into the live index when the collection's HNSW backend is SQ8 —
    /// mirror of [`Self::train_rabitq`].
    ///
    /// The int8 traversal distance is L2 over codes, so training is rejected
    /// on metrics whose ordering it cannot preserve (everything except
    /// Euclidean and Cosine — see `Sq8Codec::supports_metric`). For cosine
    /// collections the index stores unit-normalized vectors, so the quantizer
    /// trains on the same normalized form.
    fn train_sq8(
        collection: &crate::collection::Collection,
        mut vectors: Vec<Vec<f32>>,
        dim: usize,
    ) -> Result<Vec<SearchResult>> {
        let metric = collection.config().metric;
        if !crate::index::hnsw::native::Sq8Codec::supports_metric(metric) {
            return Err(Error::InvalidQuantizerConfig(format!(
                "SQ8 traversal supports euclidean and cosine metrics, not {metric:?}"
            )));
        }
        Self::normalize_training_for_cosine(collection, &mut vectors);

        let refs: Vec<&[f32]> = vectors.iter().map(Vec::as_slice).collect();
        let quantizer = crate::index::hnsw::native::ScalarQuantizer::train(&refs)
            .map_err(|e| Error::TrainingFailed(e.to_string()))?;

        quantizer
            .save(collection.data_path())
            .map_err(|e| Error::TrainingFailed(e.to_string()))?;

        // Live install: re-encodes the collection's vectors (O(n·d)) so SQ8
        // traversal activates without a restart.
        let installed = collection
            .install_sq8_quantizer(std::sync::Arc::new(quantizer))
            .map_err(|e| Error::TrainingFailed(e.to_string()))?;
        if !installed {
            tracing::debug!(
                "SQ8 trained on a non-SQ8-backend collection; takes effect at next open"
            );
        }

        Self::finalize_pq_config(collection, StorageMode::SQ8, 4)?;

        Ok(train_result_response(serde_json::json!({
            "status": "trained",
            "type": "sq8",
            "dimension": dim,
            "training_vectors": vectors.len()
        })))
    }

    /// Updates storage mode and oversampling in config, then persists it.
    fn finalize_pq_config(
        collection: &crate::collection::Collection,
        mode: StorageMode,
        oversampling: u32,
    ) -> Result<()> {
        {
            let mut cfg = collection.config_write();
            cfg.storage_mode = mode;
            cfg.pq_rescore_oversampling = Some(oversampling);
        }
        collection.save_config()
    }
}

/// Wraps a JSON metadata value into a single-element `SearchResult` vector.
fn train_result_response(metadata: serde_json::Value) -> Vec<SearchResult> {
    vec![SearchResult::new(
        crate::Point::metadata_only(0, metadata),
        0.0,
    )]
}

/// Parsed parameters for a TRAIN QUANTIZER statement.
struct TrainParams {
    m: usize,
    k: usize,
    train_type: String,
    oversampling: u32,
    sample_limit: Option<usize>,
    force: bool,
}

impl TrainParams {
    /// Extracts training parameters from a `TrainStatement`.
    #[allow(clippy::cast_possible_truncation, clippy::cast_sign_loss)]
    fn from_statement(stmt: &crate::velesql::TrainStatement) -> Self {
        use crate::velesql::WithValue;

        let m = stmt
            .params
            .get("m")
            .and_then(WithValue::as_integer)
            .map_or(0_usize, |v| v.max(0) as usize);
        let k = stmt
            .params
            .get("k")
            .and_then(WithValue::as_integer)
            .map_or(256_usize, |v| v.max(0) as usize);
        let train_type = stmt
            .params
            .get("type")
            .and_then(WithValue::as_str)
            .unwrap_or("pq")
            .to_lowercase();
        let oversampling = stmt
            .params
            .get("oversampling")
            .and_then(WithValue::as_integer)
            .map_or(4_u32, |v| v.max(0) as u32);
        let sample_limit = stmt
            .params
            .get("sample")
            .and_then(WithValue::as_integer)
            .map(|v| v.max(0) as usize);
        let force = stmt
            .params
            .get("force")
            .and_then(WithValue::as_bool)
            .unwrap_or(false);

        Self {
            m,
            k,
            train_type,
            oversampling,
            sample_limit,
            force,
        }
    }

    /// Validates basic constraints (m > 0, k > 0).
    ///
    /// Only PQ/OPQ consume `m` and `k`; `rabitq` and `sq8` derive everything
    /// from the vectors themselves, so demanding a dummy `m` for them was an
    /// incoherence, not a safeguard. Unknown types fall through — the
    /// dispatch rejects them with the full supported list.
    fn validate_basic(&self) -> Result<()> {
        if !matches!(self.train_type.as_str(), "pq" | "opq") {
            return Ok(());
        }
        if self.m == 0 {
            return Err(Error::InvalidQuantizerConfig(
                "m (num_subspaces) must be > 0".into(),
            ));
        }
        if self.k == 0 {
            return Err(Error::InvalidQuantizerConfig(
                "k (num_centroids) must be > 0".into(),
            ));
        }
        Ok(())
    }

    /// Validates dimension compatibility (dim % m == 0 for PQ/OPQ).
    fn validate_dimension(&self, dim: usize) -> Result<()> {
        if !matches!(self.train_type.as_str(), "rabitq" | "sq8") && !dim.is_multiple_of(self.m) {
            return Err(Error::InvalidQuantizerConfig(format!(
                "dimension {dim} is not divisible by m={}",
                self.m
            )));
        }
        Ok(())
    }
}