diskann_quantization/minmax/multi/
kernel.rs

1/*
2 * Copyright (c) Microsoft Corporation.
3 * Licensed under the MIT license.
4 */
5
6use crate::{
7    bits::{Representation, Unsigned},
8    minmax::MinMaxMeta,
9    multi_vector::{MatRef, MaxSimError},
10};
11
12/// An owned MinMax query with a backend selected at construction.
13#[derive(Debug)]
14pub struct MinMaxMaxSim<const B_BITS: usize>
15where
16    Unsigned: Representation<B_BITS>,
17{
18    pub(super) inner: Box<dyn MinMaxMaxSimKernel<B_BITS>>,
19}
20
21impl<const B_BITS: usize> MinMaxMaxSim<B_BITS>
22where
23    Unsigned: Representation<B_BITS>,
24{
25    pub fn nrows(&self) -> usize {
26        self.inner.nrows()
27    }
28
29    /// Overwrite scores, ignoring NaN inner products. Empty documents and query rows
30    /// with only NaN or negative-infinite inner products produce f32::MAX.
31    /// Nonempty zero-dimensional documents produce zero.
32    ///
33    /// # Errors
34    ///
35    /// Returns an error for a wrong output length or document dimension,
36    /// without modifying the output.
37    pub fn compute_max_sim(
38        &self,
39        document: MatRef<'_, MinMaxMeta<B_BITS>>,
40        scores: &mut [f32],
41    ) -> Result<(), MaxSimError> {
42        self.inner.compute_max_sim(document, scores)
43    }
44}
45
46// Closed implementation boundary: callers never instantiate the matrix backends.
47pub(super) trait MinMaxMaxSimKernel<const B_BITS: usize>:
48    Send + Sync + std::fmt::Debug + 'static
49where
50    Unsigned: Representation<B_BITS>,
51{
52    fn nrows(&self) -> usize;
53    fn compute_max_sim(
54        &self,
55        document: MatRef<'_, MinMaxMeta<B_BITS>>,
56        scores: &mut [f32],
57    ) -> Result<(), MaxSimError>;
58
59    #[cfg(test)]
60    fn inner_type_id(&self) -> std::any::TypeId {
61        std::any::TypeId::of::<Self>()
62    }
63}
64
65/// Validate shape and handle empty inputs; return whether scoring is still needed.
66pub(super) fn validate<const B_BITS: usize>(
67    rows: usize,
68    dim: usize,
69    document: MatRef<'_, MinMaxMeta<B_BITS>>,
70    scores: &mut [f32],
71) -> Result<bool, MaxSimError>
72where
73    Unsigned: Representation<B_BITS>,
74{
75    if scores.len() != rows {
76        return Err(MaxSimError::InvalidBufferLength(scores.len(), rows));
77    }
78    if document.repr().intrinsic_dim() != dim {
79        return Err(MaxSimError::UnequalDim(
80            document.repr().intrinsic_dim(),
81            dim,
82        ));
83    }
84    if rows == 0 {
85        return Ok(false);
86    }
87    if document.num_vectors() == 0 {
88        scores.fill(f32::MAX);
89        return Ok(false);
90    }
91    if dim == 0 {
92        scores.fill(0.0);
93        return Ok(false);
94    }
95    Ok(true)
96}
diskann_quantization/minmax/multi/kernel.rs

diskann_quantization/minmax/multi/
kernel.rs
