uqa_storage/diskann_index/
metric.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7use uqa_core::memory::BudgetedVec;
8
9use crate::read_control::StorageReadControl;
10use crate::vector_index::validate_vector_values_controlled;
11use crate::{StorageBackendError, StorageBackendResult};
12
13#[derive(Debug, Clone, Copy, PartialEq, Eq)]
14pub enum ExactVectorReason {
15    ZeroNorm,
16    NonFiniteNorm,
17}
18
19/// Classification applies equally to indexed vectors and queries. Exact inputs retain their raw canonical scoring path.
20#[derive(Debug)]
21pub enum NavigationInput {
22    Navigable(NavigationVector),
23    Exact(ExactVectorReason),
24}
25
26#[derive(Debug)]
27pub struct NavigationVector {
28    coordinates: BudgetedVec<f64>,
29}
30
31/// Nonnegative squared Euclidean distance, with no conversion to a score or probability carrier.
32#[derive(Debug, Clone, Copy, PartialEq, PartialOrd)]
33pub struct SquaredNavigationDistance(f64);
34
35impl SquaredNavigationDistance {
36    pub fn get(self) -> f64 {
37        self.0
38    }
39}
40
41impl NavigationInput {
42    pub fn from_raw(
43        dimensions: u32,
44        raw: &[f32],
45        control: &StorageReadControl,
46    ) -> StorageBackendResult<Self> {
47        let (canonical_norm, norm) = norms(dimensions, raw, control)?;
48        if let Some(reason) = exact_reason(canonical_norm) {
49            return Ok(Self::Exact(reason));
50        }
51        let mut coordinates = BudgetedVec::new(control.memory());
52        coordinates.reserve(raw.len())?;
53        for (offset, &value) in raw.iter().enumerate() {
54            checkpoint(offset, control)?;
55            coordinates.push(f64::from(value) / norm)?;
56        }
57        control.check()?;
58        Ok(Self::Navigable(NavigationVector { coordinates }))
59    }
60}
61
62pub(super) fn exact_reason(canonical_norm: f32) -> Option<ExactVectorReason> {
63    if canonical_norm == 0.0 {
64        Some(ExactVectorReason::ZeroNorm)
65    } else if !canonical_norm.is_finite() {
66        Some(ExactVectorReason::NonFiniteNorm)
67    } else {
68        None
69    }
70}
71
72pub(super) fn norms(
73    dimensions: u32,
74    raw: &[f32],
75    control: &StorageReadControl,
76) -> StorageBackendResult<(f32, f64)> {
77    validate_vector_values_controlled(dimensions, raw, Some(control))?;
78    if dimensions == 0 {
79        return Err(StorageBackendError::Other(
80            "DiskANN dimensions must be positive".into(),
81        ));
82    }
83    let mut canonical_squared_norm = 0.0_f32;
84    let mut squared_norm = 0.0_f64;
85    for (offset, &value) in raw.iter().enumerate() {
86        checkpoint(offset, control)?;
87        canonical_squared_norm += value * value;
88        let value = f64::from(value);
89        squared_norm += value * value;
90    }
91    Ok((canonical_squared_norm.sqrt(), squared_norm.sqrt()))
92}
93
94impl NavigationVector {
95    pub fn coordinates(&self) -> &[f64] {
96        &self.coordinates
97    }
98
99    pub(super) fn copy(&self, control: &StorageReadControl) -> StorageBackendResult<Self> {
100        control.check()?;
101        let mut coordinates = BudgetedVec::new(control.memory());
102        coordinates.reserve(self.coordinates.len())?;
103        for (offset, &value) in self.coordinates.iter().enumerate() {
104            checkpoint(offset, control)?;
105            coordinates.push(value)?;
106        }
107        control.check()?;
108        Ok(Self { coordinates })
109    }
110
111    /// Refine an already decoded physical node without allocating another normalized vector.
112    pub(super) fn squared_distance_to_raw(
113        &self,
114        raw: &[f32],
115        control: &StorageReadControl,
116    ) -> StorageBackendResult<SquaredNavigationDistance> {
117        let (canonical_norm, norm) = norms(self.coordinates.len() as u32, raw, control)?;
118        if exact_reason(canonical_norm).is_some() {
119            return Err(StorageBackendError::Other(
120                "DiskANN physical node is not navigable".into(),
121            ));
122        }
123        let mut distance = 0.0;
124        for (offset, (&query, &value)) in self.coordinates.iter().zip(raw).enumerate() {
125            checkpoint(offset, control)?;
126            let difference = query - f64::from(value) / norm;
127            distance += difference * difference;
128        }
129        control.check()?;
130        Ok(SquaredNavigationDistance(distance))
131    }
132
133    pub fn squared_distance(
134        &self,
135        other: &Self,
136        control: &StorageReadControl,
137    ) -> StorageBackendResult<SquaredNavigationDistance> {
138        control.check()?;
139        if self.coordinates.len() != other.coordinates.len() {
140            return Err(StorageBackendError::Other(
141                "DiskANN navigation dimensions differ".into(),
142            ));
143        }
144        Ok(SquaredNavigationDistance(squared_distance(
145            &self.coordinates,
146            &other.coordinates,
147            control,
148        )?))
149    }
150}
151
152pub(super) fn squared_distance(
153    left: &[f64],
154    right: &[f64],
155    control: &StorageReadControl,
156) -> StorageBackendResult<f64> {
157    debug_assert_eq!(left.len(), right.len());
158    let mut distance = 0.0;
159    for (offset, (left, right)) in left.iter().zip(right).enumerate() {
160        checkpoint(offset, control)?;
161        let difference = left - right;
162        distance += difference * difference;
163    }
164    Ok(distance)
165}
166
167pub(super) fn checkpoint(offset: usize, control: &StorageReadControl) -> StorageBackendResult<()> {
168    if offset.is_multiple_of(1024) {
169        control.check()?;
170    }
171    Ok(())
172}
173
174#[cfg(test)]
175mod tests;
uqa_storage/diskann_index/metric.rs

uqa_storage/diskann_index/
metric.rs
