diff --git a/BENCHMARKS.md b/BENCHMARKS.md index c42be84..e8ba435 100644 --- a/BENCHMARKS.md +++ b/BENCHMARKS.md @@ -243,6 +243,28 @@ results. | 10000 | 104 | 1487 | 33.8 | 13.7 | 13.9 | 0.49 | 0.51 | 2020.9 | | 100000 | 1376 | 20285 | 728.9 | 353.1 | 142.2 | 4.65 | 4.78 | 214.7 | +### After: batched bulk build (optionally parallel); deletions handled in search + +Profiling showed **90% of a build's distance evaluations are in back-link +pruning**. The bulk build now inserts in batches: plan each node's neighbours +against the graph as it stood at the start of the batch, link, then prune every +overflowing list once. That is less work even single-threaded (a node gaining +several back-links in a batch is pruned once), and with the `parallel` feature +planning and pruning run on a thread pool. The graph is deterministic and the +same with or without the feature. Parallelising *within* one insert was tried +first and gave only 1.45x on 16 cores (tasks too small). + +| build | 1K | 10K | 100K | +|---|---:|---:|---:| +| v2.4.0 | 116 ms | 1676 ms | ~21 s | +| batched | 83 ms | 1074 ms | 19.2 s | +| batched + `parallel` (16 cores) | 34 ms | 388 ms | 5.9 s | + +Recall on clustered data is unchanged or slightly better (100K, `ef = 64`: +0.984 -> 0.9945). On uniform random data it dips slightly (10K, `ef = 64`: +0.474 -> 0.444), the cost of batch members not seeing each other while +planning; batches are capped at 1/16 of the graph and 512 nodes. + ## Vector Search Latency Brute-force cosine similarity over 384-dimensional embeddings (OpenAI text-embedding-3-small size). diff --git a/CHANGELOG.md b/CHANGELOG.md index 0f54dbf..0c75b95 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,24 @@ # Changelog +## Unreleased + +### Search +- `clawhdf5-ann`: **faster index builds.** Back-link pruning is 90% of a + build's distance evaluations; the bulk build now inserts in batches and + prunes each overflowing neighbour list once per batch (10K: 1676 -> 1074 ms). + With the `parallel` feature, planning and pruning run on a thread pool (10K: + 388 ms, 100K: ~21 s -> 5.9 s on 16 cores). The graph is deterministic and + identical with or without the feature. `clawhdf5-agent`'s `parallel` feature + enables it for the agent's index and is now **on by default** (adds `rayon` + to the default dependency set; build with `--no-default-features --features + float16,hnsw` to opt out). +- `clawhdf5-ann`: `HnswIndex::search` returned fewer than `k` results — often + none — when the records nearest the query had been deleted: it collected `ef` + candidates, *then* dropped the deleted ones, *then* took `k`. Deleted nodes + are now traversed as waypoints but never occupy a result slot, so a search + returns the `k` nearest live records. Matters for any store that deletes or + supersedes memories without compacting straight away. + ## v2.4.0 (2026-09-19) ### Upgrade Notes diff --git a/CLAUDE.md b/CLAUDE.md index 2a47ad5..2dfc918 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -33,6 +33,8 @@ Cargo workspace with 16 crates under `crates/` (plus `libaec-sys`, an internal F the approximate `clawhdf5-ann` index for the vector stage (the index mirrors the cache and self-heals on drift). Build the agent with `--no-default-features --features float16` to force the exact linear cosine scan. + The agent's `parallel` feature (also default) builds the index on a thread + pool; the graph is identical with or without it. The index uses the HNSW paper's diversity heuristic for neighbour selection (plain closest-M capped recall on clustered data: 0.31 recall@10 at 100K). Its graph is saved to `.h5.ann` at each checkpoint and reloaded by `open()` diff --git a/crates/clawhdf5-agent/Cargo.toml b/crates/clawhdf5-agent/Cargo.toml index 4f2dbb4..0d44038 100644 --- a/crates/clawhdf5-agent/Cargo.toml +++ b/crates/clawhdf5-agent/Cargo.toml @@ -45,9 +45,11 @@ name = "memory_bench" harness = false [features] -default = ["float16", "hnsw"] +default = ["float16", "hnsw", "parallel"] float16 = ["half"] -parallel = ["rayon"] +# Rayon-parallel brute-force search strategies, and a parallel bulk build of +# the HNSW index (same graph, several times faster on a multi-core machine). +parallel = ["rayon", "clawhdf5-ann?/parallel"] # Compress embeddings with Zstd instead of deflate when # `MemoryConfig::compression` is on. Off by default: it links libzstd (C). zstd = ["clawhdf5/zstd"] diff --git a/crates/clawhdf5-ann/src/hnsw.rs b/crates/clawhdf5-ann/src/hnsw.rs index b9e7ce3..eb7124d 100644 --- a/crates/clawhdf5-ann/src/hnsw.rs +++ b/crates/clawhdf5-ann/src/hnsw.rs @@ -248,65 +248,61 @@ impl HnswIndex { let mut entry_point = 0; let mut ep_level = node_levels[0]; - // Insert nodes one by one - for i in 1..n { - let node_level = node_levels[i]; - let mut ep = entry_point; - - // Phase 1: greedy search from top layer down to node_level + 1 - let start_layer = ep_level; - for layer in (node_level + 1..=start_layer).rev() { - ep = greedy_closest(vectors, &graph[layer], &vectors[i], ep, metric); + // Insert in batches. Each batch is planned against the graph as it + // stood when the batch began (read-only, so the plans are independent + // and run in parallel with the `parallel` feature), then linked, then + // every neighbour list that overflowed is pruned once. Pruning is ~90% + // of a build's distance evaluations, and a node that gains several + // back-links in one batch is pruned once instead of once per link. + // + // Nodes in the same batch cannot see each other while planning, so + // batches start at one node and grow only as the graph does — a batch + // is never more than a small fraction of what is already linked. The + // result is deterministic and identical with or without `parallel`. + let mut next = 1; + while next < n { + let mut end = (next + batch_len(next)).min(n); + // A node that raises the top layer becomes the new entry point and + // changes how every later node descends: give it a batch alone. + if let Some(tall) = (next..end).find(|&i| node_levels[i] > ep_level) { + end = if tall == next { next + 1 } else { tall }; } - // Phase 2: search and connect at layers node_level down to 0 - let bottom = if node_level < start_layer { - node_level - } else { - start_layer - }; - for layer in (0..=bottom).rev() { - let max_conn = if layer == 0 { m_max0 } else { m }; + let plans = plan_batch( + vectors, + &graph, + &node_levels, + next..end, + entry_point, + ep_level, + (m, m_max0, ef_construction), + metric, + ); - let neighbors = search_layer( - vectors, - &graph[layer], - &vectors[i], - ep, - ef_construction, - metric, - ); - - 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(); - for &neighbor in &selected { - graph[layer][neighbor].push(i); - // Prune if over limit - if graph[layer][neighbor].len() > max_conn { - prune_connections( - vectors, - &mut graph[layer][neighbor], - neighbor, - max_conn, - metric, - ); + let mut overflowed: Vec<(usize, usize)> = Vec::new(); + for (offset, plan) in plans.into_iter().enumerate() { + let node = next + offset; + for (layer, selected) in plan { + let max_conn = if layer == 0 { m_max0 } else { m }; + for &neighbor in &selected { + let list = &mut graph[layer][neighbor]; + list.push(node); + if list.len() == max_conn + 1 { + overflowed.push((layer, neighbor)); + } } - } - - if !selected.is_empty() { - ep = selected[0]; + graph[layer][node] = selected; } } + prune_overflowed(vectors, &mut graph, overflowed, (m, m_max0), metric); - // Update entry point if this node has a higher level - if node_level > ep_level { - entry_point = i; - ep_level = node_level; + for (i, &level) in node_levels.iter().enumerate().take(end).skip(next) { + if level > ep_level { + entry_point = i; + ep_level = level; + } } + next = end; } Self { @@ -408,22 +404,19 @@ impl HnswIndex { ep, self.ef_construction, self.metric, + None, ); 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); - if self.graph[layer][neighbor].len() > max_conn { - prune_connections( - &self.vectors, - &mut self.graph[layer][neighbor], - neighbor, - max_conn, - self.metric, - ); - } - } + link_back( + &self.vectors, + &mut self.graph[layer], + id, + &selected, + max_conn, + self.metric, + ); if !selected.is_empty() { ep = selected[0]; } @@ -517,13 +510,21 @@ impl HnswIndex { ep = greedy_closest(&self.vectors, &self.graph[layer], query, ep, self.metric); } - // Search layer 0 with ef candidates. Deleted nodes are still traversed - // (they remain valid graph waypoints) but are filtered from the result. - let candidates = search_layer(&self.vectors, &self.graph[0], query, ep, ef, self.metric); + // Search layer 0 for the ef nearest *live* nodes. Deleted nodes are + // still traversed (they remain valid graph waypoints) but take no + // result slot, so deletions near the query don't shrink the answer. + let candidates = search_layer( + &self.vectors, + &self.graph[0], + query, + ep, + ef, + self.metric, + Some(&self.deleted), + ); candidates .into_iter() - .filter(|c| !self.deleted[c.id]) .take(k) .map(|c| (c.id, c.distance)) .collect() @@ -940,6 +941,14 @@ fn greedy_closest( } /// Search a single layer for the ef closest nodes to `query`. +/// Best-first search of one layer, returning up to `ef` nodes by ascending +/// distance. +/// +/// `skip` marks nodes that must not be *returned* (soft-deleted ones). They +/// are still traversed — a tombstone is a perfectly good waypoint — but they +/// never occupy one of the `ef` result slots. Filtering them out afterwards +/// instead meant a query whose neighbourhood had been deleted got back fewer +/// than `k` results, or none, however many live records were nearby. fn search_layer( vectors: &[Vec], layer: &[Vec], @@ -947,6 +956,7 @@ fn search_layer( ep: usize, ef: usize, metric: DistanceMetric, + skip: Option<&[bool]>, ) -> Vec { let ep_dist = compute_distance(query, &vectors[ep], metric); @@ -959,16 +969,18 @@ fn search_layer( // Max-heap of current results (furthest first) let mut results = BinaryHeap::new(); - results.push(FarCandidate { - id: ep, - distance: ep_dist, - }); + if !skip.is_some_and(|s| s[ep]) { + results.push(FarCandidate { + id: ep, + distance: ep_dist, + }); + } VISITED.with_borrow_mut(|visited| { visited.begin(vectors.len()); visited.insert(ep); search_layer_visit( - vectors, layer, query, ef, metric, visited, candidates, results, + vectors, layer, query, ef, metric, skip, visited, candidates, results, ) }) } @@ -1016,6 +1028,7 @@ fn search_layer_visit( query: &[f32], ef: usize, metric: DistanceMetric, + skip: Option<&[bool]>, visited: &mut Visited, mut candidates: BinaryHeap, mut results: BinaryHeap, @@ -1039,6 +1052,9 @@ fn search_layer_visit( id: neighbor, distance: d, }); + if skip.is_some_and(|s| s[neighbor]) { + continue; // explore through it, but never return it + } results.push(FarCandidate { id: neighbor, distance: d, @@ -1111,6 +1127,125 @@ fn select_neighbors( selected } +/// How many nodes to plan together once `linked` nodes are in the graph. +fn batch_len(linked: usize) -> usize { + (linked / 16).clamp(1, 512) +} + +/// For each node in `batch`: the neighbours to link it to on each of its +/// layers, found by searching the graph as it currently stands. +#[allow(clippy::too_many_arguments)] +fn plan_batch( + vectors: &[Vec], + graph: &[Vec>], + node_levels: &[usize], + batch: std::ops::Range, + entry_point: usize, + ep_level: usize, + (m, m_max0, ef_construction): (usize, usize, usize), + metric: DistanceMetric, +) -> Vec)>> { + let plan_one = |i: usize| -> Vec<(usize, Vec)> { + let node_level = node_levels[i]; + let mut ep = entry_point; + // Phase 1: greedy descent from the top layer down to node_level + 1. + for layer in (node_level + 1..=ep_level).rev() { + ep = greedy_closest(vectors, &graph[layer], &vectors[i], ep, metric); + } + // Phase 2: search and select on every layer the node lives on. + let mut plan = Vec::with_capacity(node_level.min(ep_level) + 1); + for layer in (0..=node_level.min(ep_level)).rev() { + let max_conn = if layer == 0 { m_max0 } else { m }; + let neighbors = search_layer( + vectors, + &graph[layer], + &vectors[i], + ep, + ef_construction, + metric, + None, + ); + let scored: Vec<(usize, f32)> = neighbors.iter().map(|c| (c.id, c.distance)).collect(); + let selected = select_neighbors(vectors, &scored, max_conn, metric); + if let Some(&closest) = selected.first() { + ep = closest; + } + plan.push((layer, selected)); + } + plan + }; + + #[cfg(feature = "parallel")] + if batch.len() >= PARALLEL_MIN { + use rayon::prelude::*; + return batch.into_par_iter().map(plan_one).collect(); + } + batch.map(plan_one).collect() +} + +/// Prune every `(layer, node)` neighbour list in `overflowed` back to its +/// limit. Each list belongs to a different node, so they are independent. +fn prune_overflowed( + vectors: &[Vec], + graph: &mut [Vec>], + overflowed: Vec<(usize, usize)>, + (m, m_max0): (usize, usize), + metric: DistanceMetric, +) { + let limit = |layer: usize| if layer == 0 { m_max0 } else { m }; + + #[cfg(feature = "parallel")] + if overflowed.len() >= PARALLEL_MIN { + use rayon::prelude::*; + let mut work: Vec<(usize, usize, Vec)> = overflowed + .into_iter() + .map(|(layer, node)| (layer, node, std::mem::take(&mut graph[layer][node]))) + .collect(); + work.par_iter_mut().for_each(|(layer, node, list)| { + prune_connections(vectors, list, *node, limit(*layer), metric); + }); + for (layer, node, list) in work { + graph[layer][node] = list; + } + return; + } + + for (layer, node) in overflowed { + prune_connections(vectors, &mut graph[layer][node], node, limit(layer), metric); + } +} + +/// Fewest independent tasks worth handing to the thread pool. +#[cfg(feature = "parallel")] +const PARALLEL_MIN: usize = 8; + +/// Add the back-link `neighbor -> new_id` for every selected neighbour, pruning +/// each list that overflows. +/// +/// Used by incremental [`HnswIndex::insert`]. (A single insert's handful of +/// prunes is too fine-grained to parallelise profitably — measured 1.45x on 16 +/// cores; bulk builds batch their pruning instead, see `prune_overflowed`.) +fn link_back( + vectors: &[Vec], + layer: &mut [Vec], + new_id: usize, + selected: &[usize], + max_conn: usize, + metric: DistanceMetric, +) { + let mut overflowed: Vec = Vec::new(); + for &neighbor in selected { + layer[neighbor].push(new_id); + if layer[neighbor].len() > max_conn { + overflowed.push(neighbor); + } + } + + for node in overflowed { + prune_connections(vectors, &mut layer[node], node, max_conn, metric); + } +} + /// Trim `node`'s neighbour list back to `max_conn` with [`select_neighbors`]. fn prune_connections( vectors: &[Vec], @@ -1384,6 +1519,58 @@ mod tests { assert!(recall >= 0.95, "incremental recall@10 = {recall}"); } + #[test] + fn deletions_near_the_query_do_not_shrink_or_degrade_results() { + let mut vectors = clustered(2040, 16, 20, 11); + let queries = vectors.split_off(2000); + let mut index = HnswIndex::build_with_metric(&vectors, 8, 40, DistanceMetric::L2); + + let mut short = 0; + let mut hits = 0; + for q in &queries { + // Delete this query's 40 nearest neighbours: more than ef, so every + // candidate a plain search collects is a tombstone. + let mut exact: Vec<(usize, f32)> = vectors + .iter() + .enumerate() + .filter(|(i, _)| !index.is_deleted(*i)) + .map(|(i, v)| (i, compute_distance(q, v, DistanceMetric::L2))) + .collect(); + exact.sort_by(|a, b| a.1.total_cmp(&b.1)); + for &(id, _) in &exact[..40] { + index.mark_deleted(id); + } + let want: Vec = exact[40..50].iter().map(|e| e.0).collect(); + + let got = index.search(q, 10, 32); + assert!(got.iter().all(|(id, _)| !index.is_deleted(*id))); + short += usize::from(got.len() < 10); + hits += got.iter().filter(|(id, _)| want.contains(id)).count(); + } + assert_eq!(short, 0, "searches returned fewer than k live results"); + let recall = hits as f64 / (10 * queries.len()) as f64; + assert!(recall >= 0.9, "recall@10 among live records = {recall}"); + } + + #[test] + fn bulk_build_is_deterministic() { + // Batched planning runs on a thread pool with the `parallel` feature; + // the graph must not depend on scheduling. (It is also the same graph + // with and without the feature: both take this exact code path.) + let vectors = clustered(2500, 16, 20, 21); + let a = HnswIndex::build_with_metric(&vectors, 8, 40, DistanceMetric::Cosine); + let b = HnswIndex::build_with_metric(&vectors, 8, 40, DistanceMetric::Cosine); + assert_eq!(a.graph_to_bytes(), b.graph_to_bytes()); + } + + #[test] + fn batches_stay_a_small_fraction_of_the_graph() { + assert_eq!(batch_len(1), 1); + assert_eq!(batch_len(15), 1); + assert_eq!(batch_len(160), 10); + assert_eq!(batch_len(1_000_000), 512); + } + #[test] fn graph_bytes_round_trip_gives_identical_searches() { let mut vectors = clustered(1260, 16, 12, 9); diff --git a/crates/clawhdf5-bench/src/bin/search_harness.rs b/crates/clawhdf5-bench/src/bin/search_harness.rs index 3c43af7..d058a60 100644 --- a/crates/clawhdf5-bench/src/bin/search_harness.rs +++ b/crates/clawhdf5-bench/src/bin/search_harness.rs @@ -485,8 +485,12 @@ fn main() { let mut json = Vec::new(); println!("## Search harness"); - for &n in sizes { - bench_ann(n, &mut json); + // `--e2e-only` skips the index benchmarks, so the end-to-end section runs + // in a process that has not already spun up a thread pool. + if !args.iter().any(|a| a == "--e2e-only") { + for &n in sizes { + bench_ann(n, &mut json); + } } if ann_only { diff --git a/scripts/ci-test.sh b/scripts/ci-test.sh index aa7146d..03367a0 100755 --- a/scripts/ci-test.sh +++ b/scripts/ci-test.sh @@ -63,6 +63,13 @@ run_step "cargo clippy (format feature matrix)" cargo clippy \ --features parallel,lz4,zstd,pcodec,fast-checksum \ -- -D warnings +# The HNSW index's parallel bulk build is feature-gated too. +run_step "cargo clippy (ann parallel)" cargo clippy \ + -p clawhdf5-ann \ + --all-targets \ + --features parallel \ + -- -D warnings + # 4. Tests (exclude clawhdf5-py) run_step "cargo test" cargo test \ --workspace \ @@ -72,6 +79,10 @@ run_step "cargo test (format feature matrix)" cargo test \ -p clawhdf5-format \ --features parallel,lz4,zstd,pcodec,fast-checksum +run_step "cargo test (ann parallel)" cargo test \ + -p clawhdf5-ann \ + --features parallel + # 5. Python interop suites. The h5py writer tests are #[ignore]d so a plain # `cargo test` stays hermetic; run them explicitly here. if python3 -c "import h5py" >/dev/null 2>&1 || [ "${CLAWHDF5_REQUIRE_INTEROP:-0}" = "1" ]; then