diskann_quantization/minmax/
vectors.rs

1/*
2 * Copyright (c) Microsoft Corporation.
3 * Licensed under the MIT license.
4 */
5
6use std::ops::{Add, Mul};
7
8use diskann_vector::{MathematicalValue, PureDistanceFunction};
9use thiserror::Error;
10
11use crate::{
12    alloc::GlobalAllocator,
13    bits::{BitSlice, Dense, Representation, Unsigned},
14    distances,
15    distances::{InnerProduct, MV},
16    meta::{self, slice},
17};
18
19/// A per-vector precomputed coefficients to help compute inner products
20/// and squared L2 distances for the MinMax quantized vectors.
21///
22/// The inner product between `X = ax * X' + bx` and `Y = ay * Y' + by` for d-dimensional
23/// vectors X and Y is:
24/// ```math
25/// <X, Y> = <ax * X' + bx, ay * Y' + by>
26///        = ax * ay * <X', Y'> + ax * <X', by> + ay * <Y', bx> + d * bx * by.
27/// ```
28/// Let us define a grouping of these terms to make it easier to understand:
29/// ```math
30///  Nx = ax * sum_i X'[i],     Ny = ay * sum_i Y'[i],
31/// ```
32/// We can then simplify the dot product calculation as follows:
33/// ```math
34/// <X, Y> = ax * ay * <X', Y'> + Nx * by + Ny * bx +  d * bx * by
35///                    --------
36///                       |
37///               Integer Dot Product
38/// ```
39///
40/// To compute the squared L2 distance,
41/// ```math
42/// |X - Y|^2 = |ax * X' + bx|^2 + |ay * Y' + by|^2 - 2 * <X, Y>
43/// ```
44/// we can re-use the computation for inner-product from above.
45#[derive(Default, Debug, Clone, Copy, PartialEq, bytemuck::Zeroable, bytemuck::Pod)]
46#[repr(C)]
47pub struct MinMaxCompensation {
48    pub dim: u32,          // - dimension
49    pub b: f32,            // - bx
50    pub n: f32,            // - Nx
51    pub a: f32,            // - ax
52    pub norm_squared: f32, // - |ax * X' + bx|^2
53}
54
55#[derive(Debug, Clone, Copy, Default)]
56pub(crate) struct Coefficients<T> {
57    pub(crate) a: T,
58    pub(crate) b: T,
59    pub(crate) n: T,
60}
61
62impl From<MinMaxCompensation> for Coefficients<f32> {
63    fn from(meta: MinMaxCompensation) -> Self {
64        Self {
65            a: meta.a,
66            b: meta.b,
67            n: meta.n,
68        }
69    }
70}
71
72impl<T: Copy + Add<Output = T> + Mul<Output = T>> Coefficients<T> {
73    /// Compensate a raw code dot product using the logical (unpadded) dimension.
74    #[inline(always)]
75    pub(crate) fn compensate_inner_product(self, other: Self, dot: T, dim: T) -> T {
76        ((self.a * other.a) * dot + self.n * other.b) + other.n * self.b + (self.b * other.b) * dim
77    }
78}
79
80const META_BYTES: usize = std::mem::size_of::<MinMaxCompensation>(); // This will be 5 * 4 = 20 bytes.
81
82/// Error type for parsing a slice of bytes as a `DataRef`
83/// and returning corresponding dimension.
84#[derive(Debug, Error, Clone, PartialEq, Eq)]
85pub enum MetaParseError {
86    #[error("Invalid size: {0}, must contain at least {META_BYTES} bytes")]
87    NotCanonical(usize),
88}
89
90impl MinMaxCompensation {
91    /// Reads the dimension from the first 4 bytes of a MinMax quantized vector's metadata.
92    ///
93    /// This function is used to extract the vector dimension from serialized MinMax quantized
94    /// vector data without fully deserializing the entire vector structure.
95    ///
96    /// # Arguments
97    /// * `bytes` - A byte slice containing the serialized MinMax vector data
98    ///
99    /// # Returns
100    /// * `Ok(dimension)` - The dimension of the vector as a `usize`
101    /// * `Err(MetaParseError::NotCanonical(size))` - If the byte slice is shorter than 20 bytes (META_BYTES)
102    ///
103    /// # Usage
104    /// Use this when you need to determine the vector dimension from serialized data before
105    /// creating a `DataRef` or allocating appropriately sized buffers for decompression.
106    #[inline(always)]
107    pub fn read_dimension(bytes: &[u8]) -> Result<usize, MetaParseError> {
108        if bytes.len() < META_BYTES {
109            return Err(MetaParseError::NotCanonical(bytes.len()));
110        }
111
112        // SAFETY: There are at least `META_BYTES` = 20 bytes in the array so this access is within bounds.
113        let dim_bytes: [u8; 4] = bytes.get(..4).map_or_else(
114            || Err(MetaParseError::NotCanonical(bytes.len())),
115            |slice| {
116                slice
117                    .try_into()
118                    .map_err(|_| MetaParseError::NotCanonical(bytes.len()))
119            },
120        )?;
121
122        let dim = u32::from_le_bytes(dim_bytes) as usize;
123
124        Ok(dim)
125    }
126}
127
128/// An owning compressed data vector
129///
130/// See: [`meta::Vector`].
131pub type Data<const NBITS: usize> = meta::Vector<NBITS, Unsigned, MinMaxCompensation, Dense>;
132
133/// A borrowed `Data` vector
134///
135/// See: [`meta::Vector`].
136pub type DataRef<'a, const NBITS: usize> =
137    meta::VectorRef<'a, NBITS, Unsigned, MinMaxCompensation, Dense>;
138
139#[derive(Debug, Error, Clone, Copy, PartialEq, Eq)]
140pub enum DecompressError {
141    #[error("expected src and dst length to be identical, instead src is {0}, and dst is {1}")]
142    LengthMismatch(usize, usize),
143}
144impl<const NBITS: usize> DataRef<'_, NBITS>
145where
146    Unsigned: Representation<NBITS>,
147{
148    /// Decompresses a MinMax quantized vector back into its original floating-point representation.
149    ///
150    /// This method reconstructs the original vector values using the stored quantization parameters
151    /// and the MinMax dequantization formula: `x = x' * a + b` and stores the result in `dst`
152    ///
153    /// # Arguments
154    ///
155    /// * `dst` - A mutable slice of `f32` values where the decompressed data will be written.
156    ///   Must have the same length as the compressed vector.
157    ///
158    /// # Returns
159    ///
160    /// * `Ok(())` - On successful decompression
161    /// * `Err(DecompressError::LengthMismatch(src_len, dst_len))` - If the destination slice
162    ///   length doesn't match the compressed vector length
163    pub fn decompress_into(&self, dst: &mut [f32]) -> Result<(), DecompressError> {
164        if dst.len() != self.len() {
165            return Err(DecompressError::LengthMismatch(self.len(), dst.len()));
166        }
167        let meta = self.meta();
168
169        // SAFETY: We checked that the length of the underlying vector is the same as
170        // as `dst` so we are guaranteed to be within bounds when accessing the vector.
171        dst.iter_mut().enumerate().for_each(|(i, d)| unsafe {
172            *d = self.vector().get_unchecked(i) as f32 * meta.a + meta.b
173        });
174        Ok(())
175    }
176}
177
178/// A mutable borrowed `Data` vector
179///
180/// See: [`meta::Vector`].
181pub type DataMutRef<'a, const NBITS: usize> =
182    meta::VectorMut<'a, NBITS, Unsigned, MinMaxCompensation, Dense>;
183
184////////////////////
185// Full Precision //
186////////////////////
187
188/// A meta struct storing the `sum` and `norm_squared` of a
189/// full query after transformation is applied to it.
190///
191/// The inner product between `X = ax * X' + bx` and `Y` for d-dimensional
192/// vectors X and Y is:
193/// ```math
194/// <X, Y> = <ax * X' + bx, Y>
195///        = ax * <X', Y> + bx * sum(Y).
196///               --------
197///                  |
198///          Integer-Float Dot Product
199/// ```
200///
201/// To compute the squared L2 distance,
202/// ```math
203/// |X - Y|^2 = |ax * X' + bx|^2 + |Y|^2 - 2 * <X', Y>
204/// ```
205#[derive(Debug, Clone, Copy, Default, bytemuck::Zeroable, bytemuck::Pod)]
206#[repr(C)]
207pub struct FullQueryMeta {
208    /// The sum of `data`.
209    pub sum: f32,
210    /// The norm of the 'data'.
211    pub norm_squared: f32,
212}
213
214/// A full precision query.
215///
216/// See: [`slice::Slice`].
217pub type FullQuery<A = GlobalAllocator> = slice::PolySlice<f32, FullQueryMeta, A>;
218
219/// A borrowed full precision query.
220///
221/// See: [`slice::SliceRef`].
222pub type FullQueryRef<'a> = slice::SliceRef<'a, f32, FullQueryMeta>;
223
224/// A mutable borrowed full precision query.
225///
226/// See: [`slice::SliceMut`].
227pub type FullQueryMut<'a> = slice::SliceMut<'a, f32, FullQueryMeta>;
228
229///////////////////////////
230// Compensated Distances //
231///////////////////////////
232#[inline(always)]
233fn kernel<const NBITS: usize, const MBITS: usize, F>(
234    x: DataRef<'_, NBITS>,
235    y: DataRef<'_, MBITS>,
236    f: F,
237) -> distances::MathematicalResult<f32>
238where
239    Unsigned: Representation<NBITS> + Representation<MBITS>,
240    InnerProduct: for<'a, 'b> PureDistanceFunction<
241            BitSlice<'a, NBITS, Unsigned>,
242            BitSlice<'b, MBITS, Unsigned>,
243            distances::MathematicalResult<u32>,
244        >,
245    F: Fn(f32, &MinMaxCompensation, &MinMaxCompensation) -> f32,
246{
247    let raw_product = InnerProduct::evaluate(x.vector(), y.vector())?;
248    let (xm, ym) = (x.meta(), y.meta());
249    let v = Coefficients::from(xm).compensate_inner_product(
250        ym.into(),
251        raw_product.into_inner() as f32,
252        x.len() as f32,
253    );
254    Ok(MV::new(f(v, &xm, &ym)))
255}
256
257pub struct MinMaxIP;
258
259impl<const NBITS: usize, const MBITS: usize>
260    PureDistanceFunction<DataRef<'_, NBITS>, DataRef<'_, MBITS>, distances::MathematicalResult<f32>>
261    for MinMaxIP
262where
263    Unsigned: Representation<NBITS> + Representation<MBITS>,
264    InnerProduct: for<'a, 'b> PureDistanceFunction<
265            BitSlice<'a, NBITS, Unsigned>,
266            BitSlice<'b, MBITS, Unsigned>,
267            distances::MathematicalResult<u32>,
268        >,
269{
270    #[inline(always)]
271    fn evaluate(
272        x: DataRef<'_, NBITS>,
273        y: DataRef<'_, MBITS>,
274    ) -> distances::MathematicalResult<f32> {
275        kernel(x, y, |v, _, _| v)
276    }
277}
278
279impl<const NBITS: usize, const MBITS: usize>
280    PureDistanceFunction<DataRef<'_, NBITS>, DataRef<'_, MBITS>, distances::Result<f32>>
281    for MinMaxIP
282where
283    Unsigned: Representation<NBITS> + Representation<MBITS>,
284    InnerProduct: for<'a, 'b> PureDistanceFunction<
285            BitSlice<'a, NBITS, Unsigned>,
286            BitSlice<'b, MBITS, Unsigned>,
287            distances::MathematicalResult<u32>,
288        >,
289{
290    #[inline(always)]
291    fn evaluate(x: DataRef<'_, NBITS>, y: DataRef<'_, MBITS>) -> distances::Result<f32> {
292        let v: distances::MathematicalResult<f32> = Self::evaluate(x, y);
293        Ok(-v?.into_inner())
294    }
295}
296
297impl<const NBITS: usize>
298    PureDistanceFunction<FullQueryRef<'_>, DataRef<'_, NBITS>, distances::MathematicalResult<f32>>
299    for MinMaxIP
300where
301    Unsigned: Representation<NBITS>,
302    InnerProduct: for<'a, 'b> PureDistanceFunction<
303            &'a [f32],
304            BitSlice<'b, NBITS, Unsigned>,
305            distances::MathematicalResult<f32>,
306        >,
307{
308    #[inline(always)]
309    fn evaluate(x: FullQueryRef<'_>, y: DataRef<'_, NBITS>) -> distances::MathematicalResult<f32> {
310        let raw_product: f32 = InnerProduct::evaluate(x.vector(), y.vector())?.into_inner();
311        Ok(MathematicalValue::new(
312            raw_product * y.meta().a + x.meta().sum * y.meta().b,
313        ))
314    }
315}
316
317impl<const NBITS: usize>
318    PureDistanceFunction<FullQueryRef<'_>, DataRef<'_, NBITS>, distances::Result<f32>> for MinMaxIP
319where
320    Unsigned: Representation<NBITS>,
321    InnerProduct: for<'a, 'b> PureDistanceFunction<
322            &'a [f32],
323            BitSlice<'b, NBITS, Unsigned>,
324            distances::MathematicalResult<f32>,
325        >,
326{
327    #[inline(always)]
328    fn evaluate(x: FullQueryRef<'_>, y: DataRef<'_, NBITS>) -> distances::Result<f32> {
329        let v: distances::MathematicalResult<f32> = Self::evaluate(x, y);
330        Ok(-v?.into_inner())
331    }
332}
333
334pub struct MinMaxL2Squared;
335
336impl<const NBITS: usize, const MBITS: usize>
337    PureDistanceFunction<DataRef<'_, NBITS>, DataRef<'_, MBITS>, distances::MathematicalResult<f32>>
338    for MinMaxL2Squared
339where
340    Unsigned: Representation<NBITS> + Representation<MBITS>,
341    InnerProduct: for<'a, 'b> PureDistanceFunction<
342            BitSlice<'a, NBITS, Unsigned>,
343            BitSlice<'b, MBITS, Unsigned>,
344            distances::MathematicalResult<u32>,
345        >,
346{
347    #[inline(always)]
348    fn evaluate(
349        x: DataRef<'_, NBITS>,
350        y: DataRef<'_, MBITS>,
351    ) -> distances::MathematicalResult<f32> {
352        kernel(x, y, |v, xm, ym| {
353            -2.0 * v + xm.norm_squared + ym.norm_squared
354        })
355    }
356}
357
358impl<const NBITS: usize, const MBITS: usize>
359    PureDistanceFunction<DataRef<'_, NBITS>, DataRef<'_, MBITS>, distances::Result<f32>>
360    for MinMaxL2Squared
361where
362    Unsigned: Representation<NBITS> + Representation<MBITS>,
363    InnerProduct: for<'a, 'b> PureDistanceFunction<
364            BitSlice<'a, NBITS, Unsigned>,
365            BitSlice<'b, MBITS, Unsigned>,
366            distances::MathematicalResult<u32>,
367        >,
368{
369    #[inline(always)]
370    fn evaluate(x: DataRef<'_, NBITS>, y: DataRef<'_, MBITS>) -> distances::Result<f32> {
371        let v: distances::MathematicalResult<f32> = Self::evaluate(x, y);
372        Ok(v?.into_inner())
373    }
374}
375
376impl<const NBITS: usize>
377    PureDistanceFunction<FullQueryRef<'_>, DataRef<'_, NBITS>, distances::MathematicalResult<f32>>
378    for MinMaxL2Squared
379where
380    Unsigned: Representation<NBITS>,
381    InnerProduct: for<'a, 'b> PureDistanceFunction<
382            &'a [f32],
383            BitSlice<'b, NBITS, Unsigned>,
384            distances::MathematicalResult<f32>,
385        >,
386{
387    #[inline(always)]
388    fn evaluate(x: FullQueryRef<'_>, y: DataRef<'_, NBITS>) -> distances::MathematicalResult<f32> {
389        let raw_product = InnerProduct::evaluate(x.vector(), y.vector())?.into_inner();
390
391        let ym = y.meta();
392        let compensated_ip = raw_product * ym.a + x.meta().sum * ym.b;
393        Ok(MV::new(
394            x.meta().norm_squared + ym.norm_squared - 2.0 * compensated_ip,
395        ))
396    }
397}
398
399impl<const NBITS: usize>
400    PureDistanceFunction<FullQueryRef<'_>, DataRef<'_, NBITS>, distances::Result<f32>>
401    for MinMaxL2Squared
402where
403    Unsigned: Representation<NBITS>,
404    InnerProduct: for<'a, 'b> PureDistanceFunction<
405            &'a [f32],
406            BitSlice<'b, NBITS, Unsigned>,
407            distances::MathematicalResult<f32>,
408        >,
409{
410    #[inline(always)]
411    fn evaluate(x: FullQueryRef<'_>, y: DataRef<'_, NBITS>) -> distances::Result<f32> {
412        let v: distances::MathematicalResult<f32> = Self::evaluate(x, y);
413        Ok(v?.into_inner())
414    }
415}
416
417///////////////////////
418// Cosine Distances //
419///////////////////////
420
421pub struct MinMaxCosine;
422
423impl<const NBITS: usize, const MBITS: usize>
424    PureDistanceFunction<DataRef<'_, NBITS>, DataRef<'_, MBITS>, distances::Result<f32>>
425    for MinMaxCosine
426where
427    Unsigned: Representation<NBITS> + Representation<MBITS>,
428    MinMaxIP: for<'a, 'b> PureDistanceFunction<
429            DataRef<'a, NBITS>,
430            DataRef<'b, MBITS>,
431            distances::MathematicalResult<f32>,
432        >,
433{
434    // 1 - <X, Y> / (|X| * |Y|)
435    #[inline(always)]
436    fn evaluate(x: DataRef<'_, NBITS>, y: DataRef<'_, MBITS>) -> distances::Result<f32> {
437        let ip: MV<f32> = MinMaxIP::evaluate(x, y)?;
438        let (xm, ym) = (x.meta(), y.meta());
439        Ok(1.0 - ip.into_inner() / (xm.norm_squared.sqrt() * ym.norm_squared.sqrt()))
440    }
441}
442
443impl<const NBITS: usize>
444    PureDistanceFunction<FullQueryRef<'_>, DataRef<'_, NBITS>, distances::Result<f32>>
445    for MinMaxCosine
446where
447    Unsigned: Representation<NBITS>,
448    MinMaxIP: for<'a, 'b> PureDistanceFunction<
449            FullQueryRef<'a>,
450            DataRef<'b, NBITS>,
451            distances::MathematicalResult<f32>,
452        >,
453{
454    #[inline(always)]
455    fn evaluate(x: FullQueryRef<'_>, y: DataRef<'_, NBITS>) -> distances::Result<f32> {
456        let ip: MathematicalValue<f32> = MinMaxIP::evaluate(x, y)?;
457        let (xm, ym) = (x.meta().norm_squared, y.meta());
458        Ok(1.0 - ip.into_inner() / (xm.sqrt() * ym.norm_squared.sqrt()))
459        // 1 - <X, Y> / (|X| * |Y|)
460    }
461}
462
463pub struct MinMaxCosineNormalized;
464
465impl<const NBITS: usize, const MBITS: usize>
466    PureDistanceFunction<DataRef<'_, NBITS>, DataRef<'_, MBITS>, distances::Result<f32>>
467    for MinMaxCosineNormalized
468where
469    Unsigned: Representation<NBITS> + Representation<MBITS>,
470    MinMaxIP: for<'a, 'b> PureDistanceFunction<
471            DataRef<'a, NBITS>,
472            DataRef<'b, MBITS>,
473            distances::MathematicalResult<f32>,
474        >,
475{
476    #[inline(always)]
477    fn evaluate(x: DataRef<'_, NBITS>, y: DataRef<'_, MBITS>) -> distances::Result<f32> {
478        let ip: MathematicalValue<f32> = MinMaxIP::evaluate(x, y)?;
479        Ok(1.0 - ip.into_inner()) // 1 - <X, Y>
480    }
481}
482
483impl<const NBITS: usize>
484    PureDistanceFunction<FullQueryRef<'_>, DataRef<'_, NBITS>, distances::Result<f32>>
485    for MinMaxCosineNormalized
486where
487    Unsigned: Representation<NBITS>,
488    MinMaxIP: for<'a, 'b> PureDistanceFunction<
489            FullQueryRef<'a>,
490            DataRef<'b, NBITS>,
491            distances::MathematicalResult<f32>,
492        >,
493{
494    #[inline(always)]
495    fn evaluate(x: FullQueryRef<'_>, y: DataRef<'_, NBITS>) -> distances::Result<f32> {
496        let ip: MathematicalValue<f32> = MinMaxIP::evaluate(x, y)?;
497        Ok(1.0 - ip.into_inner()) // 1 - <X, Y>
498    }
499}
500
501///////////
502// Tests //
503///////////
504
505#[cfg(test)]
506#[cfg(not(miri))]
507mod minmax_vector_tests {
508    use diskann_utils::Reborrow;
509    use rand::{
510        Rng, SeedableRng,
511        distr::{Distribution, Uniform},
512        rngs::StdRng,
513    };
514
515    use super::*;
516    use crate::{alloc::GlobalAllocator, scalar::bit_scale};
517
518    /// Builds a random MinMax quantized vector and its full-precision reconstruction.
519    ///
520    /// Returns `(compressed, original)` where `compressed` has its `MinMaxCompensation`
521    /// metadata fully populated and `original` is the dequantized f32 vector.
522    fn random_minmax_vector<const NBITS: usize>(
523        dim: usize,
524        rng: &mut impl Rng,
525    ) -> (Data<NBITS>, Vec<f32>)
526    where
527        Unsigned: Representation<NBITS>,
528    {
529        let mut v = Data::<NBITS>::new_boxed(dim);
530
531        let domain = Unsigned::domain_const::<NBITS>();
532        let code_dist = Uniform::new_inclusive(*domain.start(), *domain.end()).unwrap();
533
534        {
535            let mut bs = v.vector_mut();
536            for i in 0..dim {
537                bs.set(i, code_dist.sample(rng)).unwrap();
538            }
539        }
540
541        let a: f32 = Uniform::new_inclusive(0.0, 2.0).unwrap().sample(rng);
542        let b: f32 = Uniform::new_inclusive(0.0, 2.0).unwrap().sample(rng);
543
544        let original: Vec<f32> = (0..dim)
545            .map(|i| a * v.vector().get(i).unwrap() as f32 + b)
546            .collect();
547
548        let code_sum: f32 = (0..dim).map(|i| v.vector().get(i).unwrap() as f32).sum();
549        let norm_squared: f32 = original.iter().map(|x| x * x).sum();
550
551        v.set_meta(MinMaxCompensation {
552            a,
553            b,
554            n: a * code_sum,
555            norm_squared,
556            dim: dim as u32,
557        });
558
559        (v, original)
560    }
561
562    fn test_minmax_compensated_vectors<const NBITS: usize, R>(dim: usize, rng: &mut R)
563    where
564        Unsigned: Representation<NBITS>,
565        InnerProduct: for<'a, 'b> PureDistanceFunction<
566                BitSlice<'a, NBITS, Unsigned>,
567                BitSlice<'b, NBITS, Unsigned>,
568                distances::MathematicalResult<u32>,
569            >,
570        InnerProduct: for<'a, 'b> PureDistanceFunction<
571                &'a [f32],
572                BitSlice<'b, NBITS, Unsigned>,
573                distances::MathematicalResult<f32>,
574            >,
575        R: Rng,
576    {
577        assert!(dim <= bit_scale::<NBITS>() as usize);
578
579        let (v1, original1) = random_minmax_vector::<NBITS>(dim, rng);
580        let (v2, original2) = random_minmax_vector::<NBITS>(dim, rng);
581
582        let norm1_squared = v1.meta().norm_squared;
583        let norm2_squared = v2.meta().norm_squared;
584
585        // Calculate raw integer dot product
586        let expected_ip = (0..dim).map(|i| original1[i] * original2[i]).sum::<f32>();
587
588        // Test inner product with f32
589        let computed_ip_f32: distances::Result<f32> =
590            MinMaxIP::evaluate(v1.reborrow(), v2.reborrow());
591        let computed_ip_f32 = computed_ip_f32.unwrap();
592        assert!(
593            (expected_ip - (-computed_ip_f32)).abs() / expected_ip.abs() < 1e-3,
594            "Inner product (f32) failed: expected {}, got {} on dim : {}",
595            -expected_ip,
596            computed_ip_f32,
597            dim
598        );
599
600        // Expected L2 distance = |X|² + |Y|² - 2<X,Y>
601        let expected_l2 = (0..dim)
602            .map(|i| original1[i] - original2[i])
603            .map(|x| x.powf(2.0))
604            .sum::<f32>();
605
606        // Test L2 distance with f32
607        let computed_l2_f32: distances::Result<f32> =
608            MinMaxL2Squared::evaluate(v1.reborrow(), v2.reborrow());
609        let computed_l2_f32 = computed_l2_f32.unwrap();
610        assert!(
611            ((computed_l2_f32 - expected_l2).abs() / expected_l2) < 1e-3,
612            "L2 distance (f32) failed: expected {}, got {} on dim : {}",
613            expected_l2,
614            computed_l2_f32,
615            dim
616        );
617
618        let expected_cosine = 1.0 - expected_ip / (norm1_squared.sqrt() * norm2_squared.sqrt());
619
620        let computed_cosine: distances::Result<f32> =
621            MinMaxCosine::evaluate(v1.reborrow(), v2.reborrow());
622        let computed_cosine = computed_cosine.unwrap();
623
624        {
625            let passed = (computed_cosine - expected_cosine).abs() < 1e-6
626                || ((computed_cosine - expected_cosine).abs() / expected_cosine) < 1e-3;
627
628            assert!(
629                passed,
630                "Cosine distance (f32) failed: expected {}, got {} on dim : {}",
631                expected_cosine, computed_cosine, dim
632            );
633        }
634
635        let cosine_normalized: distances::Result<f32> =
636            MinMaxCosineNormalized::evaluate(v1.reborrow(), v2.reborrow());
637        let cosine_normalized = cosine_normalized.unwrap();
638        let expected_cos_normalized = 1.0 - expected_ip;
639        assert!(
640            ((expected_cos_normalized - cosine_normalized).abs() / expected_cos_normalized.abs())
641                < 1e-6,
642            "CosineNormalized distance (f32) failed: expected {}, got {} on dim : {}",
643            expected_cos_normalized,
644            cosine_normalized,
645            dim
646        );
647
648        //Calculate inner product with full precision vector
649        let mut fp_query = FullQuery::new_in(dim, GlobalAllocator).unwrap();
650        fp_query.vector_mut().copy_from_slice(&original1);
651        *fp_query.meta_mut() = FullQueryMeta {
652            norm_squared: norm1_squared,
653            sum: original1.iter().sum::<f32>(),
654        };
655
656        let fp_ip: distances::Result<f32> = MinMaxIP::evaluate(fp_query.reborrow(), v2.reborrow());
657        let fp_ip = fp_ip.unwrap();
658        assert!(
659            (expected_ip - (-fp_ip)).abs() / expected_ip.abs() < 1e-3,
660            "Inner product (f32) failed: expected {}, got {} on dim : {}",
661            -expected_ip,
662            fp_ip,
663            dim
664        );
665
666        let fp_l2: distances::Result<f32> =
667            MinMaxL2Squared::evaluate(fp_query.reborrow(), v2.reborrow());
668        let fp_l2 = fp_l2.unwrap();
669        assert!(
670            ((fp_l2 - expected_l2).abs() / expected_l2) < 1e-3,
671            "L2 distance (f32) failed: expected {}, got {} on dim : {}",
672            expected_l2,
673            computed_l2_f32,
674            dim
675        );
676
677        let fp_cosine: distances::Result<f32> =
678            MinMaxCosine::evaluate(fp_query.reborrow(), v2.reborrow());
679        let fp_cosine = fp_cosine.unwrap();
680        let diff = (fp_cosine - expected_cosine).abs();
681        assert!(
682            (diff / expected_cosine) < 1e-3 || diff <= 1e-6,
683            "Cosine distance (f32) failed: expected {}, got {} on dim : {}",
684            expected_cosine,
685            fp_cosine,
686            dim
687        );
688
689        let fp_cos_norm: distances::Result<f32> =
690            MinMaxCosineNormalized::evaluate(fp_query.reborrow(), v2.reborrow());
691        let fp_cos_norm = fp_cos_norm.unwrap();
692        assert!(
693            (((1.0 - expected_ip) - fp_cos_norm).abs() / (1.0 - expected_ip)) < 1e-3,
694            "Cosine distance (f32) failed: expected {}, got {} on dim : {}",
695            (1.0 - expected_ip),
696            fp_cos_norm,
697            dim
698        );
699
700        //Test `decompress_into` to make sure it outputs tje full-precision vector correctly.
701        let meta = v1.meta();
702        let v1_ref = DataRef::new(v1.vector(), &meta);
703        let dim = v1_ref.len();
704        let mut boxed = vec![0f32; dim + 1];
705
706        let pre = v1_ref.decompress_into(&mut boxed);
707        assert_eq!(
708            pre.unwrap_err(),
709            DecompressError::LengthMismatch(dim, dim + 1)
710        );
711        let pre = v1_ref.decompress_into(&mut boxed[..dim - 1]);
712        assert_eq!(
713            pre.unwrap_err(),
714            DecompressError::LengthMismatch(dim, dim - 1)
715        );
716        let pre = v1_ref.decompress_into(&mut boxed[..dim]);
717        assert!(pre.is_ok());
718
719        boxed
720            .iter()
721            .zip(original1.iter())
722            .for_each(|(x, y)| assert!((*x - *y).abs() <= 1e-6));
723
724        // Verify `read_dimension` is correct.
725        let mut bytes = vec![0u8; Data::canonical_bytes(dim)];
726        let mut data = DataMutRef::from_canonical_front_mut(bytes.as_mut_slice(), dim).unwrap();
727        data.set_meta(meta);
728
729        let pre = MinMaxCompensation::read_dimension(&bytes);
730        assert!(pre.is_ok());
731        let read_dim = pre.unwrap();
732        assert_eq!(read_dim, dim);
733
734        let pre = MinMaxCompensation::read_dimension(&[0_u8; 2]);
735        assert_eq!(pre.unwrap_err(), MetaParseError::NotCanonical(2));
736    }
737
738    cfg_if::cfg_if! {
739        if #[cfg(miri)] {
740            // The max dim does not need to be as high for `CompensatedVectors` because they
741            // defer their distance function implementation to `BitSlice`, which is more
742            // heavily tested.
743            const TRIALS: usize = 2;
744        } else {
745            const TRIALS: usize = 10;
746        }
747    }
748
749    macro_rules! test_minmax_compensated {
750        ($name:ident, $nbits:literal, $seed:literal) => {
751            #[test]
752            fn $name() {
753                let mut rng = StdRng::seed_from_u64($seed);
754                const MAX_DIM: usize = (bit_scale::<$nbits>() as usize);
755                for dim in 1..=MAX_DIM {
756                    for _ in 0..TRIALS {
757                        test_minmax_compensated_vectors::<$nbits, _>(dim, &mut rng);
758                    }
759                }
760            }
761        };
762    }
763    test_minmax_compensated!(unsigned_minmax_compensated_test_u1, 1, 0xa33d5658097a1c35);
764    test_minmax_compensated!(unsigned_minmax_compensated_test_u2, 2, 0xaedf3d2a223b7b77);
765    test_minmax_compensated!(unsigned_minmax_compensated_test_u4, 4, 0xf60c0c8d1aadc126);
766    test_minmax_compensated!(unsigned_minmax_compensated_test_u8, 8, 0x09fa14c42a9d7d98);
767
768    /// Test the heterogeneous MinMax kernel for N-bit queries × M-bit database vectors.
769    ///
770    /// Verifies that `kernel::<N, M, _>` produces inner-product and squared-L2
771    /// results matching the full-precision reference, for random codes and
772    /// random compensation coefficients.
773    fn test_minmax_heterogeneous_kernel<const NBITS: usize, const MBITS: usize, R>(
774        dim: usize,
775        rng: &mut R,
776    ) where
777        Unsigned: Representation<NBITS> + Representation<MBITS>,
778        InnerProduct: for<'a, 'b> PureDistanceFunction<
779                BitSlice<'a, NBITS, Unsigned>,
780                BitSlice<'b, MBITS, Unsigned>,
781                distances::MathematicalResult<u32>,
782            >,
783        R: Rng,
784    {
785        let (v_query, original1) = random_minmax_vector::<NBITS>(dim, rng);
786        let (v_data, original2) = random_minmax_vector::<MBITS>(dim, rng);
787
788        // ── Inner Product ──
789        let expected_ip: f32 = original1.iter().zip(&original2).map(|(x, y)| x * y).sum();
790        let computed_ip = kernel(v_query.reborrow(), v_data.reborrow(), |v, _, _| v)
791            .unwrap()
792            .into_inner();
793        assert!(
794            (expected_ip - computed_ip).abs() / expected_ip.abs().max(1e-10) < 1e-6,
795            "Heterogeneous IP ({},{}) failed: expected {}, got {} on dim: {}",
796            NBITS,
797            MBITS,
798            expected_ip,
799            computed_ip,
800            dim,
801        );
802    }
803
804    macro_rules! test_minmax_heterogeneous {
805        ($name:ident, $N:literal, $M:literal, $seed:literal) => {
806            #[test]
807            fn $name() {
808                let mut rng = StdRng::seed_from_u64($seed);
809                // Use the smaller bit width's scale as max dimension.
810                const MAX_DIM: usize = bit_scale::<$M>() as usize;
811                for dim in 1..=MAX_DIM {
812                    for _ in 0..TRIALS {
813                        test_minmax_heterogeneous_kernel::<$N, $M, _>(dim, &mut rng);
814                    }
815                }
816            }
817        };
818    }
819
820    test_minmax_heterogeneous!(minmax_heterogeneous_8x4, 8, 4, 0xb7c3d9e5f1a20864);
821    test_minmax_heterogeneous!(minmax_heterogeneous_8x2, 8, 2, 0x4e8f2c6a1d3b5079);
822    test_minmax_heterogeneous!(minmax_heterogeneous_8x1, 8, 1, 0x1b0f2c614d2a7141);
823}
diskann_quantization/minmax/vectors.rs

diskann_quantization/minmax/
vectors.rs
