diff --git a/crates/ruvector-diskann/examples/bench_search_hot.rs b/crates/ruvector-diskann/examples/bench_search_hot.rs new file mode 100644 index 000000000..8c06a3cd9 --- /dev/null +++ b/crates/ruvector-diskann/examples/bench_search_hot.rs @@ -0,0 +1,83 @@ +use rand::rngs::StdRng; +use rand::{Rng, SeedableRng}; +use ruvector_diskann::{DiskAnnConfig, DiskAnnIndex}; +use std::hint::black_box; +use std::time::Instant; + +const DIM: usize = 128; +const QUERY_COUNT: usize = 2_000; +const WARMUP_COUNT: usize = 100; +const K: usize = 10; + +fn random_unit_vector(rng: &mut StdRng) -> Vec { + let mut vector: Vec = (0..DIM).map(|_| rng.gen_range(-1.0..1.0)).collect(); + let norm = vector.iter().map(|value| value * value).sum::().sqrt(); + for value in &mut vector { + *value /= norm; + } + vector +} + +fn parse_n() -> usize { + let mut args = std::env::args().skip(1); + match args.next() { + Some(flag) if flag == "--n" => args + .next() + .expect("--n requires a value") + .parse() + .expect("n must be a positive integer"), + Some(value) => value.parse().expect("n must be a positive integer"), + None => 100_000, + } +} + +fn main() { + let n = parse_n(); + assert!(n > 1, "n must be greater than one"); + + let mut data_rng = StdRng::seed_from_u64(0x677D_15CA_77); + let mut index = DiskAnnIndex::new(DiskAnnConfig { + dim: DIM, + max_degree: 4, + build_beam: 8, + search_beam: 64, + alpha: 1.0, + ..Default::default() + }); + + let build_started = Instant::now(); + for id in 0..n { + index + .insert(id.to_string(), random_unit_vector(&mut data_rng)) + .expect("benchmark insert failed"); + } + index.build().expect("benchmark build failed"); + let build_seconds = build_started.elapsed().as_secs_f64(); + + let mut query_rng = StdRng::seed_from_u64(0x6770_0A11_CE); + let queries: Vec> = (0..QUERY_COUNT) + .map(|_| random_unit_vector(&mut query_rng)) + .collect(); + + for query in queries.iter().take(WARMUP_COUNT) { + black_box(index.search(black_box(query), K).expect("warmup failed")); + } + + let mut latencies_ns = Vec::with_capacity(QUERY_COUNT); + let run_started = Instant::now(); + for query in &queries { + let query_started = Instant::now(); + black_box(index.search(black_box(query), K).expect("search failed")); + latencies_ns.push(query_started.elapsed().as_nanos() as u64); + } + let elapsed = run_started.elapsed(); + + latencies_ns.sort_unstable(); + let p50_us = latencies_ns[QUERY_COUNT / 2] as f64 / 1_000.0; + let p99_us = latencies_ns[(QUERY_COUNT * 99 / 100).min(QUERY_COUNT - 1)] as f64 / 1_000.0; + let qps = QUERY_COUNT as f64 / elapsed.as_secs_f64(); + + println!( + "RESULT n={n} d={DIM} k={K} queries={QUERY_COUNT} build_s={build_seconds:.3} p50_us={p50_us:.3} p99_us={p99_us:.3} qps={qps:.3}" + ); +} diff --git a/crates/ruvector-diskann/src/distance.rs b/crates/ruvector-diskann/src/distance.rs index 5716380dc..a6d0f67ca 100644 --- a/crates/ruvector-diskann/src/distance.rs +++ b/crates/ruvector-diskann/src/distance.rs @@ -199,7 +199,30 @@ impl VisitedSet { /// Reset for a new search — O(1) via generation counter #[inline] pub fn clear(&mut self) { - self.generation += 1; + if self.generation == u64::MAX { + self.bits.fill(0); + self.gens.fill(0); + self.generation = 1; + } else { + self.generation += 1; + } + } + + /// Prepare this set for an index of `n` nodes. + /// + /// A size mismatch reinitializes the backing storage. Repeated use with the + /// same index size takes the O(1) generation-counter path in [`Self::clear`]. + #[inline] + pub(crate) fn prepare(&mut self, n: usize) { + if self.gens.len() != n { + self.bits.resize((n + 63) / 64, 0); + self.bits.fill(0); + self.gens.resize(n, 0); + self.gens.fill(0); + self.generation = 1; + } else { + self.clear(); + } } /// Mark node as visited @@ -341,6 +364,31 @@ mod tests { assert!(vs.contains(43)); } + #[test] + fn test_visited_set_generation_wrap() { + let mut vs = VisitedSet::new(100); + vs.generation = u64::MAX; + vs.insert(42); + + vs.clear(); + + assert_eq!(vs.generation, 1); + assert!(!vs.contains(42)); + } + + #[test] + fn test_visited_set_size_mismatch_reinitializes() { + let mut vs = VisitedSet::new(2); + vs.insert(1); + + vs.prepare(100); + + assert_eq!(vs.gens.len(), 100); + assert!(!vs.contains(1)); + vs.insert(99); + assert!(vs.contains(99)); + } + #[test] fn test_pq_flat_table() { // 2 subspaces, 4 centroids each (k=4 for test) diff --git a/crates/ruvector-diskann/src/graph.rs b/crates/ruvector-diskann/src/graph.rs index c8d6e5bff..d6c8f0a47 100644 --- a/crates/ruvector-diskann/src/graph.rs +++ b/crates/ruvector-diskann/src/graph.rs @@ -141,7 +141,7 @@ impl VamanaGraph { beam_width: usize, visited: &mut VisitedSet, ) -> (Vec, usize) { - visited.clear(); + visited.prepare(vectors.len()); let mut candidates = BinaryHeap::::new(); let mut best = BinaryHeap::::new(); diff --git a/crates/ruvector-diskann/src/index.rs b/crates/ruvector-diskann/src/index.rs index 587bcebec..aaa3d53b7 100644 --- a/crates/ruvector-diskann/src/index.rs +++ b/crates/ruvector-diskann/src/index.rs @@ -9,6 +9,9 @@ use std::collections::HashMap; use std::fs::{self, File}; use std::io::{BufWriter, Write}; use std::path::{Path, PathBuf}; +use std::sync::{Mutex, TryLockError}; + +const MAX_VISITED_POOL_SIZE: usize = 4; /// Search result #[derive(Debug, Clone)] @@ -70,8 +73,8 @@ pub struct DiskAnnIndex { pq_codes: Vec>, /// Whether index has been built built: bool, - /// Reusable visited set for search (avoids per-query allocation) - visited: Option, + /// Small pool of reusable visited sets for shared-reference searches. + visited_pool: Mutex>, /// Memory-mapped vector data (for large datasets) mmap: Option, } @@ -89,7 +92,7 @@ impl DiskAnnIndex { pq: None, pq_codes: Vec::new(), built: false, - visited: None, + visited_pool: Mutex::new(Vec::new()), mmap: None, } } @@ -154,8 +157,12 @@ impl DiskAnnIndex { graph.build(&self.vectors)?; self.graph = Some(graph); - // Pre-allocate visited set for search - self.visited = Some(VisitedSet::new(n)); + // Seed the single-threaded fast path. Concurrent searches allocate a + // fallback set when every pooled set is checked out. + *self + .visited_pool + .get_mut() + .unwrap_or_else(|poisoned| poisoned.into_inner()) = vec![VisitedSet::new(n)]; self.built = true; if let Some(ref path) = self.config.storage_path { @@ -167,20 +174,49 @@ impl DiskAnnIndex { /// Search for k nearest neighbors pub fn search(&self, query: &[f32], k: usize) -> Result> { - if !self.built { - return Err(DiskAnnError::NotBuilt); - } - if query.len() != self.config.dim { - return Err(DiskAnnError::DimensionMismatch { - expected: self.config.dim, - actual: query.len(), - }); + self.validate_search(query)?; + + let mut visited = match self.visited_pool.try_lock() { + Ok(mut pool) => pool + .pop() + .unwrap_or_else(|| VisitedSet::new(self.vectors.len())), + Err(TryLockError::Poisoned(poisoned)) => poisoned + .into_inner() + .pop() + .unwrap_or_else(|| VisitedSet::new(self.vectors.len())), + Err(TryLockError::WouldBlock) => VisitedSet::new(self.vectors.len()), + }; + + let result = self.search_with(query, k, &mut visited); + + let mut pool = self + .visited_pool + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + if pool.len() < MAX_VISITED_POOL_SIZE { + pool.push(visited); } + result + } + + /// Search with caller-owned visited-state storage. + /// + /// Reuse the same [`VisitedSet`] for steady-state searches without visited + /// set allocation. If its size differs from this index, the set is resized + /// and fully reset before the search. + pub fn search_with( + &self, + query: &[f32], + k: usize, + visited: &mut VisitedSet, + ) -> Result> { + self.validate_search(query)?; + let graph = self.graph.as_ref().unwrap(); let beam = self.config.search_beam.max(k); - let (candidates, _) = graph.greedy_search(&self.vectors, query, beam); + let (candidates, _) = graph.greedy_search_fast(&self.vectors, query, beam, visited); // Re-rank candidates with exact distance let mut scored: Vec<(u32, f32)> = candidates @@ -199,6 +235,20 @@ impl DiskAnnIndex { .collect()) } + fn validate_search(&self, query: &[f32]) -> Result<()> { + if !self.built { + return Err(DiskAnnError::NotBuilt); + } + if query.len() != self.config.dim { + return Err(DiskAnnError::DimensionMismatch { + expected: self.config.dim, + actual: query.len(), + }); + } + + Ok(()) + } + /// Get the number of vectors in the index pub fn count(&self) -> usize { self.vectors.len() @@ -407,7 +457,7 @@ impl DiskAnnIndex { pq, pq_codes, built: true, - visited: Some(VisitedSet::new(n)), + visited_pool: Mutex::new(vec![VisitedSet::new(n)]), mmap: Some(mmap), }) } @@ -459,6 +509,64 @@ mod tests { assert!(results[0].distance < 1e-6); // Exact match } + #[test] + fn search_paths_return_identical_ids() { + use rand::prelude::*; + + let dim = 32; + let n = 500; + let mut index = DiskAnnIndex::new(DiskAnnConfig { + dim, + max_degree: 16, + build_beam: 32, + search_beam: 32, + alpha: 1.2, + ..Default::default() + }); + index.insert_batch(random_vectors(n, dim)).unwrap(); + index.build().unwrap(); + + let mut query_rng = rand::rngs::StdRng::seed_from_u64(0x6775_EA2C); + let mut visited = VisitedSet::new(n); + + for _ in 0..100 { + let query: Vec = (0..dim).map(|_| query_rng.gen()).collect(); + let pooled_ids: Vec = index + .search(&query, 10) + .unwrap() + .into_iter() + .map(|result| result.id) + .collect(); + let caller_owned_ids: Vec = index + .search_with(&query, 10, &mut visited) + .unwrap() + .into_iter() + .map(|result| result.id) + .collect(); + + // This mirrors the pre-fix search shape: allocate a VisitedSet in + // graph.greedy_search, then perform the same exact-distance rerank. + let graph = index.graph.as_ref().unwrap(); + let beam = index.config.search_beam.max(10); + let (candidates, _) = graph.greedy_search(&index.vectors, &query, beam); + let mut scored: Vec<(u32, f32)> = candidates + .into_iter() + .map(|id| (id, l2_squared(index.vectors.get(id as usize), &query))) + .collect(); + scored.sort_unstable_by(|a, b| { + a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal) + }); + let old_path_ids: Vec = scored + .into_iter() + .take(10) + .map(|(id, _)| index.id_map[id as usize].clone()) + .collect(); + + assert_eq!(pooled_ids, old_path_ids); + assert_eq!(caller_owned_ids, old_path_ids); + } + } + #[test] fn test_diskann_with_pq() { let mut index = DiskAnnIndex::new(DiskAnnConfig { diff --git a/crates/ruvector-diskann/src/lib.rs b/crates/ruvector-diskann/src/lib.rs index 4b84ad035..cccab6ff3 100644 --- a/crates/ruvector-diskann/src/lib.rs +++ b/crates/ruvector-diskann/src/lib.rs @@ -19,6 +19,7 @@ pub mod pq; #[cfg(feature = "reuse-under-drift")] pub mod reuse; +pub use distance::VisitedSet; pub use error::{DiskAnnError, Result}; pub use index::{DiskAnnConfig, DiskAnnIndex}; pub use pq::ProductQuantizer; diff --git a/crates/ruvector-diskann/tests/search_allocations.rs b/crates/ruvector-diskann/tests/search_allocations.rs new file mode 100644 index 000000000..e09a42856 --- /dev/null +++ b/crates/ruvector-diskann/tests/search_allocations.rs @@ -0,0 +1,91 @@ +use rand::rngs::StdRng; +use rand::{Rng, SeedableRng}; +use ruvector_diskann::{DiskAnnConfig, DiskAnnIndex}; +use std::alloc::{GlobalAlloc, Layout, System}; +use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; + +struct CountingAllocator; + +static TRACKING: AtomicBool = AtomicBool::new(false); +static LARGE_ALLOCATION_THRESHOLD: AtomicUsize = AtomicUsize::new(usize::MAX); +static LARGE_ALLOCATED_BYTES: AtomicUsize = AtomicUsize::new(0); +static TOTAL_ALLOCATED_BYTES: AtomicUsize = AtomicUsize::new(0); + +unsafe impl GlobalAlloc for CountingAllocator { + unsafe fn alloc(&self, layout: Layout) -> *mut u8 { + if TRACKING.load(Ordering::Relaxed) { + TOTAL_ALLOCATED_BYTES.fetch_add(layout.size(), Ordering::Relaxed); + if layout.size() >= LARGE_ALLOCATION_THRESHOLD.load(Ordering::Relaxed) { + LARGE_ALLOCATED_BYTES.fetch_add(layout.size(), 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 TRACKING.load(Ordering::Relaxed) { + TOTAL_ALLOCATED_BYTES.fetch_add(new_size, Ordering::Relaxed); + if new_size >= LARGE_ALLOCATION_THRESHOLD.load(Ordering::Relaxed) { + LARGE_ALLOCATED_BYTES.fetch_add(new_size, Ordering::Relaxed); + } + } + System.realloc(ptr, layout, new_size) + } +} + +#[global_allocator] +static ALLOCATOR: CountingAllocator = CountingAllocator; + +#[test] +fn pooled_search_does_not_allocate_visited_set_storage() { + const N: usize = 250_000; + const DIM: usize = 8; + const SEARCHES: usize = 100; + const MAX_HOT_SEARCH_ALLOCATED_BYTES: usize = 300_000; + + let mut rng = StdRng::seed_from_u64(0x677A_110C); + let mut index = DiskAnnIndex::new(DiskAnnConfig { + dim: DIM, + max_degree: 4, + build_beam: 8, + search_beam: 32, + alpha: 1.0, + ..Default::default() + }); + for id in 0..N { + let vector = (0..DIM).map(|_| rng.gen()).collect(); + index.insert(id.to_string(), vector).unwrap(); + } + index.build().unwrap(); + + let query: Vec = (0..DIM).map(|_| rng.gen()).collect(); + index.search(&query, 10).unwrap(); + + // VisitedSet's generation vector alone requests N * sizeof(u64) bytes. + // Count every allocation at least that large while allowing the unrelated, + // much smaller candidate and result buffers used by the search itself. + let visited_generation_bytes = N * std::mem::size_of::(); + LARGE_ALLOCATION_THRESHOLD.store(visited_generation_bytes, Ordering::Relaxed); + LARGE_ALLOCATED_BYTES.store(0, Ordering::Relaxed); + TOTAL_ALLOCATED_BYTES.store(0, Ordering::Relaxed); + TRACKING.store(true, Ordering::Relaxed); + for _ in 0..SEARCHES { + std::hint::black_box(index.search(&query, 10).unwrap()); + } + TRACKING.store(false, Ordering::Relaxed); + + let allocated = LARGE_ALLOCATED_BYTES.load(Ordering::Relaxed); + let total_allocated = TOTAL_ALLOCATED_BYTES.load(Ordering::Relaxed); + assert_eq!( + allocated, 0, + "{SEARCHES} hot searches allocated {allocated} bytes in visited-set-sized blocks" + ); + assert!( + total_allocated < MAX_HOT_SEARCH_ALLOCATED_BYTES, + "{SEARCHES} hot searches allocated {total_allocated} total bytes, expected less than {MAX_HOT_SEARCH_ALLOCATED_BYTES} bytes and far below the {visited_generation_bytes}-byte generation array" + ); +}