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)]
468mod minmax_vector_tests {
469    use diskann_utils::Reborrow;
470    use rand::{
471        Rng, SeedableRng,
472        distr::{Distribution, Uniform},
473        rngs::StdRng,
474    };
475
476    use super::*;
477    use crate::{alloc::GlobalAllocator, scalar::bit_scale};
478
479    fn test_minmax_compensated_vectors<const NBITS: usize, R>(dim: usize, rng: &mut R)
480    where
481        Unsigned: Representation<NBITS>,
482        InnerProduct: for<'a, 'b> PureDistanceFunction<
483                BitSlice<'a, NBITS, Unsigned>,
484                BitSlice<'b, NBITS, Unsigned>,
485                distances::MathematicalResult<u32>,
486            >,
487        InnerProduct: for<'a, 'b> PureDistanceFunction<
488                &'a [f32],
489                BitSlice<'b, NBITS, Unsigned>,
490                distances::MathematicalResult<f32>,
491            >,
492        R: Rng,
493    {
494        assert!(dim <= bit_scale::<NBITS>() as usize);
495
496        // Create two vectors with known compensation values
497        let mut v1 = Data::<NBITS>::new_boxed(dim);
498        let mut v2 = Data::<NBITS>::new_boxed(dim);
499
500        let domain = Unsigned::domain_const::<NBITS>();
501        let code_distribution = Uniform::new_inclusive(*domain.start(), *domain.end()).unwrap();
502
503        // Set bit values
504        {
505            let mut bitslice1 = v1.vector_mut();
506            let mut bitslice2 = v2.vector_mut();
507
508            for i in 0..dim {
509                bitslice1.set(i, code_distribution.sample(rng)).unwrap();
510                bitslice2.set(i, code_distribution.sample(rng)).unwrap();
511            }
512        }
513        let a_rnd = Uniform::new_inclusive(0.0, 2.0).unwrap();
514        let b_rnd = Uniform::new_inclusive(0.0, 2.0).unwrap();
515
516        // Set compensation coefficients
517        // v1: X = a1 * X' + b1
518        // v2: Y = a2 * Y' + b2
519        let a1 = a_rnd.sample(rng);
520        let b1 = b_rnd.sample(rng);
521        let a2 = a_rnd.sample(rng);
522        let b2 = b_rnd.sample(rng);
523
524        // Calculate sum of vector elements for n values
525        let sum1: f32 = (0..dim).map(|i| v1.vector().get(i).unwrap() as f32).sum();
526        let sum2: f32 = (0..dim).map(|i| v2.vector().get(i).unwrap() as f32).sum();
527
528        // Create original full-precision vectors for reference calculations
529        let mut original1 = Vec::with_capacity(dim);
530        let mut original2 = Vec::with_capacity(dim);
531
532        // Calculate the reconstructed original vectors and their norms
533        for i in 0..dim {
534            let val1 = a1 * v1.vector().get(i).unwrap() as f32 + b1;
535            let val2 = a2 * v2.vector().get(i).unwrap() as f32 + b2;
536            original1.push(val1);
537            original2.push(val2);
538        }
539
540        // Calculate squared norms
541        let norm1_squared: f32 = original1.iter().map(|x| x * x).sum();
542        let norm2_squared: f32 = original2.iter().map(|x| x * x).sum();
543
544        // Set compensation coefficients
545        v1.set_meta(MinMaxCompensation {
546            a: a1,
547            b: b1,
548            n: a1 * sum1,
549            norm_squared: norm1_squared,
550            dim: dim as u32,
551        });
552
553        v2.set_meta(MinMaxCompensation {
554            a: a2,
555            b: b2,
556            n: a2 * sum2,
557            norm_squared: norm2_squared,
558            dim: dim as u32,
559        });
560
561        // Calculate raw integer dot product
562        let expected_ip = (0..dim).map(|i| original1[i] * original2[i]).sum::<f32>();
563
564        // Test inner product with f32
565        let computed_ip_f32: distances::Result<f32> =
566            MinMaxIP::evaluate(v1.reborrow(), v2.reborrow());
567        let computed_ip_f32 = computed_ip_f32.unwrap();
568        assert!(
569            (expected_ip - (-computed_ip_f32)).abs() / expected_ip.abs() < 1e-3,
570            "Inner product (f32) failed: expected {}, got {} on dim : {}",
571            -expected_ip,
572            computed_ip_f32,
573            dim
574        );
575
576        // Expected L2 distance = |X|² + |Y|² - 2<X,Y>
577        let expected_l2 = (0..dim)
578            .map(|i| original1[i] - original2[i])
579            .map(|x| x.powf(2.0))
580            .sum::<f32>();
581
582        // Test L2 distance with f32
583        let computed_l2_f32: distances::Result<f32> =
584            MinMaxL2Squared::evaluate(v1.reborrow(), v2.reborrow());
585        let computed_l2_f32 = computed_l2_f32.unwrap();
586        assert!(
587            ((computed_l2_f32 - expected_l2).abs() / expected_l2) < 1e-3,
588            "L2 distance (f32) failed: expected {}, got {} on dim : {}",
589            expected_l2,
590            computed_l2_f32,
591            dim
592        );
593
594        let expected_cosine = 1.0 - expected_ip / (norm1_squared.sqrt() * norm2_squared.sqrt());
595
596        let computed_cosine: distances::Result<f32> =
597            MinMaxCosine::evaluate(v1.reborrow(), v2.reborrow());
598        let computed_cosine = computed_cosine.unwrap();
599
600        {
601            let passed = (computed_cosine - expected_cosine).abs() < 1e-6
602                || ((computed_cosine - expected_cosine).abs() / expected_cosine) < 1e-3;
603
604            assert!(
605                passed,
606                "Cosine distance (f32) failed: expected {}, got {} on dim : {}",
607                expected_cosine, computed_cosine, dim
608            );
609        }
610
611        let cosine_normalized: distances::Result<f32> =
612            MinMaxCosineNormalized::evaluate(v1.reborrow(), v2.reborrow());
613        let cosine_normalized = cosine_normalized.unwrap();
614        let expected_cos_normalized = 1.0 - expected_ip;
615        assert!(
616            ((expected_cos_normalized - cosine_normalized).abs() / expected_cos_normalized.abs())
617                < 1e-6,
618            "CosineNormalized distance (f32) failed: expected {}, got {} on dim : {}",
619            expected_cos_normalized,
620            cosine_normalized,
621            dim
622        );
623
624        //Calculate inner product with full precision vector
625        let mut fp_query = FullQuery::new_in(dim, GlobalAllocator).unwrap();
626        fp_query.vector_mut().copy_from_slice(&original1);
627        *fp_query.meta_mut() = FullQueryMeta {
628            norm_squared: norm1_squared,
629            sum: original1.iter().sum::<f32>(),
630        };
631
632        let fp_ip: distances::Result<f32> = MinMaxIP::evaluate(fp_query.reborrow(), v2.reborrow());
633        let fp_ip = fp_ip.unwrap();
634        assert!(
635            (expected_ip - (-fp_ip)).abs() / expected_ip.abs() < 1e-3,
636            "Inner product (f32) failed: expected {}, got {} on dim : {}",
637            -expected_ip,
638            fp_ip,
639            dim
640        );
641
642        let fp_l2: distances::Result<f32> =
643            MinMaxL2Squared::evaluate(fp_query.reborrow(), v2.reborrow());
644        let fp_l2 = fp_l2.unwrap();
645        assert!(
646            ((fp_l2 - expected_l2).abs() / expected_l2) < 1e-3,
647            "L2 distance (f32) failed: expected {}, got {} on dim : {}",
648            expected_l2,
649            computed_l2_f32,
650            dim
651        );
652
653        let fp_cosine: distances::Result<f32> =
654            MinMaxCosine::evaluate(fp_query.reborrow(), v2.reborrow());
655        let fp_cosine = fp_cosine.unwrap();
656        let diff = (fp_cosine - expected_cosine).abs();
657        assert!(
658            (diff / expected_cosine) < 1e-3 || diff <= 1e-6,
659            "Cosine distance (f32) failed: expected {}, got {} on dim : {}",
660            expected_cosine,
661            fp_cosine,
662            dim
663        );
664
665        let fp_cos_norm: distances::Result<f32> =
666            MinMaxCosineNormalized::evaluate(fp_query.reborrow(), v2.reborrow());
667        let fp_cos_norm = fp_cos_norm.unwrap();
668        assert!(
669            (((1.0 - expected_ip) - fp_cos_norm).abs() / (1.0 - expected_ip)) < 1e-3,
670            "Cosine distance (f32) failed: expected {}, got {} on dim : {}",
671            (1.0 - expected_ip),
672            fp_cos_norm,
673            dim
674        );
675
676        //Test `decompress_into` to make sure it outputs tje full-precision vector correctly.
677        let meta = v1.meta();
678        let v1_ref = DataRef::new(v1.vector(), &meta);
679        let dim = v1_ref.len();
680        let mut boxed = vec![0f32; dim + 1];
681
682        let pre = v1_ref.decompress_into(&mut boxed);
683        assert_eq!(
684            pre.unwrap_err(),
685            DecompressError::LengthMismatch(dim, dim + 1)
686        );
687        let pre = v1_ref.decompress_into(&mut boxed[..dim - 1]);
688        assert_eq!(
689            pre.unwrap_err(),
690            DecompressError::LengthMismatch(dim, dim - 1)
691        );
692        let pre = v1_ref.decompress_into(&mut boxed[..dim]);
693        assert!(pre.is_ok());
694
695        boxed
696            .iter()
697            .zip(original1.iter())
698            .for_each(|(x, y)| assert!((*x - *y).abs() <= 1e-6));
699
700        // Verify `read_dimension` is correct.
701        let mut bytes = vec![0u8; Data::canonical_bytes(dim)];
702        let mut data = DataMutRef::from_canonical_front_mut(bytes.as_mut_slice(), dim).unwrap();
703        data.set_meta(meta);
704
705        let pre = MinMaxCompensation::read_dimension(&bytes);
706        assert!(pre.is_ok());
707        let read_dim = pre.unwrap();
708        assert_eq!(read_dim, dim);
709
710        let pre = MinMaxCompensation::read_dimension(&[0_u8; 2]);
711        assert_eq!(pre.unwrap_err(), MetaParseError::NotCanonical(2));
712    }
713
714    cfg_if::cfg_if! {
715        if #[cfg(miri)] {
716            // The max dim does not need to be as high for `CompensatedVectors` because they
717            // defer their distance function implementation to `BitSlice`, which is more
718            // heavily tested.
719            const TRIALS: usize = 2;
720        } else {
721            const TRIALS: usize = 10;
722        }
723    }
724
725    macro_rules! test_minmax_compensated {
726        ($name:ident, $nbits:literal, $seed:literal) => {
727            #[test]
728            fn $name() {
729                let mut rng = StdRng::seed_from_u64($seed);
730                for dim in 1..(bit_scale::<$nbits>() as usize) {
731                    for _ in 0..TRIALS {
732                        test_minmax_compensated_vectors::<$nbits, _>(dim, &mut rng);
733                    }
734                }
735            }
736        };
737    }
738    test_minmax_compensated!(unsigned_minmax_compensated_test_u1, 1, 0xa32d5658097a1c35);
739    test_minmax_compensated!(unsigned_minmax_compensated_test_u2, 2, 0xaedf3d2a223b7b77);
740    test_minmax_compensated!(unsigned_minmax_compensated_test_u4, 4, 0xf60c0c8d1aadc126);
741    test_minmax_compensated!(unsigned_minmax_compensated_test_u8, 8, 0x09fa14c42a9d7d98);
742}
diskann_quantization/minmax/vectors.rs

diskann_quantization/minmax/
vectors.rs
