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

uqa_storage/
vector_index.rs
