From c19199f3eb88586dd0e58df8a6e59ae9e6e9e0f8 Mon Sep 17 00:00:00 2001 From: osobh Date: Sat, 19 Sep 2026 13:43:49 -0700 Subject: [PATCH] perf(ann): batched bulk build, parallel with the `parallel` feature Profiling the build showed 90% of all distance evaluations are in back-link pruning (40.8M of 44.9M at 10K): every overflow re-runs the diversity heuristic pairwise over ~max_conn candidates. The bulk build now inserts in batches: plan every node's neighbours against the graph as it stood when the batch began (read-only, so plans are independent), link, then prune each overflowing list once. A node gaining several back-links in a batch is pruned once rather than once per link, so this is faster even single-threaded (10K: 1676 -> 1074 ms). With `parallel`, planning and pruning use rayon (10K: 388 ms; 100K: ~21 s -> 5.9 s on 16 cores). Batches start at one node and are capped at 1/16 of the linked graph and 512 nodes; a node that raises the top layer gets a batch to itself. The result is deterministic and identical with or without the feature (one code path; test compares two builds byte for byte). Parallelising within a single insert was tried first: 1.45x on 16 cores, tasks too small. Incremental insert() stays sequential. Recall on clustered data is unchanged or slightly better; uniform random data dips slightly (10K, ef=64: 0.474 -> 0.444). clawhdf5-agent's `parallel` feature now passes through to the index. Co-Authored-By: Claude Fable 5.1 --- BENCHMARKS.md | 22 +++ CHANGELOG.md | 7 + crates/clawhdf5-agent/Cargo.toml | 4 +- crates/clawhdf5-ann/src/hnsw.rs | 257 +++++++++++++++++++++++-------- 4 files changed, 225 insertions(+), 65 deletions(-) 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 2a8f8c1..8564e22 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -3,6 +3,13 @@ ## 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 + now enables it for the agent's index. - `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 diff --git a/crates/clawhdf5-agent/Cargo.toml b/crates/clawhdf5-agent/Cargo.toml index 4f2dbb4..a029583 100644 --- a/crates/clawhdf5-agent/Cargo.toml +++ b/crates/clawhdf5-agent/Cargo.toml @@ -47,7 +47,9 @@ harness = false [features] default = ["float16", "hnsw"] 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 3fc6f00..8648a6a 100644 --- a/crates/clawhdf5-ann/src/hnsw.rs +++ b/crates/clawhdf5-ann/src/hnsw.rs @@ -248,66 +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, - None, - ); - - 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 in next..end { + if node_levels[i] > ep_level { + entry_point = i; + ep_level = node_levels[i]; + } } + next = end; } Self { @@ -414,18 +409,14 @@ impl HnswIndex { 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]; } @@ -1136,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], @@ -1442,6 +1552,25 @@ mod tests { 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);