diff --git a/CHANGELOG.md b/CHANGELOG.md index 07adce5..cc08329 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,21 @@ # Changelog +## Unreleased + +### Search +- `clawhdf5-ann`: **HNSW recall fix.** Neighbours were chosen as the plain + closest-M, which on clustered data (what embeddings look like) turns each + cluster into an island: recall@10 was 0.87 / 0.67 / 0.31 at 1K / 10K / 100K + vectors and did not improve with `ef`. The index now uses the HNSW paper's + diversity heuristic (Algorithm 4 with kept pruned connections) when linking a + new node and when pruning back-links: recall@10 at `ef = 64` is 1.00 / 1.00 / + 0.98 and responds to `ef`. Builds are slower (~3.5x at 10K). Existing + persisted indexes keep their old graph until rebuilt; the agent rebuilds its + index from the cache, so stores pick this up automatically. +- `clawhdf5-bench`: new `search_harness` binary — HNSW recall@10 / QPS / latency + per `ef` against an exact scan, and end-to-end `hybrid_search` timings, on + deterministic clustered (or `--uniform`) data. Baseline in `BENCHMARKS.md`. + ## v2.3.0 (2026-09-19) ### Upgrade Notes diff --git a/crates/clawhdf5-ann/src/hnsw.rs b/crates/clawhdf5-ann/src/hnsw.rs index c8ae49e..5755f2b 100644 --- a/crates/clawhdf5-ann/src/hnsw.rs +++ b/crates/clawhdf5-ann/src/hnsw.rs @@ -252,8 +252,9 @@ impl HnswIndex { metric, ); - // Select up to m closest neighbors - let selected: Vec = neighbors.iter().take(max_conn).map(|c| c.id).collect(); + let scored: Vec<(usize, f32)> = + neighbors.iter().map(|c| (c.id, c.distance)).collect(); + let selected = select_neighbors(vectors, &scored, max_conn, metric); // Add bidirectional connections graph[layer][i] = selected.clone(); @@ -382,7 +383,8 @@ impl HnswIndex { self.ef_construction, self.metric, ); - let selected: Vec = neighbors.iter().take(max_conn).map(|c| c.id).collect(); + let scored: Vec<(usize, f32)> = neighbors.iter().map(|c| (c.id, c.distance)).collect(); + let selected = select_neighbors(&self.vectors, &scored, max_conn, self.metric); self.graph[layer][id] = selected.clone(); for &neighbor in &selected { self.graph[layer][neighbor].push(id); @@ -828,7 +830,52 @@ fn search_layer( result } -/// Prune connections for a node to keep only the closest `max_conn` neighbors. +/// Choose up to `max_conn` neighbours for a node from `candidates` (sorted by +/// ascending distance to that node) — the HNSW paper's Algorithm 4 with +/// `keepPrunedConnections`. +/// +/// Taking the plain `max_conn` closest is what breaks the graph on clustered +/// data: every link of a node inside a tight cluster goes to that same cluster, +/// so clusters become islands that a search entering elsewhere can never +/// reach, however large `ef` is. Instead a candidate is accepted only if it is +/// closer to the node than to every neighbour already accepted, which spreads +/// links across directions and keeps the long edges that join clusters. Any +/// remaining slots are then filled with the closest rejected candidates, so a +/// node is never left under-connected. +fn select_neighbors( + vectors: &[Vec], + candidates: &[(usize, f32)], + max_conn: usize, + metric: DistanceMetric, +) -> Vec { + if candidates.len() <= max_conn { + return candidates.iter().map(|&(id, _)| id).collect(); + } + let mut selected: Vec = Vec::with_capacity(max_conn); + let mut rejected: Vec = Vec::new(); + for &(id, dist_to_node) in candidates { + if selected.len() >= max_conn { + break; + } + let diverse = selected + .iter() + .all(|&s| compute_distance(&vectors[id], &vectors[s], metric) > dist_to_node); + if diverse { + selected.push(id); + } else { + rejected.push(id); + } + } + for id in rejected { + if selected.len() >= max_conn { + break; + } + selected.push(id); + } + selected +} + +/// Trim `node`'s neighbour list back to `max_conn` with [`select_neighbors`]. fn prune_connections( vectors: &[Vec], neighbors: &mut Vec, @@ -839,22 +886,12 @@ fn prune_connections( if neighbors.len() <= max_conn { return; } - #[cfg(feature = "parallel")] - let mut scored: Vec<(usize, f32)> = { - use rayon::prelude::*; - neighbors - .par_iter() - .map(|&n| (n, compute_distance(&vectors[node], &vectors[n], metric))) - .collect() - }; - #[cfg(not(feature = "parallel"))] let mut scored: Vec<(usize, f32)> = neighbors .iter() .map(|&n| (n, compute_distance(&vectors[node], &vectors[n], metric))) .collect(); - scored.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal)); - scored.truncate(max_conn); - *neighbors = scored.into_iter().map(|(id, _)| id).collect(); + scored.sort_by(|a, b| a.1.total_cmp(&b.1).then(a.0.cmp(&b.0))); + *neighbors = select_neighbors(vectors, &scored, max_conn, metric); } // --------------------------------------------------------------------------- @@ -1043,6 +1080,104 @@ fn get_attr_string(attrs: &[(String, AttrValue)], name: &str) -> Result Vec> { + let mut state = seed; + let mut next = move || { + 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)) >> 40) as f32 / (1u64 << 24) as f32 - 0.5 + }; + let centres: Vec> = (0..clusters) + .map(|_| (0..dim).map(|_| next() * 10.0).collect()) + .collect(); + (0..n) + .map(|i| { + centres[i % clusters] + .iter() + .map(|c| c + next() * 0.5) + .collect() + }) + .collect() + } + + fn recall_at_10( + index: &HnswIndex, + vectors: &[Vec], + queries: &[Vec], + ef: usize, + ) -> f64 { + let mut hits = 0; + for q in queries { + let mut exact: Vec<(usize, f32)> = vectors + .iter() + .enumerate() + .map(|(i, v)| (i, compute_distance(q, v, DistanceMetric::L2))) + .collect(); + exact.sort_by(|a, b| a.1.total_cmp(&b.1)); + let want: Vec = exact[..10].iter().map(|e| e.0).collect(); + hits += index + .search(q, 10, ef) + .iter() + .filter(|(id, _)| want.contains(id)) + .count(); + } + hits as f64 / (10 * queries.len()) as f64 + } + + #[test] + fn clustered_data_keeps_high_recall() { + // Data and queries come from the same clusters: one draw, split. + let mut vectors = clustered(3060, 24, 30, 1); + let queries = vectors.split_off(3000); + let built = HnswIndex::build_with_metric(&vectors, 8, 40, DistanceMetric::L2); + let recall = recall_at_10(&built, &vectors, &queries, 64); + assert!(recall >= 0.95, "bulk build recall@10 = {recall}"); + + // Incremental inserts go through the same neighbour selection. + let mut incremental = HnswIndex::new(8, 40, DistanceMetric::L2); + for v in &vectors { + incremental.insert(v.clone()); + } + let recall = recall_at_10(&incremental, &vectors, &queries, 64); + assert!(recall >= 0.95, "incremental recall@10 = {recall}"); + } + + #[test] + fn select_neighbors_prefers_diverse_directions_and_fills_up() { + // Node at the origin. Three candidates bunched together on the right, + // one on the left. With room for two, plain closest-M would take two + // from the bunch and lose the only link leftwards. + let vectors = vec![ + vec![0.0, 0.0], // 0: the node + vec![1.0, 0.0], // 1 + vec![1.1, 0.0], // 2 + vec![1.2, 0.0], // 3 + vec![-2.0, 0.0], // 4 + ]; + let scored: Vec<(usize, f32)> = (1..5) + .map(|i| { + ( + i, + compute_distance(&vectors[0], &vectors[i], DistanceMetric::L2), + ) + }) + .collect(); + assert_eq!( + select_neighbors(&vectors, &scored, 2, DistanceMetric::L2), + [1, 4] + ); + // Spare capacity is filled with the closest rejected candidates. + assert_eq!( + select_neighbors(&vectors, &scored, 3, DistanceMetric::L2), + [1, 4, 2] + ); + } + fn make_random_vectors(n: usize, dim: usize, seed: u64) -> Vec> { let mut vectors = Vec::with_capacity(n); let mut state = seed; diff --git a/crates/clawhdf5-bench/src/bin/search_harness.rs b/crates/clawhdf5-bench/src/bin/search_harness.rs index 0db6931..d6e5d19 100644 --- a/crates/clawhdf5-bench/src/bin/search_harness.rs +++ b/crates/clawhdf5-bench/src/bin/search_harness.rs @@ -18,6 +18,7 @@ //! cargo run --release -p clawhdf5-bench --bin search_harness # 1K, 10K //! cargo run --release -p clawhdf5-bench --bin search_harness -- --full # + 100K //! cargo run --release -p clawhdf5-bench --bin search_harness -- --json out.json +//! cargo run --release -p clawhdf5-bench --bin search_harness -- --ann-only --uniform //! ``` use std::time::{Duration, Instant}; @@ -78,8 +79,26 @@ struct Dataset { query_cluster: Vec, } +/// `--uniform`: isotropic random unit vectors instead of clusters. Not a +/// realistic workload, but a useful second distribution — a recall problem +/// that appears only on clustered data points at graph connectivity. +static UNIFORM: std::sync::atomic::AtomicBool = std::sync::atomic::AtomicBool::new(false); + fn make_dataset(n: usize, seed: u64) -> Dataset { let mut rng = Rng(seed); + if UNIFORM.load(std::sync::atomic::Ordering::Relaxed) { + let random_unit = |rng: &mut Rng| { + let mut v: Vec = (0..DIM).map(|_| rng.gauss()).collect(); + normalize(&mut v); + v + }; + return Dataset { + vectors: (0..n).map(|_| random_unit(&mut rng)).collect(), + queries: (0..N_QUERIES).map(|_| random_unit(&mut rng)).collect(), + cluster_of: vec![0; n], + query_cluster: vec![0; N_QUERIES], + }; + } let n_clusters = (n / 100).clamp(8, 512); let centres: Vec> = (0..n_clusters) .map(|_| { @@ -344,6 +363,11 @@ fn bench_end_to_end(n: usize, json: &mut Vec) { fn main() { let args: Vec = std::env::args().skip(1).collect(); let full = args.iter().any(|a| a == "--full"); + let ann_only = args.iter().any(|a| a == "--ann-only"); + if args.iter().any(|a| a == "--uniform") { + UNIFORM.store(true, std::sync::atomic::Ordering::Relaxed); + println!("(uniform random data)"); + } let json_path = args .iter() .position(|a| a == "--json") @@ -365,6 +389,9 @@ fn main() { bench_ann(n, &mut json); } + if ann_only { + return; + } println!("\n### End to end: `HDF5Memory::hybrid_search` (k = {K}, weights 0.7 / 0.3)\n"); println!( "| N | ingest ms | checkpoint ms | open ms | first query ms | p50 ms | p99 ms | QPS |"