diff --git a/DESIGN.md b/DESIGN.md index 66dff99..46a27ab 100644 --- a/DESIGN.md +++ b/DESIGN.md @@ -210,6 +210,60 @@ search(sliceHandles, queryVector: Float32Array, k, ef, filter?): Promise<{ids, d - Auto-ef / auto-efC read the node count from the header high-water minus freelist length — same semantics as today, minus the #2182 inflation (freed ids return to the pool). +**Upper-layer descent is a beam, not hill climbing** (`beam_descend`, `DESCENT_EF = 16`). The +textbook width-1 descent halts at the first upper-layer node no neighbor improves on. On a +clustered corpus that local minimum can sit in the wrong basin, and layer-0 adjacency is +intra-basin, so the layer-0 beam has no uphill edge with which to leave — the query's true +nearest neighbor is then unreachable at *any* ef, and raising ef only expands the wrong basin. +Measured on the `tests/concurrent.rs` corpus (8 000 nodes, 64-d, self-query every node at +ef 256, insertion order fixed by seed): width 1 loses 125 nodes over 200 builds, width 4 loses +20 over 200, width 8 loses 8 over 700, width 16 loses 0 over 700. Cost at 50 000 × 768-d: +visits/query +17 % to +27 %, p50 +0.06 ms flat (0.15 → 0.21 ms at ef 16, 0.21 → 0.28 at ef 64, +0.46 → 0.47 at ef 512 — the descent is a fixed cost, so it hurts most where ef is small), build +throughput -20 %. recall@10 improves below ef 128 (0.844 → 0.903 at ef 16, 0.983 → 1.000 at +ef 64) and is unchanged above. + +`beam_descend` is shared by the read and write paths deliberately: insert must route through +the same graph its queries will, or nodes get their neighbors chosen from a basin searches +never reach. The width is a compile-time constant rather than a parameter because it is a +correctness floor, not a recall/latency dial — `ef` is the dial. + +Three facts worth keeping when working on this. + +The trap is a property of graph *shape*, not of concurrency: it reproduces single-threaded from +a fixed insertion permutation, and concurrency only shuffles that permutation. It also needs the +full corpus — no seed reproduces it at 32 dims, or at 2 000 / 4 000 nodes, so a shrunken repro +is not evidence of a fix. `descent_width_sweep` in `tests/concurrent.rs` (ignored by default) is +the harness behind the table above. + +Read and write descent widths must match. Measured over 200 builds per cell: width 1 both sides +loses 125 nodes, read-only widening loses 37, **write-only widening loses 245 — worse than +either**, and both sides widened loses 0. `insert` seeds each level's `search_layer` from the +descent's landing point, so a graph wired under one routing policy and queried under another is +less navigable than one where they agree. This is also the upgrade story: an existing plane file +read by a new binary is the read-only row, improved but not repaired until its nodes are +re-inserted. + +`HNSW_SWEEP_READ_EF` sets the sweep's query-side width; the build side is whatever `DESCENT_EF` +is compiled as, so the four cells are two runs per value of the constant: + +```text +HNSW_SWEEP_SEEDS=200 HNSW_SWEEP_READ_EF=1 cargo test --release --test concurrent \ + descent_width_sweep -- --ignored --nocapture +HNSW_SWEEP_SEEDS=200 HNSW_SWEEP_READ_EF=16 cargo test --release --test concurrent \ + descent_width_sweep -- --ignored --nocapture +``` + +Do not add a per-level visit cap to the descent without re-measuring. The obvious ceiling, +`DESCENT_EF * UPPER_CAP` = 1024, is already exceeded by ordinary queries: the worst of 3 000 +random queries visits 788 nodes at level 1 on a 50 000-node graph and 1 044 on a 500 000-node +one. A cap that binds silently degrades recall, which is the defect this exists to fix. What +bounds the pathological case instead is `search_layer`'s strict `d < worst`: with every distance +tied — a zero query ties them all at exactly 1.0 — a full result set never admits another +candidate, so the descent drains after `ef` expansions per level (measured 608 visits at 50 000 +nodes, 990 at 500 000). `a_tied_distance_descent_stops_at_its_visit_cap` fails if that `<` is +ever relaxed. + **Filtering** (predicate-aware / ACORN, `filteredSearch = true` today): 1. **Bitset fast path.** RBAC allow-lists and companion-condition candidate sets are computed diff --git a/src/insert.rs b/src/insert.rs index f20f84e..6e64a2f 100644 --- a/src/insert.rs +++ b/src/insert.rs @@ -7,7 +7,7 @@ use crate::distance::{quantize_int8, Query}; use crate::format::{NO_ID, NO_UPPER}; use crate::graph::Graph; -use crate::search::{greedy_descend, search_layer, SearchScratch, SearchStats}; +use crate::search::{beam_descend, search_layer, SearchScratch, SearchStats, DESCENT_EF}; pub struct InsertParams { pub m: usize, // base connection count (JS M, default 16) @@ -229,17 +229,38 @@ pub fn insert( return Err(InsertError::Wedged); }; let top = level.min(entry_level as u8); - let (mut ep, mut ep_dist) = - greedy_descend(graph, &query, entry_id, entry_dist, entry_level, top as u32, &mut stats); + let (mut ep, mut ep_dist) = beam_descend( + graph, + &query, + entry_id, + entry_dist, + entry_level, + top as u32, + DESCENT_EF, + scratch, + &mut stats, + ); // Per-level connection lists for the new node, selection-ordered. let mut connections: Vec> = vec![Vec::new(); level as usize + 1]; let mut nbuf: Vec = Vec::new(); + let mut neighbors: Vec<(u32, f32)> = Vec::new(); for l in (0..=top).rev() { scratch_begin(graph, scratch); - let mut neighbors = - search_layer(graph, &query, ep, ep_dist, params.ef_construction, l, scratch, &mut stats, None, u64::MAX); + search_layer( + graph, + &query, + ep, + ep_dist, + params.ef_construction, + l, + scratch, + &mut stats, + None, + u64::MAX, + &mut neighbors, + ); neighbors.truncate(m << 1); if let Some(&(best, best_d)) = neighbors.first() { ep = best; diff --git a/src/search.rs b/src/search.rs index 9e8ca8b..adc5ab9 100644 --- a/src/search.rs +++ b/src/search.rs @@ -3,7 +3,7 @@ //! array; neighbor ids stream through a reusable scratch buffer. use crate::distance::Query; -use crate::format::NO_ID; +use crate::format::{MAX_UPPER_LEVELS, NO_ID}; use crate::graph::Graph; use std::cmp::Ordering as CmpOrdering; use std::collections::BinaryHeap; @@ -49,11 +49,21 @@ pub struct SearchScratch { visited: Vec, epoch: u32, neighbors: Vec, + candidates: BinaryHeap, + results: BinaryHeap, + descent_out: Vec<(u32, f32)>, } impl SearchScratch { pub fn new() -> Self { - SearchScratch { visited: Vec::new(), epoch: 0, neighbors: Vec::new() } + SearchScratch { + visited: Vec::new(), + epoch: 0, + neighbors: Vec::new(), + candidates: BinaryHeap::new(), + results: BinaryHeap::new(), + descent_out: Vec::new(), + } } pub fn begin_public(&mut self, capacity: u64) { @@ -109,10 +119,14 @@ fn bit_allowed(filter: Option<&[u8]>, id: u32) -> bool { } } -/// Beam search within one layer, starting from `entry`. Level 0 reads slot adjacency; -/// upper levels read the resident upper map. Returns (id, distance) ascending by distance. +/// Beam search within one layer, starting from `entry`. Level 0 reads slot adjacency; upper +/// levels read the resident upper map. Fills `out` with (id, distance) ascending by distance. /// Assumes scratch.begin() was called for this query; entry is marked visited here. /// +/// `out` is the caller's so the descent can reuse one buffer across all levels. It is filled by +/// one `extend` off an exact-size drain, so it takes at most one allocation sized to the results +/// actually found — never to `ef`, which arrives unvalidated from the NAPI boundary. +/// /// `filter`: optional allow-bitset over node ids (bit i of byte i>>3). Filtered-out nodes /// are traversed (their edges route) but excluded from results — ACORN-style — with /// `visit_budget` bounding total visits so a selective filter terminates. @@ -127,18 +141,21 @@ pub fn search_layer( stats: &mut SearchStats, filter: Option<&[u8]>, visit_budget: u64, -) -> Vec<(u32, f32)> { - let mut candidates = BinaryHeap::new(); - let mut results: BinaryHeap = BinaryHeap::new(); + out: &mut Vec<(u32, f32)>, +) { + // take() the scratch buffers to sidestep the double-borrow of scratch + let mut candidates = std::mem::take(&mut scratch.candidates); + let mut results = std::mem::take(&mut scratch.results); + let mut nbuf = std::mem::take(&mut scratch.neighbors); + candidates.clear(); + results.clear(); + scratch.visit(entry); candidates.push(Candidate { distance: entry_dist, id: entry }); if bit_allowed(filter, entry) { results.push(Result_ { distance: entry_dist, id: entry }); } - // take() the scratch neighbor buffer to sidestep the double-borrow of scratch - let mut nbuf = std::mem::take(&mut scratch.neighbors); - while let Some(c) = candidates.pop() { let worst = results.peek().map(|r| r.distance).unwrap_or(f32::INFINITY); if results.len() >= ef && c.distance > worst { @@ -177,45 +194,52 @@ pub fn search_layer( } } } - scratch.neighbors = nbuf; - - let mut out: Vec<(u32, f32)> = results.into_iter().map(|r| (r.id, r.distance)).collect(); + out.clear(); + out.extend(results.drain().map(|r| (r.id, r.distance))); out.sort_by(|a, b| a.1.total_cmp(&b.1)); - out + + candidates.clear(); + scratch.neighbors = nbuf; + scratch.candidates = candidates; + scratch.results = results; } -/// Greedy single-candidate descent through upper layers from `from_level` down to -/// `to_level` (exclusive lower bound handled by caller loops). Returns improved entry. -pub fn greedy_descend( +/// Beam width at every upper level of the descent. A width of 1 is hill climbing, which halts in +/// the first basin no neighbor improves on; layer-0 adjacency is intra-basin, so a query that +/// lands in the wrong one has no uphill edge out at any ef. See DESIGN.md §7 for the width sweep. +pub const DESCENT_EF: usize = 16; + +/// Beam descent through upper layers from `from_level` down to `to_level` (exclusive), each level +/// a width-`ef` beam seeded by the level above's best. Returns the entry for the caller's +/// layer-`to_level` search, which the caller must `begin()` a fresh epoch for. +/// +/// A node reachable at several levels must be expandable at each, so the epoch rolls per level. +pub fn beam_descend( graph: &Graph, query: &Query, mut current: u32, mut current_dist: f32, from_level: u32, to_level: u32, + ef: usize, + scratch: &mut SearchScratch, stats: &mut SearchStats, ) -> (u32, f32) { - let mut nbuf: Vec = Vec::new(); - let mut level = from_level; + let mut found = std::mem::take(&mut scratch.descent_out); + let mut level = from_level.min(MAX_UPPER_LEVELS as u32); while level > to_level { - let mut improved = true; - while improved { - improved = false; - graph.upper_neighbors_into(current, level.min(255) as u8, &mut nbuf); - for i in 0..nbuf.len() { - let nid = nbuf[i]; - if let Some(d) = graph.distance_to(nid, query) { - stats.visits += 1; - if d < current_dist { - current = nid; - current_dist = d; - improved = true; - } - } - } + scratch.begin(graph.file.id_high_water()); + search_layer( + graph, query, current, current_dist, ef, level as u8, scratch, stats, None, u64::MAX, &mut found, + ); + // search_layer admits the entry itself, so `first` is never worse than what went in + if let Some(&(id, d)) = found.first() { + current = id; + current_dist = d; } level -= 1; } + scratch.descent_out = found; (current, current_dist) } @@ -255,7 +279,7 @@ fn resolve_entry(graph: &Graph, query: &Query, stats: &mut SearchStats) -> Optio Some((id, level as u32, d)) } -/// Full search: greedy descent through upper layers, then beam at layer 0. +/// Full search: beam descent through upper layers, then the layer-0 beam at `ef`. pub fn search( graph: &Graph, query: &Query, @@ -267,9 +291,11 @@ pub fn search( let Some((entry_id, entry_level, entry_dist)) = resolve_entry(graph, query, &mut stats) else { return (Vec::new(), stats); }; + let (ep, ep_dist) = + beam_descend(graph, query, entry_id, entry_dist, entry_level, 0, DESCENT_EF, scratch, &mut stats); scratch.begin(graph.file.id_high_water()); - let (ep, ep_dist) = greedy_descend(graph, query, entry_id, entry_dist, entry_level, 0, &mut stats); - let mut out = search_layer(graph, query, ep, ep_dist, ef, 0, scratch, &mut stats, None, u64::MAX); + let mut out = Vec::new(); + search_layer(graph, query, ep, ep_dist, ef, 0, scratch, &mut stats, None, u64::MAX, &mut out); out.truncate(k); (out, stats) } @@ -289,10 +315,16 @@ pub fn search_filtered( let Some((entry_id, entry_level, entry_dist)) = resolve_entry(graph, query, &mut stats) else { return (Vec::new(), stats); }; + let (ep, ep_dist) = + beam_descend(graph, query, entry_id, entry_dist, entry_level, 0, DESCENT_EF, scratch, &mut stats); scratch.begin_public(graph.file.id_high_water()); - let (ep, ep_dist) = greedy_descend(graph, query, entry_id, entry_dist, entry_level, 0, &mut stats); - let budget = if filter.is_some() { (ef * filter_expansion) as u64 } else { u64::MAX }; - let mut out = search_layer(graph, query, ep, ep_dist, ef, 0, scratch, &mut stats, filter, budget); + let budget = if filter.is_some() { + stats.visits.saturating_add((ef * filter_expansion) as u64) + } else { + u64::MAX + }; + let mut out = Vec::new(); + search_layer(graph, query, ep, ep_dist, ef, 0, scratch, &mut stats, filter, budget, &mut out); out.truncate(k); (out, stats) } @@ -332,8 +364,10 @@ pub fn search_predicated( let Some((entry_id, entry_level, entry_dist)) = resolve_entry(graph, query, &mut stats) else { return (Vec::new(), stats); }; + let (ep, ep_dist) = + beam_descend(graph, query, entry_id, entry_dist, entry_level, 0, DESCENT_EF, scratch, &mut stats); scratch.begin_public(graph.file.id_high_water()); - let (ep, ep_dist) = greedy_descend(graph, query, entry_id, entry_dist, entry_level, 0, &mut stats); + let layer0_budget = stats.visits.saturating_add(visit_budget); use std::collections::HashMap; let mut verdicts: HashMap = HashMap::new(); @@ -341,15 +375,19 @@ pub fn search_predicated( let mut batch: Vec = Vec::new(); let mut outstanding = 0usize; - let mut candidates = BinaryHeap::new(); - let mut results: BinaryHeap = BinaryHeap::new(); + // this path hand-rolls the layer-0 beam because admission waits on a verdict rather than + // being decided at expansion time + let mut candidates = std::mem::take(&mut scratch.candidates); + let mut results = std::mem::take(&mut scratch.results); + let mut nbuf = std::mem::take(&mut scratch.neighbors); + candidates.clear(); + results.clear(); + scratch.visit(ep); candidates.push(Candidate { distance: ep_dist, id: ep }); speculative.push((ep, ep_dist)); batch.push(ep); - let mut nbuf = std::mem::take(&mut scratch.neighbors); - // guarded on `outstanding` rather than draining until the channel is empty: with a blocking // receive the unguarded form pays another full timeout after the last verdict lands, on every // filtered query @@ -387,7 +425,7 @@ pub fn search_predicated( if results.len() >= ef && c.distance > worst { break; } - if stats.visits >= visit_budget { + if stats.visits >= layer0_budget { break; } if graph.neighbors_into(c.id, &mut nbuf).is_none() { @@ -435,12 +473,188 @@ pub fn search_predicated( false }); - let mut out: Vec<(u32, f32)> = results.into_iter().map(|r| (r.id, r.distance)).collect(); + let mut out: Vec<(u32, f32)> = results.drain().map(|r| (r.id, r.distance)).collect(); out.sort_by(|a, b| a.1.total_cmp(&b.1)); out.truncate(k); + + candidates.clear(); + scratch.candidates = candidates; + scratch.results = results; (out, stats) } +#[cfg(test)] +mod descent_tests { + use super::*; + use crate::format::UPPER_CAP; + use crate::insert::{insert, InsertParams}; + use crate::PlaneFile; + + /// The descent takes no caller-supplied visit budget. What bounds it under tied distances is + /// `search_layer`'s strict `d < worst`: a full result set admits no tied candidate, so the + /// beam drains after `ef` expansions. Relaxing that to `<=` lets a zero query — which ties + /// every cosine distance at 1.0, and any caller can send one — walk the whole upper level. + #[test] + fn a_tied_distance_descent_stops_at_its_visit_cap() { + let dims = 16; + let n = 6_000u32; + let path = std::env::temp_dir().join(format!("hnsw-tied-{}.hnsw", std::process::id())); + let _ = std::fs::remove_file(&path); + let graph = Graph::new(PlaneFile::create(&path, dims, 16, n as u64 + 1024).expect("create")); + let params = InsertParams::default(); + let mut scratch = SearchScratch::new(); + for i in 0..n { + let v: Vec = (0..dims).map(|d| ((i as f32 * 0.31 + d as f32) * 0.7).sin()).collect(); + insert(&graph, &v, ¶ms, &mut scratch).expect("insert"); + } + + let (entry_id, entry_level) = graph.file.entry_point(); + assert!(entry_level >= 1, "precondition: the graph has an upper level to descend"); + let query = Query::new(vec![0.0f32; dims]); + let entry_dist = graph.distance_to(entry_id, &query).expect("the entry point is live"); + assert_eq!(entry_dist, 1.0, "precondition: a zero query ties every stored vector at 1.0"); + + let ef = 2usize; // one level, so the bound under test is per-level rather than a sum + let mut stats = SearchStats { visits: 0 }; + beam_descend(&graph, &query, entry_id, entry_dist, 1, 0, ef, &mut scratch, &mut stats); + + let ceiling = (ef * UPPER_CAP) as u64; + assert!( + stats.visits <= ceiling, + "the tied-distance descent visited {} nodes at level 1 against a ceiling of {ceiling} — \ +level 1 holds roughly {} nodes, and a beam that pushed tied candidates would walk all of them", + stats.visits, + n / 16 + ); + let _ = std::fs::remove_file(&path); + } + + /// `filter_expansion` bounds layer 0, so the budget has to start counting where the descent + /// left off. Measured against a counter the descent has already advanced, a descent costing + /// more than the whole budget leaves layer 0 unable to expand even one candidate, and the + /// search returns its entry point instead of a result set — silently, since a filtered search + /// is allowed to return short. + #[test] + fn the_filter_budget_bounds_layer_zero_not_the_descent() { + let dims = 16; + let n = 6_000u32; + let path = std::env::temp_dir().join(format!("hnsw-budget-{}.hnsw", std::process::id())); + let _ = std::fs::remove_file(&path); + let graph = Graph::new(PlaneFile::create(&path, dims, 16, n as u64 + 1024).expect("create")); + let params = InsertParams::default(); + let mut scratch = SearchScratch::new(); + for i in 0..n { + let v: Vec = (0..dims).map(|d| ((i as f32 * 0.31 + d as f32) * 0.7).sin()).collect(); + insert(&graph, &v, ¶ms, &mut scratch).expect("insert"); + } + + let query = Query::new((0..dims).map(|d| ((97.0f32 * 0.31 + d as f32) * 0.7).sin()).collect()); + let (entry_id, entry_level) = graph.file.entry_point(); + let entry_dist = graph.distance_to(entry_id, &query).expect("the entry point is live"); + let mut descent = SearchStats { visits: 0 }; + beam_descend(&graph, &query, entry_id, entry_dist, entry_level, 0, DESCENT_EF, &mut scratch, &mut descent); + + // a budget deliberately far below what the descent spends + let (ef, filter_expansion) = (16usize, 1usize); + assert!( + descent.visits > (ef * filter_expansion) as u64, + "precondition: the descent ({} visits) must cost more than the whole budget ({})", + descent.visits, + ef * filter_expansion + ); + let allow = vec![0xffu8; (n as usize).div_ceil(8)]; + let (hits, stats) = + search_filtered(&graph, &query, 10, ef, Some(&allow), filter_expansion, &mut scratch); + + assert_eq!(hits.len(), 10, "layer 0 got no budget of its own: {hits:?}"); + assert!( + stats.visits > descent.visits, + "layer 0 expanded nothing beyond the descent ({} total vs {} for the descent alone)", + stats.visits, + descent.visits + ); + let _ = std::fs::remove_file(&path); + } + + /// `ef` arrives from an unvalidated `u32` at the NAPI boundary. Sizing any allocation by it + /// turns a caller's typo into a request for tens of gigabytes, which `handle_alloc_error` + /// answers by aborting — uncatchable from JS, and it takes every in-flight query with it. + #[test] + fn an_absurd_ef_answers_instead_of_reserving_by_it() { + let dims = 8; + let n = 64u32; + let path = std::env::temp_dir().join(format!("hnsw-absurdef-{}.hnsw", std::process::id())); + let _ = std::fs::remove_file(&path); + let graph = Graph::new(PlaneFile::create(&path, dims, 8, n as u64 + 16).expect("create")); + let params = InsertParams::default(); + let mut scratch = SearchScratch::new(); + for i in 0..n { + let v: Vec = (0..dims).map(|d| ((i as f32 * 0.31 + d as f32) * 0.7).sin()).collect(); + insert(&graph, &v, ¶ms, &mut scratch).expect("insert"); + } + + let query = Query::new((0..dims).map(|d| ((7.0f32 * 0.31 + d as f32) * 0.7).sin()).collect()); + let (hits, _) = search(&graph, &query, 5, u32::MAX as usize, &mut scratch); + assert_eq!(hits.len(), 5, "an absurd ef must still answer from a {n}-node plane"); + let _ = std::fs::remove_file(&path); + } + + /// The same contract on the predicate path, which carries its own layer-0 loop: a host budget + /// smaller than the descent must still buy layer-0 visits. `predicate_tests` all pass + /// `64 * 24`, orders above what a descent costs, so they hold either way. + #[test] + fn the_predicate_visit_budget_bounds_layer_zero_not_the_descent() { + let dims = 16; + let n = 6_000u32; + let path = std::env::temp_dir().join(format!("hnsw-predbudget-{}.hnsw", std::process::id())); + let _ = std::fs::remove_file(&path); + let graph = Graph::new(PlaneFile::create(&path, dims, 16, n as u64 + 1024).expect("create")); + let params = InsertParams::default(); + let mut scratch = SearchScratch::new(); + for i in 0..n { + let v: Vec = (0..dims).map(|d| ((i as f32 * 0.31 + d as f32) * 0.7).sin()).collect(); + insert(&graph, &v, ¶ms, &mut scratch).expect("insert"); + } + + let query = Query::new((0..dims).map(|d| ((97.0f32 * 0.31 + d as f32) * 0.7).sin()).collect()); + let (entry_id, entry_level) = graph.file.entry_point(); + let entry_dist = graph.distance_to(entry_id, &query).expect("the entry point is live"); + let mut descent = SearchStats { visits: 0 }; + beam_descend(&graph, &query, entry_id, entry_dist, entry_level, 0, DESCENT_EF, &mut scratch, &mut descent); + + let budget = 64u64; + assert!( + descent.visits > budget, + "precondition: the descent ({} visits) must cost more than the whole budget ({budget})", + descent.visits + ); + + let (req_tx, req_rx) = std::sync::mpsc::channel::>(); + let (res_tx, res_rx) = std::sync::mpsc::channel::<(Vec, Vec)>(); + let worker = std::thread::spawn(move || { + while let Ok(ids) = req_rx.recv() { + let verdicts = vec![1u8; ids.len()]; + if res_tx.send((ids, verdicts)).is_err() { + break; + } + } + }); + let mut pipe = PredicatePipe { dispatch: Box::new(move |ids| req_tx.send(ids).is_ok()), rx: res_rx }; + let (hits, stats) = search_predicated(&graph, &query, 10, 16, &mut pipe, budget, &mut scratch); + + assert_eq!(hits.len(), 10, "layer 0 got no budget of its own: {hits:?}"); + assert!( + stats.visits > descent.visits, + "layer 0 expanded nothing beyond the descent ({} total vs {} for the descent alone)", + stats.visits, + descent.visits + ); + drop(pipe); + worker.join().unwrap(); + let _ = std::fs::remove_file(&path); + } +} + #[cfg(test)] mod predicate_tests { use super::*; diff --git a/tests/allocation.rs b/tests/allocation.rs new file mode 100644 index 0000000..41210bb --- /dev/null +++ b/tests/allocation.rs @@ -0,0 +1,125 @@ +//! The module header calls the search path zero-copy with a reusable scratch. This binary is the +//! only place that can hold a counting global allocator, so it owns the one test that checks it: +//! per-query allocations must not scale with the graph's height. A descent that allocated per +//! upper level would pass every other test in the suite. + +use hnsw_plane::distance::Query; +use hnsw_plane::insert::{insert, InsertParams}; +use hnsw_plane::search::{search, SearchScratch}; +use hnsw_plane::{Graph, PlaneFile}; +use std::alloc::{GlobalAlloc, Layout, System}; +use std::cell::Cell; +use std::sync::atomic::{AtomicUsize, Ordering}; + +static ALLOCATIONS: AtomicUsize = AtomicUsize::new(0); + +thread_local! { + /// A global allocator sees every thread in the binary, the test harness's own included, so a + /// process-wide flag would fold their allocations into the measurement. Scoping it per thread + /// keeps the count to the one doing the searching. `const` init on a `Cell` so that reading it + /// neither lazily initializes nor registers a destructor — either would allocate, and + /// allocating inside the allocator recurses. + static COUNTING: Cell = const { Cell::new(false) }; +} + +#[inline] +fn counting_here() -> bool { + COUNTING.try_with(|c| c.get()).unwrap_or(false) +} + +struct CountingAllocator; + +unsafe impl GlobalAlloc for CountingAllocator { + unsafe fn alloc(&self, layout: Layout) -> *mut u8 { + if counting_here() { + ALLOCATIONS.fetch_add(1, Ordering::Relaxed); + } + System.alloc(layout) + } + unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) { + System.dealloc(ptr, layout) + } + unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 { + if counting_here() { + ALLOCATIONS.fetch_add(1, Ordering::Relaxed); + } + System.realloc(ptr, layout, new_size) + } +} + +#[global_allocator] +static ALLOCATOR: CountingAllocator = CountingAllocator; + +/// Allocations attributable to one `search` on a warmed scratch, averaged over `queries`. +fn allocations_per_search(graph: &Graph, dims: usize, scratch: &mut SearchScratch, queries: usize) -> f64 { + let vector = |i: usize| -> Vec { + (0..dims).map(|d| ((i as f32 * 0.11 + d as f32) * 0.9).sin()).collect() + }; + for i in 0..8 { + let _ = search(graph, &Query::new(vector(i)), 10, 64, scratch); + } + let queries_prepared: Vec = (0..queries).map(|i| Query::new(vector(i + 100))).collect(); + + ALLOCATIONS.store(0, Ordering::Relaxed); + COUNTING.with(|c| c.set(true)); + for query in &queries_prepared { + let (hits, _) = search(graph, query, 10, 64, scratch); + std::hint::black_box(hits); + } + COUNTING.with(|c| c.set(false)); + ALLOCATIONS.load(Ordering::Relaxed) as f64 / queries as f64 +} + +fn build(path: &std::path::Path, n: u32, dims: usize) -> Graph { + let _ = std::fs::remove_file(path); + let graph = Graph::new(PlaneFile::create(path, dims, 16, n as u64 + 1024).expect("create")); + // only the graph's height matters here, not its quality + let params = InsertParams { ef_construction: 24, ..InsertParams::default() }; + let mut scratch = SearchScratch::new(); + for i in 0..n { + let v: Vec = (0..dims).map(|d| ((i as f32 * 0.31 + d as f32) * 0.7).sin()).collect(); + insert(&graph, &v, ¶ms, &mut scratch).expect("insert"); + } + graph +} + +/// A taller graph means more upper levels to descend. If the descent allocated per level, the +/// per-query count would rise with height; it must not. +#[test] +fn a_search_does_not_allocate_per_upper_level() { + let dims = 16; + let pid = std::process::id(); + let shallow_path = std::env::temp_dir().join(format!("hnsw-alloc-shallow-{pid}.hnsw")); + let tall_path = std::env::temp_dir().join(format!("hnsw-alloc-tall-{pid}.hnsw")); + + let shallow = build(&shallow_path, 2_000, dims); + let tall = build(&tall_path, 60_000, dims); + let shallow_levels = shallow.file.entry_point().1; + let tall_levels = tall.file.entry_point().1; + assert!( + tall_levels >= shallow_levels + 2, + "precondition: the two graphs must differ in height ({shallow_levels} vs {tall_levels}). \ +This is derived from `ml` and `level_for`; if either was tuned, re-pick the two node counts \ +rather than reading this as an allocation regression" + ); + + let mut scratch = SearchScratch::new(); + let shallow_allocs = allocations_per_search(&shallow, dims, &mut scratch, 200); + let tall_allocs = allocations_per_search(&tall, dims, &mut scratch, 200); + + assert!( + tall_allocs <= shallow_allocs + 0.5, + "search allocates per upper level: {shallow_allocs:.2}/query at {shallow_levels} levels \ +vs {tall_allocs:.2}/query at {tall_levels}" + ); + assert!( + tall_allocs <= 2.0, + "search allocates {tall_allocs:.2} times per query on a warmed scratch; the result vector \ +should be the only one" + ); + + drop(shallow); + drop(tall); + let _ = std::fs::remove_file(&shallow_path); + let _ = std::fs::remove_file(&tall_path); +} diff --git a/tests/concurrent.rs b/tests/concurrent.rs index e3699e2..5040859 100644 --- a/tests/concurrent.rs +++ b/tests/concurrent.rs @@ -1,9 +1,11 @@ //! Concurrent-write torture: writers insert while readers search; then verify the graph is -//! coherent (every stored vector findable, edge lists within cap, freelist reuse works). +//! coherent (every stored vector findable, edge lists within cap, freelist reuse works). The +//! deterministic descent tests at the bottom share `vector_for`: its clustered corpus is what +//! makes them bite. use hnsw_plane::distance::Query; use hnsw_plane::insert::{insert, InsertParams}; -use hnsw_plane::search::{search, SearchScratch}; +use hnsw_plane::search::{beam_descend, search, search_layer, SearchScratch, SearchStats, DESCENT_EF}; use hnsw_plane::{Graph, PlaneFile}; use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::{Arc, Barrier}; @@ -172,3 +174,149 @@ fn racing_first_inserts_all_stay_reachable() { let _ = std::fs::remove_file(&path); } } + +/// Fisher-Yates over a xorshift stream: one seed names one exact graph, with no thread +/// interleaving in it. +fn insertion_order(n: u32, seed: u64) -> Vec { + let mut order: Vec = (0..n).collect(); + let mut s = seed.wrapping_mul(0x9e37_79b9_7f4a_7c15) | 1; + for i in (1..order.len()).rev() { + s ^= s << 13; + s ^= s >> 7; + s ^= s << 17; + order.swap(i, (s % (i as u64 + 1)) as usize); + } + order +} + +/// `concurrent_insert_search`'s reachability assertion with the concurrency removed. Concurrency +/// only shuffles the insertion order, which this fixes outright, so what remains is the descent: +/// a width-1 hill climb halts in the first basin no neighbor improves on, and layer-0 adjacency +/// is intra-basin, leaving the query's own vector unreachable at any ef. Sampling every 97th node +/// as the test above does catches that ~1.5% of the time; these two seeds catch it every time, +/// losing 55 and 22 of their 8000 self-queries on a width-1 descent. +#[test] +fn a_descent_that_traps_at_a_local_minimum_still_reaches_the_true_neighborhood() { + let dims = 64; + let n = 8_000u32; + for &seed in &[57u64, 240] { + let path = std::env::temp_dir().join(format!("hnsw-descent-{}-{seed}.hnsw", std::process::id())); + let _ = std::fs::remove_file(&path); + let graph = Graph::new(PlaneFile::create(&path, dims, 32, n as u64 + 1024).expect("create")); + let params = InsertParams::default(); + let mut scratch = SearchScratch::new(); + + let inserted: Vec<(u32, u32)> = insertion_order(n, seed) + .into_iter() + .map(|index| { + let v = vector_for(index, dims); + (index, insert(&graph, &v, ¶ms, &mut scratch).expect("insert")) + }) + .collect(); + + let mut misses = Vec::new(); + for &(index, id) in &inserted { + let query = Query::new(vector_for(index, dims)); + let (results, _) = search(&graph, &query, 10, 256, &mut scratch); + if !results.iter().any(|&(rid, _)| rid == id) { + misses.push((index, id)); + } + } + assert!( + misses.is_empty(), + "seed {seed}: {} of {n} nodes are their own true nearest neighbor but unreachable \ +from the entry point — the descent stranded the search. First few (corpus index, node id): {:?}", + misses.len(), + &misses[..misses.len().min(8)] + ); + + drop(graph); + let _ = std::fs::remove_file(&path); + } +} + +/// `search` at an explicit descent width — what the sweep varies, since `search` itself reads the +/// `DESCENT_EF` constant. A single-threaded build always has a live entry point, so this skips +/// `search`'s dead-entry repair and is otherwise the same sequence. +fn search_at_descent_width( + graph: &Graph, + query: &Query, + k: usize, + ef: usize, + descent_ef: usize, + scratch: &mut SearchScratch, +) -> Vec<(u32, f32)> { + let (entry_id, entry_level) = graph.file.entry_point(); + let Some(entry_dist) = graph.distance_to(entry_id, query) else { + return Vec::new(); + }; + let mut stats = SearchStats { visits: 0 }; + let (ep, ep_dist) = + beam_descend(graph, query, entry_id, entry_dist, entry_level, 0, descent_ef, scratch, &mut stats); + scratch.begin_public(graph.file.id_high_water()); + let mut out = Vec::new(); + search_layer(graph, query, ep, ep_dist, ef, 0, scratch, &mut stats, None, u64::MAX, &mut out); + out.truncate(k); + out +} + +/// The measurement behind DESIGN.md's descent-width table, kept runnable so the numbers can be +/// re-derived when M, ml or the prune policy changes. Ignored: it reports rather than asserts, +/// and a full sweep is minutes of CPU. +/// +/// `HNSW_SWEEP_READ_EF` sets the query-side descent width; the build side is whatever +/// `DESCENT_EF` is compiled as. Varying them independently is what separates a routing defect +/// from a construction one — DESIGN.md §7 records that matrix. +/// +/// ```text +/// HNSW_SWEEP_SEEDS=700 HNSW_SWEEP_READ_EF=8 cargo test --release --test concurrent \ +/// descent_width_sweep -- --ignored --nocapture +/// ``` +#[test] +#[ignore] +fn descent_width_sweep() { + let dims = 64; + let n = 8_000u32; + let seeds: u64 = std::env::var("HNSW_SWEEP_SEEDS").ok().and_then(|v| v.parse().ok()).unwrap_or(50); + // 0 would silently mean width 1, since the entry is admitted before any cap check + let read_ef: usize = std::env::var("HNSW_SWEEP_READ_EF") + .ok() + .and_then(|v| v.parse().ok()) + .map(|v: usize| v.max(1)) + .unwrap_or(DESCENT_EF); + let mut total = 0usize; + let mut bad = 0usize; + for seed in 0..seeds { + let path = std::env::temp_dir().join(format!("hnsw-sweep-{}-{seed}.hnsw", std::process::id())); + let _ = std::fs::remove_file(&path); + let graph = Graph::new(PlaneFile::create(&path, dims, 32, n as u64 + 1024).expect("create")); + let params = InsertParams::default(); + let mut scratch = SearchScratch::new(); + let inserted: Vec<(u32, u32)> = insertion_order(n, seed) + .into_iter() + .map(|index| { + let v = vector_for(index, dims); + (index, insert(&graph, &v, ¶ms, &mut scratch).expect("insert")) + }) + .collect(); + let misses = inserted + .iter() + .filter(|&&(index, id)| { + let q = Query::new(vector_for(index, dims)); + let results = search_at_descent_width(&graph, &q, 10, 256, read_ef, &mut scratch); + !results.iter().any(|&(rid, _)| rid == id) + }) + .count(); + if misses > 0 { + println!("seed {seed}: {misses} misses"); + bad += 1; + } + total += misses; + drop(graph); + let _ = std::fs::remove_file(&path); + } + println!( + "descent width sweep (build {DESCENT_EF} / read {read_ef}): {total} misses over {seeds} \ +builds of {n}, {bad} builds affected" + ); +}