diskann_bftree/
quant.rs

1/*
2 * Copyright (c) Microsoft Corporation.
3 * Licensed under the MIT license.
4 */
5
6//! Bf-Tree quant vector provider.
7
8use crate::{AccessError, VectorError, VectorUnavailable};
9use bf_tree::{BfTree, Config};
10use diskann::{error::IntoANNResult, utils::VectorRepr, ANNError, ANNResult};
11use diskann_quantization::{
12    alloc::{GlobalAllocator, Poly, ScopedAllocator},
13    spherical::iface::{
14        DistanceComputer, Opaque, OpaqueMut, Quantizer, QueryComputer, QueryLayout,
15    },
16};
17use diskann_vector::PreprocessedDistanceFunction;
18
19use super::ConfigError;
20use crate::{bftree_insert, TestCallCount};
21
22pub struct QuantQueryComputer(QueryComputer<GlobalAllocator>);
23
24impl QuantQueryComputer {
25    pub(crate) fn evaluate(&self, x: &[u8]) -> ANNResult<f32> {
26        match self.0.evaluate_similarity(Opaque::new(x)) {
27            Ok(distance) => Ok(distance),
28            Err(err) => Err(ANNError::new(diskann::ANNErrorKind::IndexError, err)),
29        }
30    }
31}
32
33pub struct QuantVectorProvider {
34    quant_vector_index: BfTree,
35    pub(crate) quantizer: Poly<dyn Quantizer>,
36    pub(super) num_get_calls: TestCallCount,
37}
38
39impl QuantVectorProvider {
40    pub fn new_with_config(quantizer: Poly<dyn Quantizer>, config: Config) -> ANNResult<Self> {
41        crate::validate_record_size(
42            "quant_vector_provider",
43            &config,
44            std::mem::size_of::<usize>(),
45            quantizer.bytes(),
46        )?;
47
48        let quant_vector_index = BfTree::with_config(config, None).map_err(ConfigError)?;
49
50        Ok(Self {
51            quant_vector_index,
52            quantizer,
53            num_get_calls: TestCallCount::default(),
54        })
55    }
56
57    /// Access the BfTree config
58    pub(crate) fn config(&self) -> &Config {
59        self.quant_vector_index.config()
60    }
61
62    /// Access the underlying BfTree
63    pub(crate) fn bftree(&self) -> &BfTree {
64        &self.quant_vector_index
65    }
66
67    /// Create a new instance from an existing BfTree (for loading from snapshot)
68    ///
69    pub(crate) fn new_from_bftree(
70        quantizer: Poly<dyn Quantizer>,
71        quant_vector_index: BfTree,
72    ) -> Self {
73        Self {
74            quant_vector_index,
75            quantizer,
76            num_get_calls: TestCallCount::default(),
77        }
78    }
79
80    pub(crate) fn delete_vector(&self, i: usize) {
81        let key = bytemuck::bytes_of(&i);
82        self.quant_vector_index.delete(key);
83    }
84
85    /// Return the dimension of the full-precision data associated with this provider
86    pub fn full_dim(&self) -> usize {
87        self.quantizer.full_dim()
88    }
89
90    /// Create a query computer for the provided query vector
91    pub fn query_computer<T>(&self, query: &[T]) -> ANNResult<QuantQueryComputer>
92    where
93        T: VectorRepr,
94    {
95        let query_f32 = T::as_f32(query).into_ann_result()?;
96        let inner = self
97            .quantizer
98            .fused_query_computer(
99                &query_f32,
100                QueryLayout::FullPrecision,
101                true,
102                GlobalAllocator,
103                ScopedAllocator::global(),
104            )
105            .map_err(|e| ANNError::log_sq_error(e))?;
106        Ok(QuantQueryComputer(inner))
107    }
108
109    /// Create a distance computer for the underlying schema
110    pub fn distance_computer(&self) -> ANNResult<DistanceComputer> {
111        self.quantizer
112            .distance_computer(GlobalAllocator)
113            .map_err(|e| ANNError::log_sq_error(e))
114    }
115
116    pub(crate) fn get_vector_into(&self, i: usize, buffer: &mut [u8]) -> Result<(), AccessError> {
117        use diskann::ANNErrorKind;
118        use thiserror::Error;
119
120        let expected = self.quantizer.bytes();
121        if buffer.len() != expected {
122            #[derive(Debug, Error)]
123            #[error("expected a buffer with dim {0}, instead got {1}")]
124            struct WrongDim(usize, usize);
125
126            return Err(AccessError::Error(ANNError::new(
127                ANNErrorKind::IndexError,
128                WrongDim(expected, buffer.len()),
129            )));
130        }
131
132        self.num_get_calls.increment();
133        match self.quant_vector_index.read(bytemuck::bytes_of(&i), buffer) {
134            bf_tree::LeafReadResult::Found(read_size) => {
135                if read_size as usize != expected {
136                    return Err(AccessError::Error(ANNError::log_index_error(format!(
137                        "The bf-tree entry for vector id {} is marked as found but has size {} instead of the expected size {}",
138                        i, read_size, expected,
139                    ))));
140                }
141            }
142            bf_tree::LeafReadResult::Deleted => {
143                return Err(AccessError::Transient(VectorUnavailable {
144                    id: i,
145                    err: VectorError::Deleted,
146                }));
147            }
148            bf_tree::LeafReadResult::InvalidKey => {
149                return Err(AccessError::Error(ANNError::log_index_error(format!(
150                    "The bf-tree entry for vector id {} is marked as invalid",
151                    i,
152                ))));
153            }
154            bf_tree::LeafReadResult::NotFound => {
155                return Err(AccessError::Transient(VectorUnavailable {
156                    id: i,
157                    err: VectorError::NotFound,
158                }));
159            }
160        };
161
162        Ok(())
163    }
164
165    /// Return the quant vector at index `i`
166    #[cfg(test)]
167    pub(crate) fn get_vector_sync(&self, i: usize) -> Result<Vec<u8>, AccessError> {
168        let mut value = vec![0u8; self.quantizer.bytes()];
169        self.get_vector_into(i, &mut value)?;
170        Ok(value)
171    }
172
173    /// Compress the vector, `v`, and set the compressed quant vector with Id, `i`, to it
174    ///
175    /// Errors if:
176    ///
177    /// * `v.dim() != self.full_dim()`: The slice must have the proper length.
178    /// * PQ compression encounters an error (such as the presence of `NaN`s).
179    pub(crate) fn set_vector_sync<T>(&self, i: usize, v: &[T]) -> ANNResult<()>
180    where
181        T: Copy + VectorRepr,
182    {
183        let vf32: &[f32] = &T::as_f32(v).into_ann_result()?;
184
185        if vf32.len() != self.full_dim() {
186            return Err(ANNError::log_dimension_mismatch_error(
187                "Vector f32 dimension is not equal to the expected dimension.".to_string(),
188            ));
189        }
190
191        // Serialize the key into a byte string, &[u8]
192        let key = bytemuck::bytes_of(&i);
193
194        let dim = self.quantizer.bytes();
195        let quant_vector = &mut vec![0u8; dim];
196        self.quantizer
197            .compress(
198                vf32,
199                OpaqueMut::new(quant_vector),
200                ScopedAllocator::global(),
201            )
202            .map_err(|e| ANNError::log_sq_error(e))?;
203
204        bftree_insert(&self.quant_vector_index, key, quant_vector)?;
205
206        Ok(())
207    }
208
209    /// Set the quant vector with Id, `i`, to `v`
210    ///
211    /// Errors if:
212    ///
213    /// * `v.len() != self.pq_chunks()`: `v` must have the right length.
214    #[cfg(test)]
215    pub(crate) fn set_quant_vector(&self, i: usize, v: &[u8]) -> ANNResult<()> {
216        if v.len() != self.quantizer.bytes() {
217            return Err(ANNError::log_index_error(
218                "Vector dimension is not equal to the expected dimension.",
219            ));
220        }
221
222        // Update pq vector with id = i to v
223        let key = bytemuck::bytes_of(&i);
224
225        bftree_insert(&self.quant_vector_index, key, v)?;
226
227        Ok(())
228    }
229}
230
231/// Train a spherical quantizer on simple data and return it as a `Poly<dyn Quantizer>`.
232#[cfg(test)]
233pub(crate) fn create_test_quantizer(dim: usize) -> Poly<dyn Quantizer> {
234    use diskann_quantization::{
235        algorithms::TransformKind,
236        alloc::poly,
237        spherical::{iface, PreScale, SphericalQuantizer, SupportedMetric},
238    };
239    use diskann_utils::views::Init;
240    use diskann_utils::views::Matrix;
241    use rand::{rngs::StdRng, SeedableRng};
242
243    // Create training data with spread-out values.
244    let nrows = 8;
245    let mut counter = 0.0f32;
246    let data = Matrix::new(
247        Init(move || {
248            counter += 0.5;
249            counter
250        }),
251        nrows,
252        dim,
253    );
254
255    let mut rng = StdRng::seed_from_u64(42);
256    let quantizer = SphericalQuantizer::train(
257        data.as_view(),
258        TransformKind::Null,
259        SupportedMetric::SquaredL2,
260        PreScale::None,
261        &mut rng,
262        GlobalAllocator,
263    )
264    .unwrap();
265
266    let imp = iface::Impl::<1>::new(quantizer).unwrap();
267    poly!(Quantizer, imp, GlobalAllocator).unwrap()
268}
269
270///////////
271// Tests //
272///////////
273/// These unit tests target the functionality of Bf-Tree quant vector provider alone
274#[cfg(test)]
275mod tests {
276    use std::sync::Arc;
277
278    use diskann::ANNErrorKind;
279    use diskann_quantization::spherical::iface::Opaque;
280    use diskann_vector::DistanceFunction;
281    use tokio::task::JoinSet;
282
283    use super::*;
284
285    /// Test edge cases of the Bf-Tree quant vector provider
286    #[tokio::test]
287    async fn common_errors() {
288        let dim = 5;
289        let quantizer = create_test_quantizer(dim);
290        let quant_bytes = quantizer.bytes();
291
292        let bf_tree_config = Config::default();
293        let provider = QuantVectorProvider::new_with_config(quantizer, bf_tree_config).unwrap();
294
295        // try to set an out of bounds vector
296        let result = provider.set_quant_vector(20, &[]).unwrap_err();
297        assert_eq!(result.kind(), ANNErrorKind::IndexError);
298
299        // try to set an out of bounds vector via set_vector_sync
300        let result = provider.set_vector_sync::<f32>(20, &[]).unwrap_err();
301        assert_eq!(result.kind(), ANNErrorKind::DimensionMismatchError);
302
303        // try to set a quant vector with the wrong dimension
304        let result = provider.set_quant_vector(0, &[]).unwrap_err();
305        assert_eq!(result.kind(), ANNErrorKind::IndexError);
306
307        // verify expected quant vector byte count
308        assert_eq!(quant_bytes, provider.quantizer.bytes());
309    }
310
311    fn create_test_provider() -> QuantVectorProvider {
312        let dim = 2;
313
314        let quantizer = create_test_quantizer(dim);
315
316        let bf_tree_config = Config::default();
317        let provider = QuantVectorProvider::new_with_config(quantizer, bf_tree_config).unwrap();
318
319        assert_eq!(provider.full_dim(), dim);
320
321        // Set vectors.
322        provider.set_vector_sync(0, &[-1.5, -1.5]).unwrap();
323        provider.set_vector_sync(1, &[-0.5, -0.5]).unwrap();
324        provider.set_vector_sync(2, &[0.5, 0.5]).unwrap();
325        provider.set_vector_sync(3, &[1.5, 1.5]).unwrap();
326        provider.set_vector_sync(4, &[2.5, 2.5]).unwrap();
327        provider
328    }
329
330    /// Test the distance computation functions of the provider
331    #[tokio::test]
332    async fn test_similarity_function() {
333        let provider = create_test_provider();
334        let quant_bytes = provider.quantizer.bytes();
335
336        // Verify compressed vectors are the expected size.
337        for i in 0..5 {
338            let v = provider.get_vector_sync(i).unwrap();
339            assert_eq!(v.len(), quant_bytes);
340        }
341
342        // Error checking.
343        assert!(provider.set_vector_sync(2, &[0.0]).is_err());
344
345        // Query Computer — verify it returns finite distances.
346        let c = provider.query_computer(&[-0.5f32, -0.5]).unwrap();
347        let dist = c.evaluate(&provider.get_vector_sync(3).unwrap()).unwrap();
348        assert!(dist.is_finite(), "query distance should be finite");
349
350        // Distance Computer — verify distances between compressed vectors are finite
351        // and that identical vectors produce zero distance.
352        let d = provider.distance_computer().unwrap();
353        let v0 = provider.get_vector_sync(0).unwrap();
354        let v3 = provider.get_vector_sync(3).unwrap();
355        let dist = d
356            .evaluate_similarity(Opaque::new(&v0), Opaque::new(&v3))
357            .unwrap();
358        assert!(dist.is_finite(), "distance should be finite");
359
360        // Same vector should have small self-distance (may not be exactly zero
361        // due to quantization loss, especially at low bit-widths).
362        let self_dist = d
363            .evaluate_similarity(Opaque::new(&v0), Opaque::new(&v0))
364            .unwrap();
365        assert!(
366            self_dist.abs() < 1.0,
367            "self-distance should be small, got {}",
368            self_dist
369        );
370    }
371
372    /// Test the interleaved and parallel traversal of the Bf-Tree
373    /// by invoking the async accessors of the quant vector provider
374    #[tokio::test(flavor = "multi_thread", worker_threads = 5)]
375    async fn test_parallel_tree_traversal() {
376        let dim = 2;
377        let quantizer = create_test_quantizer(dim);
378
379        let bf_tree_config = Config::default();
380        let provider =
381            Arc::new(QuantVectorProvider::new_with_config(quantizer, bf_tree_config).unwrap());
382        let mut set = JoinSet::new();
383        for i in 0..11 {
384            let vector = vec![i as f32, (i + 1) as f32];
385            let provider_clone = Arc::clone(&provider);
386            set.spawn(async move { provider_clone.set_vector_sync(i as usize, &vector).unwrap() });
387        }
388
389        while let Some(res) = set.join_next().await {
390            res.unwrap();
391        }
392
393        // Verify that each vector was stored and can be retrieved with the correct size.
394        let quant_bytes = provider.quantizer.bytes();
395        let mut expected_buf = vec![0u8; quant_bytes];
396
397        for i in 0..11 {
398            let stored = provider.get_vector_sync(i).unwrap();
399            assert_eq!(stored.len(), quant_bytes);
400
401            // Compress the same input again and verify we get the same output
402            // (spherical compression is deterministic).
403            provider
404                .quantizer
405                .compress(
406                    &[i as f32, (i + 1) as f32],
407                    OpaqueMut::new(&mut expected_buf),
408                    ScopedAllocator::global(),
409                )
410                .unwrap();
411            assert_eq!(stored, expected_buf);
412        }
413    }
414}
diskann_bftree/quant.rs

diskann_bftree/
quant.rs
