tidaldb/tidal/tests/m12p1_vector_recall.rs

#![allow(clippy::unwrap_used, clippy::cast_precision_loss)]
//! m12p1 — pure k-NN recall probe (`TidalDb::vector_search_items`).
//!
//! The recall harness (`tidal-stress --verify-recall`) measures ANN recall@k by
//! comparing the engine's nearest-neighbor result against a brute-force cosine
//! ground truth. This test exercises the engine-side surface that harness hits
//! — the raw k-NN probe — and proves it returns the true nearest neighbors
//! (recall == 1.0) at a corpus size where the slot index is exact brute-force,
//! so any later HNSW recall shortfall is attributable to the index, not the
//! plumbing.
//!
//! # UAT Scenario
//!
//! ```
//! Given: A db with an Item "content" embedding slot and N indexed vectors
//! When:  db.vector_search_items(q, 10, None)
//! Then:  Returns the 10 items nearest q by cosine, closest-first
//! And:   The set equals a brute-force cosine top-10 ground truth (recall 1.0)
//! ```

use std::{collections::HashMap, time::Duration};

use tidaldb::{
    TidalDb,
    schema::{DecaySpec, EntityId, EntityKind, SchemaBuilder, Window},
};

const DIM: usize = 16;
// Below `USEARCH_MIN_VECTORS` (10k) the slot is an exact BruteForceIndex, so the
// probe's recall against a brute-force ground truth must be exactly 1.0 — this
// isolates the plumbing from HNSW approximation, which the live harness measures.
const N: u64 = 500;

/// `SplitMix64` — a tiny deterministic generator so the corpus is reproducible
/// without a dev-dependency on a seeded RNG. Same scheme the live harness uses.
const fn splitmix64(state: &mut u64) -> u64 {
    *state = state.wrapping_add(0x9E37_79B9_7F4A_7C15);
    let mut z = *state;
    z = (z ^ (z >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9);
    z = (z ^ (z >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB);
    z ^ (z >> 31)
}

/// A deterministic, non-zero vector keyed by entity id (each component in
/// `[-0.5, 0.5)`). The engine L2-normalizes on write, so cosine order over these
/// raw vectors matches the engine's L2-on-normalized order.
fn vector_for(id: u64) -> Vec<f32> {
    let mut state = id.wrapping_mul(0x2545_F491_4F6C_DD1D).wrapping_add(1);
    (0..DIM)
        .map(|_| {
            let bits = (splitmix64(&mut state) >> 40) as u32; // 24 random bits
            (bits as f32 / f32::from(1u16 << 12) / 4096.0) - 0.5
        })
        .collect()
}

fn cosine(a: &[f32], b: &[f32]) -> f32 {
    let dot: f32 = a.iter().zip(b).map(|(x, y)| x * y).sum();
    let na: f32 = a.iter().map(|x| x * x).sum::<f32>().sqrt();
    let nb: f32 = b.iter().map(|x| x * x).sum::<f32>().sqrt();
    if na == 0.0 || nb == 0.0 {
        0.0
    } else {
        dot / (na * nb)
    }
}

/// Brute-force cosine top-`k` ground truth over the deterministic corpus.
fn ground_truth(query: &[f32], k: usize) -> Vec<u64> {
    let mut scored: Vec<(u64, f32)> = (1..=N)
        .map(|id| (id, cosine(query, &vector_for(id))))
        .collect();
    scored.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap());
    scored.into_iter().take(k).map(|(id, _)| id).collect()
}

fn build_db() -> TidalDb {
    let mut builder = SchemaBuilder::new();
    // A schema must declare at least one signal; the probe ignores signals
    // entirely, but the engine requires one to open.
    let _ = builder
        .signal(
            "view",
            EntityKind::Item,
            DecaySpec::Exponential {
                half_life: Duration::from_secs(7 * 24 * 3600),
            },
        )
        .windows(&[Window::TwentyFourHours])
        .velocity(false)
        .add();
    builder.embedding_slot("content", EntityKind::Item, DIM);
    let schema = builder.build().unwrap();
    let db = TidalDb::builder()
        .ephemeral()
        .with_schema(schema)
        .open()
        .unwrap();
    for id in 1..=N {
        // An item must exist for its embedding; metadata is irrelevant to k-NN.
        db.write_item_with_metadata(EntityId::new(id), &HashMap::new())
            .unwrap();
        db.write_item_embedding(EntityId::new(id), &vector_for(id))
            .unwrap();
    }
    db
}

#[test]
fn vector_search_returns_exact_nearest_neighbors() {
    let db = build_db();

    // Query with item 42's own vector: it must come back first, ~zero distance.
    let q = vector_for(42);
    let results = db.vector_search_items(&q, 10, None).unwrap();
    assert_eq!(results.len(), 10, "k=10 nearest requested");
    assert_eq!(
        results[0].id, 42,
        "an item's own vector is its nearest neighbor"
    );
    // ~0 modulo F16 quantization (the slot's default), which perturbs a stored
    // unit vector by ~1e-3 — far below any other item's distance.
    assert!(
        results[0].distance <= 0.01,
        "self-distance must be ~0, got {}",
        results[0].distance
    );

    // Results are ordered closest-first (ascending L2 distance).
    for w in results.windows(2) {
        assert!(
            w[0].distance <= w[1].distance,
            "results must be sorted by ascending distance"
        );
    }

    // Recall@10 vs an INDEPENDENT brute-force cosine ground truth. The exhaustive
    // index is exact, but it scores in F16 (and computes L2-on-normalized where the
    // oracle computes cosine), so a single item at the k=10 boundary may swap —
    // hence `>= 9`, not `== 10`. A real recall miss (HNSW approximation) shows up
    // as a much larger shortfall, which the live harness measures at scale.
    let truth: std::collections::HashSet<u64> = ground_truth(&q, 10).into_iter().collect();
    let got: std::collections::HashSet<u64> = results.iter().map(|r| r.id).collect();
    let hits = got.intersection(&truth).count();
    assert!(
        hits >= 9,
        "exact index must achieve recall@10 ~ 1.0, got {hits}/10"
    );
}

#[test]
fn vector_search_accepts_ef_search_override() {
    let db = build_db();
    let q = vector_for(7);
    // Both the slot default (None) and an explicit override resolve and return
    // the same exact nearest set on a brute-force slot (ef is ignored there).
    let a = db.vector_search_items(&q, 5, None).unwrap();
    let b = db.vector_search_items(&q, 5, Some(64)).unwrap();
    let ids_a: Vec<u64> = a.iter().map(|r| r.id).collect();
    let ids_b: Vec<u64> = b.iter().map(|r| r.id).collect();
    assert_eq!(ids_a, ids_b);
    assert_eq!(ids_a[0], 7);
}

#[test]
fn vector_search_rejects_dimension_mismatch() {
    let db = build_db();
    let wrong = vec![0.1_f32; DIM + 1];
    let err = db.vector_search_items(&wrong, 10, None);
    assert!(
        err.is_err(),
        "a query vector of the wrong dimension must error, not silently return"
    );
}