diskann_quantization/minmax/
quantizer.rs

1/*
2 * Copyright (c) Microsoft Corporation.
3 * Licensed under the MIT license.
4 */
5
6#[cfg(feature = "flatbuffers")]
7use thiserror::Error;
8
9use super::vectors::{DataMutRef, FullQueryMut, MinMaxCompensation, MinMaxIP, MinMaxL2Squared};
10
11use crate::{
12    AsFunctor, CompressInto,
13    algorithms::Transform,
14    alloc::{GlobalAllocator, ScopedAllocator},
15    bits::{Representation, Unsigned},
16    minmax::{MinMaxCosine, MinMaxCosineNormalized, vectors::FullQueryMeta},
17    num::Positive,
18    scalar::{InputContainsNaN, bit_scale},
19};
20#[cfg(feature = "flatbuffers")]
21use crate::{
22    algorithms::transforms::TransformError,
23    alloc::{Allocator, AllocatorError, Poly},
24    flatbuffers as fb,
25};
26
27/// The base number of bytes to allocate when attempting to serialize a quantizer.
28#[cfg(all(not(test), feature = "flatbuffers"))]
29const DEFAULT_SERIALIZED_BYTES: usize = 1024;
30
31// When testing, use a small value so we trigger the reallocation logic.
32#[cfg(all(test, feature = "flatbuffers"))]
33const DEFAULT_SERIALIZED_BYTES: usize = 1;
34
35/// Recall that from the module-level documentation, MinMaxQuantizer, quantizes X
36/// into `n` bit vectors as follows  -
37/// ```math
38/// X' = round((X - s) * (2^n - 1) / c).clamp(0, 2^n - 1))
39/// ```
40/// where `s` is a shift value and `c` is a scaling parameter computed from the range of values.
41///
42/// For most bit widths (>1), given a positive scaling parameter `grid_scale : f32`,
43/// these are computed as:
44/// ```math
45/// - m = (max_i X[i] + min_i X[i]) / 2.0
46/// - w = max_i X[i] - min_i X[i]
47///
48/// - s = m - w * grid_scale
49/// - c = 2 * w * grid_scale
50///
51/// where `grid_scale` is an input to the quantizer.
52/// ```
53/// For 1-bit quantization, to avoid outliers, `s` and `c` are derived differently:
54/// - Values are first split into two groups: those below and above the mean.
55/// - `s` is the average of values below the mean.
56/// - `c` is the difference between the average of values above the mean and `s`.
57///
58/// See [`MinMaxCompensation`] for notation.
59/// We have then that
60/// ```math
61/// X = X' * (c / (2^n - 1)) + s
62///          --------------    -
63///                 |          |
64///                ax          bx
65/// ```
66#[derive(Debug)]
67#[cfg_attr(test, derive(PartialEq))]
68pub struct MinMaxQuantizer {
69    /// Support for different strategies of pre-transforming vectors before applying compression.
70    /// See [`Transform`] for more details on supported types. The input dimension of vectors
71    /// to the quantizer is derived from `transform.input_dim()`.
72    transform: Transform<GlobalAllocator>,
73
74    /// Scaling parameter used to scale the range (min, max) in order to avoid outliers.
75    /// The input must be a positive value. In general, any value between [0.8, 1] does well.
76    grid_scale: Positive<f32>,
77}
78
79impl MinMaxQuantizer {
80    /// Instantiates a new quantizer with specific transform.
81    pub fn new(transform: Transform<GlobalAllocator>, grid_scale: Positive<f32>) -> Self {
82        Self {
83            transform,
84            grid_scale,
85        }
86    }
87
88    /// Input dimension of vectors to quantizer.
89    pub fn dim(&self) -> usize {
90        self.transform.input_dim()
91    }
92
93    /// Output dimension of vectors after applying transform.
94    ///
95    /// Output storage vectors should use this dimension instead of `self.dim()` because
96    /// in general, the output dim **may** be different from the input dimension.
97    pub fn output_dim(&self) -> usize {
98        self.transform.output_dim()
99    }
100
101    /// Outputs the minimum and maximum value of the range of values
102    /// for an input vector `vec`. The function cases based on the
103    /// intended number of bits `NBITS` per dimension.
104    ///
105    /// * `1-bit` - In order to avoid outlier values, the range
106    ///   is defined by taking the values larger and smaller than
107    ///   the numeric mean, and then taking the respective means of
108    ///   each of these sets as the `max` and `min`.
109    ///
110    /// * `N-bits` - Computes the `min` and `max` of the vector values.
111    ///
112    /// # Returns
113    ///
114    /// * `(m - w * g, m + w * g)` - the lower and upper end of the range, where,
115    ///   `m = (max + min) / 2.0`, `w = (max - min) / 2.0`, and `g = self.grid_scale`.
116    fn get_range<const NBITS: usize>(&self, vec: &[f32]) -> (f32, f32) {
117        let (min, max) = match NBITS {
118            1 => {
119                let (mut min, mut min_count) = (0.0f32, 0.0f32);
120                let (mut max, mut max_count) = (0.0f32, 0.0f32);
121
122                let mean = vec.iter().sum::<f32>() / (vec.len() as f32);
123
124                vec.iter().for_each(|x| {
125                    let m = f32::from((*x < mean) as u8);
126                    min += m * x;
127                    min_count += m;
128                    max += (1.0 - m) * x;
129                    max_count += 1.0 - m;
130                });
131
132                ((min / min_count).min(mean), (max / max_count).max(mean))
133            }
134            _ => {
135                vec // Using `f32::NAN` since [`core::f32::min`] and `max` output the other value if one of them is NAN .
136                    .iter()
137                    .fold((f32::NAN, f32::NAN), |(cmin, cmax), &e| {
138                        (cmin.min(e), cmax.max(e))
139                    })
140            }
141        };
142
143        let width = (max - min) / 2.0;
144        let mid = min + width;
145
146        (
147            mid - width * self.grid_scale.into_inner(),
148            mid + width * self.grid_scale.into_inner(),
149        )
150    }
151
152    fn compress<const NBITS: usize, T>(
153        &self,
154        from: &[T],
155        mut into: DataMutRef<'_, NBITS>,
156    ) -> Result<L2Loss, InputContainsNaN>
157    where
158        T: Copy + Into<f32>,
159        Unsigned: Representation<NBITS>,
160    {
161        let mut into_vec = into.vector_mut();
162
163        assert_eq!(from.len(), self.dim());
164        assert_eq!(self.output_dim(), into_vec.len());
165
166        let domain = Unsigned::domain_const::<NBITS>();
167        let domain_min = *domain.start() as f32;
168        let domain_max = *domain.end() as f32;
169
170        let mut vec = vec![f32::default(); self.output_dim()];
171
172        // We know vec.len() == self.output_dim() and `from.len() == self.dim`
173        #[expect(clippy::unwrap_used)]
174        self.transform
175            .transform_into(
176                &mut vec,
177                &from.iter().map(|&x| x.into()).collect::<Vec<f32>>(),
178                ScopedAllocator::global(),
179            )
180            .unwrap();
181
182        let (min, max) = self.get_range::<NBITS>(&vec);
183
184        let inverse_scale = (max - min).max(1e-8) / bit_scale::<NBITS>(); // To avoid NaN. This is ONLY possible if the vector is all the same value.
185        let mut norm_squared: f32 = 0.0;
186        let mut code_sum: f32 = 0.0;
187        let mut loss: f32 = 0.0;
188
189        let mut nan_check = false;
190
191        vec.iter().enumerate().for_each(|(i, &v)| {
192            nan_check |= v.is_nan();
193
194            let code = ((v - min) / inverse_scale)
195                .clamp(domain_min, domain_max)
196                .round();
197
198            let v_r = (code * inverse_scale) + min; // reconstructed value for `v`.
199            norm_squared += v_r * v_r;
200            code_sum += code;
201            loss += (v_r - v).powi(2);
202
203            //SAFETY: we checked that the lengths of `from` and `into_vec` are the same.
204            unsafe {
205                into_vec.set_unchecked(i, code as u8);
206            }
207        });
208
209        let meta = MinMaxCompensation {
210            dim: self.output_dim() as u32,
211            b: min,
212            a: inverse_scale,
213            n: inverse_scale * code_sum,
214            norm_squared,
215        };
216
217        into.set_meta(meta);
218
219        if nan_check {
220            Err(InputContainsNaN)
221        } else {
222            Ok(match Positive::new(loss) {
223                Ok(p) => L2Loss::Positive(p),
224                Err(_) => L2Loss::Zero,
225            })
226        }
227    }
228
229    /// Attempt to deserialize a FlatBuffer `fb::minmax::Quantizer` (as produced by [`MinMaxQuantizer::serialize`]) into [`MinMaxQuantizer`].
230    #[cfg(feature = "flatbuffers")]
231    #[cfg_attr(docsrs, doc(cfg(feature = "flatbuffers")))]
232    pub fn try_deserialize(data: &[u8]) -> Result<Self, DeserializationError> {
233        // Check that this is one of the known identifiers.
234        if !fb::minmax::quantizer_buffer_has_identifier(data) {
235            return Err(DeserializationError::InvalidIdentifier);
236        }
237
238        let root = fb::minmax::root_as_quantizer(data)?;
239
240        Ok(Self::new(
241            Transform::try_unpack(GlobalAllocator, root.transform())?,
242            Positive::new(root.grid_scale())
243                .map_err(|_| DeserializationError::GridScaleNotPositive)?,
244        ))
245    }
246
247    #[cfg(feature = "flatbuffers")]
248    pub fn serialize<A>(&self, allocator: A) -> Result<Poly<[u8], A>, AllocatorError>
249    where
250        A: Allocator + std::panic::UnwindSafe,
251    {
252        use flatbuffers::FlatBufferBuilder;
253
254        let mut buf = FlatBufferBuilder::new_in(Poly::broadcast(
255            0u8,
256            DEFAULT_SERIALIZED_BYTES,
257            allocator.clone(),
258        )?);
259
260        let transform = &self.transform;
261
262        let (root, mut buf) = match std::panic::catch_unwind(move || {
263            let offset = transform.pack(&mut buf);
264
265            let root = fb::minmax::Quantizer::create(
266                &mut buf,
267                &fb::minmax::QuantizerArgs {
268                    transform: Some(offset),
269                    grid_scale: self.grid_scale.into_inner(),
270                },
271            );
272            (root, buf)
273        }) {
274            Ok(ret) => ret,
275            Err(err) => match err.downcast_ref::<String>() {
276                Some(msg) => {
277                    if msg.contains("AllocatorError") {
278                        return Err(AllocatorError);
279                    } else {
280                        std::panic::resume_unwind(err);
281                    }
282                }
283                None => std::panic::resume_unwind(err),
284            },
285        };
286
287        // Finish serializing and then copy out the finished data into a newly allocated buffer.
288        fb::minmax::finish_quantizer_buffer(&mut buf, root);
289        Poly::from_iter(buf.finished_data().iter().copied(), allocator)
290    }
291}
292
293////////////////
294// Flatbuffer //
295////////////////
296
297#[cfg(feature = "flatbuffers")]
298#[cfg_attr(docsrs, doc(cfg(feature = "flatbuffers")))]
299#[derive(Debug, Clone, Error)]
300#[non_exhaustive]
301pub enum DeserializationError {
302    #[error("unhandled file identifier in flatbuffer")]
303    InvalidIdentifier,
304    #[error(transparent)]
305    InvalidFlatBuffer(#[from] flatbuffers::InvalidFlatbuffer),
306    #[error(transparent)]
307    AllocatorError(#[from] AllocatorError),
308    #[error(transparent)]
309    TransformError(#[from] TransformError),
310    #[error("grid-scale is missing or is not positive")]
311    GridScaleNotPositive,
312}
313
314/////////////////
315// Compression //
316/////////////////
317
318/// A struct defining euclidean loss from quantization.
319///
320/// For an input vector `x` and its representation `x'`,
321/// this is supposed to store `||x - x'||^2`.
322#[derive(Clone, Copy, Debug)]
323pub enum L2Loss {
324    Zero,
325    Positive(Positive<f32>),
326}
327
328impl L2Loss {
329    /// Euclidean loss as a `f32` value
330    pub fn as_f32(&self) -> f32 {
331        match self {
332            L2Loss::Zero => 0.0,
333            L2Loss::Positive(p) => p.into_inner(),
334        }
335    }
336}
337
338impl<const NBITS: usize, T> CompressInto<&[T], DataMutRef<'_, NBITS>> for MinMaxQuantizer
339where
340    T: Copy + Into<f32>,
341    Unsigned: Representation<NBITS>,
342{
343    type Error = InputContainsNaN;
344
345    type Output = L2Loss;
346
347    /// Compress the input vector `from` into a mut ref of Data `to`.
348    ///
349    /// This method computes and stores the compensation coefficients required for computing
350    /// distances correctly.
351    ///
352    /// # Error
353    ///
354    /// Returns an error if the input contains `NaN`.
355    ///
356    /// # Panics
357    ///
358    /// Panics if:
359    /// * `from.len() != self.dim()`: Vector to be compressed must have the same
360    ///   dimensionality as the quantizer.
361    /// * `to.vector().len() != self.output_dim()`: Compressed vector must have the same dimensionality
362    ///   as the quantizer.
363    fn compress_into(&self, from: &[T], to: DataMutRef<'_, NBITS>) -> Result<L2Loss, Self::Error> {
364        self.compress::<NBITS, T>(from, to)
365    }
366}
367
368impl<'a, T> CompressInto<&[T], FullQueryMut<'a>> for MinMaxQuantizer
369where
370    T: Copy + Into<f32>,
371{
372    type Error = InputContainsNaN;
373
374    type Output = ();
375
376    /// Compress the input vector `from` into a [`FullQueryMut`] `to`.
377    ///
378    /// This method simply applies the transformation to the input without
379    /// any compression.
380    ///
381    /// # Error
382    ///
383    /// Returns an error if the input contains `NaN`.
384    ///
385    /// # Panics
386    ///
387    /// Panics if:
388    /// * `from.len() != self.dim()`: Vector to be compressed must have the same
389    ///   dimensionality as the quantizer.
390    /// * `to.len() != self.output_dim()`: Compressed vector must have the same dimensionality
391    ///   as the quantizer.
392    fn compress_into(&self, from: &[T], mut to: FullQueryMut<'a>) -> Result<(), Self::Error> {
393        assert_eq!(from.len(), self.dim());
394        assert_eq!(self.output_dim(), to.len());
395
396        // Transform the input vector and return error if it contains NaN
397        let from: Vec<f32> = from.iter().map(|&x| x.into()).collect();
398        if from.iter().any(|x| x.is_nan()) {
399            return Err(InputContainsNaN);
400        }
401
402        // We know vec.len() == self.output_dim() and `from.len() == self.dim`
403        #[expect(clippy::unwrap_used)]
404        self.transform
405            .transform_into(to.vector_mut(), &from, ScopedAllocator::global())
406            .unwrap();
407
408        let norm_squared = to.vector().iter().map(|x| *x * *x).sum::<f32>();
409        let sum = to.vector().iter().sum::<f32>();
410
411        *to.meta_mut() = FullQueryMeta { norm_squared, sum };
412
413        Ok(())
414    }
415}
416
417///////////////////////
418// Distance Functors //
419///////////////////////
420
421macro_rules! impl_functor {
422    ($dist:ident) => {
423        impl AsFunctor<$dist> for MinMaxQuantizer {
424            // no need to do any work here.
425            fn as_functor(&self) -> $dist {
426                $dist
427            }
428        }
429    };
430}
431
432impl_functor!(MinMaxIP);
433impl_functor!(MinMaxL2Squared);
434impl_functor!(MinMaxCosine);
435impl_functor!(MinMaxCosineNormalized);
436
437///////////
438// Tests //
439///////////
440#[cfg(test)]
441#[cfg(not(miri))]
442mod minmax_quantizer_tests {
443    use std::num::NonZeroUsize;
444
445    use diskann_utils::{Reborrow, ReborrowMut};
446    use diskann_vector::{PureDistanceFunction, distance::SquaredL2};
447    use rand::{
448        SeedableRng,
449        distr::{Distribution, Uniform},
450        rngs::StdRng,
451    };
452
453    use super::*;
454    use crate::{
455        algorithms::transforms::NullTransform,
456        alloc::GlobalAllocator,
457        minmax::vectors::{Data, DataRef, FullQuery, FullQueryMut},
458    };
459
460    fn reconstruct_minmax<const NBITS: usize>(v: DataRef<'_, NBITS>) -> Vec<f32>
461    where
462        Unsigned: Representation<NBITS>,
463    {
464        (0..v.len())
465            .map(|i| {
466                let m = v.meta();
467                v.vector().get(i).unwrap() as f32 * m.a + m.b
468            })
469            .collect()
470    }
471
472    fn test_quantizer_encoding_random<const NBITS: usize>(
473        dim: usize,
474        rng: &mut StdRng,
475        relative_err: f32,
476        scale: f32,
477    ) where
478        Unsigned: Representation<NBITS>,
479        MinMaxQuantizer: for<'a, 'b> CompressInto<&'a [f32], DataMutRef<'b, NBITS>, Output = L2Loss>
480            + for<'a, 'b> CompressInto<&'a [f32], FullQueryMut<'b>, Output = ()>,
481    {
482        let distribution = Uniform::new_inclusive::<f32, f32>(-1.0, 1.0).unwrap();
483
484        let quantizer = MinMaxQuantizer::new(
485            Transform::Null(NullTransform::new(NonZeroUsize::new(dim).unwrap())),
486            Positive::new(scale).unwrap(),
487        );
488
489        assert_eq!(quantizer.dim(), dim);
490
491        let vector: Vec<f32> = distribution.sample_iter(rng).take(dim).collect();
492
493        let mut encoded = Data::new_boxed(dim);
494        let loss = quantizer
495            .compress_into(&*vector, encoded.reborrow_mut())
496            .unwrap();
497
498        let reconstructed = reconstruct_minmax::<NBITS>(encoded.reborrow());
499        assert_eq!(reconstructed.len(), dim);
500
501        let reconstruction_error: f32 = SquaredL2::evaluate(&*vector, &*reconstructed);
502        let norm = vector.iter().map(|x| x * x).sum::<f32>();
503        assert!(
504            (reconstruction_error / norm) <= relative_err,
505            "Expected vector : {:?} to be reconstructed within error {} but instead got : {:?}, with error {} for dim : {}",
506            vector,
507            relative_err,
508            reconstructed,
509            reconstruction_error / norm,
510            dim,
511        );
512
513        assert!((loss.as_f32() - reconstruction_error) <= 1e-4);
514
515        let expected_code_sum = (0..dim)
516            .map(|i| encoded.vector().get(i).unwrap() as f32)
517            .sum::<f32>();
518        let code_sum = encoded.reborrow().meta().n / encoded.reborrow().meta().a;
519        assert!(
520            (code_sum - expected_code_sum).abs() <= 2e-5 * (dim as f32),
521            "Encoded vector with dim : {dim} is {:?}, got error : {} for vector : {:?}",
522            encoded.reborrow(),
523            (code_sum - expected_code_sum).abs(),
524            vector,
525        );
526        let recon_norm_sq = reconstructed.iter().map(|x| x * x).sum::<f32>();
527        assert!((encoded.reborrow().meta().norm_squared - recon_norm_sq).abs() <= 1e-3);
528
529        // FullQuery
530        let mut f = FullQuery::new_in(dim, GlobalAllocator).unwrap();
531        quantizer
532            .compress_into(vector.as_slice(), f.reborrow_mut())
533            .unwrap();
534
535        f.vector()
536            .iter()
537            .enumerate()
538            .zip(vector.iter())
539            .for_each(|((i, x), y)| {
540                assert!(
541                    (*x - *y).abs() < 1e-10,
542                    "Full Query did not compress dimension {i} with value {} correctly, got {} instead.",
543                    *y,
544                    *x,
545                )
546            });
547
548        assert!(
549            (f.meta().norm_squared - norm).abs() < 1e-10,
550            "Full Query norm in meta should be {norm} but instead got {}",
551            f.meta().norm_squared
552        );
553
554        let sum = vector.iter().sum::<f32>();
555        assert!(
556            (f.meta().sum - sum) < 1e-10,
557            "Full Query norm in meta should be {sum} but instead got {}",
558            f.meta().sum
559        );
560    }
561
562    cfg_if::cfg_if! {
563        if #[cfg(miri)] {
564            // The max dim does not need to be as high for `CompensatedVectors` because they
565            // defer their distance function implementation to `BitSlice`, which is more
566            // heavily tested.
567            const TRIALS: usize = 2;
568        } else {
569            const TRIALS: usize = 10;
570        }
571    }
572
573    macro_rules! test_minmax_quantizer_encoding {
574        ($name:ident, $dim:literal, $nbits:literal, $seed:literal, $err:expr) => {
575            #[test]
576            fn $name() {
577                let mut rng = StdRng::seed_from_u64($seed);
578                let scales = [1.0, 1.1, 0.9];
579                for (s, e) in scales.iter().zip($err) {
580                    for d in 10..$dim {
581                        for _ in 0..TRIALS {
582                            test_quantizer_encoding_random::<$nbits>(d, &mut rng, e, *s);
583                        }
584                    }
585                }
586            }
587        };
588    }
589    test_minmax_quantizer_encoding!(
590        test_minmax_encoding_1bit,
591        100,
592        1,
593        0xa32d5658097a1c35,
594        vec![0.5, 0.5, 0.5]
595    );
596    test_minmax_quantizer_encoding!(
597        test_minmax_encoding_2bit,
598        100,
599        2,
600        0xf60c0c8d1aadc126,
601        vec![0.5, 0.5, 0.5]
602    );
603    test_minmax_quantizer_encoding!(
604        test_minmax_encoding_4bit,
605        100,
606        4,
607        0x09fa14c42a9d7d98,
608        vec![1.0e-2, 1.0e-2, 3.0e-2]
609    );
610    test_minmax_quantizer_encoding!(
611        test_minmax_encoding_8bit,
612        100,
613        8,
614        0xaedf3d2a223b7b77,
615        vec![2.0e-3, 2.0e-3, 7.0e-3]
616    );
617
618    macro_rules! expand_to_bitrates {
619        ($name:ident, $func:ident) => {
620            #[test]
621            fn $name() {
622                $func::<1>();
623                $func::<2>();
624                $func::<4>();
625                $func::<8>();
626            }
627        };
628    }
629
630    /// Tests the edge case where min == max but both are non-zero.
631    fn test_all_same_value_vector<const NBITS: usize>()
632    where
633        Unsigned: Representation<NBITS>,
634        MinMaxQuantizer:
635            for<'a, 'b> CompressInto<&'a [f32], DataMutRef<'b, NBITS>, Output = L2Loss>,
636    {
637        let dim = 30;
638        let quantizer = MinMaxQuantizer::new(
639            Transform::Null(NullTransform::new(NonZeroUsize::new(dim).unwrap())),
640            Positive::new(1.0).unwrap(),
641        );
642        let constant_value = 42.5f32;
643        let vector = vec![constant_value; dim];
644
645        let mut encoded = Data::new_boxed(dim);
646        let result = quantizer.compress_into(&vector, encoded.reborrow_mut());
647
648        assert!(
649            result.is_ok(),
650            "Constant-value vector should compress successfully"
651        );
652
653        assert!(result.unwrap().as_f32().abs() <= 1e-6);
654
655        // Reconstruction should yield the original constant value (approximately)
656        let reconstructed = reconstruct_minmax(encoded.reborrow());
657        for &val in &reconstructed {
658            assert!(
659                (val - constant_value).abs() < 1e-3,
660                "Reconstructed value {} should be close to original {}. Compressed vector is {:?}",
661                val,
662                constant_value,
663                encoded.meta(),
664            );
665        }
666    }
667
668    /// This tests boundary conditions in the quantization logic.
669    fn test_two_distinct_values<const NBITS: usize>()
670    where
671        Unsigned: Representation<NBITS>,
672        MinMaxQuantizer:
673            for<'a, 'b> CompressInto<&'a [f32], DataMutRef<'b, NBITS>, Output = L2Loss>,
674    {
675        let dim = 20;
676        let quantizer = MinMaxQuantizer::new(
677            Transform::Null(NullTransform::new(NonZeroUsize::new(dim).unwrap())),
678            Positive::new(1.0).unwrap(),
679        );
680
681        let val1 = -10.0f32;
682        let val2 = 15.0f32;
683        let mut vector = vec![val1; dim];
684        // Make half the vector the second value
685        for i in vector.iter_mut().skip(dim) {
686            *i = val2;
687        }
688
689        let mut encoded = Data::new_boxed(dim);
690        let result = quantizer.compress_into(&vector, encoded.reborrow_mut());
691
692        assert!(
693            result.is_ok(),
694            "Two-value vector should compress successfully"
695        );
696
697        assert!(result.unwrap().as_f32().abs() <= 1e-6);
698
699        // Verify that only two distinct codes are used
700        let mut codes_used = std::collections::HashSet::new();
701        for i in 0..dim {
702            codes_used.insert(encoded.vector().get(i).unwrap());
703        }
704
705        // For most bit widths, we should see exactly 2 codes (min and max of domain)
706        if NBITS > 1 {
707            assert!(
708                codes_used.len() <= 2,
709                "Should use at most 2 distinct codes for 2-value input, but used: {:?}",
710                codes_used
711            );
712        }
713
714        // Verify reconstruction maintains the two-value structure approximately
715        let reconstructed = reconstruct_minmax(encoded.reborrow());
716        for ((i, val), v) in reconstructed.into_iter().enumerate().zip(&vector) {
717            // Round to nearest 0.1 to account for quantization error
718            assert!(
719                (val - v).abs() < 1e-4,
720                "Reconstructed value in dim : {i} is {val}, when it should be {v}."
721            );
722        }
723    }
724
725    /// Verifies that NaN values in the input cause the expected error but
726    /// dimension in meta is correctly set.
727    fn test_nan_input_error<const NBITS: usize>()
728    where
729        Unsigned: Representation<NBITS>,
730        MinMaxQuantizer:
731            for<'a, 'b> CompressInto<&'a [f32], DataMutRef<'b, NBITS>, Output = L2Loss>,
732    {
733        let dim = 100;
734        let quantizer = MinMaxQuantizer::new(
735            Transform::Null(NullTransform::new(NonZeroUsize::new(dim).unwrap())),
736            Positive::new(1.0).unwrap(),
737        );
738
739        // Test vector with NaN in the middle.
740        let mut vector_nan = vec![1.0f32; dim];
741        vector_nan[33] = f32::NAN;
742        let mut encoded = Data::new_boxed(dim);
743        let result = quantizer.compress_into(&vector_nan, encoded.reborrow_mut());
744        assert!(result.is_err(), "Vector with NaN should cause an error");
745
746        let meta = encoded.meta();
747        assert_eq!(meta.dim as usize, dim);
748    }
749
750    expand_to_bitrates!(all_same_values_vector, test_all_same_value_vector);
751    expand_to_bitrates!(two_distinct_values, test_two_distinct_values);
752    expand_to_bitrates!(nan_input_error, test_nan_input_error);
753
754    /// Verifies that providing a vector with wrong dimensionality causes a panic.
755    #[test]
756    #[should_panic(expected = "assertion `left == right` failed\n  left: 15\n right: 10")]
757    fn test_dimension_mismatch_panic()
758    where
759        Unsigned: Representation<8>,
760        MinMaxQuantizer: for<'a, 'b> CompressInto<&'a [f32], DataMutRef<'b, 8>, Output = L2Loss>,
761    {
762        let expected_dim = 10;
763        let quantizer = MinMaxQuantizer::new(
764            Transform::Null(NullTransform::new(NonZeroUsize::new(expected_dim).unwrap())),
765            Positive::new(1.0).unwrap(),
766        );
767
768        // Provide vector with wrong dimension
769        let wrong_vector = vec![1.0f32; expected_dim + 5]; // Too many dimensions
770        let mut encoded = Data::new_boxed(expected_dim);
771
772        // This should panic due to assertion in compress_into
773        let _ = quantizer.compress_into(&wrong_vector, encoded.reborrow_mut());
774    }
775
776    #[cfg(feature = "flatbuffers")]
777    mod serialization {
778        use std::sync::{
779            Arc,
780            atomic::{AtomicBool, Ordering},
781        };
782
783        use super::*;
784        use crate::{
785            algorithms::{TransformKind, transforms::TargetDim},
786            alloc::{AllocatorCore, AllocatorError},
787        };
788
789        /// Build a quantizer backed by a random `DoubleHadamard` transform so that
790        /// serialization has to round-trip non-trivial transform state (the random
791        /// sign flips), not just the scalar grid-scale.
792        fn make_hadamard_quantizer(dim: usize, grid_scale: f32, seed: u64) -> MinMaxQuantizer {
793            let mut rng = StdRng::seed_from_u64(seed);
794            let transform = Transform::new(
795                TransformKind::DoubleHadamard {
796                    target_dim: TargetDim::Same,
797                },
798                NonZeroUsize::new(dim).unwrap(),
799                Some(&mut rng),
800                GlobalAllocator,
801            )
802            .unwrap();
803
804            MinMaxQuantizer::new(transform, Positive::new(grid_scale).unwrap())
805        }
806
807        fn test_roundtrip_inner<const NBITS: usize>(seed: u64)
808        where
809            Unsigned: Representation<NBITS>,
810            MinMaxQuantizer:
811                for<'a, 'b> CompressInto<&'a [f32], DataMutRef<'b, NBITS>, Output = L2Loss>,
812        {
813            let dim = 16;
814            let quantizer = make_hadamard_quantizer(dim, 0.9, seed);
815
816            let serialized = quantizer.serialize(GlobalAllocator).unwrap();
817            let deserialized = MinMaxQuantizer::try_deserialize(&serialized).unwrap();
818
819            assert_eq!(deserialized, quantizer);
820        }
821
822        macro_rules! roundtrip_test {
823            ($name:ident, $nbits:literal, $seed:literal) => {
824                #[test]
825                fn $name() {
826                    test_roundtrip_inner::<$nbits>($seed);
827                }
828            };
829        }
830        roundtrip_test!(serialize_roundtrip_1bit, 1, 0x1b3f_8c21_6d90_a54e);
831        roundtrip_test!(serialize_roundtrip_2bit, 2, 0x77c2_e0a4_35bd_1f08);
832        roundtrip_test!(serialize_roundtrip_4bit, 4, 0x0a9d_4471_c8e2_9b6f);
833        roundtrip_test!(serialize_roundtrip_8bit, 8, 0xd4e1_5aa9_3f7c_2210);
834
835        /// Deserializing a buffer without the expected file identifier is rejected
836        /// (rather than panicking or silently misparsing).
837        #[test]
838        fn deserialize_rejects_invalid_identifier() {
839            // A zero-filled buffer long enough to hold an identifier slot has the wrong
840            // identifier and must be rejected.
841            assert!(matches!(
842                MinMaxQuantizer::try_deserialize(&[0u8; 32]),
843                Err(DeserializationError::InvalidIdentifier),
844            ));
845
846            // Corrupting the file identifier of an otherwise-valid buffer is also rejected.
847            let quantizer = make_hadamard_quantizer(16, 1.0, 0x51ed_9c33);
848            let mut serialized = quantizer.serialize(GlobalAllocator).unwrap().to_vec();
849            // The file identifier occupies bytes 4..8 of a (non size-prefixed) flatbuffer.
850            serialized[4] ^= 0xff;
851            assert!(matches!(
852                MinMaxQuantizer::try_deserialize(&serialized),
853                Err(DeserializationError::InvalidIdentifier),
854            ));
855        }
856
857        /// An allocator that succeeds on its first allocation but fails on every
858        /// subsequent one, used to exercise the reallocation path in `serialize`.
859        #[derive(Debug, Clone)]
860        struct FlakyAllocator {
861            have_allocated: Arc<AtomicBool>,
862        }
863
864        impl FlakyAllocator {
865            fn new(have_allocated: Arc<AtomicBool>) -> Self {
866                Self { have_allocated }
867            }
868        }
869
870        // SAFETY: This wraps `GlobalAllocator` and only changes when allocation fails.
871        unsafe impl AllocatorCore for FlakyAllocator {
872            fn allocate(
873                &self,
874                layout: std::alloc::Layout,
875            ) -> Result<std::ptr::NonNull<[u8]>, AllocatorError> {
876                if self.have_allocated.swap(true, Ordering::Relaxed) {
877                    Err(AllocatorError)
878                } else {
879                    GlobalAllocator.allocate(layout)
880                }
881            }
882
883            unsafe fn deallocate(&self, ptr: std::ptr::NonNull<[u8]>, layout: std::alloc::Layout) {
884                // SAFETY: Inherited from caller.
885                unsafe { GlobalAllocator.deallocate(ptr, layout) }
886            }
887        }
888
889        /// Serialization must surface an [`AllocatorError`] instead of panicking when the
890        /// buffer needs to grow but reallocation fails. Under `cfg(test)` the initial
891        /// buffer is intentionally tiny, so growth (and thus the second allocation) always
892        /// happens.
893        #[test]
894        fn serialize_allocation_failure_does_not_panic() {
895            let quantizer = make_hadamard_quantizer(16, 1.0, 0x7a11_36bd);
896            let have_allocated = Arc::new(AtomicBool::new(false));
897            let _: AllocatorError = quantizer
898                .serialize(FlakyAllocator::new(have_allocated.clone()))
899                .unwrap_err();
900            assert!(have_allocated.load(Ordering::Relaxed));
901        }
902    }
903}
diskann_quantization/minmax/quantizer.rs

diskann_quantization/minmax/
quantizer.rs
