src/vann/hnsw.rs
//! `hnsw` — a **vendored, dependency-free HNSW** (Hierarchical Navigable Small
//! World, Malkov & Yashunin 2016) approximate nearest-neighbour index, built to
//! sit behind the same [`super::VectorIndex`] seam as the exact brute-force flat
//! scan.
//!
//! ## Why vendored (and not `hnsw_rs` / `instant-distance`)
//! The nordisk ROOT LAW is **gatling-only** parallelism — `rayon` /
//! `thread::spawn` (incl. transitively in any dependency) are forbidden. The two
//! obvious embeddable Rust HNSW crates both break it:
//! - `hnsw_rs` (hnswlib-rs) parallelises inserts with `rayon`.
//! - `instant-distance` pulls `rayon` into its `Builder::build`.
//!
//! So this is a minimal HNSW written directly onto `std` (no new dep at all): the
//! crate that gains the code (`znippy-zoomies`) stays `cargo tree | grep rayon`
//! **empty**, and any parallelism we add later goes through this crate's own
//! [`crate::gatling_forkjoin`] pool, never rayon. The build here is
//! single-threaded (graph mutation is inherently serial); the win is on the
//! **query** side — O(log n) graph descent instead of the O(n) flat scan.
//!
//! ## Correctness contract
//! The flat [`super::VectorIndex`] remains the exact **oracle**: HNSW is
//! *approximate*, so its top-`k` is validated against the flat top-`k` by a
//! **recall floor** (see the tests + the `nornir.mimir.search.vec_hnsw_*`
//! benches). Vectors are L2-normalized (the flat index already stores them so),
//! and the graph metric is `dist = 1 − cosine`, so a returned score is the plain
//! cosine `1 − dist` — identical units to [`super::VectorIndex::search`].
//!
//! ## Recall depends on the CORPUS, not just on `ef_search` — read this before
//! ## "fixing" a recall number
//! A graph ANN index answers queries by walking short edges between *mutually
//! near* points. How well that works is a property of the data, not of the code:
//!
//! | corpus (n=10k, k=10) | mean pairwise cosine *within* the true top-10 | recall@10 @ `ef_search`=64 |
//! |-----------------------------------------|-----------------------------------------------|----------------------------|
//! | real `jina-v2-base-code` 768-dim chunks | **0.658** | **0.998** |
//! | i.i.d. uniform random 128-dim | **0.082** | **0.733** |
//!
//! On i.i.d. uniform vectors the *k* true nearest neighbours of a query are
//! mutually near-orthogonal — measured on a 10k×128 uniform corpus, only **2.6%**
//! of the true top-10 appear in each other's exact 32-NN lists. No bounded-degree
//! proximity graph can hop between them, so recall has to be bought with search
//! *breadth* (`ef_search`), and the breadth needed grows with `n`. That is the
//! textbook curse of dimensionality (intrinsic dimensionality = 128), not a
//! defect; on such a corpus this index is also **slower than the flat scan**, so
//! it is a regime where an ANN index would never be used in the first place.
//!
//! Measured `ef_search` → recall@10 on i.i.d. uniform 128-dim (oden, 32 cores):
//!
//! | `ef_search` | 16 | 32 | 64 | 128 | 256 | 512 | 1024 | 2048 |
//! |-------------|------|------|------|------|------|------|------|------|
//! | n=10k | .339 | .523 | .733 | .906 | .981 | **1.000** | 1.000 | 1.000 |
//! | n=100k | .103 | .153 | .266 | .438 | .628 | .841 | .948 | **.991** |
//!
//! Real corpus (`jina-v2-base-code`, 768-dim, held-out queries): 0.974 @ ef=16,
//! 0.998 @ ef=64 at n=10k; 1.000 at n=39,559.
//!
//! **The structural invariant to test against is the convergence, not a single
//! number**: for a correctly built, connected graph, recall@k → 1.0 as
//! `ef_search` → n, on *any* corpus. A graph with broken links, a mis-signed
//! metric, unreachable nodes or a bad level assignment plateaus below 1.0 no
//! matter how wide the search gets. [`tests::recall_converges_to_exact_as_ef_grows`]
//! is that guard; [`tests::ci_recall_floor_realistic_corpus`] is the cheap
//! always-on floor on a realistic (clustered) corpus.
use std::cmp::{Ordering, Reverse};
use std::collections::{BinaryHeap, HashSet};
use super::{DotFn, normalized, select_dot_kernel};
/// Tunables for an HNSW build + query. Defaults target **recall@10 ≥ ~0.95** on
/// code-embedding-shaped corpora while keeping the single-threaded build cheap
/// enough for a 100k×128 corpus in the LIGHT bench.
#[derive(Debug, Clone, Copy)]
pub struct HnswParams {
/// Max neighbours per node on the upper layers (`M`). Layer 0 uses `2*M`.
pub m: usize,
/// Candidate list width during construction (`efConstruction`). Larger =
/// better graph quality (recall) at higher build cost.
pub ef_construction: usize,
/// Candidate list width during search (`efSearch`); the query raises it to
/// `max(ef_search, k)`. Larger = higher recall, slower query.
pub ef_search: usize,
/// Seed for the deterministic level-assignment RNG, so a build is
/// reproducible (same graph → same benches/tests every run).
pub seed: u64,
}
impl Default for HnswParams {
fn default() -> Self {
Self {
m: 16,
ef_construction: 128,
ef_search: 64,
seed: 0x5EED_1234_ABCD_0001,
}
}
}
/// A vendored HNSW index over L2-normalized `f32` vectors of a fixed
/// dimensionality, keyed by stable `u64` ids (parallel to the stored rows).
pub struct HnswIndex {
dim: usize,
/// External ids, parallel to the rows of [`Self::data`] (node index order).
ids: Vec<u64>,
/// Row-major, L2-normalized vectors: `ids.len() * dim` floats.
data: Vec<f32>,
/// `graph[node][layer]` = neighbour node indices of `node` at `layer`. A node
/// present up to `node_level[node]` has `node_level[node] + 1` layers.
graph: Vec<Vec<Vec<u32>>>,
/// Top layer each node participates in.
node_level: Vec<usize>,
/// Entry point (the node with the current max level).
entry: u32,
/// Highest layer any node reaches.
max_level: usize,
m: usize,
m_max0: usize,
ef_search: usize,
}
impl HnswIndex {
/// Build an HNSW over `ids.len()` vectors. `data` is the row-major
/// **already-L2-normalized** matrix (`ids.len() * dim` floats), in `ids`
/// order — exactly what [`super::VectorIndex`] stores, so the caller hands
/// its stored matrix straight in. Single-threaded (graph mutation is serial).
pub fn build(dim: usize, ids: Vec<u64>, data: Vec<f32>, params: HnswParams) -> Self {
let n = ids.len();
debug_assert_eq!(data.len(), n * dim, "data must be n*dim floats");
let kernel = select_dot_kernel();
let m = params.m.max(2);
let m_max0 = m * 2;
let m_l = 1.0 / (m as f64).ln();
let mut rng = SplitMix64::new(params.seed);
let mut idx = HnswIndex {
dim,
ids,
data,
graph: Vec::with_capacity(n),
node_level: Vec::with_capacity(n),
entry: 0,
max_level: 0,
m,
m_max0,
ef_search: params.ef_search.max(1),
};
for node in 0..n {
let level = assign_level(&mut rng, m_l);
idx.node_level.push(level);
idx.graph.push((0..=level).map(|_| Vec::new()).collect());
if node == 0 {
idx.entry = 0;
idx.max_level = level;
continue;
}
idx.insert(node as u32, level, params.ef_construction, kernel);
}
idx
}
/// Number of stored vectors.
pub fn len(&self) -> usize {
self.ids.len()
}
/// True when empty.
/// **Resident heap bytes of the built graph.**
///
/// An HNSW is not a small annotation on the vectors it indexes: it keeps its
/// OWN copy of the `n * dim` float matrix (see [`Self::build`]), so building
/// one roughly DOUBLES the memory of whatever already held those vectors —
/// and it is built lazily, after admission, inside whatever cache holds the
/// owner. A byte budget that never re-reads this is blind to the larger half
/// of its own entry.
///
/// Counts allocated capacity, and walks the per-node neighbour lists (the
/// graph is the term that varies with `m` and with level distribution, so it
/// is measured rather than estimated).
pub fn resident_bytes(&self) -> usize {
let vecs = self.data.capacity() * std::mem::size_of::<f32>()
+ self.ids.capacity() * std::mem::size_of::<u64>();
let mut graph = self.graph.capacity() * std::mem::size_of::<Vec<Vec<u32>>>();
for node in &self.graph {
graph += node.capacity() * std::mem::size_of::<Vec<u32>>();
for layer in node {
graph += layer.capacity() * std::mem::size_of::<u32>();
}
}
let levels = self.node_level.capacity() * std::mem::size_of::<usize>();
std::mem::size_of::<Self>() + vecs + graph + levels
}
pub fn is_empty(&self) -> bool {
self.ids.is_empty()
}
/// Dimensionality.
pub fn dim(&self) -> usize {
self.dim
}
/// Approximate top-`k` nearest ids to `query`, best-first, as `(id, score)`
/// where `score` is cosine similarity in `[-1, 1]` (higher = closer) — the
/// SAME units and ordering [`super::VectorIndex::search`] returns, so this is
/// a drop-in for the exact scan (approximate: validated by the recall floor).
///
/// # Panics
/// If `query.len() != dim`.
pub fn search(&self, query: &[f32], k: usize) -> Vec<(u64, f32)> {
self.search_with_ef(query, k, self.ef_search)
}
/// The `ef_search` this index was built with — the default breadth
/// [`Self::search`] uses.
pub fn ef_search(&self) -> usize {
self.ef_search
}
/// Like [`Self::search`], but with an explicit candidate-list width `ef`
/// for THIS query, overriding the built-in [`Self::ef_search`].
///
/// `ef` is a pure **query-time** knob: it trades latency for recall and
/// touches no graph state, so a caller that needs a higher recall floor on a
/// hard corpus must NOT have to rebuild the index (which is single-threaded
/// and costs ~80 s for 100k×128). Recall rises monotonically with `ef` and
/// reaches the exact answer at `ef = n`; see the module docs for the measured
/// curve.
///
/// # Panics
/// If `query.len() != dim`.
pub fn search_with_ef(&self, query: &[f32], k: usize, ef: usize) -> Vec<(u64, f32)> {
assert_eq!(
query.len(),
self.dim,
"query dim {} != index dim {}",
query.len(),
self.dim
);
if self.ids.is_empty() || k == 0 {
return Vec::new();
}
let kernel = select_dot_kernel();
let qn = normalized(query);
let ef = ef.max(k);
// Greedy descent from the top layer with ef=1, then a wide ef search at
// layer 0 (the standard HNSW K-NN-SEARCH).
let mut ep = vec![self.entry];
let mut lc = self.max_level;
while lc > 0 {
let w = self.search_layer(&qn, &ep, 1, lc, kernel);
if let Some(&(_, best)) = w.first() {
ep = vec![best];
}
lc -= 1;
}
let w = self.search_layer(&qn, &ep, ef, 0, kernel);
w.into_iter()
.take(k)
.map(|(d, node)| (self.ids[node as usize], 1.0 - d))
.collect()
}
// ----- internals ---------------------------------------------------------
#[inline]
fn row(&self, i: usize) -> &[f32] {
&self.data[i * self.dim..(i + 1) * self.dim]
}
/// `dist = 1 − cosine` between the query-vector `target` and node `i`.
#[inline]
fn dist_to(&self, target: &[f32], i: u32, kernel: DotFn) -> f32 {
// SAFETY: `kernel` was chosen by `select_dot_kernel` to match a CPU
// feature confirmed present at runtime (or the always-sound scalar
// fallback); `target` and the row are both `self.dim` long.
1.0 - unsafe { kernel(target, self.row(i as usize)) }
}
/// `dist = 1 − cosine` between two stored nodes.
#[inline]
fn dist_nodes(&self, a: u32, b: u32, kernel: DotFn) -> f32 {
// SAFETY: as `dist_to`; both rows are `self.dim` long.
1.0 - unsafe { kernel(self.row(a as usize), self.row(b as usize)) }
}
/// Insert node `node` (already present in `data`/`graph` with empty edges) at
/// `level`, wiring its bidirectional links per the HNSW INSERT algorithm.
fn insert(&mut self, node: u32, level: usize, ef_c: usize, kernel: DotFn) {
// The new node's own vector is the search target. Clone it so the later
// `&mut self.graph` mutations don't alias the `&self.data` borrow.
let target = self.row(node as usize).to_vec();
let mut ep = vec![self.entry];
// Descend the layers ABOVE `level` greedily (ef = 1) to reach a good ep.
let mut lc = self.max_level;
while lc > level {
let w = self.search_layer(&target, &ep, 1, lc, kernel);
if let Some(&(_, best)) = w.first() {
ep = vec![best];
}
lc -= 1;
}
// From min(max_level, level) down to 0: find efC candidates, select
// neighbours (heuristic), link both ways, prune over-full neighbours.
let mut lc = level.min(self.max_level);
loop {
let w = self.search_layer(&target, &ep, ef_c, lc, kernel);
let m_at = if lc == 0 { self.m_max0 } else { self.m };
let selected = self.select_neighbors(&w, m_at, kernel);
self.graph[node as usize][lc] = selected.clone();
for &nb in &selected {
self.graph[nb as usize][lc].push(node);
let m_nb = if lc == 0 { self.m_max0 } else { self.m };
if self.graph[nb as usize][lc].len() > m_nb {
// Re-select nb's connections down to m_nb (keeps degree bounded
// while preserving the best-navigable edges).
let conns: Vec<(f32, u32)> = self.graph[nb as usize][lc]
.iter()
.map(|&e| (self.dist_nodes(nb, e, kernel), e))
.collect();
let pruned = self.select_neighbors(&conns, m_nb, kernel);
self.graph[nb as usize][lc] = pruned;
}
}
ep = w.iter().map(|&(_, n)| n).collect();
if lc == 0 {
break;
}
lc -= 1;
}
if level > self.max_level {
self.max_level = level;
self.entry = node;
}
}
/// HNSW SEARCH-LAYER: from entry points `ep`, greedily explore `layer` keeping
/// the `ef` closest to `target` seen. Returns them sorted **ascending by
/// distance** (closest first) as `(dist, node)`.
fn search_layer(
&self,
target: &[f32],
ep: &[u32],
ef: usize,
layer: usize,
kernel: DotFn,
) -> Vec<(f32, u32)> {
let mut visited: HashSet<u32> = HashSet::with_capacity(ef * 8);
// candidates: min-dist on top (explore closest first).
let mut candidates: BinaryHeap<Reverse<Ord32>> = BinaryHeap::new();
// results: max-dist on top (so we can evict the worst when over ef).
let mut results: BinaryHeap<Ord32> = BinaryHeap::new();
for &e in ep {
if visited.insert(e) {
let d = self.dist_to(target, e, kernel);
candidates.push(Reverse(Ord32(d, e)));
results.push(Ord32(d, e));
}
}
while let Some(Reverse(Ord32(cd, c))) = candidates.pop() {
let worst = results.peek().map(|x| x.0).unwrap_or(f32::INFINITY);
if cd > worst && results.len() >= ef {
break;
}
// `layer` is always <= node_level[c] for any node reachable here, so
// this index is in bounds by construction.
for &nb in &self.graph[c as usize][layer] {
if visited.insert(nb) {
let d = self.dist_to(target, nb, kernel);
let worst = results.peek().map(|x| x.0).unwrap_or(f32::INFINITY);
if d < worst || results.len() < ef {
candidates.push(Reverse(Ord32(d, nb)));
results.push(Ord32(d, nb));
if results.len() > ef {
results.pop();
}
}
}
}
}
let mut out: Vec<(f32, u32)> = results.into_iter().map(|Ord32(d, n)| (d, n)).collect();
out.sort_by(|a, b| a.0.total_cmp(&b.0).then(a.1.cmp(&b.1)));
out
}
/// HNSW neighbour selection (the simple heuristic, Malkov & Yashunin
/// Algorithm 4): keep a candidate only if it is closer to the target than to
/// every already-selected neighbour (spreads edges around the query, the key
/// to navigability). Falls back to nearest-remaining to top up to `m` so node
/// degree — and thus graph connectivity/recall — is not starved.
fn select_neighbors(&self, candidates: &[(f32, u32)], m: usize, kernel: DotFn) -> Vec<u32> {
let mut c = candidates.to_vec();
c.sort_by(|a, b| a.0.total_cmp(&b.0).then(a.1.cmp(&b.1)));
let mut result: Vec<u32> = Vec::with_capacity(m);
for &(d_q, e) in &c {
if result.len() >= m {
break;
}
let good = result.iter().all(|&r| self.dist_nodes(e, r, kernel) >= d_q);
if good {
result.push(e);
}
}
if result.len() < m {
for &(_, e) in &c {
if result.len() >= m {
break;
}
if !result.contains(&e) {
result.push(e);
}
}
}
result
}
}
/// Ordered `f32` (with an id tie-break) so distances live in a `BinaryHeap`.
/// `total_cmp` keeps the order well-defined even for NaN.
#[derive(Clone, Copy)]
struct Ord32(f32, u32);
impl PartialEq for Ord32 {
fn eq(&self, other: &Self) -> bool {
self.0.to_bits() == other.0.to_bits() && self.1 == other.1
}
}
impl Eq for Ord32 {}
impl Ord for Ord32 {
fn cmp(&self, other: &Self) -> Ordering {
self.0.total_cmp(&other.0).then(self.1.cmp(&other.1))
}
}
impl PartialOrd for Ord32 {
fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
Some(self.cmp(other))
}
}
/// The hard ceiling on an assigned level. Layer `l` holds ~`n / m^l` nodes, so
/// even `m = 2` puts 2^32 nodes above layer 32 — unreachable. The cap exists
/// because `assign_level`'s `-ln(U)` is unbounded as `U → 0`: a single unlucky
/// draw (`U` denormal) would otherwise mint a 255-layer node, allocating 256
/// neighbour vectors and 255 pointless greedy descents on EVERY subsequent
/// query, since that node also becomes the entry point.
const MAX_LEVEL_CAP: usize = 32;
/// Assign an HNSW level: `floor(-ln(U) * mL)`, `U ∈ (0,1]`. The geometric decay
/// gives the O(log n) layer structure.
fn assign_level(rng: &mut SplitMix64, m_l: f64) -> usize {
let mut r = rng.next_f64();
if r <= 0.0 {
r = f64::MIN_POSITIVE;
}
((-(r.ln()) * m_l).floor() as usize).min(MAX_LEVEL_CAP)
}
/// A tiny deterministic PRNG (splitmix64) — no `rand` dep, reproducible builds.
struct SplitMix64(u64);
impl SplitMix64 {
fn new(seed: u64) -> Self {
Self(seed)
}
fn next_u64(&mut self) -> u64 {
self.0 = self.0.wrapping_add(0x9E37_79B9_7F4A_7C15);
let mut z = self.0;
z = (z ^ (z >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9);
z = (z ^ (z >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB);
z ^ (z >> 31)
}
/// Uniform in `[0, 1)`.
fn next_f64(&mut self) -> f64 {
(self.next_u64() >> 11) as f64 / (1u64 << 53) as f64
}
/// Standard normal, Box–Muller (the second variate is discarded — this is a
/// fixture generator, not a hot loop).
fn next_gauss(&mut self) -> f64 {
let u1 = self.next_f64().max(f64::MIN_POSITIVE);
let u2 = self.next_f64();
(-2.0 * u1.ln()).sqrt() * (std::f64::consts::TAU * u2).cos()
}
}
/// The default cluster spread for [`clustered_row`], chosen so the mean pairwise
/// cosine *within* a cluster is ≈ `1/(1+spread²)` ≈ **0.66** — the value measured
/// on nornir's real `jina-v2-base-code` corpus (39,559 × 768) for the true top-10
/// of a query. See the module docs.
pub const REALISTIC_SPREAD: f32 = 0.72;
/// Generate row `row` of a **deterministic corpus with realistic embedding
/// geometry**: `clusters` random centroids, each row a centroid plus Gaussian
/// noise scaled by `spread`, L2-normalized.
///
/// Why this and not i.i.d. uniform vectors: uniform vectors have *intrinsic
/// dimensionality `dim`*, so a query's true k nearest neighbours are mutually
/// near-orthogonal and no proximity graph can walk between them — recall@k then
/// measures the curse of dimensionality rather than the health of the index (see
/// the module docs for the measured contrast). Real embedding corpora are
/// strongly clustered; a recall floor is only meaningful against that shape.
///
/// Rows are independent given `(seed, row)`, so held-out **queries** are just
/// rows drawn from beyond the corpus range — in-distribution but not members.
pub fn clustered_row(seed: u64, row: u64, dim: usize, clusters: u64, spread: f32) -> Vec<f32> {
let clusters = clusters.max(1);
let mut rr = SplitMix64::new(seed ^ row.wrapping_mul(0x9E37_79B9_7F4A_7C15));
let c = rr.next_u64() % clusters;
let mut rc = SplitMix64::new(
seed.wrapping_add(0x00C0_FFEE)
.wrapping_add(c.wrapping_mul(0x1234_5678_9ABC_DEF1)),
);
let v: Vec<f32> = (0..dim)
.map(|_| (rc.next_gauss() + spread as f64 * rr.next_gauss()) as f32)
.collect();
normalized(&v)
}
/// A whole `n × dim` row-major corpus from [`clustered_row`], with the cluster
/// count defaulted to `n/50` (≈50 rows per cluster, so a top-10 sits comfortably
/// inside one). Returns the L2-normalized matrix ready for
/// [`super::VectorIndex::add`].
pub fn clustered_corpus(seed: u64, n: usize, dim: usize) -> Vec<f32> {
let clusters = ((n / 50) as u64).max(4);
let mut out = Vec::with_capacity(n * dim);
for row in 0..n as u64 {
out.extend_from_slice(&clustered_row(seed, row, dim, clusters, REALISTIC_SPREAD));
}
out
}
#[cfg(test)]
mod tests {
use super::super::VectorIndex;
use super::*;
/// Well-spread deterministic pseudo-random vector (NOT the smooth near-tie
/// shape) so recall@k is a meaningful metric — near-duplicate corpora make
/// the "true top-k" ill-defined.
fn rand_vec(seed: u64, dim: usize) -> Vec<f32> {
let mut s = seed
.wrapping_mul(0x9E37_79B9_7F4A_7C15)
.wrapping_add(0xABCD_1234);
(0..dim)
.map(|_| {
s ^= s << 13;
s ^= s >> 7;
s ^= s << 17;
((s >> 11) as f64 / (1u64 << 53) as f64 * 2.0 - 1.0) as f32
})
.collect()
}
fn recall_at_k(flat: &VectorIndex, hnsw: &HnswIndex, queries: &[Vec<f32>], k: usize) -> f64 {
let mut hit = 0usize;
let mut tot = 0usize;
for q in queries {
let oracle: std::collections::HashSet<u64> =
flat.search(q, k).into_iter().map(|(id, _)| id).collect();
let got = hnsw.search(q, k);
for (id, _) in got {
if oracle.contains(&id) {
hit += 1;
}
}
tot += oracle.len();
}
hit as f64 / tot.max(1) as f64
}
#[test]
fn hnsw_recall_floor_vs_flat_oracle() {
let dim = 64;
let n = 3_000usize;
let mut flat = VectorIndex::new(dim).unwrap();
let mut data = Vec::with_capacity(n * dim);
let mut ids = Vec::with_capacity(n);
for i in 0..n as u64 {
let v = rand_vec(i, dim);
data.extend_from_slice(&v);
ids.push(i);
}
flat.add(&data, &ids).unwrap();
let hnsw = flat.build_hnsw(HnswParams::default());
assert_eq!(hnsw.len(), n);
let queries: Vec<Vec<f32>> = (0..100).map(|q| rand_vec(1_000_000 + q, dim)).collect();
let recall = recall_at_k(&flat, &hnsw, &queries, 10);
assert!(
recall >= 0.90,
"HNSW recall@10 {recall:.3} fell below the 0.90 floor vs the flat oracle"
);
}
#[test]
fn hnsw_finds_exact_planted_match() {
// A corpus of random vectors plus the query itself planted at a known id —
// its self-cosine is 1.0, so it MUST be the rank-1 hit.
let dim = 48;
let n = 2_000usize;
let mut flat = VectorIndex::new(dim).unwrap();
let mut data = Vec::with_capacity(n * dim);
let mut ids = Vec::with_capacity(n);
// Seed far outside the corpus id range so the planted vector is unique
// (a colliding seed would create a legitimate rank-1 tie).
let planted = rand_vec(5_000_000, dim);
for i in 0..n as u64 {
if i == 1234 {
data.extend_from_slice(&planted);
} else {
data.extend_from_slice(&rand_vec(i, dim));
}
ids.push(i);
}
flat.add(&data, &ids).unwrap();
let hnsw = flat.build_hnsw(HnswParams::default());
let hits = hnsw.search(&planted, 1);
assert_eq!(hits[0].0, 1234, "planted self-match must be rank 1");
assert!(
(hits[0].1 - 1.0).abs() < 1e-3,
"self-cosine ~1.0, got {}",
hits[0].1
);
}
/// Build a flat oracle + an HNSW over the same rows.
fn corpus(n: usize, dim: usize, rows: impl Fn(u64) -> Vec<f32>) -> (VectorIndex, HnswIndex) {
let mut flat = VectorIndex::new(dim).unwrap();
let mut data = Vec::with_capacity(n * dim);
let mut ids = Vec::with_capacity(n);
for i in 0..n as u64 {
data.extend_from_slice(&rows(i));
ids.push(i);
}
flat.add(&data, &ids).unwrap();
let h = flat.build_hnsw(HnswParams::default());
(flat, h)
}
fn recall_ef(
flat: &VectorIndex,
h: &HnswIndex,
queries: &[Vec<f32>],
k: usize,
ef: usize,
) -> f64 {
let mut hit = 0usize;
let mut tot = 0usize;
for q in queries {
let oracle: std::collections::HashSet<u64> =
flat.search(q, k).into_iter().map(|(id, _)| id).collect();
for (id, _) in h.search_with_ef(q, k, ef) {
if oracle.contains(&id) {
hit += 1;
}
}
tot += oracle.len();
}
hit as f64 / tot.max(1) as f64
}
/// **THE CI RECALL GUARD** — cheap, always-on, no `#[ignore]`.
///
/// 1k vectors with the cluster geometry of a real embedding corpus (see
/// [`clustered_corpus`]), 64 held-out in-distribution queries, ground truth
/// from the exact flat oracle. Asserts `recall@10 ≥ 0.90` at the shipped
/// default `ef_search`. Sub-second, so it runs on every `cargo test`, not
/// only in the heavy `nornir.mimir.search.vec_hnsw_*` bench arms.
///
/// Verified red-when-broken: truncating every layer-0 adjacency list to its
/// first 4 entries drops this to ~0.31.
#[test]
fn ci_recall_floor_realistic_corpus() {
let (dim, n, k, clusters) = (128usize, 1_000usize, 10usize, 20u64);
let (flat, h) = corpus(n, dim, |i| {
clustered_row(0xC0DE_1234, i, dim, clusters, REALISTIC_SPREAD)
});
// Held-out queries: rows drawn from beyond the corpus range, so they are
// in-distribution but not members (no free self-match).
let queries: Vec<Vec<f32>> = (0..64u64)
.map(|q| clustered_row(0xC0DE_1234, 1_000_000 + q, dim, clusters, REALISTIC_SPREAD))
.collect();
let recall = recall_ef(&flat, &h, &queries, k, HnswParams::default().ef_search);
assert!(
recall >= 0.90,
"CI recall floor: HNSW recall@{k} {recall:.4} < 0.90 on a realistic \
(clustered) 1k×{dim} corpus at the default ef_search — the index is broken. \
Do NOT fix this by raising ef_search; see the module docs."
);
}
/// **THE STRUCTURAL INVARIANT** — corpus-independent, so it cannot be argued
/// away as "hard data".
///
/// For a correctly built, fully connected graph with a consistent metric,
/// widening the search must converge on the exact answer: `recall@k → 1.0` as
/// `ef_search → n`. Broken bidirectional links, unreachable nodes, a
/// mis-signed distance or a bad level assignment all plateau *below* 1.0 no
/// matter how wide the search gets.
///
/// Run on the ADVERSARIAL i.i.d.-uniform corpus on purpose: recall at the
/// default `ef_search` is only ~0.73 there (curse of dimensionality, see the
/// module docs), yet a healthy graph still reaches exactly 1.0 by ef=512.
#[test]
fn recall_converges_to_exact_as_ef_grows() {
let (dim, n, k) = (128usize, 10_000usize, 10usize);
let (flat, h) = corpus(n, dim, |i| rand_vec(i, dim));
let queries: Vec<Vec<f32>> = (0..32).map(|q| rand_vec(9_000_000 + q, dim)).collect();
let mut prev = 0.0f64;
for &ef in &[64usize, 128, 256, 512] {
let r = recall_ef(&flat, &h, &queries, k, ef);
assert!(
r >= prev - 1e-9,
"recall@{k} must be monotone in ef_search: {prev:.4} at the previous \
step, {r:.4} at ef={ef}"
);
prev = r;
}
assert!(
prev >= 0.999,
"recall@{k} only reached {prev:.4} at ef=512 on {n} rows — a correct, \
connected graph converges to the exact answer as ef grows, so this is a \
GRAPH defect (links/levels/metric), not a tuning problem"
);
}
/// Every node must be reachable from the entry point at layer 0, and no node
/// may be left with an empty adjacency list. An index that silently strands
/// part of the corpus loses that part's recall permanently — invisibly,
/// because the query still returns k results.
#[test]
fn layer0_graph_is_fully_connected() {
let (dim, n) = (64usize, 3_000usize);
let (_flat, h) = corpus(n, dim, |i| {
clustered_row(0xFEED, i, dim, 60, REALISTIC_SPREAD)
});
assert!(
(0..n).all(|i| !h.graph[i][0].is_empty()),
"some node has no layer-0 neighbours — it can never be retrieved"
);
let mut seen = vec![false; n];
let mut stack = vec![h.entry as usize];
seen[h.entry as usize] = true;
let mut reached = 1usize;
while let Some(c) = stack.pop() {
for &nb in &h.graph[c][0] {
if !seen[nb as usize] {
seen[nb as usize] = true;
reached += 1;
stack.push(nb as usize);
}
}
}
assert_eq!(
reached, n,
"only {reached}/{n} nodes reachable from the entry point at layer 0"
);
}
/// `search_with_ef` must be a pure query-time knob: same graph, wider search,
/// never fewer correct hits — and identical to `search` at the built-in ef.
#[test]
fn ef_is_a_query_time_knob_not_a_rebuild() {
let (dim, n, k) = (64usize, 2_000usize, 10usize);
let (_flat, h) = corpus(n, dim, |i| {
clustered_row(0xBEEF, i, dim, 40, REALISTIC_SPREAD)
});
let q = clustered_row(0xBEEF, 999_999, dim, 40, REALISTIC_SPREAD);
assert_eq!(h.search(&q, k), h.search_with_ef(&q, k, h.ef_search()));
// Widening must not lose the rank-1 hit.
let wide = h.search_with_ef(&q, k, 512);
assert_eq!(wide.len(), k);
assert!(wide[0].1 >= h.search(&q, k)[0].1 - 1e-6);
}
/// A single unlucky `assign_level` draw must not mint a 255-layer node (it
/// would become the entry point and add 255 useless descents per query).
#[test]
fn level_assignment_is_capped() {
let mut rng = SplitMix64::new(1);
let m_l = 1.0 / (16f64).ln();
for _ in 0..10_000 {
assert!(assign_level(&mut rng, m_l) <= MAX_LEVEL_CAP);
}
// The degenerate draw itself.
let mut zero = SplitMix64::new(0);
zero.0 = 0u64.wrapping_sub(0x9E37_79B9_7F4A_7C15); // next_u64() == 0 → next_f64() == 0.0
assert_eq!(assign_level(&mut zero, m_l), MAX_LEVEL_CAP);
}
#[test]
fn empty_and_k0_are_safe() {
let idx = HnswIndex::build(8, vec![], vec![], HnswParams::default());
assert!(idx.is_empty());
assert!(idx.search(&[0.0; 8], 5).is_empty());
let mut flat = VectorIndex::new(8).unwrap();
flat.add(&[1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], &[9])
.unwrap();
let h = flat.build_hnsw(HnswParams::default());
assert!(
h.search(&[1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], 0)
.is_empty()
);
}
}