uqa_storage/vector_index/
retained.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! Exact readers retain canonical vectors and their allowance independently of the live session.
8
9use std::sync::Arc;
10
11use uqa_core::{
12    memory::{Budgeted, BudgetedVec},
13    DocId, Payload, PostingEntry, PostingList,
14};
15
16use crate::read_control::StorageReadControl;
17use crate::vector_index::{
18    cosine_similarity, select_top_k_scored, validate_vector_values, VectorIndex,
19};
20use crate::{StorageBackendError, StorageBackendResult};
21use sha2::{Digest, Sha256};
22
23mod builder;
24mod canonical;
25pub use builder::RetainedVectorIndexBuilder;
26
27pub(crate) type VectorEntries = Vec<(DocId, u32, Vec<f32>)>;
28
29/// An immutable exact index whose vector buffers and read workspace share the original retention allowance.
30#[derive(Clone)]
31pub struct RetainedVectorIndex {
32    entries: Arc<Budgeted<VectorEntries>>,
33    dimensions: u32,
34    index_kind: &'static str,
35    control: StorageReadControl,
36    fingerprint: [u8; 32],
37}
38
39impl RetainedVectorIndex {
40    pub(crate) fn from_entries(
41        entries: Budgeted<VectorEntries>,
42        dimensions: u32,
43        index_kind: &'static str,
44        control: &StorageReadControl,
45    ) -> StorageBackendResult<Self> {
46        control.check()?;
47        let mut digest = Sha256::new();
48        digest.update(b"uqa-retained-vector-values-v1\0");
49        digest.update(dimensions.to_le_bytes());
50        for (document, ordinal, vector) in entries.iter() {
51            control.check()?;
52            digest.update(document.to_le_bytes());
53            digest.update(ordinal.to_le_bytes());
54            digest.update((vector.len() as u64).to_le_bytes());
55            for chunk in vector.chunks(1024) {
56                control.check()?;
57                for value in chunk {
58                    digest.update(value.to_bits().to_le_bytes());
59                }
60            }
61        }
62        Ok(Self {
63            entries: entries.into_shared()?,
64            dimensions,
65            index_kind,
66            control: control.clone(),
67            fingerprint: digest.finalize().into(),
68        })
69    }
70
71    fn scores(
72        &self,
73        query: &[f32],
74        threshold: Option<f32>,
75    ) -> StorageBackendResult<BudgetedVec<(DocId, f32)>> {
76        let mut scores = BudgetedVec::new(self.control.memory());
77        for vectors in self.entries.chunk_by(|a, b| a.0 == b.0) {
78            self.control.check()?;
79            let mut best = cosine_similarity(query, &vectors[0].2);
80            for (_, _, vector) in &vectors[1..] {
81                self.control.check()?;
82                let score = cosine_similarity(query, vector);
83                if score.total_cmp(&best).is_gt() {
84                    best = score;
85                }
86            }
87            if threshold.is_none_or(|threshold| best >= threshold) {
88                scores.push((vectors[0].0, best))?;
89            }
90        }
91        Ok(scores)
92    }
93
94    fn postings(&self, scores: &[(DocId, f32)]) -> StorageBackendResult<PostingList> {
95        let mut postings = BudgetedVec::new(self.control.memory());
96        postings.reserve(scores.len())?;
97        for &(doc_id, score) in scores {
98            self.control.check()?;
99            postings.push(PostingEntry::new(
100                doc_id,
101                Payload::with_score(f64::from(score)),
102            ))?;
103        }
104        let (postings, _memory) = postings.into_parts();
105        Ok(PostingList::from_sorted_unchecked(postings))
106    }
107}
108
109impl VectorIndex for RetainedVectorIndex {
110    fn contains_document(&self, doc_id: DocId) -> StorageBackendResult<bool> {
111        self.control.check()?;
112        Ok(self
113            .entries
114            .binary_search_by_key(&doc_id, |entry| entry.0)
115            .is_ok())
116    }
117
118    fn dimensions(&self) -> u32 {
119        self.dimensions
120    }
121    fn index_kind(&self) -> &'static str {
122        self.index_kind
123    }
124    fn add(&mut self, _: DocId, _: Vec<f32>) -> StorageBackendResult<()> {
125        Err(read_only())
126    }
127    fn add_many(&mut self, _: DocId, _: Vec<Vec<f32>>) -> StorageBackendResult<()> {
128        Err(read_only())
129    }
130    fn delete(&mut self, _: DocId) -> StorageBackendResult<()> {
131        Err(read_only())
132    }
133    fn clear(&mut self) -> StorageBackendResult<()> {
134        Err(read_only())
135    }
136    fn initialize(&mut self) -> StorageBackendResult<()> {
137        Err(read_only())
138    }
139    fn search_knn(&self, query: &[f32], k: usize) -> StorageBackendResult<PostingList> {
140        self.control.check()?;
141        validate_vector_values(self.dimensions, query)?;
142        if k == 0 {
143            return Ok(PostingList::new());
144        }
145        // The tuple drops score storage before its reservation on every return path.
146        let mut scores = self.scores(query, None)?.into_parts();
147        select_top_k_scored(&mut scores.0, k);
148        scores.0.sort_unstable_by_key(|(doc_id, _)| *doc_id);
149        self.postings(&scores.0)
150    }
151    fn search_threshold(&self, query: &[f32], threshold: f32) -> StorageBackendResult<PostingList> {
152        self.control.check()?;
153        validate_vector_values(self.dimensions, query)?;
154        if !threshold.is_finite() {
155            return Err(StorageBackendError::Other(format!(
156                "vector similarity threshold must be finite, got {threshold}"
157            )));
158        }
159        let scores = self.scores(query, Some(threshold))?;
160        self.postings(&scores)
161    }
162    fn count(&self) -> StorageBackendResult<usize> {
163        self.control.check()?;
164        Ok(self.entries.len())
165    }
166    fn snapshot(&self) -> StorageBackendResult<Arc<dyn VectorIndex>> {
167        Ok(Arc::new(self.clone()))
168    }
169
170    fn vector_read_snapshot(
171        &self,
172        control: &StorageReadControl,
173    ) -> StorageBackendResult<Option<super::VectorReadSnapshot>> {
174        self.control.check()?;
175        control.check()?;
176        Ok(Some(
177            Budgeted::new(self.clone(), self.control.memory().empty_reservation()).into_shared()?,
178        ))
179    }
180}
181
182fn read_only() -> StorageBackendError {
183    StorageBackendError::Other("cannot write a retained vector snapshot".into())
184}
185
186#[cfg(test)]
187mod tests;
uqa_storage/vector_index/retained.rs

uqa_storage/vector_index/
retained.rs
