oxirs_vec/

index_dispatcher.rs

//! Runtime index dispatcher (wraps the optimizer brain with concrete indices).
//!
//! This is the public entry point for the cost-based optimizer described in
//! the `0.3.0` "Advanced query optimization" milestone:
//!
//! * The dispatcher owns one instance of each enabled index family
//!   (HNSW / IVF / LSH / PQ).
//! * It delegates "which family should I pick?" to
//!   [`crate::optimizer::OptimizerDispatcher`] (a cost-model + stats brain).
//! * It executes the query against the picked family, observes the actual
//!   latency and result count, feeds an observation back to the brain, and
//!   re-issues against the next-best family if the **observed recall** trips
//!   the dispatcher's recall threshold.
//!
//! The split between "brain" (in `optimizer/`) and "wrapper" (this file)
//! mirrors the split between `query_planning.rs` (brain-only) and
//! `dynamic_index_selector.rs` (wrapper) but uses the new cost-model formulas
//! and the persisted `QueryStats`.
//!
//! Persistence: the dispatcher periodically saves its [`QueryStats`] to
//! disk for online learning across process restarts.  Use
//! [`IndexDispatcher::with_stats_path`] to enable persistence.

use crate::hnsw::{HnswConfig, HnswIndex};
use crate::ivf::{IvfConfig, IvfIndex};
use crate::lsh::{LshConfig, LshIndex};
use crate::optimizer::cost_model::{CostModel, IndexFamily, IndexParameters, WorkloadProfile};
use crate::optimizer::index_dispatcher::{DispatchPlan, DispatcherConfig, OptimizerDispatcher};
use crate::optimizer::query_stats::{QueryObservation, QueryStats};
use crate::pq::{PQConfig, PQIndex};
use crate::{Vector, VectorIndex};
use anyhow::{anyhow, Result};
use std::path::PathBuf;
use std::sync::{Arc, RwLock};
use std::time::Instant;
use tracing::{debug, warn};

/// User-facing configuration for the runtime index dispatcher.
#[derive(Debug, Clone)]
pub struct IndexDispatcherConfig {
    /// Cost-model formula parameters.
    pub parameters: IndexParameters,
    /// Brain-level dispatcher config (recall threshold, weight refresh, etc.).
    pub dispatcher: DispatcherConfig,
    /// HNSW config used when the dispatcher chooses the HNSW family.
    pub hnsw_config: HnswConfig,
    /// IVF config used when the dispatcher chooses the IVF family.
    pub ivf_config: IvfConfig,
    /// LSH config used when the dispatcher chooses the LSH family.
    pub lsh_config: LshConfig,
    /// PQ config used when the dispatcher chooses the PQ family.
    pub pq_config: PQConfig,
    /// File on disk where [`QueryStats`] is persisted (`None` disables it).
    pub stats_path: Option<PathBuf>,
    /// Save the stats file after this many observations (rate-limit IO).
    pub stats_save_interval: u64,
}

impl Default for IndexDispatcherConfig {
    fn default() -> Self {
        Self {
            parameters: IndexParameters::default(),
            dispatcher: DispatcherConfig::default(),
            hnsw_config: HnswConfig::default(),
            ivf_config: IvfConfig::default(),
            lsh_config: LshConfig::default(),
            pq_config: PQConfig::default(),
            stats_path: None,
            stats_save_interval: 256,
        }
    }
}

/// One row produced by [`IndexDispatcher::search_knn_with_plan`].
#[derive(Debug, Clone)]
pub struct DispatchedSearch {
    /// k-NN result list returned by the chosen index.
    pub results: Vec<(String, f32)>,
    /// Family that ultimately served the query (after fallbacks).
    pub served_by: IndexFamily,
    /// Plan generated by the optimizer brain at dispatch time.
    pub plan: DispatchPlan,
    /// Number of fallback re-issues performed (0 = primary served).
    pub fallback_attempts: usize,
    /// Wall-clock latency in microseconds for the total dispatch (incl. retries).
    pub latency_us: f64,
}

/// Runtime index dispatcher.
///
/// Workflow:
///
/// ```ignore
/// let mut d = IndexDispatcher::new(config)?;
/// d.insert("v1".into(), vector)?;   // buffered
/// d.insert("v2".into(), vector)?;
/// // ...
/// d.build()?;                        // trains IVF + PQ, materialises HNSW + LSH
/// let res = d.search_knn(&query, 10)?;
/// ```
///
/// Insert-buffering is required because IVF and PQ both need a training
/// pass over a representative sample before any vectors can be inserted.
/// HNSW and LSH have no training phase but are also built during
/// [`Self::build`] for symmetry.
pub struct IndexDispatcher {
    config: IndexDispatcherConfig,
    brain: Arc<RwLock<OptimizerDispatcher>>,
    /// Buffered insertions waiting for `build()` to materialise indices.
    pending: Vec<(String, Vector)>,
    hnsw: Option<HnswIndex>,
    ivf: Option<IvfIndex>,
    lsh: Option<LshIndex>,
    pq: Option<PQIndex>,
    /// Current vector count across the dispatcher (used for workload profile).
    vector_count: usize,
    /// Current vector dimensionality.
    vector_dim: usize,
    /// `true` once `build()` has been called and indices are searchable.
    is_built: bool,
}

impl IndexDispatcher {
    /// Build a fresh dispatcher.  No indices are created yet — call
    /// [`Self::insert`] then [`Self::build`] (or rely on lazy build inside
    /// [`Self::insert`] for in-memory indices that don't need a build phase).
    pub fn new(config: IndexDispatcherConfig) -> Result<Self> {
        let cost_model = CostModel::new(config.parameters.clone(), Default::default());

        // Restore stats from disk if a path is configured and the file exists.
        let stats = if let Some(path) = &config.stats_path {
            if path.exists() {
                match QueryStats::load(path) {
                    Ok(s) => {
                        debug!(
                            "IndexDispatcher: loaded {} observations from {:?}",
                            s.total_observations, path
                        );
                        s
                    }
                    Err(e) => {
                        warn!(
                            "IndexDispatcher: failed to load stats from {:?}: {} — starting fresh",
                            path, e
                        );
                        QueryStats::default()
                    }
                }
            } else {
                QueryStats::default()
            }
        } else {
            QueryStats::default()
        };

        let brain = OptimizerDispatcher::new(cost_model, stats, config.dispatcher.clone());

        Ok(Self {
            config,
            brain: Arc::new(RwLock::new(brain)),
            pending: Vec::new(),
            hnsw: None,
            ivf: None,
            lsh: None,
            pq: None,
            vector_count: 0,
            vector_dim: 0,
            is_built: false,
        })
    }

    /// Convenience: dispatcher with default config and a stats file path.
    pub fn with_stats_path(stats_path: PathBuf) -> Result<Self> {
        let config = IndexDispatcherConfig {
            stats_path: Some(stats_path),
            ..Default::default()
        };
        Self::new(config)
    }

    /// Buffer a vector for insertion.  Indices are not built until
    /// [`Self::build`] is called.
    ///
    /// After `build()` completes, additional inserts go directly into the
    /// already-trained HNSW and LSH indices; IVF and PQ skip them since they
    /// would require re-training.
    pub fn insert(&mut self, uri: String, vector: Vector) -> Result<()> {
        if self.vector_dim == 0 {
            self.vector_dim = vector.dimensions;
        } else if vector.dimensions != self.vector_dim {
            return Err(anyhow!(
                "IndexDispatcher::insert: dim mismatch (have {}, got {})",
                self.vector_dim,
                vector.dimensions
            ));
        }

        if !self.is_built {
            self.pending.push((uri, vector));
            return Ok(());
        }

        // Post-build inserts: only HNSW and LSH (no training needed).
        if let Some(hnsw) = &mut self.hnsw {
            hnsw.insert(uri.clone(), vector.clone())?;
        }
        if let Some(lsh) = &mut self.lsh {
            lsh.insert(uri, vector)?;
        }
        self.vector_count += 1;
        Ok(())
    }

    /// Materialise every enabled family from the buffered inserts.
    ///
    /// IVF and PQ are trained on the buffered set (or a 10 000-vector
    /// sample, whichever is smaller); HNSW and LSH are bulk-inserted.
    pub fn build(&mut self) -> Result<()> {
        if self.is_built {
            return Ok(());
        }
        if self.pending.is_empty() {
            // Nothing to build — leave indices `None`; search will error.
            self.is_built = true;
            return Ok(());
        }

        // ── HNSW ────────────────────────────────────────────────────────────
        let mut hnsw = HnswIndex::new(self.config.hnsw_config.clone())
            .map_err(|e| anyhow!("IndexDispatcher::build: HnswIndex::new failed: {}", e))?;
        for (uri, v) in &self.pending {
            hnsw.insert(uri.clone(), v.clone())?;
        }
        self.hnsw = Some(hnsw);

        // ── LSH ─────────────────────────────────────────────────────────────
        let mut lsh = LshIndex::new(self.config.lsh_config.clone());
        for (uri, v) in &self.pending {
            lsh.insert(uri.clone(), v.clone())?;
        }
        self.lsh = Some(lsh);

        // ── IVF (must train first, then insert) ─────────────────────────────
        // Cap training set at 10k samples to keep build time bounded.
        let sample_size = self.pending.len().min(10_000);
        let training_set: Vec<Vector> = self
            .pending
            .iter()
            .take(sample_size)
            .map(|(_, v)| v.clone())
            .collect();

        let mut ivf_config = self.config.ivf_config.clone();
        // Keep n_clusters ≤ training set size, otherwise k-means cannot run.
        if ivf_config.n_clusters > sample_size {
            ivf_config.n_clusters = sample_size.max(1);
        }
        let mut ivf = IvfIndex::new(ivf_config)?;
        ivf.train(&training_set)?;
        for (uri, v) in &self.pending {
            ivf.insert(uri.clone(), v.clone())?;
        }
        self.ivf = Some(ivf);

        // ── PQ (must train first, then insert) ──────────────────────────────
        let pq_dim = self.pending[0].1.dimensions;
        let mut pq_config = self.config.pq_config.clone();
        // PQ requires `n_subquantizers` to divide `dim` exactly; clamp to
        // the largest divisor ≤ the configured value.
        if pq_dim % pq_config.n_subquantizers != 0 {
            // Find largest k ≤ configured that divides pq_dim.
            let mut k = pq_config.n_subquantizers.min(pq_dim).max(1);
            while k > 1 && pq_dim % k != 0 {
                k -= 1;
            }
            pq_config.n_subquantizers = k;
        }
        let mut pq = PQIndex::new(pq_config);
        pq.train(&training_set)?;
        for (uri, v) in &self.pending {
            pq.insert(uri.clone(), v.clone())?;
        }
        self.pq = Some(pq);

        self.vector_count = self.pending.len();
        self.pending.clear();
        self.is_built = true;
        Ok(())
    }

    /// `true` once [`Self::build`] has been invoked.
    pub fn is_built(&self) -> bool {
        self.is_built
    }

    /// Number of vectors indexed.
    pub fn len(&self) -> usize {
        self.vector_count
    }

    /// `true` when no vectors are indexed.
    pub fn is_empty(&self) -> bool {
        self.vector_count == 0
    }

    /// Search the dispatcher with default workload (`k`, density=1.0,
    /// recall=brain config).  This is the simple entry point.
    pub fn search_knn(&self, query: &Vector, k: usize) -> Result<Vec<(String, f32)>> {
        let dispatched = self.search_knn_with_plan(query, k, 1.0)?;
        Ok(dispatched.results)
    }

    /// Search with explicit query density (filter selectivity).  Returns
    /// the full [`DispatchedSearch`] for observability.
    pub fn search_knn_with_plan(
        &self,
        query: &Vector,
        k: usize,
        query_density: f32,
    ) -> Result<DispatchedSearch> {
        if !self.is_built {
            return Err(anyhow!(
                "IndexDispatcher::search_knn: dispatcher must be built first (call build())"
            ));
        }
        if self.vector_count == 0 {
            return Err(anyhow!("IndexDispatcher::search_knn: no vectors indexed"));
        }

        let workload = WorkloadProfile::new(
            self.vector_count,
            self.vector_dim,
            self.config.dispatcher.recall_fallback_threshold,
        )
        .with_query_density(query_density)
        .with_k(k);

        let plan = {
            let brain = self
                .brain
                .read()
                .map_err(|_| anyhow!("dispatcher brain RwLock poisoned"))?;
            brain
                .pick_plan(&workload)
                .map_err(|e| anyhow!("dispatcher brain failed to plan: {}", e))?
        };

        let start = Instant::now();
        let mut current = plan.primary;
        let mut fallback_attempts = 0;
        let mut results = self.execute_with_family(current, query, k)?;

        // Apply fallback policy: if results came up empty *and* the brain
        // says we have fallback budget, retry with the next-best family.
        let max_fallbacks = self.config.dispatcher.max_fallbacks;
        let observed_recall_proxy = if results.is_empty() { 0.0 } else { 1.0 };
        if max_fallbacks > 0 && results.is_empty() {
            for next_family in plan.fallbacks.iter().map(|e| e.family).take(max_fallbacks) {
                fallback_attempts += 1;
                tracing::info!(
                    "IndexDispatcher: empty result from {:?}, falling back to {:?}",
                    current,
                    next_family
                );
                current = next_family;
                results = self.execute_with_family(current, query, k)?;
                if !results.is_empty() {
                    break;
                }
            }
        }

        let elapsed_us = start.elapsed().as_secs_f64() * 1_000_000.0;

        // Record observation for online learning.
        let observation = QueryObservation::new(
            current,
            !results.is_empty(),
            elapsed_us,
            // We don't have ground truth; record the recall *proxy* (1.0 if hit, 0.0 else)
            // so the stats reflect operational hit rate.
            Some(observed_recall_proxy),
            plan.primary_cost,
        );
        {
            let mut brain = self
                .brain
                .write()
                .map_err(|_| anyhow!("dispatcher brain RwLock poisoned"))?;
            let refreshed = brain.record_observation(observation);
            if refreshed {
                debug!("IndexDispatcher: refreshed cost-model weights");
            }
        }

        // Optionally persist stats.
        self.maybe_persist_stats()?;

        Ok(DispatchedSearch {
            results,
            served_by: current,
            plan,
            fallback_attempts,
            latency_us: elapsed_us,
        })
    }

    /// Force-flush stats to disk now (no-op if no path configured).
    pub fn flush_stats(&self) -> Result<()> {
        if let Some(path) = &self.config.stats_path {
            let brain = self
                .brain
                .read()
                .map_err(|_| anyhow!("dispatcher brain RwLock poisoned"))?;
            brain.stats().save(path)?;
        }
        Ok(())
    }

    /// Snapshot of internal observation counts for diagnostics.
    pub fn observation_count(&self) -> Result<u64> {
        let brain = self
            .brain
            .read()
            .map_err(|_| anyhow!("dispatcher brain RwLock poisoned"))?;
        Ok(brain.stats().total_observations)
    }

    fn execute_with_family(
        &self,
        family: IndexFamily,
        query: &Vector,
        k: usize,
    ) -> Result<Vec<(String, f32)>> {
        match family {
            IndexFamily::Hnsw => self
                .hnsw
                .as_ref()
                .map(|i| i.search_knn(query, k))
                .unwrap_or_else(|| Err(anyhow!("HNSW family not built"))),
            IndexFamily::Ivf => self
                .ivf
                .as_ref()
                .map(|i| i.search_knn(query, k))
                .unwrap_or_else(|| Err(anyhow!("IVF family not built"))),
            IndexFamily::Lsh => self
                .lsh
                .as_ref()
                .map(|i| i.search_knn(query, k))
                .unwrap_or_else(|| Err(anyhow!("LSH family not built"))),
            IndexFamily::Pq => self
                .pq
                .as_ref()
                .map(|i| i.search_knn(query, k))
                .unwrap_or_else(|| Err(anyhow!("PQ family not built"))),
        }
    }

    fn maybe_persist_stats(&self) -> Result<()> {
        if self.config.stats_path.is_none() {
            return Ok(());
        }
        // Saves are throttled to once per `stats_save_interval` observations.
        // We approximate this by checking the brain's observation counter.
        let brain = self
            .brain
            .read()
            .map_err(|_| anyhow!("dispatcher brain RwLock poisoned"))?;
        let total = brain.stats().total_observations;
        let interval = self.config.stats_save_interval.max(1);
        if total % interval == 0 && total > 0 {
            if let Some(path) = &self.config.stats_path {
                if let Err(e) = brain.stats().save(path) {
                    warn!("IndexDispatcher: stats save to {:?} failed: {}", path, e);
                }
            }
        }
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::env::temp_dir;

    fn unique_stats_path() -> PathBuf {
        let stamp = std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .map(|d| d.as_nanos())
            .unwrap_or(0);
        let mut p = temp_dir();
        p.push(format!("oxirs_vec_dispatcher_{}.json", stamp));
        p
    }

    fn random_vec(seed: u64, dim: usize) -> Vector {
        let mut state = seed.wrapping_mul(2654435769).wrapping_add(0x9E37_79B9);
        let mut values = Vec::with_capacity(dim);
        for _ in 0..dim {
            state = state.wrapping_mul(6364136223846793005).wrapping_add(1);
            let f = (state as f32) / (u64::MAX as f32);
            values.push((f - 0.5) * 2.0);
        }
        Vector::new(values)
    }

    /// Build a dispatcher config tuned for small test datasets — IVF needs
    /// fewer clusters than the default 256 when working with <50 vectors.
    fn small_test_config() -> IndexDispatcherConfig {
        let mut cfg = IndexDispatcherConfig::default();
        // ivf_config is itself a Default-initialised struct; mutating the
        // public fields here is appropriate — `clippy::field_reassign_with_default`
        // does not apply when we're modifying a sub-field of an
        // already-extracted variable.
        cfg.ivf_config.n_clusters = 4;
        cfg.ivf_config.n_probes = 2;
        cfg
    }

    #[test]
    fn dispatcher_can_insert_and_search() -> Result<()> {
        let mut d = IndexDispatcher::new(small_test_config())?;
        for i in 0..50 {
            d.insert(format!("v{}", i), random_vec(i as u64 + 1, 32))?;
        }
        d.build()?;
        let q = random_vec(1, 32);
        let results = d.search_knn(&q, 5)?;
        // Some non-empty subset of 5 results expected.
        assert!(!results.is_empty());
        Ok(())
    }

    #[test]
    fn search_with_plan_returns_metadata() -> Result<()> {
        let mut d = IndexDispatcher::new(small_test_config())?;
        for i in 0..50 {
            d.insert(format!("v{}", i), random_vec(i as u64 + 1, 16))?;
        }
        d.build()?;
        let q = random_vec(2, 16);
        let dispatched = d.search_knn_with_plan(&q, 4, 1.0)?;
        assert!(dispatched.latency_us >= 0.0);
        // plan.primary must equal served_by when no fallback.
        assert_eq!(dispatched.plan.primary, dispatched.served_by);
        Ok(())
    }

    #[test]
    fn dispatcher_persists_stats_to_disk() -> Result<()> {
        let path = unique_stats_path();
        let mut config = small_test_config();
        config.stats_path = Some(path.clone());
        config.stats_save_interval = 1; // Save after every observation
        let mut d = IndexDispatcher::new(config)?;
        for i in 0..16 {
            d.insert(format!("v{}", i), random_vec(i as u64 + 1, 8))?;
        }
        d.build()?;
        let q = random_vec(99, 8);
        let _ = d.search_knn(&q, 3)?;
        assert!(path.exists(), "stats file must be created");
        let loaded = QueryStats::load(&path)?;
        assert!(loaded.total_observations >= 1);
        let _ = std::fs::remove_file(&path);
        Ok(())
    }

    #[test]
    fn search_on_unbuilt_dispatcher_errors() -> Result<()> {
        let d = IndexDispatcher::new(IndexDispatcherConfig::default())?;
        let q = random_vec(1, 4);
        let res = d.search_knn(&q, 3);
        assert!(res.is_err());
        Ok(())
    }

    #[test]
    fn mismatched_dim_errors() -> Result<()> {
        let mut d = IndexDispatcher::new(IndexDispatcherConfig::default())?;
        d.insert("v1".into(), random_vec(1, 8))?;
        let res = d.insert("v2".into(), random_vec(2, 16));
        assert!(res.is_err(), "dim mismatch must error");
        Ok(())
    }

    #[test]
    fn flush_stats_no_op_without_path() -> Result<()> {
        let d = IndexDispatcher::new(IndexDispatcherConfig::default())?;
        d.flush_stats()?;
        Ok(())
    }

    #[test]
    fn flush_stats_writes_file_when_path_set() -> Result<()> {
        let path = unique_stats_path();
        let d = IndexDispatcher::with_stats_path(path.clone())?;
        d.flush_stats()?;
        assert!(path.exists(), "flush must create the file");
        let _ = std::fs::remove_file(&path);
        Ok(())
    }

    #[test]
    fn restart_loads_previous_stats() -> Result<()> {
        let path = unique_stats_path();
        // First run: produce some observations.
        {
            let mut config = small_test_config();
            config.stats_path = Some(path.clone());
            config.stats_save_interval = 1;
            let mut d = IndexDispatcher::new(config)?;
            for i in 0..16 {
                d.insert(format!("v{}", i), random_vec(i as u64 + 1, 4))?;
            }
            d.build()?;
            let q = random_vec(99, 4);
            let _ = d.search_knn(&q, 2)?;
            d.flush_stats()?;
        }
        // Second run: should reload observations.
        let mut config = small_test_config();
        config.stats_path = Some(path.clone());
        let d2 = IndexDispatcher::new(config)?;
        let n = d2.observation_count()?;
        assert!(n >= 1, "second run must load at least 1 observation");
        let _ = std::fs::remove_file(&path);
        Ok(())
    }

    #[test]
    fn build_is_idempotent() -> Result<()> {
        let mut d = IndexDispatcher::new(small_test_config())?;
        for i in 0..16 {
            d.insert(format!("v{}", i), random_vec(i as u64 + 1, 8))?;
        }
        d.build()?;
        // Calling build again is a no-op.
        d.build()?;
        assert!(d.is_built());
        Ok(())
    }

    #[test]
    fn post_build_inserts_go_to_hnsw_and_lsh() -> Result<()> {
        let mut d = IndexDispatcher::new(small_test_config())?;
        for i in 0..16 {
            d.insert(format!("v{}", i), random_vec(i as u64 + 1, 8))?;
        }
        d.build()?;
        let pre = d.len();
        d.insert("late".into(), random_vec(999, 8))?;
        assert_eq!(d.len(), pre + 1);
        Ok(())
    }
}