diff --git a/core/src/diskann.rs b/core/src/diskann.rs index b668b59..2b0063e 100644 --- a/core/src/diskann.rs +++ b/core/src/diskann.rs @@ -352,10 +352,12 @@ impl DiskAnnIndex { let pq_build_distances = if self.build_params.build_distance == DiskAnnBuildDistance::ProductQuantized { + // The shared m * ksub * ksub table plus each worker's m * ksub copy + // of the rows for its current search's query. self.pq .m() .checked_mul(self.pq.ksub()) - .and_then(|value| value.checked_mul(self.pq.ksub())) + .and_then(|value| value.checked_mul(self.pq.ksub().checked_add(workers)?)) .and_then(|value| value.checked_mul(size_of::())) .ok_or_else(|| invalid_input("DiskANN PQ build-distance table size overflows"))? } else { @@ -986,6 +988,37 @@ mod tests { assert!(row_id_encoding_scratch_bytes(usize::MAX).is_err()); } + #[test] + fn diskann_memory_estimate_counts_pq_build_query_rows_per_worker() { + rayon::ThreadPoolBuilder::new() + .num_threads(12) + .build() + .unwrap() + .install(|| { + let estimate = |build_distance| { + let mut index = DiskAnnIndex::new( + 16, + MetricType::L2, + 4, + DiskAnnBuildParams { + build_distance, + ..DiskAnnBuildParams::default() + }, + ); + index.ids = (0..1_024).collect(); + index.vectors = vec![0.0; 1_024 * 16]; + index.estimate_build_memory_bytes().unwrap() + }; + let (m, ksub, workers) = (4, 256, 12); + + assert_eq!( + estimate(DiskAnnBuildDistance::ProductQuantized) + - estimate(DiskAnnBuildDistance::FullPrecision), + (m * ksub * ksub + workers * m * ksub) * size_of::() + ); + }); + } + #[test] fn diskann_memory_budget_automatically_selects_overlapping_shards() { // Pin the worker count so this test exercises the same memory-plan diff --git a/core/src/vamana.rs b/core/src/vamana.rs index f64ec30..228d0fa 100644 --- a/core/src/vamana.rs +++ b/core/src/vamana.rs @@ -1244,15 +1244,21 @@ enum BuildSearchDistance<'a> { } impl BuildSearchDistance<'_> { + fn prepare_query(&self, query_node: usize, table: &mut Vec) { + if let Self::ProductQuantized(distance) = self { + distance.prepare_query(query_node, table); + } + } + #[inline] - fn between(&self, left: usize, right: usize) -> f32 { + fn to_query(&self, node: usize, query_node: usize, table: &[f32]) -> f32 { match self { Self::FullPrecision { vectors, dimension, metric, - } => distance_between(vectors, *dimension, left, right, *metric), - Self::ProductQuantized(distance) => distance.between(left, right), + } => distance_between(vectors, *dimension, node, query_node, *metric), + Self::ProductQuantized(distance) => distance.to_query(node, table), } } } @@ -1332,6 +1338,32 @@ impl<'a> PqBuildDistance<'a> { } } + // A build search measures every candidate against one fixed node, so it + // only needs the m table rows picked by that node's codes. Copying them + // into one m * ksub buffer replaces lookups that stride by ksub floats + // through the whole m * ksub * ksub table. + fn prepare_query(&self, query_node: usize, table: &mut Vec) { + table.clear(); + table.reserve_exact(self.m * self.ksub); + for sub in 0..self.m { + let row = (sub * self.ksub + self.code(query_node, sub)) * self.ksub; + table.extend_from_slice(&self.centroid_distances[row..row + self.ksub]); + } + } + + // `new` writes both halves of each subspace table from one value, so the + // copied rows hold the bits of the columns `between` reads, and the sum + // keeps its subspace order: the result equals `between(node, query_node)`. + #[inline] + fn to_query(&self, node: usize, table: &[f32]) -> f32 { + let mut distance = 0.0; + for sub in 0..self.m { + distance += table[sub * self.ksub + self.code(node, sub)]; + } + distance + } + + #[cfg(test)] #[inline] fn between(&self, left: usize, right: usize) -> f32 { let mut distance = 0.0; @@ -1712,6 +1744,7 @@ pub(crate) struct GreedySearchScratch { prune_unique: Vec, prune_pool: Vec, prune_selected: Vec, + query_distances: Vec, #[cfg(test)] peak_retained: usize, #[cfg(test)] @@ -1749,6 +1782,7 @@ impl GreedySearchScratch { prune_unique: Vec::with_capacity(search_list_size.saturating_add(max_degree)), prune_pool: Vec::with_capacity(search_list_size.saturating_add(max_degree)), prune_selected: Vec::with_capacity(max_degree), + query_distances: Vec::new(), #[cfg(test)] peak_retained: 0, #[cfg(test)] @@ -1984,11 +2018,16 @@ impl ParallelVamanaBuilder<'_> { scratch: &mut GreedySearchScratch, ) { scratch.begin_search(); + let mut query_distances = std::mem::take(&mut scratch.query_distances); + self.search_distance + .prepare_query(query_node, &mut query_distances); let entry = self.entry_node as usize; scratch.insert_candidate( ScoredNode { id: self.entry_node, - distance: self.search_distance.between(entry, query_node), + distance: self + .search_distance + .to_query(entry, query_node, &query_distances), }, search_list_size, ); @@ -2006,13 +2045,18 @@ impl ParallelVamanaBuilder<'_> { scratch.insert_candidate( ScoredNode { id: neighbor as u32, - distance: self.search_distance.between(neighbor, query_node), + distance: self.search_distance.to_query( + neighbor, + query_node, + &query_distances, + ), }, search_list_size, ); } scratch.neighbor_buffer.clear(); } + scratch.query_distances = query_distances; } fn insert_reverse_edges( @@ -2548,6 +2592,113 @@ mod tests { )); } + #[test] + fn vamana_pq_build_query_rows_match_symmetric_table_bits() { + // `prepare_query` copies the rows picked by the query's codes, which + // hold the columns `between` reads only while the table is symmetric. + let (dimension, m, count) = (10, 3, 24); + for nbits in [4, 8] { + for metric in [MetricType::L2, MetricType::InnerProduct] { + let mut rng = StdRng::seed_from_u64(nbits as u64); + let mut pq = ProductQuantizer::with_nbits_balanced(dimension, m, nbits); + pq.set_centroids( + (0..dimension * pq.ksub()) + .map(|_| rng.gen_range(-2.0f32..2.0)) + .collect(), + ); + let mut codes = (0..count * pq.code_size()) + .map(|_| rng.gen::()) + .collect::>(); + // Pin the first and last centroid of every subspace. + codes[..pq.code_size()].fill(0x00); + codes[pq.code_size()..2 * pq.code_size()].fill(0xff); + let distance = PqBuildDistance::new(&pq, &codes, count, metric).unwrap(); + let ksub = pq.ksub(); + let table = &distance.centroid_distances; + + for sub in 0..m { + for left in 0..ksub { + for right in 0..ksub { + assert_eq!( + table[(sub * ksub + left) * ksub + right].to_bits(), + table[(sub * ksub + right) * ksub + left].to_bits(), + "nbits={nbits} metric={metric:?} sub={sub} codes=({left}, {right})" + ); + } + } + } + // One buffer serves every query, as a worker's scratch does. + let mut rows = Vec::new(); + distance.prepare_query(0, &mut rows); + let buffer = rows.as_ptr(); + for query in 0..count { + distance.prepare_query(query, &mut rows); + assert_eq!(rows.len(), m * ksub); + assert_eq!(rows.as_ptr(), buffer); + for node in 0..count { + assert_eq!( + distance.to_query(node, &rows).to_bits(), + distance.between(node, query).to_bits(), + "nbits={nbits} metric={metric:?} nodes=({node}, {query})" + ); + } + } + } + } + } + + #[test] + fn vamana_pq_greedy_search_scores_candidates_against_each_query() { + let (dimension, m, count, degree) = (8, 4, 96, 8); + let mut rng = StdRng::seed_from_u64(5); + let mut pq = ProductQuantizer::with_nbits_balanced(dimension, m, 8); + pq.set_centroids( + (0..dimension * pq.ksub()) + .map(|_| rng.gen_range(-1.0f32..1.0)) + .collect(), + ); + let codes = (0..count * pq.code_size()) + .map(|_| rng.gen::()) + .collect::>(); + let vectors = vec![0.0; count * dimension]; + let reference = PqBuildDistance::new(&pq, &codes, count, MetricType::L2).unwrap(); + let adjacency = ParallelAdjacency::new(count, degree, |node| { + (1..=degree) + .map(|step| ((node + step * 7) % count) as u32) + .collect::>() + }); + let builder = ParallelVamanaBuilder { + vectors: &vectors, + dimension, + metric: MetricType::L2, + entry_node: 0, + adjacency, + search_distance: BuildSearchDistance::ProductQuantized( + PqBuildDistance::new(&pq, &codes, count, MetricType::L2).unwrap(), + ), + }; + let mut scratch = GreedySearchScratch::new(count, degree, 16); + builder.greedy_search(5, 16, &mut scratch); + let rows = scratch.query_distances.as_ptr(); + + // Searches reuse one scratch, as in `run_pass`, so a stale query table + // would score later candidates against an earlier node. + for query in [5, 41, 77, 5] { + builder.greedy_search(query, 16, &mut scratch); + assert_eq!(scratch.query_distances.len(), m * pq.ksub()); + assert_eq!(scratch.query_distances.as_ptr(), rows); + assert!(!scratch.results.is_empty()); + for candidate in &scratch.results { + assert_eq!( + candidate.distance.to_bits(), + reference.between(candidate.id as usize, query).to_bits(), + "query={query} candidate={}", + candidate.id + ); + } + } + } + #[test] fn vamana_pq_guided_build_keeps_exact_pruned_graph_bounded_and_reachable() { let dimension = 8; diff --git a/docs/diskann.html b/docs/diskann.html index cf7ef12..b855e90 100644 --- a/docs/diskann.html +++ b/docs/diskann.html @@ -240,7 +240,7 @@

Capacity model

Cold-adjacency LRUMemory/local latency tiers cap it at 64 MiB. Remote/object-store latency tiers may use all remaining budget needed by the cold adjacency section after the raw-vector split; all cases are clipped to section size and total budget.Shared with retained batch workers; split across 16 shards for production-sized budgets. Adaptive adjacency payloadAt most 4 × total_edges, usually less after BFS-remapped delta-varint encoding, plus page-tail paddingActual graph degrees, ID locality, and diskann.max-degree. PQ codebook4 × 2^pq.bits × dimensionDimension and pq.bits; independent of pq.m. - Build peakRaw vectors + IDs + row-ID encoding scratch + codes + codebook + row-ID order + adjacency locators + max(graph-build, in-place graph-remap). PQ training keeps a deterministic reservoir of at most 50,000 rows, reduced automatically so the retained sample, optional cosine-normalized copy, codebook, and parallel KMeans scratch fit the build budget. The low-level Rust DiskAnnIndex::train returns an error when even the minimum one-worker plan cannot fit. If the parallel graph estimate exceeds the budget, overlapping coarse shards are built and merged one at a time; the final compact adjacency remains approximately 4 × N × R + 2 × N.N, dimension, pq.m, pq.bits, R, Lbuild, memory budget, and Rayon worker count. + Build peakRaw vectors + IDs + row-ID encoding scratch + codes + codebook + PQ build-distance tables (4 × pq.m × K × (K + workers) with K = 2^pq.bits, PQ-guided construction only) + row-ID order + adjacency locators + max(graph-build, in-place graph-remap). PQ training keeps a deterministic reservoir of at most 50,000 rows, reduced automatically so the retained sample, optional cosine-normalized copy, codebook, and parallel KMeans scratch fit the build budget. The low-level Rust DiskAnnIndex::train returns an error when even the minimum one-worker plan cannot fit. If the parallel graph estimate exceeds the budget, overlapping coarse shards are built and merged one at a time; the final compact adjacency remains approximately 4 × N × R + 2 × N.N, dimension, pq.m, pq.bits, R, Lbuild, memory budget, and Rayon worker count. Query-local adjacencyAt most 8 MiB of retained cold-window allocation capacity/references after each beam round, plus the active beamInternal safety bound; stats report peak bytes and evictions. Raw-vector cache and query buffersOne shared immutable raw-vector LRU, plus up to 8 MiB of reusable/local-reference capacity per active worker under a 64 MiB aggregate worker allowance. Memory/local latency tiers cap the shared LRU at 64 MiB; remote/object-store latency tiers may use the remaining Reader budget up to the vector-section size.The shared LRU is clipped to the total Reader budget. Session scratch and borrowed references persist for reuse across batch calls without duplicating payload bytes. Graph candidate heapsL retained candidates plus at most 2 × L live frontier entries; contiguous capacity is reused between queriesl_search and graph worker count; evicted lazy entries are periodically compacted.