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 canonical;
21mod collection;
22mod config;
23mod cosine;
24mod memory_snapshot;
25pub mod query;
26pub use query::statistics::{DiskANNExecutionRoute, DiskANNExecutionStats, VectorQueryResult};
27pub(crate) mod retained;
28
29pub(crate) use canonical::{
30    canonical_vectors_equal, copy_vector, ordinal_count, selected_fingerprint,
31};
32pub use canonical::{decode_vector_bytes, SelectedVectorRead, VectorRead, VectorReadSnapshot};
33pub use collection::{
34    RetainedVectorIndexesBuilder, VectorIndexSource, VectorIndexes, VectorIndexesIter,
35};
36pub use config::{
37    DiskANNAlpha, DiskANNIndexParams, HNSWIndexParams, IVFIndexParams, VectorIndexOpenMode,
38    VectorIndexSpec,
39};
40pub use cosine::cosine_similarity;
41pub(crate) use cosine::cosine_similarity_controlled;
42pub use retained::{RetainedVectorIndex, RetainedVectorIndexBuilder};
43
44#[cfg(test)]
45mod physical_snapshots;
46
47pub fn validate_vector_values(dimensions: u32, vector: &[f32]) -> StorageBackendResult<()> {
48    validate_vector_values_controlled(dimensions, vector, None)
49}
50
51/// Validate canonical dimensions and finite coordinates with periodic cancellation checks.
52pub fn validate_vector_values_controlled(
53    dimensions: u32,
54    vector: &[f32],
55    control: Option<&crate::read_control::StorageReadControl>,
56) -> StorageBackendResult<()> {
57    if let Some(control) = control {
58        control.check()?;
59    }
60    let dimensions = usize::try_from(dimensions).map_err(|_| {
61        StorageBackendError::Other(format!(
62            "vector dimension {dimensions} exceeds the platform usize range"
63        ))
64    })?;
65    if vector.len() != dimensions {
66        return Err(StorageBackendError::Other(format!(
67            "vector dimension mismatch: expected {dimensions}, got {}",
68            vector.len()
69        )));
70    }
71    for (index, value) in vector.iter().copied().enumerate() {
72        if index.is_multiple_of(1024) {
73            if let Some(control) = control {
74                control.check()?;
75            }
76        }
77        if !value.is_finite() {
78            return Err(StorageBackendError::Other(format!(
79                "vector component {index} must be finite, got {value}"
80            )));
81        }
82    }
83    Ok(())
84}
85
86fn checked_vector_count(counts: impl IntoIterator<Item = usize>) -> StorageBackendResult<usize> {
87    counts.into_iter().try_fold(0_usize, |total, count| {
88        total
89            .checked_add(count)
90            .ok_or_else(|| StorageBackendError::Other("vector count overflow".into()))
91    })
92}
93
94pub(crate) fn validate_threshold(threshold: f32) -> StorageBackendResult<()> {
95    if threshold.is_finite() {
96        Ok(())
97    } else {
98        Err(StorageBackendError::Other(format!(
99            "vector similarity threshold must be finite, got {threshold}"
100        )))
101    }
102}
103
104pub fn select_top_k_scored(scored: &mut Vec<(DocId, f32)>, k: usize) {
105    if scored.len() > k {
106        scored.select_nth_unstable_by(k, |a, b| b.1.total_cmp(&a.1).then_with(|| a.0.cmp(&b.0)));
107        scored.truncate(k);
108    }
109}
110
111/// Collapse tensor-vector scores to the best score for each document without
112/// allocating one tree node per candidate. The final sort performed by each
113/// caller restores the posting-list invariant after top-k selection.
114pub(crate) fn deduplicate_scored_values(scored: &mut [(DocId, f32)]) -> usize {
115    if scored.len() < 2 {
116        return scored.len();
117    }
118    scored.sort_unstable_by_key(|(doc_id, _)| *doc_id);
119    let mut write = 1;
120    for read in 1..scored.len() {
121        let (doc_id, score) = scored[read];
122        if scored[write - 1].0 == doc_id {
123            scored[write - 1].1 = scored[write - 1].1.max(score);
124        } else {
125            scored[write] = (doc_id, score);
126            write += 1;
127        }
128    }
129    write
130}
131
132pub(crate) fn vector_norm(vector: &[f32]) -> f32 {
133    let mut squared_norm = 0.0_f32;
134    for value in vector {
135        squared_norm += value * value;
136    }
137    squared_norm.sqrt()
138}
139
140/// Cosine similarity when both vector norms were computed once outside the
141/// candidate loop. This preserves the raw-vector score while avoiding two
142/// norm reductions and two square roots for every candidate.
143pub(crate) fn cosine_similarity_with_norms(a: &[f32], b: &[f32], norm_a: f32, norm_b: f32) -> f32 {
144    if a.len() != b.len() || a.is_empty() || norm_a == 0.0 || norm_b == 0.0 {
145        return 0.0;
146    }
147    let mut dot = 0.0_f32;
148    for (x, y) in a.iter().zip(b) {
149        dot += x * y;
150    }
151    dot / (norm_a * norm_b)
152}
153
154pub trait VectorIndex: Send + Sync {
155    fn dimensions(&self) -> u32;
156    fn index_kind(&self) -> &'static str {
157        "vector"
158    }
159    /// Metadata from the actual selected physical/canonical view, without executing a logical vector query. Unsupported methods return None.
160    fn diskann_query_metadata(
161        &self,
162        control: &crate::read_control::StorageReadControl,
163    ) -> StorageBackendResult<Option<crate::diskann_index::DiskANNQueryMetadata>> {
164        control.check()?;
165        Ok(None)
166    }
167    fn add(&mut self, doc_id: DocId, vector: Vec<f32>) -> StorageBackendResult<()>;
168    fn add_many(&mut self, doc_id: DocId, vectors: Vec<Vec<f32>>) -> StorageBackendResult<()>;
169    fn delete(&mut self, doc_id: DocId) -> StorageBackendResult<()>;
170    fn clear(&mut self) -> StorageBackendResult<()>;
171    fn search_knn(&self, query: &[f32], k: usize) -> StorageBackendResult<PostingList>;
172    fn search_threshold(&self, query: &[f32], threshold: f32) -> StorageBackendResult<PostingList>;
173    /// 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.
174    fn search_knn_with_control(
175        &self,
176        query: &[f32],
177        k: usize,
178        control: &crate::read_control::StorageReadControl,
179    ) -> StorageBackendResult<PostingList> {
180        control.check()?;
181        self.search_knn(query, k)
182    }
183    /// Controlled threshold search keeps the same result and caller-owned posting boundary as ordinary search.
184    fn search_threshold_with_control(
185        &self,
186        query: &[f32],
187        threshold: f32,
188        control: &crate::read_control::StorageReadControl,
189    ) -> StorageBackendResult<PostingList> {
190        control.check()?;
191        self.search_threshold(query, threshold)
192    }
193    /// Run one KNN invocation and retain its measured work when supported. An optional invoking control preserves the existing controlled-search boundary; absence preserves ordinary provider dispatch. Unsupported providers execute once and return no report.
194    fn search_knn_with_statistics(
195        &self,
196        query: &[f32],
197        k: usize,
198        control: Option<&crate::read_control::StorageReadControl>,
199    ) -> StorageBackendResult<VectorQueryResult> {
200        match control {
201            Some(control) => self.search_knn_with_control(query, k, control),
202            None => self.search_knn(query, k),
203        }
204        .map(VectorQueryResult::unmeasured)
205    }
206    /// Return exact threshold results and optional work from their one invocation. Report support must not introduce an additional scan or weaken the original provider's controls.
207    fn search_threshold_with_statistics(
208        &self,
209        query: &[f32],
210        threshold: f32,
211        control: Option<&crate::read_control::StorageReadControl>,
212    ) -> StorageBackendResult<VectorQueryResult> {
213        match control {
214            Some(control) => self.search_threshold_with_control(query, threshold, control),
215            None => self.search_threshold(query, threshold),
216        }
217        .map(VectorQueryResult::unmeasured)
218    }
219    fn count(&self) -> StorageBackendResult<usize>;
220
221    /// 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.
222    fn contains_document(&self, _doc_id: DocId) -> StorageBackendResult<bool> {
223        Err(StorageBackendError::Other(
224            "canonical vector membership is not supported by this backend".into(),
225        ))
226    }
227
228    /// Build any auxiliary physical metadata required by this index from its
229    /// current vector contents. Brute-force indexes need no extra work;
230    /// persistent IVF and HNSW implementations use this during explicit
231    /// index creation, while restore paths deliberately skip it.
232    fn initialize(&mut self) -> StorageBackendResult<()> {
233        Ok(())
234    }
235
236    /// Read-only handle suitable for an `ExecutionContext`.
237    fn snapshot(&self) -> StorageBackendResult<Arc<dyn VectorIndex>>;
238
239    /// 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.
240    fn snapshot_with_control(
241        &self,
242        control: &crate::read_control::StorageReadControl,
243    ) -> StorageBackendResult<Arc<dyn VectorIndex>> {
244        control.check()?;
245        self.snapshot()
246    }
247
248    /// Borrow an opaque retained `DiskANN` canonical source for evaluated document replacements. Capturing metadata does not execute or observe a vector query; other physical methods leave this capability absent.
249    fn diskann_read_snapshot(
250        &self,
251        control: &crate::read_control::StorageReadControl,
252    ) -> StorageBackendResult<Option<crate::diskann_index::DiskANNReadSnapshot>> {
253        control.check()?;
254        Ok(None)
255    }
256
257    /// Retain the actual canonical values independently of physical mutation provenance. Capturing this source is metadata access, not a logical vector query.
258    fn vector_read_snapshot(
259        &self,
260        control: &crate::read_control::StorageReadControl,
261    ) -> StorageBackendResult<Option<VectorReadSnapshot>> {
262        self.diskann_read_snapshot(control)?
263            .map(|source| {
264                uqa_core::memory::Budgeted::new(source, control.memory().empty_reservation())
265                    .into_shared()
266                    .map(|source| source as VectorReadSnapshot)
267                    .map_err(Into::into)
268            })
269            .transpose()
270    }
271
272    /// Check whether a SQL replacement can retain this document's canonical vectors. Unsupported providers return false and keep the normal replacement path. Live memory indexes override this to borrow their canonical data without cloning the entire index.
273    fn matches_document_vectors(
274        &self,
275        document: DocId,
276        vectors: &[Vec<f32>],
277        control: &crate::read_control::StorageReadControl,
278    ) -> StorageBackendResult<bool> {
279        self.vector_read_snapshot(control)?
280            .map_or(Ok(false), |source| {
281                canonical_vectors_equal(source.as_ref(), document, vectors, control)
282            })
283    }
284
285    /// Keep this selected physical `DiskANN` generation while reading a different fixed canonical field view. The query owner verifies raw coverage without inventing mutation origins.
286    fn snapshot_with_vector_read(
287        &self,
288        _source: VectorReadSnapshot,
289        control: &crate::read_control::StorageReadControl,
290    ) -> StorageBackendResult<Option<Arc<dyn VectorIndex>>> {
291        control.check()?;
292        Ok(None)
293    }
294
295    /// Retain this physical `DiskANN` generation with the selected actual canonical replacements. Wrappers preserve their original read observation and retention boundary on the returned index. Unsupported methods return None without changing physical selection.
296    fn snapshot_with_diskann_changes(
297        &self,
298        _changes: &crate::diskann_index::DiskANNReadChanges,
299        control: &crate::read_control::StorageReadControl,
300    ) -> StorageBackendResult<Option<Arc<dyn VectorIndex>>> {
301        control.check()?;
302        Ok(None)
303    }
304
305    /// Independent writable copy used by in-memory engine rollback. The
306    /// default keeps third-party and persistent implementations source
307    /// compatible; only indexes hosted by a memory engine need to support it.
308    fn writable_snapshot(&self) -> StorageBackendResult<Box<dyn VectorIndex>> {
309        Err(StorageBackendError::Other(
310            "writable vector-index snapshots are not supported by this backend".into(),
311        ))
312    }
313}
314
315#[derive(Debug, Clone)]
316pub struct MemoryVectorIndex {
317    dimensions: u32,
318    vectors: BTreeMap<DocId, Vec<Vec<f32>>>,
319}
320
321impl MemoryVectorIndex {
322    pub fn new(dimensions: u32) -> Self {
323        Self {
324            dimensions,
325            vectors: BTreeMap::new(),
326        }
327    }
328
329    pub fn vectors(&self) -> &BTreeMap<DocId, Vec<Vec<f32>>> {
330        &self.vectors
331    }
332}
333
334impl VectorIndex for MemoryVectorIndex {
335    fn contains_document(&self, doc_id: DocId) -> StorageBackendResult<bool> {
336        Ok(self.vectors.contains_key(&doc_id))
337    }
338
339    fn dimensions(&self) -> u32 {
340        self.dimensions
341    }
342
343    fn index_kind(&self) -> &'static str {
344        "memory-bruteforce"
345    }
346
347    fn add(&mut self, doc_id: DocId, vector: Vec<f32>) -> StorageBackendResult<()> {
348        self.validate_dimensions(&vector)?;
349        self.vectors.insert(doc_id, vec![vector]);
350        Ok(())
351    }
352
353    fn add_many(&mut self, doc_id: DocId, vectors: Vec<Vec<f32>>) -> StorageBackendResult<()> {
354        for vector in &vectors {
355            self.validate_dimensions(vector)?;
356        }
357        if vectors.is_empty() {
358            self.vectors.remove(&doc_id);
359        } else {
360            self.vectors.insert(doc_id, vectors);
361        }
362        Ok(())
363    }
364
365    fn delete(&mut self, doc_id: DocId) -> StorageBackendResult<()> {
366        self.vectors.remove(&doc_id);
367        Ok(())
368    }
369
370    fn clear(&mut self) -> StorageBackendResult<()> {
371        self.vectors.clear();
372        Ok(())
373    }
374
375    /// Brute-force top-k by cosine similarity, descending.
376    fn search_knn(&self, query: &[f32], k: usize) -> StorageBackendResult<PostingList> {
377        self.validate_dimensions(query)?;
378        if k == 0 || self.vectors.is_empty() {
379            return Ok(PostingList::new());
380        }
381        let mut scored: Vec<(DocId, f32)> = self
382            .vectors
383            .iter()
384            .filter_map(|(&doc_id, vectors)| best_vector_score(query, vectors).map(|s| (doc_id, s)))
385            .collect();
386        select_top_k_scored(&mut scored, k);
387        // The output of `top_k` is re-sorted by doc_id ascending so the
388        // posting list invariant holds; the score lives in the payload.
389        scored.sort_by_key(|(id, _)| *id);
390        let entries = scored
391            .into_iter()
392            .map(|(doc_id, sim)| PostingEntry::new(doc_id, Payload::with_score(f64::from(sim))))
393            .collect::<Vec<_>>();
394        Ok(PostingList::from_sorted_unchecked(entries))
395    }
396
397    /// Brute-force threshold scan: keep all docs with `cosine >= threshold`.
398    fn search_threshold(&self, query: &[f32], threshold: f32) -> StorageBackendResult<PostingList> {
399        self.validate_dimensions(query)?;
400        validate_threshold(threshold)?;
401        let mut entries: Vec<PostingEntry> = self
402            .vectors
403            .iter()
404            .filter_map(|(&doc_id, vectors)| {
405                let sim = best_vector_score(query, vectors)?;
406                if sim >= threshold {
407                    Some(PostingEntry::new(
408                        doc_id,
409                        Payload::with_score(f64::from(sim)),
410                    ))
411                } else {
412                    None
413                }
414            })
415            .collect();
416        // The BTreeMap iteration is already doc_id-ascending, so the
417        // filter preserves the invariant.
418        entries.sort_by_key(|e| e.doc_id);
419        Ok(PostingList::from_sorted_unchecked(entries))
420    }
421
422    fn count(&self) -> StorageBackendResult<usize> {
423        checked_vector_count(self.vectors.values().map(Vec::len))
424    }
425
426    fn snapshot(&self) -> StorageBackendResult<Arc<dyn VectorIndex>> {
427        Ok(Arc::new(self.clone()))
428    }
429
430    fn snapshot_with_control(
431        &self,
432        control: &crate::read_control::StorageReadControl,
433    ) -> StorageBackendResult<Arc<dyn VectorIndex>> {
434        memory_snapshot::capture(self, control)
435    }
436
437    fn writable_snapshot(&self) -> StorageBackendResult<Box<dyn VectorIndex>> {
438        Ok(Box::new(self.clone()))
439    }
440}
441
442impl MemoryVectorIndex {
443    fn validate_dimensions(&self, vector: &[f32]) -> StorageBackendResult<()> {
444        validate_vector_values(self.dimensions, vector)
445    }
446}
447
448fn best_vector_score(query: &[f32], vectors: &[Vec<f32>]) -> Option<f32> {
449    vectors
450        .iter()
451        .map(|vector| cosine_similarity(query, vector))
452        .max_by(f32::total_cmp)
453}
454
455#[cfg(test)]
456mod tests {
457    use super::*;
458
459    #[test]
460    fn canonical_membership_tracks_tensor_replacement_and_retained_snapshots() {
461        let indexes: [Box<dyn VectorIndex>; 3] = [
462            Box::new(MemoryVectorIndex::new(2)),
463            Box::new(crate::IVFIndex::new(2)),
464            Box::new(crate::HNSWIndex::new(2)),
465        ];
466        for mut index in indexes {
467            assert!(!index.contains_document(1).unwrap());
468            index
469                .add_many(1, vec![vec![1.0, 0.0], vec![0.0, 1.0]])
470                .unwrap();
471            let retained = index.snapshot().unwrap();
472            assert!(index.contains_document(1).unwrap());
473            assert!(!index.contains_document(2).unwrap());
474            index.add_many(1, Vec::new()).unwrap();
475            assert!(!index.contains_document(1).unwrap());
476            assert!(retained.contains_document(1).unwrap());
477            index.add(1, vec![0.5, 0.5]).unwrap();
478            index.delete(1).unwrap();
479            assert!(!index.contains_document(1).unwrap());
480        }
481    }
482
483    fn approx_eq(a: f32, b: f32, eps: f32) {
484        assert!((a - b).abs() < eps, "expected {a} ~ {b} within {eps}");
485    }
486
487    #[test]
488    fn cosine_identity_is_one() {
489        let v = vec![1.0, 2.0, 3.0];
490        approx_eq(cosine_similarity(&v, &v), 1.0, 1e-6);
491    }
492
493    #[test]
494    fn vector_count_overflow_is_reported() {
495        let error = checked_vector_count([usize::MAX, 1]).unwrap_err();
496        assert!(error.to_string().contains("vector count overflow"));
497    }
498
499    #[test]
500    fn cosine_orthogonal_is_zero() {
501        let a = vec![1.0, 0.0];
502        let b = vec![0.0, 1.0];
503        approx_eq(cosine_similarity(&a, &b), 0.0, 1e-6);
504    }
505
506    #[test]
507    fn cosine_zero_norm_is_zero() {
508        let a = vec![0.0, 0.0];
509        let b = vec![1.0, 1.0];
510        approx_eq(cosine_similarity(&a, &b), 0.0, 1e-6);
511    }
512
513    #[test]
514    fn knn_orders_by_similarity_descending_then_doc_id() {
515        let mut idx = MemoryVectorIndex::new(2);
516        idx.add(1, vec![1.0, 0.0]).unwrap();
517        idx.add(2, vec![0.5, 0.5]).unwrap();
518        idx.add(3, vec![0.0, 1.0]).unwrap();
519        let pl = idx.search_knn(&[1.0, 0.0], 2).unwrap();
520        let docs: Vec<_> = pl.iter().map(|e| e.doc_id).collect();
521        // posting list is doc_id-sorted but the top-2 should be {1, 2}.
522        assert_eq!(docs, vec![1, 2]);
523        let entry1 = pl.get_entry(1).unwrap();
524        let entry2 = pl.get_entry(2).unwrap();
525        assert!(entry1.payload.score > entry2.payload.score);
526    }
527
528    #[test]
529    fn partial_top_k_keeps_deterministic_doc_id_ties() {
530        let mut scored = vec![(10, 0.5), (3, 0.9), (1, 0.9), (8, 0.7), (2, 0.1)];
531        select_top_k_scored(&mut scored, 2);
532        scored.sort_by_key(|(doc_id, _)| *doc_id);
533        assert_eq!(scored, vec![(1, 0.9), (3, 0.9)]);
534    }
535
536    #[test]
537    fn precomputed_norm_cosine_matches_reference_bits() {
538        let a = [0.25, -3.0, 1.5, 8.0];
539        let b = [2.0, 0.75, -4.0, 0.5];
540        let expected = cosine_similarity(&a, &b);
541        let actual = cosine_similarity_with_norms(&a, &b, vector_norm(&a), vector_norm(&b));
542        assert_eq!(actual.to_bits(), expected.to_bits());
543    }
544
545    #[test]
546    fn score_deduplication_keeps_best_tensor_vector() {
547        let mut scored = vec![(7, 0.3), (2, 0.8), (7, 0.9), (2, 0.4), (9, -0.2)];
548        let count = deduplicate_scored_values(&mut scored);
549        scored.truncate(count);
550        assert_eq!(scored, vec![(2, 0.8), (7, 0.9), (9, -0.2)]);
551    }
552
553    #[test]
554    fn threshold_filters_below_cutoff() {
555        let mut idx = MemoryVectorIndex::new(2);
556        idx.add(1, vec![1.0, 0.0]).unwrap();
557        idx.add(2, vec![0.5, 0.5]).unwrap();
558        idx.add(3, vec![0.0, 1.0]).unwrap();
559        let pl = idx.search_threshold(&[1.0, 0.0], 0.5).unwrap();
560        let docs: Vec<_> = pl.iter().map(|e| e.doc_id).collect();
561        assert_eq!(docs, vec![1, 2]);
562    }
563
564    #[test]
565    fn delete_removes_vector() {
566        let mut idx = MemoryVectorIndex::new(2);
567        idx.add(1, vec![1.0, 0.0]).unwrap();
568        idx.delete(1).unwrap();
569        assert_eq!(idx.count().unwrap(), 0);
570    }
571
572    #[test]
573    fn non_finite_vectors_queries_and_thresholds_are_errors() {
574        let mut idx = MemoryVectorIndex::new(2);
575        assert!(idx.add(1, vec![f32::NAN, 0.0]).is_err());
576        idx.add(1, vec![1.0, 0.0]).unwrap();
577        assert!(idx.search_knn(&[f32::INFINITY, 0.0], 1).is_err());
578        assert!(idx
579            .search_threshold(&[1.0, 0.0], f32::NEG_INFINITY)
580            .is_err());
581    }
582}
uqa_storage/vector_index.rs

uqa_storage/
vector_index.rs
