perf(vamana): look up PQ build distances in a per-search table

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::<f32>()))
                 .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::<f32>()
+                );
+            });
+    }
+
     #[test]
     fn diskann_memory_budget_automatically_selects_overlapping_shards() {
         // Pin the worker count so this test exercises the same memory-plan

core/src/vamana.rs

@@ -1244,15 +1244,21 @@ enum BuildSearchDistance<'a> {
 }
 
 impl BuildSearchDistance<'_> {
+    fn prepare_query(&self, query_node: usize, table: &mut Vec<f32>) {
+        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<f32>) {
+        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<u32>,
     prune_pool: Vec<ScoredNode>,
     prune_selected: Vec<u32>,
+    query_distances: Vec<f32>,
     #[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::<u8>())
+                    .collect::<Vec<_>>();
+                // 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::<u8>())
+            .collect::<Vec<_>>();
+        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::<Vec<_>>()
+        });
+        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;

docs/diskann.html

@@ -240,7 +240,7 @@ <h2>Capacity model</h2>
             <tr><td>Cold-adjacency LRU</td><td>Memory/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.</td><td>Shared with retained batch workers; split across 16 shards for production-sized budgets.</td></tr>
             <tr><td>Adaptive adjacency payload</td><td>At most <code>4 × total_edges</code>, usually less after BFS-remapped delta-varint encoding, plus page-tail padding</td><td>Actual graph degrees, ID locality, and <code>diskann.max-degree</code>.</td></tr>
             <tr><td>PQ codebook</td><td><code>4 × 2^pq.bits × dimension</code></td><td>Dimension and <code>pq.bits</code>; independent of <code>pq.m</code>.</td></tr>
-            <tr><td>Build peak</td><td>Raw 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 <code>DiskAnnIndex::train</code> 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 <code>4 × N × R + 2 × N</code>.</td><td><code>N</code>, <code>dimension</code>, <code>pq.m</code>, <code>pq.bits</code>, <code>R</code>, <code>Lbuild</code>, memory budget, and Rayon worker count.</td></tr>
+            <tr><td>Build peak</td><td>Raw vectors + IDs + row-ID encoding scratch + codes + codebook + PQ build-distance tables (<code>4 × pq.m × K × (K + workers)</code> with <code>K = 2^pq.bits</code>, 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 <code>DiskAnnIndex::train</code> 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 <code>4 × N × R + 2 × N</code>.</td><td><code>N</code>, <code>dimension</code>, <code>pq.m</code>, <code>pq.bits</code>, <code>R</code>, <code>Lbuild</code>, memory budget, and Rayon worker count.</td></tr>
             <tr><td>Query-local adjacency</td><td>At most 8 MiB of retained cold-window allocation capacity/references after each beam round, plus the active beam</td><td>Internal safety bound; stats report peak bytes and evictions.</td></tr>
             <tr><td>Raw-vector cache and query buffers</td><td>One 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.</td><td>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.</td></tr>
             <tr><td>Graph candidate heaps</td><td><code>L</code> retained candidates plus at most <code>2 × L</code> live frontier entries; contiguous capacity is reused between queries</td><td><code>l_search</code> and graph worker count; evicted lazy entries are periodically compacted.</td></tr>