weavatrix_search_vector/quantized/
compact.rs

1use super::storage::{QuantizedStorage, binary_word};
2use crate::config::{DistanceMetric, IndexConfig};
3use crate::error::SearchError;
4use crate::hit::SearchHit;
5use crate::hnsw::VectorIndex;
6use crate::parallel;
7use crate::vector::{Candidate, squared_norm};
8use std::collections::BinaryHeap;
9
10/// Compact vector representation.
11#[derive(Debug, Clone, Copy, PartialEq, Eq)]
12#[non_exhaustive]
13pub enum QuantizationKind {
14    BFloat16,
15    Float16,
16    Float8E4M3,
17    Int8,
18    Binary,
19}
20
21/// Exact scan over a selected compact representation.
22///
23/// `Int8` and `Binary` require cosine distance. Floating-point formats support
24/// every [`DistanceMetric`]. Returned compact distances can be re-scored with
25/// [`Self::search_rerank`].
26#[derive(Debug)]
27pub struct QuantizedIndex {
28    pub(crate) config: IndexConfig,
29    pub(crate) kind: QuantizationKind,
30    pub(crate) keys: Vec<u64>,
31    pub(crate) storage: QuantizedStorage,
32    pub(crate) squared_norms: Vec<f32>,
33    pub(crate) words_per_vector: usize,
34}
35
36impl QuantizedIndex {
37    /// Builds a compact exact-scan index.
38    ///
39    /// # Errors
40    ///
41    /// Returns typed config, vector, key, capacity, or allocation errors.
42    pub fn build(
43        config: IndexConfig,
44        kind: QuantizationKind,
45        vectors: &[(u64, &[f32])],
46    ) -> Result<Self, SearchError> {
47        config.validate()?;
48        if matches!(kind, QuantizationKind::Int8 | QuantizationKind::Binary)
49            && config.metric != DistanceMetric::Cosine
50        {
51            return Err(SearchError::InvalidConfig(
52                "int8 and binary quantization require cosine distance",
53            ));
54        }
55        let mut order = (0..vectors.len()).collect::<Vec<_>>();
56        order.sort_unstable_by_key(|index| vectors[*index].0);
57        for pair in order.windows(2) {
58            if vectors[pair[0]].0 == vectors[pair[1]].0 {
59                return Err(SearchError::DuplicateKey(vectors[pair[0]].0));
60            }
61        }
62        let elements = config
63            .dimensions
64            .checked_mul(vectors.len())
65            .ok_or(SearchError::CapacityOverflow)?;
66        let words_per_vector = config.dimensions.div_ceil(64);
67        let mut keys = Vec::new();
68        keys.try_reserve_exact(vectors.len())
69            .map_err(|_| SearchError::AllocationFailed)?;
70        let mut squared_norms = Vec::new();
71        squared_norms
72            .try_reserve_exact(vectors.len())
73            .map_err(|_| SearchError::AllocationFailed)?;
74        let mut storage = QuantizedStorage::with_capacity(
75            kind,
76            elements,
77            words_per_vector
78                .checked_mul(vectors.len())
79                .ok_or(SearchError::CapacityOverflow)?,
80        )?;
81        for source in order {
82            let (key, vector) = vectors[source];
83            if vector.len() != config.dimensions {
84                return Err(SearchError::DimensionMismatch {
85                    expected: config.dimensions,
86                    actual: vector.len(),
87                    vector: Some(source),
88                });
89            }
90            let norm_squared = squared_norm(vector, Some(source))?;
91            if config.metric == DistanceMetric::Cosine && norm_squared == 0.0 {
92                return Err(SearchError::ZeroVector {
93                    vector: Some(source),
94                });
95            }
96            let inverse = if config.metric == DistanceMetric::Cosine {
97                norm_squared.sqrt().recip()
98            } else {
99                1.0
100            };
101            keys.push(key);
102            let quantized_norm = storage.push_vector(vector, inverse, words_per_vector)?;
103            squared_norms.push(quantized_norm);
104        }
105        Ok(Self {
106            config,
107            kind,
108            keys,
109            storage,
110            squared_norms,
111            words_per_vector,
112        })
113    }
114
115    #[must_use]
116    pub fn len(&self) -> usize {
117        self.keys.len()
118    }
119
120    #[must_use]
121    pub fn is_empty(&self) -> bool {
122        self.keys.is_empty()
123    }
124
125    #[must_use]
126    pub const fn dimensions(&self) -> usize {
127        self.config.dimensions
128    }
129
130    #[must_use]
131    pub const fn kind(&self) -> QuantizationKind {
132        self.kind
133    }
134
135    #[must_use]
136    pub const fn config(&self) -> &IndexConfig {
137        &self.config
138    }
139
140    #[must_use]
141    pub fn estimated_memory_bytes(&self) -> usize {
142        self.keys
143            .capacity()
144            .saturating_mul(std::mem::size_of::<u64>())
145            .saturating_add(
146                self.squared_norms
147                    .capacity()
148                    .saturating_mul(std::mem::size_of::<f32>()),
149            )
150            .saturating_add(self.storage.estimated_bytes())
151    }
152
153    /// Searches all compact vectors.
154    ///
155    /// # Errors
156    ///
157    /// Returns a typed query validation or allocation error.
158    pub fn search(&self, query: &[f32], count: usize) -> Result<Vec<SearchHit>, SearchError> {
159        if query.len() != self.dimensions() {
160            return Err(SearchError::DimensionMismatch {
161                expected: self.dimensions(),
162                actual: query.len(),
163                vector: None,
164            });
165        }
166        let query_squared_norm = squared_norm(query, None)?;
167        if self.config.metric == DistanceMetric::Cosine && query_squared_norm == 0.0 {
168            return Err(SearchError::ZeroVector { vector: None });
169        }
170        let limit = count.min(self.len());
171        if limit == 0 {
172            return Ok(Vec::new());
173        }
174        let mut best = BinaryHeap::with_capacity(limit);
175        for index in 0..self.len() {
176            let distance = self.compact_distance(index, query, query_squared_norm);
177            let candidate = Candidate::new(distance, index);
178            if best.len() < limit {
179                best.push(candidate);
180            } else if best
181                .peek()
182                .is_some_and(|worst| candidate.cmp(worst).is_lt())
183            {
184                best.pop();
185                best.push(candidate);
186            }
187        }
188        let mut candidates = best.into_vec();
189        candidates.sort_unstable();
190        Ok(candidates
191            .into_iter()
192            .map(|candidate| SearchHit {
193                key: self.keys[candidate.index()],
194                distance: candidate.distance,
195            })
196            .collect())
197    }
198
199    /// Searches independent queries with bounded workers.
200    ///
201    /// # Errors
202    ///
203    /// Returns the first input or worker error.
204    pub fn search_batch(
205        &self,
206        queries: &[&[f32]],
207        count: usize,
208    ) -> Result<Vec<Vec<SearchHit>>, SearchError> {
209        parallel::search_batch(queries, self.config.query_threads, |query| {
210            self.search(query, count)
211        })
212    }
213
214    /// Retrieves compact candidates and re-scores them against an f32 index.
215    ///
216    /// # Errors
217    ///
218    /// Returns a typed dimension, query, key, or allocation error.
219    pub fn search_rerank(
220        &self,
221        oracle: &VectorIndex,
222        query: &[f32],
223        count: usize,
224        candidates: usize,
225    ) -> Result<Vec<SearchHit>, SearchError> {
226        if oracle.dimensions() != self.dimensions() || oracle.config().metric != self.config.metric
227        {
228            return Err(SearchError::InvalidConfig(
229                "rerank oracle config differs from quantized index",
230            ));
231        }
232        let candidate_count = candidates.max(count).min(self.len());
233        let compact = self.search(query, candidate_count)?;
234        oracle.search_exact_filtered(query, count, |key| compact.iter().any(|hit| hit.key == key))
235    }
236
237    #[allow(clippy::cast_precision_loss)]
238    fn compact_distance(&self, index: usize, query: &[f32], query_squared_norm: f32) -> f32 {
239        match &self.storage {
240            QuantizedStorage::Binary(words) => {
241                let start = index * self.words_per_vector;
242                let mut different = 0_u32;
243                for word_index in 0..self.words_per_vector {
244                    let query_word = binary_word(query, word_index);
245                    different = different
246                        .saturating_add((words[start + word_index] ^ query_word).count_ones());
247                }
248                2.0 * different as f32 / self.dimensions() as f32
249            }
250            storage => {
251                let start = index * self.dimensions();
252                let mut dot = 0.0_f32;
253                for (dimension, query_value) in query.iter().copied().enumerate() {
254                    dot += storage.value(start + dimension) * query_value;
255                }
256                match self.config.metric {
257                    DistanceMetric::Cosine => {
258                        let denominator = (self.squared_norms[index] * query_squared_norm).sqrt();
259                        (1.0 - dot / denominator).clamp(0.0, 2.0)
260                    }
261                    DistanceMetric::Dot => -dot,
262                    DistanceMetric::SquaredEuclidean => self.squared_norms[index]
263                        .mul_add(1.0, query_squared_norm - 2.0 * dot)
264                        .max(0.0),
265                }
266            }
267        }
268    }
269}
weavatrix_search_vector/quantized/compact.rs

weavatrix_search_vector/quantized/
compact.rs
