velesdb-core 7.0.0

//! HNSW index parameters and search quality profiles.
//!
//! This module contains configuration types for tuning HNSW index
//! performance and search quality.

use crate::quantization::StorageMode;
use serde::{Deserialize, Serialize};

/// Default VAMANA alpha for neighbor diversification (1.2).
///
/// Used as serde default when deserializing configs that predate the alpha field.
const fn default_alpha() -> f32 {
    1.2
}

/// HNSW index parameters for tuning performance and recall.
///
/// Use [`HnswParams::auto`] for automatic tuning based on vector dimension,
/// or create custom parameters for specific workloads.
///
/// # Alpha (VAMANA diversification)
///
/// The `alpha` parameter controls neighbor selection diversification.
/// Values > 1.0 bias toward graph navigability (more diverse neighbors),
/// while 1.0 selects purely by proximity. The default (1.2) follows the
/// VAMANA paper recommendation (Subramanya et al., 2019).
#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)]
pub struct HnswParams {
    /// Number of bi-directional links per node (M parameter).
    /// Higher = better recall, more memory, slower insert.
    pub max_connections: usize,
    /// Size of dynamic candidate list during construction.
    /// Higher = better recall, slower indexing.
    pub ef_construction: usize,
    /// Initial capacity (grows automatically if exceeded).
    pub max_elements: usize,
    /// Vector storage mode (Full, SQ8, Binary, PQ, or `RaBitQ`).
    ///
    /// `RaBitQ` selects the binary-traversal HNSW backend (32x bandwidth
    /// reduction during graph traversal, f32 re-ranking), persisted and
    /// restored across reopens. PQ enables the ADC rescore path once a
    /// codebook is trained (`TRAIN QUANTIZER`). SQ8/Binary currently keep
    /// quantized representations as an **additional** cache that no
    /// production search path reads — selecting them adds memory rather
    /// than reducing it; the 4x-reduction figure applies only to a backend
    /// that stores quantized vectors *instead of* f32.
    #[serde(default)]
    pub storage_mode: StorageMode,
    /// VAMANA alpha for neighbor diversification (default: 1.2).
    ///
    /// Alpha > 1.0 biases `select_neighbors` toward diversity over proximity,
    /// producing a more navigable graph with better recall. Set to 1.0 for
    /// strict nearest-neighbor selection.
    #[serde(default = "default_alpha")]
    pub alpha: f32,
}

// Eq is sound because `alpha` is finite in every supported path: the preset
// constructors and `default_alpha()` produce 1.2, JSON deserialization cannot
// carry NaN/Inf, and any `alpha` arriving from untrusted public input
// (REST/Python/Tauri) is rejected by [`HnswParams::validate`] before the params
// reach the engine. Embedders that build the struct directly via its public
// fields are responsible for passing a finite value, as with any `pub` float.
impl Eq for HnswParams {}

impl Default for HnswParams {
    fn default() -> Self {
        Self::auto(768)
    }
}

impl HnswParams {
    /// Creates optimized parameters based on vector dimension.
    ///
    /// These defaults are tuned for datasets up to 100K vectors with high recall targets.
    /// For larger datasets, use [`HnswParams::for_dataset_size`].
    #[must_use]
    pub fn auto(dimension: usize) -> Self {
        match dimension {
            0..=256 => Self {
                max_connections: 24,
                ef_construction: 300,
                max_elements: 100_000,
                storage_mode: StorageMode::Full,
                alpha: default_alpha(),
            },
            // 257+ dimensions: aggressive params targeting high recall
            _ => Self {
                max_connections: 32,
                ef_construction: 400,
                max_elements: 100_000,
                storage_mode: StorageMode::Full,
                alpha: default_alpha(),
            },
        }
    }

    /// Builds creation-time parameters for `dimension`, applying the
    /// configured [`[hnsw]`](crate::config::HnswConfig) section on top of
    /// [`HnswParams::auto`].
    ///
    /// This is the **single mapping point** between the `[hnsw]` TOML table
    /// and the engine, the same role `RuntimeLimits::from_config` plays for
    /// `[limits]`, so the field correspondence cannot drift across the
    /// creation paths that consult it.
    ///
    /// A section field left `None` leaves the auto-tuned value untouched —
    /// a default [`HnswConfig`](crate::config::HnswConfig) therefore produces
    /// exactly `HnswParams::auto(dimension)`, byte for byte. Wiring the
    /// section is a no-op for anyone who did not set it.
    ///
    /// `hnsw.max_layers` has no counterpart here: the HNSW layer count is
    /// drawn per node by the level generator and no engine path caps it, so
    /// the knob stays inert and [`VelesConfig::validate`] keeps warning about
    /// it (issue #2087).
    ///
    /// `storage_mode` is deliberately not sourced from config either: it is a
    /// per-collection property, and every constructor overwrites this field
    /// with the collection's own storage mode (see
    /// `Collection::build_hnsw_index`).
    ///
    /// [`VelesConfig::validate`]: crate::config::VelesConfig::validate
    #[must_use]
    pub fn from_config(dimension: usize, hnsw: &crate::config::HnswConfig) -> Self {
        let mut params = Self::auto(dimension);
        if let Some(m) = hnsw.m {
            params.max_connections = m;
        }
        if let Some(ef) = hnsw.ef_construction {
            params.ef_construction = ef;
        }
        params
    }

    /// Creates parameters optimized for a specific dataset size.
    ///
    /// Targets high recall up to 1M vectors under benchmark-calibrated settings.
    ///
    /// # Parameters by Scale
    ///
    /// | Dataset Size | M | `ef_construction` | Target Recall |
    /// |--------------|---|-------------------|---------------|
    /// | ≤10K | 32 | 200 | ≥98% |
    /// | ≤100K | 64 | 800 | ≥95% |
    /// | ≤500K | 96 | 1200 | ≥95% |
    /// | ≤1M | 128 | 1600 | ≥95% |
    #[must_use]
    pub fn for_dataset_size(dimension: usize, expected_vectors: usize) -> Self {
        let (m_low, ef_low, m_high, ef_high, max_elems) = match expected_vectors {
            0..=10_000 => (24, 200, 32, 400, 20_000),
            10_001..=100_000 => (64, 800, 128, 1600, 150_000),
            100_001..=500_000 => (96, 1200, 128, 2000, 750_000),
            _ => (64, 800, 128, 1600, 1_500_000),
        };
        let (m, ef) = if dimension <= 256 {
            (m_low, ef_low)
        } else {
            (m_high, ef_high)
        };
        Self {
            max_connections: m,
            ef_construction: ef,
            max_elements: max_elems,
            storage_mode: StorageMode::Full,
            alpha: default_alpha(),
        }
    }

    /// Creates parameters optimized for large datasets (100K+ vectors).
    ///
    /// Higher M and `ef_construction` ensure good recall at scale.
    /// For 1M+ vectors, use [`HnswParams::for_dataset_size`] instead.
    #[must_use]
    pub fn large_dataset(dimension: usize) -> Self {
        Self::for_dataset_size(dimension, 500_000)
    }

    /// Creates parameters for 1 million vectors with high-recall target settings.
    ///
    /// Based on `OpenSearch` 2025 research: M=128, `ef_construction`=1600.
    #[must_use]
    pub fn million_scale(dimension: usize) -> Self {
        Self::for_dataset_size(dimension, 1_000_000)
    }

    /// Creates fast parameters optimized for insertion speed.
    /// Lower recall but faster indexing. Best for small datasets (<10K).
    #[must_use]
    pub fn fast() -> Self {
        Self {
            max_connections: 16,
            ef_construction: 150,
            max_elements: 100_000,
            storage_mode: StorageMode::Full,
            alpha: default_alpha(),
        }
    }

    /// Creates turbo parameters for maximum insert throughput.
    ///
    /// **Goal**: faster inserts than `auto` params (measured in `docs/BENCHMARKS.md`)
    ///
    /// # Trade-offs
    ///
    /// - **Recall**: lower than with standard params (not measured)
    /// - **Best for**: Bulk loading, development, benchmarking
    /// - **Not recommended for**: Production search workloads
    ///
    /// # Parameters
    ///
    /// - `M = 12`: Minimal connections for fast graph construction
    /// - `ef_construction = 100`: Low expansion factor
    ///
    /// After bulk loading, consider rebuilding with higher params for production.
    #[must_use]
    pub fn turbo() -> Self {
        Self {
            max_connections: 12,
            ef_construction: 100,
            max_elements: 100_000,
            storage_mode: StorageMode::Full,
            alpha: default_alpha(),
        }
    }

    /// Creates parameters optimized for high recall.
    #[must_use]
    pub fn high_recall(dimension: usize) -> Self {
        let base = Self::auto(dimension);
        Self {
            max_connections: base.max_connections + 8,
            ef_construction: base.ef_construction + 200,
            ..base
        }
    }

    /// Creates parameters optimized for maximum recall.
    #[must_use]
    pub fn max_recall(dimension: usize) -> Self {
        match dimension {
            0..=256 => Self {
                max_connections: 32,
                ef_construction: 500,
                max_elements: 100_000,
                storage_mode: StorageMode::Full,
                alpha: default_alpha(),
            },
            257..=768 => Self {
                max_connections: 48,
                ef_construction: 800,
                max_elements: 100_000,
                storage_mode: StorageMode::Full,
                alpha: default_alpha(),
            },
            _ => Self {
                max_connections: 64,
                ef_construction: 1000,
                max_elements: 100_000,
                storage_mode: StorageMode::Full,
                alpha: default_alpha(),
            },
        }
    }

    /// Creates parameters optimized for fast indexing.
    #[must_use]
    pub fn fast_indexing(dimension: usize) -> Self {
        let base = Self::auto(dimension);
        Self {
            max_connections: (base.max_connections / 2).max(8),
            ef_construction: base.ef_construction / 2,
            ..base
        }
    }

    /// Creates custom parameters with default alpha (1.2).
    #[must_use]
    pub const fn custom(
        max_connections: usize,
        ef_construction: usize,
        max_elements: usize,
    ) -> Self {
        Self {
            max_connections,
            ef_construction,
            max_elements,
            storage_mode: StorageMode::Full,
            alpha: 1.2,
        }
    }

    /// Creates parameters with SQ8 quantization for 4x memory reduction.
    ///
    /// # Memory Savings
    ///
    /// | Dimension | Full (f32) | SQ8 (u8) | Memory reduction |
    /// |-----------|------------|----------|----------|
    /// | 768 | 3 KB | 776 B | 4x |
    /// | 1536 | 6 KB | 1.5 KB | 4x |
    #[must_use]
    pub fn with_sq8(dimension: usize) -> Self {
        let mut params = Self::auto(dimension);
        params.storage_mode = StorageMode::SQ8;
        params
    }

    /// Creates parameters with binary quantization for 32x memory reduction.
    /// Best for edge/IoT devices with limited RAM.
    #[must_use]
    pub fn with_binary(dimension: usize) -> Self {
        let mut params = Self::auto(dimension);
        params.storage_mode = StorageMode::Binary;
        params
    }

    /// Returns a copy of these parameters with a custom alpha value.
    ///
    /// Alpha controls VAMANA neighbor diversification:
    /// - `1.0`: strict nearest-neighbor selection (no diversification)
    /// - `1.2`: recommended default (VAMANA paper)
    /// - `>1.2`: more diverse neighbors, better recall, slower insert
    ///
    /// # Example
    ///
    /// ```
    /// use velesdb_core::HnswParams;
    ///
    /// let params = HnswParams::auto(768).with_alpha(1.0);
    /// assert!((params.alpha - 1.0).abs() < f32::EPSILON);
    /// ```
    #[must_use]
    pub const fn with_alpha(self, alpha: f32) -> Self {
        Self { alpha, ..self }
    }

    /// Validates parameters that can originate from untrusted public input
    /// (REST, Python, Tauri) before they reach the engine.
    ///
    /// `alpha` must be finite and `>= 1.0` — the valid VAMANA range. Values
    /// below 1.0 under-diversify the graph and silently degrade recall, while
    /// non-finite values (`NaN`/`inf`) would break the [`Eq`] invariant. Valid
    /// alpha (the default 1.2 and anything `>= 1.0`) passes unchanged, so this
    /// is a no-op on every well-formed configuration.
    ///
    /// # Errors
    ///
    /// Returns [`crate::error::Error::Config`] when `alpha` is out of range.
    pub fn validate(&self) -> crate::error::Result<()> {
        if !self.alpha.is_finite() || self.alpha < 1.0 {
            return Err(crate::error::Error::Config(format!(
                "hnsw alpha must be finite and >= 1.0 (VAMANA range), got {}",
                self.alpha
            )));
        }
        Ok(())
    }
}

/// Search quality profile controlling the recall/latency tradeoff.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
#[non_exhaustive]
pub enum SearchQuality {
    /// Fast search with `ef_search=96`: 97.4% recall@10 in `recall_benchmark`
    /// (10K random 128-D points), lowest latency.
    Fast,
    /// Balanced search with `ef_search=160`: 99.8% recall@10 in
    /// `recall_benchmark`, production default.
    #[default]
    Balanced,
    /// Accurate search with `ef_search=512`: 100% recall@10 in
    /// `recall_benchmark`, 0.98 on SIFT1M's 1M (`docs/BENCHMARKS.md`).
    Accurate,
    /// Exhaustive unless the index's exact-distance features are off: every
    /// stored vector is scored, so it returns the exact top-k under the index's
    /// own distance, ties aside. This variant leaves the graph —
    /// `try_search_special_quality` routes it straight to `search_brute_force`
    /// before `ef_search` is ever consulted, so [`Self::ef_search`]'s
    /// `4096.max(k * 100)` is not the number this mode runs at. The scan is
    /// O(n / cores). An index built with those features off (`new_fast_insert`,
    /// `with_params_full(.., false)`) still keeps its vectors in the graph, but
    /// there `search_brute_force` falls back to a graph search at `Accurate`'s
    /// scaled ef.
    ///
    /// **Guarded, and the guard is the reason to read this.** A collection
    /// larger than `limits.max_perfect_mode_vectors` (default 500 000) makes
    /// the search fail with [`crate::error::Error::GuardRail`] rather than run
    /// a linear scan over it. Verified in
    /// `perfect_quality_is_refused_above_the_configured_cap`.
    ///
    /// This paragraph said the opposite until #2238 — "tunes the HNSW graph's
    /// effort and is not exhaustive". Its ~0.9994 figure at 1M was measured
    /// through this very scan, on `HnswIndex` directly, which has no cap
    /// (#1225); a collection that size is refused. The paragraph was read off
    /// `ef_search()` without checking which arm runs.
    Perfect,
    /// Custom `ef_search` value.
    Custom(usize),
    /// Adaptive `ef_search` that starts low and doubles once if the query is
    /// "hard".
    ///
    /// Uses a two-phase approach:
    /// 1. Search at `max(min_ef, k)`
    /// 2. If the result spread (the first-to-last score gap over the lower
    ///    score's distance from the metric's floor, or over the smaller
    ///    absolute score on an unbounded metric) reaches 2.0, search once more
    ///    at twice its ef, capped at `max_ef`, when that exceeds the first ef
    ///
    /// Easy queries (dense cluster hits) stop after phase 1. No recorded run
    /// measures its latency or recall against a fixed ef yet (#2266).
    ///
    /// # Phase 2 is out of reach on Cosine
    ///
    /// The spread is divided by the lower score's distance from the metric's
    /// floor, and Cosine's floor is `-1`. So `spread = (first - last) /
    /// (last + 1)`, and reaching 2.0 with a perfect top hit needs
    /// `last <= -1/3`: the **k-th** result must be anti-correlated with the
    /// query. Ordinary data does not do that — in 256 dimensions the best of
    /// 3,000 random vectors sits near `+0.17`, not below `-0.33`.
    ///
    /// Measured on one 3,000 x 256 corpus built to be hard (clusters smaller
    /// than `k`, so every top-`k` tail falls into the far background),
    /// phase 1 at `min_ef = 32`, `k = 10`:
    ///
    /// | metric | first | last | spread | queries over 2.0 |
    /// | --- | --- | --- | --- | --- |
    /// | Cosine | 0.9988 | 0.16 to 0.18 | 0.70 to 0.72 | 0 / 40 |
    /// | Euclidean | 0.049 | 1.28 to 1.29 | 24.9 to 25.1 | 40 / 40 |
    ///
    /// On a Cosine collection this variant is therefore a single pass at
    /// `min_ef`, whatever `max_ef` says, and so is [`Self::AutoTune`], which
    /// delegates to the same algorithm. Whether the threshold or the Cosine
    /// baseline should change is a behaviour question, open on #2266: moving
    /// either one trades latency for recall on every Cosine query, and no run
    /// measures that trade yet.
    Adaptive {
        /// Minimum `ef_search` (starting point). Default: 32.
        min_ef: usize,
        /// Maximum `ef_search` (cap). Default: 512.
        max_ef: usize,
    },
    /// Auto-tuned adaptive search based on collection statistics.
    ///
    /// Computes `min_ef` / `max_ef` from the collection's current size and
    /// vector dimension, then delegates to the same two-phase adaptive
    /// algorithm used by [`SearchQuality::Adaptive`]. `min_ef` grows by tiers:
    ///
    /// - up to 1K vectors: `k * 2`
    /// - 1K–10K: `k * 4`
    /// - 10K–100K: `k * 8`
    /// - more than 100K: `k * 12`
    /// - above 512 dimensions: times 1.5
    ///
    /// `min_ef` is at least `k`, and `max_ef` is `4 * min_ef`. Some search paths
    /// run one pass, scan exactly or ignore the mode instead: the Search Modes
    /// guide names the main ones ("When the two phases run"), and #2268
    /// tracks the rest. No recorded run measures its latency or recall (#2266).
    ///
    /// # Example
    ///
    /// ```rust,ignore
    /// use velesdb_core::SearchQuality;
    /// let results = index.search_with_quality(&query, 10, SearchQuality::AutoTune);
    /// ```
    AutoTune,
}

impl SearchQuality {
    /// Returns the `ef_search` value for this quality profile.
    ///
    /// # Large-scale optimization (v0.9+)
    ///
    /// - **Accurate**: 512 base (was 256), scales with k×16 to hold recall at 100K+
    /// - **Perfect**: 4096 base, scaled with k×100 — the ef a graph traversal
    ///   would use. With its exact-distance features on, `HnswIndex` does not
    ///   traverse for `Perfect`: it scans exhaustively (#2238), so this is not
    ///   what `Perfect` costs there.
    /// - **Adaptive**: returns `min_ef` (first phase); caller handles second phase
    #[must_use]
    pub fn ef_search(&self, k: usize) -> usize {
        match self {
            Self::Fast => 96.max(k * 3),
            // AutoTune: resolved at search time via auto_ef_range(); fallback
            // to Balanced for contexts that call ef_search() without collection info.
            Self::Balanced | Self::AutoTune => 160.max(k * 5),
            // Increased from 256 to 512 for better recall at 100K+ scale
            Self::Accurate => 512.max(k * 16),
            // The traversal ef for Perfect-class recall, for a caller that walks the
            // graph; `HnswIndex` scans exhaustively instead (#2238) unless its
            // exact-distance features are off.
            Self::Perfect => 4096.max(k * 100),
            Self::Custom(ef) => (*ef).max(k),
            // Adaptive: start with min_ef (first phase)
            Self::Adaptive { min_ef, .. } => (*min_ef).max(k),
        }
    }

    /// Returns `ef_search` scaled to dataset size.
    ///
    /// For datasets > 10K, the HNSW graph is deeper and requires a slightly
    /// larger search budget. Uses sqrt-based scaling (gentler than log2) to
    /// avoid excessive latency when the graph is already well-constructed
    /// with appropriate M and `ef_construction` parameters.
    ///
    /// Capped at 2x the base ef to prevent latency explosion.
    #[must_use]
    #[allow(
        clippy::cast_possible_truncation,
        clippy::cast_sign_loss,
        clippy::cast_precision_loss
    )]
    pub fn ef_search_for_scale(&self, k: usize, dataset_size: usize) -> usize {
        let base = self.ef_search(k);
        if dataset_size <= 10_000 {
            return base;
        }
        // sqrt scaling: gentle increase that doesn't destroy latency.
        // At 100K: sqrt(10) ≈ 3.16 → capped at 2.0 → ef * 2.
        // At 1M: sqrt(100) = 10 → capped at 2.0 → ef * 2.
        // Reason: dataset_size and base are small enough that f64 is exact.
        let ratio = dataset_size as f64 / 10_000.0;
        let scale = ratio.sqrt().min(2.0);
        let scaled = (base as f64 * scale) as usize;
        scaled.min(base * 2)
    }

    /// Returns `true` if this quality profile uses two-phase adaptive search.
    #[must_use]
    pub const fn is_adaptive(&self) -> bool {
        matches!(self, Self::Adaptive { .. } | Self::AutoTune)
    }

    /// Returns the maximum ef for adaptive search, or `None` for fixed profiles.
    #[must_use]
    pub const fn adaptive_max_ef(&self) -> Option<usize> {
        match self {
            Self::Adaptive { max_ef, .. } => Some(*max_ef),
            _ => None,
        }
    }
}