uqa_storage/
vector_index.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! Vector index abstraction and an in-memory brute-force fallback.
8//!
9//! Operators (`KNNOperator`, `VectorSimilarityOperator`,
10//! `QueryPoolVectorScoreOperator`) depend only on this trait. IVF and HNSW
11//! backends slot in by implementing the same surface.
12
13use std::collections::BTreeMap;
14use std::sync::Arc;
15
16use uqa_core::{DocId, Payload, PostingEntry, PostingList};
17
18use crate::{StorageBackendError, StorageBackendResult};
19
20mod collection;
21mod config;
22mod memory_snapshot;
23pub mod query;
24pub(crate) mod retained;
25
26pub use collection::{
27    RetainedVectorIndexesBuilder, VectorIndexSource, VectorIndexes, VectorIndexesIter,
28};
29pub use config::{HNSWIndexParams, IVFIndexParams, VectorIndexOpenMode, VectorIndexSpec};
30pub use retained::{RetainedVectorIndex, RetainedVectorIndexBuilder};
31
32#[cfg(test)]
33mod physical_snapshots;
34
35pub fn validate_vector_values(dimensions: u32, vector: &[f32]) -> StorageBackendResult<()> {
36    let dimensions = usize::try_from(dimensions).map_err(|_| {
37        StorageBackendError::Other(format!(
38            "vector dimension {dimensions} exceeds the platform usize range"
39        ))
40    })?;
41    if vector.len() != dimensions {
42        return Err(StorageBackendError::Other(format!(
43            "vector dimension mismatch: expected {dimensions}, got {}",
44            vector.len()
45        )));
46    }
47    if let Some((index, value)) = vector
48        .iter()
49        .copied()
50        .enumerate()
51        .find(|(_, value)| !value.is_finite())
52    {
53        return Err(StorageBackendError::Other(format!(
54            "vector component {index} must be finite, got {value}"
55        )));
56    }
57    Ok(())
58}
59
60fn checked_vector_count(counts: impl IntoIterator<Item = usize>) -> StorageBackendResult<usize> {
61    counts.into_iter().try_fold(0_usize, |total, count| {
62        total
63            .checked_add(count)
64            .ok_or_else(|| StorageBackendError::Other("vector count overflow".into()))
65    })
66}
67
68fn validate_threshold(threshold: f32) -> StorageBackendResult<()> {
69    if threshold.is_finite() {
70        Ok(())
71    } else {
72        Err(StorageBackendError::Other(format!(
73            "vector similarity threshold must be finite, got {threshold}"
74        )))
75    }
76}
77
78pub fn select_top_k_scored(scored: &mut Vec<(DocId, f32)>, k: usize) {
79    if scored.len() > k {
80        scored.select_nth_unstable_by(k, |a, b| b.1.total_cmp(&a.1).then_with(|| a.0.cmp(&b.0)));
81        scored.truncate(k);
82    }
83}
84
85/// Collapse tensor-vector scores to the best score for each document without
86/// allocating one tree node per candidate. The final sort performed by each
87/// caller restores the posting-list invariant after top-k selection.
88pub(crate) fn deduplicate_scored_values(scored: &mut [(DocId, f32)]) -> usize {
89    if scored.len() < 2 {
90        return scored.len();
91    }
92    scored.sort_unstable_by_key(|(doc_id, _)| *doc_id);
93    let mut write = 1;
94    for read in 1..scored.len() {
95        let (doc_id, score) = scored[read];
96        if scored[write - 1].0 == doc_id {
97            scored[write - 1].1 = scored[write - 1].1.max(score);
98        } else {
99            scored[write] = (doc_id, score);
100            write += 1;
101        }
102    }
103    write
104}
105
106pub(crate) fn vector_norm(vector: &[f32]) -> f32 {
107    let mut squared_norm = 0.0_f32;
108    for value in vector {
109        squared_norm += value * value;
110    }
111    squared_norm.sqrt()
112}
113
114/// Cosine similarity when both vector norms were computed once outside the
115/// candidate loop. This preserves the raw-vector score while avoiding two
116/// norm reductions and two square roots for every candidate.
117pub(crate) fn cosine_similarity_with_norms(a: &[f32], b: &[f32], norm_a: f32, norm_b: f32) -> f32 {
118    if a.len() != b.len() || a.is_empty() || norm_a == 0.0 || norm_b == 0.0 {
119        return 0.0;
120    }
121    let mut dot = 0.0_f32;
122    for (x, y) in a.iter().zip(b) {
123        dot += x * y;
124    }
125    dot / (norm_a * norm_b)
126}
127
128/// Cosine similarity between two equal-length vectors. Returns `0.0` when
129/// either vector has zero norm or the dimensions differ.
130///
131/// Arithmetic stays in `f32` end-to-end so the result is bit-equal to
132/// the reference `NumPy` implementation (`np.dot(q, v) / (||q|| * ||v||)`
133/// over `float32` arrays).
134pub fn cosine_similarity(a: &[f32], b: &[f32]) -> f32 {
135    if a.len() != b.len() || a.is_empty() {
136        return 0.0;
137    }
138    let mut dot = 0.0f32;
139    let mut norm_a = 0.0f32;
140    let mut norm_b = 0.0f32;
141    for (x, y) in a.iter().zip(b.iter()) {
142        dot += x * y;
143        norm_a += x * x;
144        norm_b += y * y;
145    }
146    if norm_a == 0.0 || norm_b == 0.0 {
147        return 0.0;
148    }
149    dot / (norm_a.sqrt() * norm_b.sqrt())
150}
151
152pub trait VectorIndex: Send + Sync {
153    fn dimensions(&self) -> u32;
154    fn index_kind(&self) -> &'static str {
155        "vector"
156    }
157    fn add(&mut self, doc_id: DocId, vector: Vec<f32>) -> StorageBackendResult<()>;
158    fn add_many(&mut self, doc_id: DocId, vectors: Vec<Vec<f32>>) -> StorageBackendResult<()>;
159    fn delete(&mut self, doc_id: DocId) -> StorageBackendResult<()>;
160    fn clear(&mut self) -> StorageBackendResult<()>;
161    fn search_knn(&self, query: &[f32], k: usize) -> StorageBackendResult<PostingList>;
162    fn search_threshold(&self, query: &[f32], threshold: f32) -> StorageBackendResult<PostingList>;
163    /// Request a controlled search. The default checks cancellation and delegates to the provider's existing read boundary; in-memory physical algorithms override it to charge query workspace to the supplied allowance.
164    fn search_knn_with_control(
165        &self,
166        query: &[f32],
167        k: usize,
168        control: &crate::read_control::StorageReadControl,
169    ) -> StorageBackendResult<PostingList> {
170        control.check()?;
171        self.search_knn(query, k)
172    }
173    /// Controlled threshold search keeps the same result and caller-owned posting boundary as ordinary search.
174    fn search_threshold_with_control(
175        &self,
176        query: &[f32],
177        threshold: f32,
178        control: &crate::read_control::StorageReadControl,
179    ) -> StorageBackendResult<PostingList> {
180        control.check()?;
181        self.search_threshold(query, threshold)
182    }
183    fn count(&self) -> StorageBackendResult<usize>;
184
185    /// Test canonical membership at this handle's visibility boundary without searching or materializing the corpus. Empty tensor replacements have no membership. This is a maintenance probe, not a logical query observation.
186    fn contains_document(&self, _doc_id: DocId) -> StorageBackendResult<bool> {
187        Err(StorageBackendError::Other(
188            "canonical vector membership is not supported by this backend".into(),
189        ))
190    }
191
192    /// Build any auxiliary physical metadata required by this index from its
193    /// current vector contents. Brute-force indexes need no extra work;
194    /// persistent IVF and HNSW implementations use this during explicit
195    /// index creation, while restore paths deliberately skip it.
196    fn initialize(&mut self) -> StorageBackendResult<()> {
197        Ok(())
198    }
199
200    /// Read-only handle suitable for an `ExecutionContext`.
201    fn snapshot(&self) -> StorageBackendResult<Arc<dyn VectorIndex>>;
202
203    /// Capture using the caller's retention allowance and cancellation. Owners that copy index data override this method to reserve construction before allocation. The default preserves the provider's existing snapshot and its independently retained control.
204    fn snapshot_with_control(
205        &self,
206        control: &crate::read_control::StorageReadControl,
207    ) -> StorageBackendResult<Arc<dyn VectorIndex>> {
208        control.check()?;
209        self.snapshot()
210    }
211
212    /// Independent writable copy used by in-memory engine rollback. The
213    /// default keeps third-party and persistent implementations source
214    /// compatible; only indexes hosted by a memory engine need to support it.
215    fn writable_snapshot(&self) -> StorageBackendResult<Box<dyn VectorIndex>> {
216        Err(StorageBackendError::Other(
217            "writable vector-index snapshots are not supported by this backend".into(),
218        ))
219    }
220}
221
222#[derive(Debug, Clone)]
223pub struct MemoryVectorIndex {
224    dimensions: u32,
225    vectors: BTreeMap<DocId, Vec<Vec<f32>>>,
226}
227
228impl MemoryVectorIndex {
229    pub fn new(dimensions: u32) -> Self {
230        Self {
231            dimensions,
232            vectors: BTreeMap::new(),
233        }
234    }
235
236    pub fn vectors(&self) -> &BTreeMap<DocId, Vec<Vec<f32>>> {
237        &self.vectors
238    }
239}
240
241impl VectorIndex for MemoryVectorIndex {
242    fn contains_document(&self, doc_id: DocId) -> StorageBackendResult<bool> {
243        Ok(self.vectors.contains_key(&doc_id))
244    }
245
246    fn dimensions(&self) -> u32 {
247        self.dimensions
248    }
249
250    fn index_kind(&self) -> &'static str {
251        "memory-bruteforce"
252    }
253
254    fn add(&mut self, doc_id: DocId, vector: Vec<f32>) -> StorageBackendResult<()> {
255        self.validate_dimensions(&vector)?;
256        self.vectors.insert(doc_id, vec![vector]);
257        Ok(())
258    }
259
260    fn add_many(&mut self, doc_id: DocId, vectors: Vec<Vec<f32>>) -> StorageBackendResult<()> {
261        for vector in &vectors {
262            self.validate_dimensions(vector)?;
263        }
264        if vectors.is_empty() {
265            self.vectors.remove(&doc_id);
266        } else {
267            self.vectors.insert(doc_id, vectors);
268        }
269        Ok(())
270    }
271
272    fn delete(&mut self, doc_id: DocId) -> StorageBackendResult<()> {
273        self.vectors.remove(&doc_id);
274        Ok(())
275    }
276
277    fn clear(&mut self) -> StorageBackendResult<()> {
278        self.vectors.clear();
279        Ok(())
280    }
281
282    /// Brute-force top-k by cosine similarity, descending.
283    fn search_knn(&self, query: &[f32], k: usize) -> StorageBackendResult<PostingList> {
284        self.validate_dimensions(query)?;
285        if k == 0 || self.vectors.is_empty() {
286            return Ok(PostingList::new());
287        }
288        let mut scored: Vec<(DocId, f32)> = self
289            .vectors
290            .iter()
291            .filter_map(|(&doc_id, vectors)| best_vector_score(query, vectors).map(|s| (doc_id, s)))
292            .collect();
293        select_top_k_scored(&mut scored, k);
294        // The output of `top_k` is re-sorted by doc_id ascending so the
295        // posting list invariant holds; the score lives in the payload.
296        scored.sort_by_key(|(id, _)| *id);
297        let entries = scored
298            .into_iter()
299            .map(|(doc_id, sim)| PostingEntry::new(doc_id, Payload::with_score(f64::from(sim))))
300            .collect::<Vec<_>>();
301        Ok(PostingList::from_sorted_unchecked(entries))
302    }
303
304    /// Brute-force threshold scan: keep all docs with `cosine >= threshold`.
305    fn search_threshold(&self, query: &[f32], threshold: f32) -> StorageBackendResult<PostingList> {
306        self.validate_dimensions(query)?;
307        validate_threshold(threshold)?;
308        let mut entries: Vec<PostingEntry> = self
309            .vectors
310            .iter()
311            .filter_map(|(&doc_id, vectors)| {
312                let sim = best_vector_score(query, vectors)?;
313                if sim >= threshold {
314                    Some(PostingEntry::new(
315                        doc_id,
316                        Payload::with_score(f64::from(sim)),
317                    ))
318                } else {
319                    None
320                }
321            })
322            .collect();
323        // The BTreeMap iteration is already doc_id-ascending, so the
324        // filter preserves the invariant.
325        entries.sort_by_key(|e| e.doc_id);
326        Ok(PostingList::from_sorted_unchecked(entries))
327    }
328
329    fn count(&self) -> StorageBackendResult<usize> {
330        checked_vector_count(self.vectors.values().map(Vec::len))
331    }
332
333    fn snapshot(&self) -> StorageBackendResult<Arc<dyn VectorIndex>> {
334        Ok(Arc::new(self.clone()))
335    }
336
337    fn snapshot_with_control(
338        &self,
339        control: &crate::read_control::StorageReadControl,
340    ) -> StorageBackendResult<Arc<dyn VectorIndex>> {
341        memory_snapshot::capture(self, control)
342    }
343
344    fn writable_snapshot(&self) -> StorageBackendResult<Box<dyn VectorIndex>> {
345        Ok(Box::new(self.clone()))
346    }
347}
348
349impl MemoryVectorIndex {
350    fn validate_dimensions(&self, vector: &[f32]) -> StorageBackendResult<()> {
351        validate_vector_values(self.dimensions, vector)
352    }
353}
354
355fn best_vector_score(query: &[f32], vectors: &[Vec<f32>]) -> Option<f32> {
356    vectors
357        .iter()
358        .map(|vector| cosine_similarity(query, vector))
359        .max_by(f32::total_cmp)
360}
361
362#[cfg(test)]
363mod tests {
364    use super::*;
365
366    #[test]
367    fn canonical_membership_tracks_tensor_replacement_and_retained_snapshots() {
368        let indexes: [Box<dyn VectorIndex>; 3] = [
369            Box::new(MemoryVectorIndex::new(2)),
370            Box::new(crate::IVFIndex::new(2)),
371            Box::new(crate::HNSWIndex::new(2)),
372        ];
373        for mut index in indexes {
374            assert!(!index.contains_document(1).unwrap());
375            index
376                .add_many(1, vec![vec![1.0, 0.0], vec![0.0, 1.0]])
377                .unwrap();
378            let retained = index.snapshot().unwrap();
379            assert!(index.contains_document(1).unwrap());
380            assert!(!index.contains_document(2).unwrap());
381            index.add_many(1, Vec::new()).unwrap();
382            assert!(!index.contains_document(1).unwrap());
383            assert!(retained.contains_document(1).unwrap());
384            index.add(1, vec![0.5, 0.5]).unwrap();
385            index.delete(1).unwrap();
386            assert!(!index.contains_document(1).unwrap());
387        }
388    }
389
390    fn approx_eq(a: f32, b: f32, eps: f32) {
391        assert!((a - b).abs() < eps, "expected {a} ~ {b} within {eps}");
392    }
393
394    #[test]
395    fn cosine_identity_is_one() {
396        let v = vec![1.0, 2.0, 3.0];
397        approx_eq(cosine_similarity(&v, &v), 1.0, 1e-6);
398    }
399
400    #[test]
401    fn vector_count_overflow_is_reported() {
402        let error = checked_vector_count([usize::MAX, 1]).unwrap_err();
403        assert!(error.to_string().contains("vector count overflow"));
404    }
405
406    #[test]
407    fn cosine_orthogonal_is_zero() {
408        let a = vec![1.0, 0.0];
409        let b = vec![0.0, 1.0];
410        approx_eq(cosine_similarity(&a, &b), 0.0, 1e-6);
411    }
412
413    #[test]
414    fn cosine_zero_norm_is_zero() {
415        let a = vec![0.0, 0.0];
416        let b = vec![1.0, 1.0];
417        approx_eq(cosine_similarity(&a, &b), 0.0, 1e-6);
418    }
419
420    #[test]
421    fn knn_orders_by_similarity_descending_then_doc_id() {
422        let mut idx = MemoryVectorIndex::new(2);
423        idx.add(1, vec![1.0, 0.0]).unwrap();
424        idx.add(2, vec![0.5, 0.5]).unwrap();
425        idx.add(3, vec![0.0, 1.0]).unwrap();
426        let pl = idx.search_knn(&[1.0, 0.0], 2).unwrap();
427        let docs: Vec<_> = pl.iter().map(|e| e.doc_id).collect();
428        // posting list is doc_id-sorted but the top-2 should be {1, 2}.
429        assert_eq!(docs, vec![1, 2]);
430        let entry1 = pl.get_entry(1).unwrap();
431        let entry2 = pl.get_entry(2).unwrap();
432        assert!(entry1.payload.score > entry2.payload.score);
433    }
434
435    #[test]
436    fn partial_top_k_keeps_deterministic_doc_id_ties() {
437        let mut scored = vec![(10, 0.5), (3, 0.9), (1, 0.9), (8, 0.7), (2, 0.1)];
438        select_top_k_scored(&mut scored, 2);
439        scored.sort_by_key(|(doc_id, _)| *doc_id);
440        assert_eq!(scored, vec![(1, 0.9), (3, 0.9)]);
441    }
442
443    #[test]
444    fn precomputed_norm_cosine_matches_reference_bits() {
445        let a = [0.25, -3.0, 1.5, 8.0];
446        let b = [2.0, 0.75, -4.0, 0.5];
447        let expected = cosine_similarity(&a, &b);
448        let actual = cosine_similarity_with_norms(&a, &b, vector_norm(&a), vector_norm(&b));
449        assert_eq!(actual.to_bits(), expected.to_bits());
450    }
451
452    #[test]
453    fn score_deduplication_keeps_best_tensor_vector() {
454        let mut scored = vec![(7, 0.3), (2, 0.8), (7, 0.9), (2, 0.4), (9, -0.2)];
455        let count = deduplicate_scored_values(&mut scored);
456        scored.truncate(count);
457        assert_eq!(scored, vec![(2, 0.8), (7, 0.9), (9, -0.2)]);
458    }
459
460    #[test]
461    fn threshold_filters_below_cutoff() {
462        let mut idx = MemoryVectorIndex::new(2);
463        idx.add(1, vec![1.0, 0.0]).unwrap();
464        idx.add(2, vec![0.5, 0.5]).unwrap();
465        idx.add(3, vec![0.0, 1.0]).unwrap();
466        let pl = idx.search_threshold(&[1.0, 0.0], 0.5).unwrap();
467        let docs: Vec<_> = pl.iter().map(|e| e.doc_id).collect();
468        assert_eq!(docs, vec![1, 2]);
469    }
470
471    #[test]
472    fn delete_removes_vector() {
473        let mut idx = MemoryVectorIndex::new(2);
474        idx.add(1, vec![1.0, 0.0]).unwrap();
475        idx.delete(1).unwrap();
476        assert_eq!(idx.count().unwrap(), 0);
477    }
478
479    #[test]
480    fn non_finite_vectors_queries_and_thresholds_are_errors() {
481        let mut idx = MemoryVectorIndex::new(2);
482        assert!(idx.add(1, vec![f32::NAN, 0.0]).is_err());
483        idx.add(1, vec![1.0, 0.0]).unwrap();
484        assert!(idx.search_knn(&[f32::INFINITY, 0.0], 1).is_err());
485        assert!(idx
486            .search_threshold(&[1.0, 0.0], f32::NEG_INFINITY)
487            .is_err());
488    }
489}
uqa_storage/vector_index.rs

uqa_storage/
vector_index.rs
