Merge fix/hnsw-deleted-topk: live-only search results, batched parallel index build
CI / test (push) Failing after 2s

Co-Authored-By: Claude Fable 5.1 <[email protected]>
This commit is contained in:
osobh
2026-09-19 13:52:24 -07:00
co-authored by Claude Fable 5.1
7 changed files with 323 additions and 76 deletions
+22
View File
@@ -243,6 +243,28 @@ results.
| 10000 | 104 | 1487 | 33.8 | 13.7 | 13.9 | 0.49 | 0.51 | 2020.9 | | 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 | | 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 ## Vector Search Latency
Brute-force cosine similarity over 384-dimensional embeddings (OpenAI text-embedding-3-small size). Brute-force cosine similarity over 384-dimensional embeddings (OpenAI text-embedding-3-small size).
+19
View File
@@ -1,5 +1,24 @@
# Changelog # 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) ## v2.4.0 (2026-09-19)
### Upgrade Notes ### Upgrade Notes
+2
View File
@@ -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 approximate `clawhdf5-ann` index for the vector stage (the index mirrors
the cache and self-heals on drift). Build the agent with the cache and self-heals on drift). Build the agent with
`--no-default-features --features float16` to force the exact linear cosine scan. `--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 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 (plain closest-M capped recall on clustered data: 0.31 recall@10 at 100K). Its
graph is saved to `<store>.h5.ann` at each checkpoint and reloaded by `open()` graph is saved to `<store>.h5.ann` at each checkpoint and reloaded by `open()`
+4 -2
View File
@@ -45,9 +45,11 @@ name = "memory_bench"
harness = false harness = false
[features] [features]
default = ["float16", "hnsw"] default = ["float16", "hnsw", "parallel"]
float16 = ["half"] 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 # Compress embeddings with Zstd instead of deflate when
# `MemoryConfig::compression` is on. Off by default: it links libzstd (C). # `MemoryConfig::compression` is on. Off by default: it links libzstd (C).
zstd = ["clawhdf5/zstd"] zstd = ["clawhdf5/zstd"]
+259 -72
View File
@@ -248,65 +248,61 @@ impl HnswIndex {
let mut entry_point = 0; let mut entry_point = 0;
let mut ep_level = node_levels[0]; let mut ep_level = node_levels[0];
// Insert nodes one by one // Insert in batches. Each batch is planned against the graph as it
for i in 1..n { // stood when the batch began (read-only, so the plans are independent
let node_level = node_levels[i]; // and run in parallel with the `parallel` feature), then linked, then
let mut ep = entry_point; // every neighbour list that overflowed is pruned once. Pruning is ~90%
// of a build's distance evaluations, and a node that gains several
// Phase 1: greedy search from top layer down to node_level + 1 // back-links in one batch is pruned once instead of once per link.
let start_layer = ep_level; //
for layer in (node_level + 1..=start_layer).rev() { // Nodes in the same batch cannot see each other while planning, so
ep = greedy_closest(vectors, &graph[layer], &vectors[i], ep, metric); // 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 plans = plan_batch(
let bottom = if node_level < start_layer { vectors,
node_level &graph,
} else { &node_levels,
start_layer next..end,
}; entry_point,
for layer in (0..=bottom).rev() { ep_level,
let max_conn = if layer == 0 { m_max0 } else { m }; (m, m_max0, ef_construction),
metric,
);
let neighbors = search_layer( let mut overflowed: Vec<(usize, usize)> = Vec::new();
vectors, for (offset, plan) in plans.into_iter().enumerate() {
&graph[layer], let node = next + offset;
&vectors[i], for (layer, selected) in plan {
ep, let max_conn = if layer == 0 { m_max0 } else { m };
ef_construction, for &neighbor in &selected {
metric, let list = &mut graph[layer][neighbor];
); list.push(node);
if list.len() == max_conn + 1 {
let scored: Vec<(usize, f32)> = overflowed.push((layer, neighbor));
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,
);
} }
} graph[layer][node] = selected;
if !selected.is_empty() {
ep = selected[0];
} }
} }
prune_overflowed(vectors, &mut graph, overflowed, (m, m_max0), metric);
// Update entry point if this node has a higher level for (i, &level) in node_levels.iter().enumerate().take(end).skip(next) {
if node_level > ep_level { if level > ep_level {
entry_point = i; entry_point = i;
ep_level = node_level; ep_level = level;
}
} }
next = end;
} }
Self { Self {
@@ -408,22 +404,19 @@ impl HnswIndex {
ep, ep,
self.ef_construction, self.ef_construction,
self.metric, self.metric,
None,
); );
let scored: Vec<(usize, f32)> = neighbors.iter().map(|c| (c.id, c.distance)).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); let selected = select_neighbors(&self.vectors, &scored, max_conn, self.metric);
self.graph[layer][id] = selected.clone(); self.graph[layer][id] = selected.clone();
for &neighbor in &selected { link_back(
self.graph[layer][neighbor].push(id); &self.vectors,
if self.graph[layer][neighbor].len() > max_conn { &mut self.graph[layer],
prune_connections( id,
&self.vectors, &selected,
&mut self.graph[layer][neighbor], max_conn,
neighbor, self.metric,
max_conn, );
self.metric,
);
}
}
if !selected.is_empty() { if !selected.is_empty() {
ep = selected[0]; ep = selected[0];
} }
@@ -517,13 +510,21 @@ impl HnswIndex {
ep = greedy_closest(&self.vectors, &self.graph[layer], query, ep, self.metric); ep = greedy_closest(&self.vectors, &self.graph[layer], query, ep, self.metric);
} }
// Search layer 0 with ef candidates. Deleted nodes are still traversed // Search layer 0 for the ef nearest *live* nodes. Deleted nodes are
// (they remain valid graph waypoints) but are filtered from the result. // still traversed (they remain valid graph waypoints) but take no
let candidates = search_layer(&self.vectors, &self.graph[0], query, ep, ef, self.metric); // 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 candidates
.into_iter() .into_iter()
.filter(|c| !self.deleted[c.id])
.take(k) .take(k)
.map(|c| (c.id, c.distance)) .map(|c| (c.id, c.distance))
.collect() .collect()
@@ -940,6 +941,14 @@ fn greedy_closest(
} }
/// Search a single layer for the ef closest nodes to `query`. /// 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( fn search_layer(
vectors: &[Vec<f32>], vectors: &[Vec<f32>],
layer: &[Vec<usize>], layer: &[Vec<usize>],
@@ -947,6 +956,7 @@ fn search_layer(
ep: usize, ep: usize,
ef: usize, ef: usize,
metric: DistanceMetric, metric: DistanceMetric,
skip: Option<&[bool]>,
) -> Vec<Candidate> { ) -> Vec<Candidate> {
let ep_dist = compute_distance(query, &vectors[ep], metric); let ep_dist = compute_distance(query, &vectors[ep], metric);
@@ -959,16 +969,18 @@ fn search_layer(
// Max-heap of current results (furthest first) // Max-heap of current results (furthest first)
let mut results = BinaryHeap::new(); let mut results = BinaryHeap::new();
results.push(FarCandidate { if !skip.is_some_and(|s| s[ep]) {
id: ep, results.push(FarCandidate {
distance: ep_dist, id: ep,
}); distance: ep_dist,
});
}
VISITED.with_borrow_mut(|visited| { VISITED.with_borrow_mut(|visited| {
visited.begin(vectors.len()); visited.begin(vectors.len());
visited.insert(ep); visited.insert(ep);
search_layer_visit( 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], query: &[f32],
ef: usize, ef: usize,
metric: DistanceMetric, metric: DistanceMetric,
skip: Option<&[bool]>,
visited: &mut Visited, visited: &mut Visited,
mut candidates: BinaryHeap<Candidate>, mut candidates: BinaryHeap<Candidate>,
mut results: BinaryHeap<FarCandidate>, mut results: BinaryHeap<FarCandidate>,
@@ -1039,6 +1052,9 @@ fn search_layer_visit(
id: neighbor, id: neighbor,
distance: d, distance: d,
}); });
if skip.is_some_and(|s| s[neighbor]) {
continue; // explore through it, but never return it
}
results.push(FarCandidate { results.push(FarCandidate {
id: neighbor, id: neighbor,
distance: d, distance: d,
@@ -1111,6 +1127,125 @@ fn select_neighbors(
selected 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`]. /// Trim `node`'s neighbour list back to `max_conn` with [`select_neighbors`].
fn prune_connections( fn prune_connections(
vectors: &[Vec<f32>], vectors: &[Vec<f32>],
@@ -1384,6 +1519,58 @@ mod tests {
assert!(recall >= 0.95, "incremental recall@10 = {recall}"); 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<usize> = 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] #[test]
fn graph_bytes_round_trip_gives_identical_searches() { fn graph_bytes_round_trip_gives_identical_searches() {
let mut vectors = clustered(1260, 16, 12, 9); let mut vectors = clustered(1260, 16, 12, 9);
@@ -485,8 +485,12 @@ fn main() {
let mut json = Vec::new(); let mut json = Vec::new();
println!("## Search harness"); println!("## Search harness");
for &n in sizes { // `--e2e-only` skips the index benchmarks, so the end-to-end section runs
bench_ann(n, &mut json); // 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 { if ann_only {
+11
View File
@@ -63,6 +63,13 @@ run_step "cargo clippy (format feature matrix)" cargo clippy \
--features parallel,lz4,zstd,pcodec,fast-checksum \ --features parallel,lz4,zstd,pcodec,fast-checksum \
-- -D warnings -- -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) # 4. Tests (exclude clawhdf5-py)
run_step "cargo test" cargo test \ run_step "cargo test" cargo test \
--workspace \ --workspace \
@@ -72,6 +79,10 @@ run_step "cargo test (format feature matrix)" cargo test \
-p clawhdf5-format \ -p clawhdf5-format \
--features parallel,lz4,zstd,pcodec,fast-checksum --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 # 5. Python interop suites. The h5py writer tests are #[ignore]d so a plain
# `cargo test` stays hermetic; run them explicitly here. # `cargo test` stays hermetic; run them explicitly here.
if python3 -c "import h5py" >/dev/null 2>&1 || [ "${CLAWHDF5_REQUIRE_INTEROP:-0}" = "1" ]; then if python3 -c "import h5py" >/dev/null 2>&1 || [ "${CLAWHDF5_REQUIRE_INTEROP:-0}" = "1" ]; then