diskann_quantization/minmax/
vectors.rs

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

diskann_quantization/minmax/
vectors.rs
