quantumclaw/clawhdf5
3
1
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 1 Wiki Activity

perf(ann): batched bulk build, parallel with the parallel feature

Browse Source 
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 <[email protected]>
This commit is contained in:
osobh
2026-09-19 13:43:49 -07:00
co-authored by Claude Fable 5.1
parent 41db450c92
commit c19199f3eb
4 changed files with 225 additions and 65 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/clawhdf5-ann/src/hnsw.rs
+193 -64
View File
@@ -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<f32>],
graph: &[Vec<Vec<usize>>],
node_levels: &[usize],
batch: std::ops::Range<usize>,
entry_point: usize,
ep_level: usize,
(m, m_max0, ef_construction): (usize, usize, usize),
metric: DistanceMetric,
) -> Vec<Vec<(usize, Vec<usize>)>> {
let plan_one = |i: usize| -> Vec<(usize, Vec<usize>)> {
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<f32>],
graph: &mut [Vec<Vec<usize>>],
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<usize>)> = 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<f32>],
layer: &mut [Vec<usize>],
new_id: usize,
selected: &[usize],
max_conn: usize,
metric: DistanceMetric,
) {
let mut overflowed: Vec<usize> = 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<f32>],
@@ -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);