diskann_disk/search/pq/
pq_scratch.rs

1/*
2 * Copyright (c) Microsoft Corporation.
3 * Licensed under the MIT license.
4 */
5//! Aligned allocator
6
7use diskann::{ANNError, ANNResult};
8
9use diskann_quantization::alloc::{AlignedAllocator, Poly};
10
11#[derive(Debug)]
12/// PQ scratch
13pub struct PQScratch {
14    /// Aligned pq table distance scratch, the length must be at least [256 * NCHUNKS]. 256 is the number of PQ centroids.
15    /// This is used to store the distance between each chunk in the query vector to each centroid, which is why the length is num of centroids * num of chunks
16    pub aligned_pqtable_dist_scratch: Poly<[f32], AlignedAllocator>,
17
18    /// Aligned dist scratch, must be at least diskann MAX_DEGREE
19    /// This is used to temporarily save the pq distance between query vector to the candidate vectors.
20    pub aligned_dist_scratch: Poly<[f32], AlignedAllocator>,
21
22    /// Aligned pq coord scratch, must be at least [N_CHUNKS * MAX_DEGREE]
23    /// This is used to store the pq coordinates of the candidate vectors.
24    pub aligned_pq_coord_scratch: Poly<[u8], AlignedAllocator>,
25
26    /// Query scratch buffer stored as `f32`, sized by the PQ table's logical dimension.
27    /// `set` populates it from a caller-provided `&[f32]`; `PQTable::preprocess_query` can
28    /// then rotate or otherwise preprocess it.
29    pub query_scratch: Vec<f32>,
30}
31
32impl PQScratch {
33    /// Create a new pq scratch.
34    ///
35    /// `dim` is the PQ table's logical dimension (`PQData::get_dim()`); the
36    /// internal `query_scratch` buffer is sized to exactly this many `f32` slots.
37    pub fn new(
38        graph_degree: usize,
39        dim: usize,
40        num_pq_chunks: usize,
41        num_centers: usize,
42    ) -> ANNResult<Self> {
43        let aligned_pq_coord_scratch =
44            Poly::broadcast(0u8, graph_degree * num_pq_chunks, AlignedAllocator::A128)
45                .map_err(ANNError::log_index_error)?;
46        let aligned_pqtable_dist_scratch =
47            Poly::broadcast(0f32, num_centers * num_pq_chunks, AlignedAllocator::A128)
48                .map_err(ANNError::log_index_error)?;
49        let aligned_dist_scratch = Poly::broadcast(0f32, graph_degree, AlignedAllocator::A128)
50            .map_err(ANNError::log_index_error)?;
51        let query_scratch = vec![0.0f32; dim];
52
53        Ok(Self {
54            aligned_pqtable_dist_scratch,
55            aligned_dist_scratch,
56            aligned_pq_coord_scratch,
57            query_scratch,
58        })
59    }
60
61    /// Copy `query` into `query_scratch`.
62    ///
63    /// `query` must already be in full-precision `f32` representation; quantized
64    /// inputs (e.g. `MinMaxElement`) should be decoded via `VectorRepr::as_f32`
65    /// at the caller boundary before invoking this method.
66    ///
67    /// Returns `DimensionMismatchError` if `query.len() != query_scratch.len()`.
68    pub fn set(&mut self, query: &[f32]) -> ANNResult<()> {
69        let dim = self.query_scratch.len();
70        if query.len() != dim {
71            return Err(ANNError::log_dimension_mismatch_error(format!(
72                "PQScratch::set: expected query of length {dim}, got {}",
73                query.len()
74            )));
75        }
76        self.query_scratch.copy_from_slice(query);
77        Ok(())
78    }
79
80    /// Return the largest number of PQ vectors whose distances can be computed using this
81    /// scratch data structure.
82    pub(crate) fn max_vectors(&self) -> usize {
83        self.aligned_dist_scratch.len()
84    }
85}
86
87#[cfg(test)]
88mod tests {
89    use diskann_quantization::num::PowerOfTwo;
90    use rstest::rstest;
91
92    use super::PQScratch;
93
94    #[rstest]
95    #[case(512, 8, 128, 256)] // default test case
96    #[case(59, 16, 37, 41)] // not multiple of 256
97    fn test_pq_scratch(
98        #[case] graph_degree: usize,
99        #[case] dim: usize,
100        #[case] num_pq_chunks: usize,
101        #[case] num_centers: usize,
102    ) {
103        let mut pq_scratch: PQScratch =
104            PQScratch::new(graph_degree, dim, num_pq_chunks, num_centers).unwrap();
105
106        assert_eq!(
107            (pq_scratch.aligned_pqtable_dist_scratch.as_ptr() as usize) % PowerOfTwo::V128.raw(),
108            0
109        );
110        assert_eq!(
111            (pq_scratch.aligned_dist_scratch.as_ptr() as usize) % PowerOfTwo::V128.raw(),
112            0
113        );
114        assert_eq!(
115            (pq_scratch.aligned_pq_coord_scratch.as_ptr() as usize) % PowerOfTwo::V128.raw(),
116            0
117        );
118
119        assert_eq!(pq_scratch.max_vectors(), graph_degree);
120
121        // Test set() method
122        let query: Vec<f32> = (1..=dim).map(|i| i as f32).collect();
123        pq_scratch.set(&query).unwrap();
124
125        (0..query.len()).for_each(|i| {
126            assert_eq!(pq_scratch.query_scratch[i], query[i]);
127        });
128    }
129
130    #[test]
131    fn test_pq_scratch_set_rejects_short_query() {
132        let dim = 16;
133        let mut pq_scratch = PQScratch::new(64, dim, 4, 256).unwrap();
134
135        // Query shorter than dim should fail
136        let short_query: Vec<f32> = (1..dim).map(|i| i as f32).collect(); // dim-1 elements
137        let err = pq_scratch.set(&short_query).unwrap_err();
138        assert_eq!(err.kind(), diskann::ANNErrorKind::DimensionMismatchError);
139        assert!(err.to_string().contains("expected query of length"));
140    }
141
142    #[test]
143    fn test_pq_scratch_set_rejects_oversized_query() {
144        let dim = 8;
145        let mut pq_scratch = PQScratch::new(64, dim, 4, 256).unwrap();
146
147        // Query longer than dim should fail
148        let long_query: Vec<f32> = (1..=dim + 10).map(|i| i as f32).collect();
149        let err = pq_scratch.set(&long_query).unwrap_err();
150        assert_eq!(err.kind(), diskann::ANNErrorKind::DimensionMismatchError);
151        assert!(err.to_string().contains("expected query of length"));
152    }
153}
diskann_disk/search/pq/pq_scratch.rs

diskann_disk/search/pq/
pq_scratch.rs
