1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
/*
* Copyright (c) Microsoft Corporation.
* Licensed under the MIT license.
*/
//! Utilities for supporting full precision, quantized, and mixed distance computations.
/// A temporary adaptor to promote the error types for fallible distance functions into
/// panics until DiskANN gets proper support for such fallible functions.
#[derive(Debug, Clone)]
pub struct UnwrapErr<T, E>(T, std::marker::PhantomData<E>);
impl<T, E> UnwrapErr<T, E> {
pub fn new(v: T) -> Self {
Self(v, std::marker::PhantomData)
}
pub fn into_inner(self) -> T {
self.0
}
}
impl<T, E> std::ops::Deref for UnwrapErr<T, E> {
type Target = T;
fn deref(&self) -> &T {
&self.0
}
}
impl<T, E> std::ops::DerefMut for UnwrapErr<T, E> {
fn deref_mut(&mut self) -> &mut T {
&mut self.0
}
}
impl<A, B, T, E> diskann_vector::DistanceFunction<A, B, f32> for UnwrapErr<T, E>
where
T: diskann_vector::DistanceFunction<A, B, Result<f32, E>>,
E: std::fmt::Debug,
{
fn evaluate_similarity(&self, a: A, b: B) -> f32 {
// Lint: We don't quite have full support for fallible distance functions.
#[expect(clippy::unwrap_used)]
self.0.evaluate_similarity(a, b).unwrap()
}
}
impl<A, T, E> diskann_vector::PreprocessedDistanceFunction<A, f32> for UnwrapErr<T, E>
where
T: diskann_vector::PreprocessedDistanceFunction<A, Result<f32, E>>,
E: std::fmt::Debug,
{
fn evaluate_similarity(&self, a: A) -> f32 {
// Lint: We don't quite have full support for fallible distance functions.
#[expect(clippy::unwrap_used)]
self.0.evaluate_similarity(a).unwrap()
}
}
pub mod pq {
//! Support for hybrid data types for full-precision and PQ compressed vectors.
//!
//! During hybrid PQ pruning, each candidate is represented as either a full-precision
//! vector or a PQ-compressed code. The [`Hybrid`] enum captures this duality, and the
//! remaining types adapt it to the [`workingset`](diskann::graph::workingset) framework.
use std::sync::Arc;
use diskann::utils::VectorRepr;
use diskann_utils::Reborrow;
use diskann_vector::DistanceFunction;
use crate::model::pq::{self, FixedChunkPQTable};
/// An element that is either a full-precision vector or a PQ-compressed code.
///
/// During hybrid pruning, the closest candidates receive full-precision vectors for
/// accurate distance computation, while the remaining candidates use cheaper PQ codes.
/// The [`HybridComputer`] dispatches to the appropriate distance function based on
/// which variant each operand is.
pub enum Hybrid<F, Q> {
Full(F),
Quant(Q),
}
impl<F, Q> Hybrid<F, Q> {
pub fn is_full(&self) -> bool {
matches!(self, Self::Full(_))
}
}
// NOTE: This definition always maps slices to the full-precision type and is used
// for zero-copy multi-insert compatibility.
impl<'a, F, Q> From<&'a [F]> for Hybrid<&'a [F], &'a [Q]> {
fn from(slice: &'a [F]) -> Self {
Self::Full(slice)
}
}
impl<'short, F, Q> Reborrow<'short> for Hybrid<F, Q>
where
F: Reborrow<'short>,
Q: Reborrow<'short>,
{
type Target = Hybrid<F::Target, Q::Target>;
fn reborrow(&'short self) -> Self::Target {
match self {
Self::Full(v) => Hybrid::Full(v.reborrow()),
Self::Quant(v) => Hybrid::Quant(v.reborrow()),
}
}
}
/// Distance computer that handles mixed full-precision and PQ-compressed operands.
///
/// When both operands are full-precision, the native distance function is used. When
/// at least one is quantized, the PQ distance table is used instead. Mixed pairs
/// (full vs quant) convert the full-precision side to `f32` for the PQ lookup.
pub struct HybridComputer<T>
where
T: VectorRepr,
{
quant: pq::distance::DistanceComputer<Arc<FixedChunkPQTable>>,
full: T::Distance,
}
impl<T> HybridComputer<T>
where
T: VectorRepr,
{
pub fn new(
quant: pq::distance::DistanceComputer<Arc<FixedChunkPQTable>>,
full: T::Distance,
) -> Self {
Self { quant, full }
}
}
/// The implementation of `DistanceFunction` for the hybrid computer.
impl<T> DistanceFunction<Hybrid<&[T], &[u8]>, Hybrid<&[T], &[u8]>, f32> for HybridComputer<T>
where
T: VectorRepr,
{
#[inline(always)]
fn evaluate_similarity(&self, x: Hybrid<&[T], &[u8]>, y: Hybrid<&[T], &[u8]>) -> f32 {
match x {
Hybrid::Full(x) => match y {
Hybrid::Full(y) => self.full.evaluate_similarity(x, y),
Hybrid::Quant(y) => {
// SAFETY: This can only panic when T = `MinMaxElement` and the underlying slice is ill-defined.
// we are ok with panicking in distance functions for now.
#[allow(clippy::unwrap_used)]
self.quant.evaluate_similarity(&*T::as_f32(x).unwrap(), y)
}
},
Hybrid::Quant(x) => match y {
Hybrid::Full(y) => {
// SAFETY: This can only panic when T = `MinMaxElement` and the underlying slice is ill-defined.
// we are ok with panicking in distance functions for now.
#[allow(clippy::unwrap_used)]
self.quant.evaluate_similarity(&*T::as_f32(y).unwrap(), x)
}
Hybrid::Quant(y) => self.quant.evaluate_similarity(x, y),
},
}
}
}
}