tidaldb/tidal/benches/vector.rs

#![allow(clippy::unwrap_used)]

//! Criterion benchmarks for the vector index subsystem.
//!
//! Measures ANN search latency across the selectivity spectrum: unfiltered,
//! filtered (20%), widened filtered (5%, ef=400), and high-selectivity
//! filtered (0.5%). Also benchmarks recall@100, single insert, and single
//! delete. Calls the `VectorIndex` search API directly.
//!
//! All setup (index construction, vector insertion) is done OUTSIDE the
//! `b.iter()` closure. Only the search/insert/delete call is measured.

use std::collections::HashSet;

use criterion::{Criterion, black_box, criterion_group, criterion_main};
use rand::Rng;
use tidaldb::storage::vector::{
    BruteForceIndex, DistanceMetric, QuantizationLevel, UsearchIndex, VectorId, VectorIndex,
    VectorIndexConfig,
};

// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------

/// Generate a random unit vector of the given dimensionality.
fn random_unit_vector(dim: usize, rng: &mut impl Rng) -> Vec<f32> {
    let v: Vec<f32> = (0..dim)
        .map(|_| {
            let x: f32 = rng.random();
            x - 0.5
        })
        .collect();
    let norm: f32 = v.iter().map(|x| x * x).sum::<f32>().sqrt();
    if norm < f32::EPSILON {
        // Degenerate case: return a unit vector along the first axis.
        let mut fallback = vec![0.0_f32; dim];
        fallback[0] = 1.0;
        return fallback;
    }
    v.iter().map(|x| x / norm).collect()
}

/// Build a brute-force index with `n` random unit vectors of dimension `dim`.
fn build_brute_index(n: u64, dim: usize) -> BruteForceIndex {
    let config = VectorIndexConfig {
        dimensions: dim,
        metric: DistanceMetric::L2,
        quantization: QuantizationLevel::F32,
        connectivity: 16,
        ef_construction: 200,
        ef_search: 200,
    };
    let index = BruteForceIndex::new(config);
    let mut rng = rand::rng();
    for id in 0..n {
        let vec = random_unit_vector(dim, &mut rng);
        index.insert(id, &vec).unwrap();
    }
    index
}

// ---------------------------------------------------------------------------
// Benchmarks
// ---------------------------------------------------------------------------

/// Benchmark: unfiltered ANN search over 10K vectors, dim=128, k=100.
/// Measures baseline search latency without any filter overhead.
fn bench_ann_search_unfiltered(c: &mut Criterion) {
    let dim = 128;
    let n = 10_000_u64;
    let index = build_brute_index(n, dim);

    let mut rng = rand::rng();
    let query = random_unit_vector(dim, &mut rng);

    c.bench_function("ann_search_unfiltered_10k", |b| {
        b.iter(|| {
            index
                .search(black_box(&query), black_box(100), black_box(200))
                .unwrap()
        });
    });
}

/// Benchmark: filtered ANN search with 20% selectivity (in-graph filter).
/// 10K vectors, dim=128, k=100.
fn bench_ann_search_filtered_20pct(c: &mut Criterion) {
    let dim = 128;
    let n = 10_000_u64;
    let index = build_brute_index(n, dim);

    let mut rng = rand::rng();
    let query = random_unit_vector(dim, &mut rng);

    // ~20% selectivity: IDs 0..1999 pass (20% of 10K).
    let filter = |id: VectorId| id < 2000;

    c.bench_function("ann_search_filtered_20pct_10k", |b| {
        b.iter(|| {
            index
                .filtered_search(black_box(&query), black_box(100), black_box(200), &filter)
                .unwrap()
        });
    });
}

/// Benchmark: filtered ANN search with 5% selectivity (widened filter, ef=400).
/// 10K vectors, dim=128, k=100.
fn bench_ann_search_filtered_5pct(c: &mut Criterion) {
    let dim = 128;
    let n = 10_000_u64;
    let index = build_brute_index(n, dim);

    let mut rng = rand::rng();
    let query = random_unit_vector(dim, &mut rng);

    // ~5% selectivity: IDs 0..499 pass (5% of 10K). Widened beam (ef=400).
    let filter = |id: VectorId| id < 500;

    c.bench_function("ann_search_filtered_5pct_10k", |b| {
        b.iter(|| {
            index
                .filtered_search(black_box(&query), black_box(100), black_box(400), &filter)
                .unwrap()
        });
    });
}

/// Benchmark: pre-filter brute-force search with 0.5% selectivity.
/// 10K vectors, dim=128, k=100. Uses a separate brute-force index.
fn bench_ann_search_brute_force(c: &mut Criterion) {
    let dim = 128;
    let n = 10_000_u64;
    let index = build_brute_index(n, dim);

    let mut rng = rand::rng();
    let query = random_unit_vector(dim, &mut rng);

    // ~0.5% selectivity: IDs 0..49 pass (0.5% of 10K).
    let filter = |id: VectorId| id < 50;

    c.bench_function("ann_search_brute_force_10k", |b| {
        b.iter(|| {
            index
                .filtered_search(black_box(&query), black_box(100), black_box(200), &filter)
                .unwrap()
        });
    });
}

/// Benchmark: recall@100 measurement.
/// Builds a 10K brute-force index, runs search, and compares against
/// ground truth (which for brute-force is exact). This benchmarks the
/// search + comparison loop to establish a baseline measurement cost.
fn bench_ann_recall_at_100(c: &mut Criterion) {
    let dim = 128;
    let n = 10_000_u64;
    let k = 100;
    let index = build_brute_index(n, dim);

    let mut rng = rand::rng();
    let query = random_unit_vector(dim, &mut rng);

    // Pre-compute ground truth.
    let ground_truth = index.search(&query, k, 200).unwrap();
    let gt_ids: Vec<VectorId> = ground_truth.iter().map(|r| r.id).collect();

    c.bench_function("ann_recall_at_100_10k", |b| {
        b.iter(|| {
            let results = index
                .search(black_box(&query), black_box(k), black_box(200))
                .unwrap();
            let result_ids: Vec<VectorId> = results.iter().map(|r| r.id).collect();

            // Compute recall: fraction of ground truth IDs found in results.
            let hits = result_ids.iter().filter(|id| gt_ids.contains(id)).count();
            #[allow(clippy::cast_precision_loss)]
            let recall = hits as f64 / gt_ids.len() as f64;
            black_box(recall)
        });
    });
}

/// Benchmark: HNSW **recall@10 at the production shape** (1536D, F16) vs an exact
/// brute-force ground truth over the SAME 10K vectors.
///
/// This extends the 128D `bench_ann_recall_at_100` to the production embedding
/// width (m12p1): unlike the brute-force variants above (which are exact, so
/// recall is trivially 1.0 and the bench measures only the search+compare cost),
/// this builds a real `UsearchIndex` (HNSW) and measures the approximation's
/// recall against a `BruteForceIndex` oracle on each iteration. It is a LOCAL
/// micro-tripwire at one shape; the authoritative recall@10 + true p99 across
/// 100k/1M is `tidal-stress --verify-recall` (open-loop, real server).
fn bench_ann_recall_at_10_1536d(c: &mut Criterion) {
    let dim = 1536;
    let n = 10_000_u64;
    let recall_k = 10;

    // Generate the corpus ONCE and insert the identical vectors into both indexes
    // so the brute-force result is a true ground truth for the HNSW result.
    let mut rng = rand::rng();
    let vectors: Vec<Vec<f32>> = (0..n).map(|_| random_unit_vector(dim, &mut rng)).collect();

    let brute = BruteForceIndex::new(VectorIndexConfig {
        dimensions: dim,
        metric: DistanceMetric::L2,
        quantization: QuantizationLevel::F32,
        connectivity: 16,
        ef_construction: 400,
        ef_search: 200,
    });
    // Production HNSW posture: M=16, ef_construction=400, F16 quantization.
    let hnsw = UsearchIndex::new(VectorIndexConfig {
        dimensions: dim,
        metric: DistanceMetric::L2,
        quantization: QuantizationLevel::F16,
        connectivity: 16,
        ef_construction: 400,
        ef_search: 200,
    })
    .unwrap();
    hnsw.reserve(n as usize).unwrap();
    for (id, v) in vectors.iter().enumerate() {
        brute.insert(id as VectorId, v).unwrap();
        hnsw.insert(id as VectorId, v).unwrap();
    }

    let query = random_unit_vector(dim, &mut rng);
    // Exact top-`recall_k` ground truth from the brute-force index.
    let gt: HashSet<VectorId> = brute
        .search(&query, recall_k, 200)
        .unwrap()
        .iter()
        .map(|r| r.id)
        .collect();

    c.bench_function("ann_recall_at_10_1536d_10k", |b| {
        b.iter(|| {
            let results = hnsw
                .search(black_box(&query), black_box(recall_k), black_box(200))
                .unwrap();
            let hits = results.iter().filter(|r| gt.contains(&r.id)).count();
            #[allow(clippy::cast_precision_loss)]
            let recall = hits as f64 / recall_k as f64;
            black_box(recall)
        });
    });
}

/// Benchmark: single vector insert into a pre-filled 10K index.
fn bench_ann_insert_single(c: &mut Criterion) {
    let dim = 128;
    let n = 10_000_u64;
    let index = build_brute_index(n, dim);

    let mut rng = rand::rng();
    let vec = random_unit_vector(dim, &mut rng);

    // Use an ID outside the pre-filled range to avoid replacement overhead.
    let mut next_id = n;

    c.bench_function("ann_insert_single_10k", |b| {
        b.iter(|| {
            index.insert(black_box(next_id), black_box(&vec)).unwrap();
            next_id += 1;
        });
    });
}

/// Benchmark: single vector delete from a pre-filled 10K index.
/// After each delete, re-inserts the vector so the bench remains iterable.
fn bench_ann_delete_single(c: &mut Criterion) {
    let dim = 128;
    let n = 10_000_u64;
    let index = build_brute_index(n, dim);

    let mut rng = rand::rng();
    let vec = random_unit_vector(dim, &mut rng);

    // Target a fixed ID for delete/reinsert cycle.
    let target_id = 0_u64;

    c.bench_function("ann_delete_single_10k", |b| {
        b.iter(|| {
            index.delete(black_box(target_id)).unwrap();
            // Re-insert so the next iteration can delete it again.
            index.insert(black_box(target_id), black_box(&vec)).unwrap();
        });
    });
}

// ---------------------------------------------------------------------------
// Criterion group + main
// ---------------------------------------------------------------------------

criterion_group!(
    benches,
    bench_ann_search_unfiltered,
    bench_ann_search_filtered_20pct,
    bench_ann_search_filtered_5pct,
    bench_ann_search_brute_force,
    bench_ann_recall_at_100,
    bench_ann_recall_at_10_1536d,
    bench_ann_insert_single,
    bench_ann_delete_single,
);
criterion_main!(benches);