Skip to main content

uqa_storage/read_only_snapshot/
vectors.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7use std::sync::Arc;
8
9use uqa_core::{DocId, PostingList};
10
11use crate::read_control::StorageReadControl;
12use crate::{StorageBackendResult, VectorIndex};
13
14use super::{read_only_error, ReadOnlySnapshot};
15
16impl<T: VectorIndex + ?Sized> ReadOnlySnapshot<T> {
17    /// Select query control for an already captured physical value. Cache owners bind each returned reader separately; nested readers preserve the first selected allowance and cancellation token.
18    pub fn with_vector_read_control(
19        mut self,
20        control: &StorageReadControl,
21    ) -> StorageBackendResult<Self> {
22        control.check()?;
23        if self.2.is_none() {
24            self.2 = Some(control.clone());
25        }
26        Ok(self)
27    }
28}
29
30impl<T: VectorIndex + ?Sized + 'static> VectorIndex for ReadOnlySnapshot<T> {
31    fn dimensions(&self) -> u32 {
32        self.0.dimensions()
33    }
34
35    fn index_kind(&self) -> &'static str {
36        self.0.index_kind()
37    }
38    fn diskann_query_metadata(
39        &self,
40        control: &StorageReadControl,
41    ) -> StorageBackendResult<Option<crate::diskann_index::DiskANNQueryMetadata>> {
42        control.check()?;
43        let original = self.2.as_ref().unwrap_or(control);
44        original.check()?;
45        let metadata = self.0.diskann_query_metadata(original)?;
46        control.check()?;
47        Ok(metadata)
48    }
49
50    fn add(&mut self, _doc_id: DocId, _vector: Vec<f32>) -> StorageBackendResult<()> {
51        Err(read_only_error())
52    }
53
54    fn add_many(&mut self, _doc_id: DocId, _vectors: Vec<Vec<f32>>) -> StorageBackendResult<()> {
55        Err(read_only_error())
56    }
57
58    fn delete(&mut self, _doc_id: DocId) -> StorageBackendResult<()> {
59        Err(read_only_error())
60    }
61
62    fn clear(&mut self) -> StorageBackendResult<()> {
63        Err(read_only_error())
64    }
65
66    fn search_knn(&self, query: &[f32], k: usize) -> StorageBackendResult<PostingList> {
67        match &self.2 {
68            Some(control) => self.0.search_knn_with_control(query, k, control),
69            None => self.0.search_knn(query, k),
70        }
71    }
72
73    fn search_threshold(&self, query: &[f32], threshold: f32) -> StorageBackendResult<PostingList> {
74        match &self.2 {
75            Some(control) => self
76                .0
77                .search_threshold_with_control(query, threshold, control),
78            None => self.0.search_threshold(query, threshold),
79        }
80    }
81
82    fn search_knn_with_control(
83        &self,
84        query: &[f32],
85        k: usize,
86        control: &StorageReadControl,
87    ) -> StorageBackendResult<PostingList> {
88        control.check()?;
89        if self.2.is_some() {
90            self.search_knn(query, k)
91        } else {
92            self.0.search_knn_with_control(query, k, control)
93        }
94    }
95
96    fn search_threshold_with_control(
97        &self,
98        query: &[f32],
99        threshold: f32,
100        control: &StorageReadControl,
101    ) -> StorageBackendResult<PostingList> {
102        control.check()?;
103        if self.2.is_some() {
104            self.search_threshold(query, threshold)
105        } else {
106            self.0
107                .search_threshold_with_control(query, threshold, control)
108        }
109    }
110
111    fn search_knn_with_statistics(
112        &self,
113        query: &[f32],
114        k: usize,
115        control: Option<&StorageReadControl>,
116    ) -> StorageBackendResult<crate::vector_index::VectorQueryResult> {
117        if let Some(control) = control {
118            control.check()?;
119        }
120        let result = self
121            .0
122            .search_knn_with_statistics(query, k, self.2.as_ref().or(control))?;
123        if let Some(control) = control {
124            control.check()?;
125        }
126        Ok(result)
127    }
128
129    fn search_threshold_with_statistics(
130        &self,
131        query: &[f32],
132        threshold: f32,
133        control: Option<&StorageReadControl>,
134    ) -> StorageBackendResult<crate::vector_index::VectorQueryResult> {
135        if let Some(control) = control {
136            control.check()?;
137        }
138        let result = self.0.search_threshold_with_statistics(
139            query,
140            threshold,
141            self.2.as_ref().or(control),
142        )?;
143        if let Some(control) = control {
144            control.check()?;
145        }
146        Ok(result)
147    }
148
149    fn count(&self) -> StorageBackendResult<usize> {
150        if let Some(control) = &self.2 {
151            control.check()?;
152        }
153        self.0.count()
154    }
155
156    fn contains_document(&self, doc_id: DocId) -> StorageBackendResult<bool> {
157        if let Some(control) = &self.2 {
158            control.check()?;
159        }
160        self.0.contains_document(doc_id)
161    }
162
163    fn initialize(&mut self) -> StorageBackendResult<()> {
164        Err(read_only_error())
165    }
166
167    fn snapshot(&self) -> StorageBackendResult<Arc<dyn VectorIndex>> {
168        Ok(Arc::new(self.clone()))
169    }
170
171    fn vector_read_snapshot(
172        &self,
173        control: &StorageReadControl,
174    ) -> StorageBackendResult<Option<crate::vector_index::VectorReadSnapshot>> {
175        control.check()?;
176        let original = self.2.as_ref().unwrap_or(control);
177        original.check()?;
178        self.3
179            .as_ref()
180            .map(Arc::clone)
181            .map_or_else(
182                || self.0.vector_read_snapshot(original),
183                |source| Ok(Some(source)),
184            )?
185            .map(|source| {
186                let source = ReadOnlySnapshot(source, self.1.clone(), Some(original.clone()), None);
187                uqa_core::memory::Budgeted::new(source, original.memory().empty_reservation())
188                    .into_shared()
189                    .map(|source| source as crate::vector_index::VectorReadSnapshot)
190                    .map_err(Into::into)
191            })
192            .transpose()
193    }
194
195    fn snapshot_with_vector_read(
196        &self,
197        source: crate::vector_index::VectorReadSnapshot,
198        control: &StorageReadControl,
199    ) -> StorageBackendResult<Option<Arc<dyn VectorIndex>>> {
200        control.check()?;
201        let original = self.2.as_ref().unwrap_or(control);
202        original.check()?;
203        Ok(self
204            .0
205            .snapshot_with_vector_read(source, original)?
206            .map(|index| {
207                Arc::new(ReadOnlySnapshot(
208                    index,
209                    self.1.clone(),
210                    self.2.clone(),
211                    None,
212                )) as Arc<dyn VectorIndex>
213            }))
214    }
215
216    fn diskann_read_snapshot(
217        &self,
218        control: &StorageReadControl,
219    ) -> StorageBackendResult<Option<crate::diskann_index::DiskANNReadSnapshot>> {
220        control.check()?;
221        if let Some(original) = &self.2 {
222            original.check()?;
223        }
224        self.0
225            .diskann_read_snapshot(self.2.as_ref().unwrap_or(control))?
226            .map(|source| source.with_guard(self.1.clone(), self.2.clone()))
227            .transpose()
228    }
229
230    fn snapshot_with_diskann_changes(
231        &self,
232        changes: &crate::diskann_index::DiskANNReadChanges,
233        control: &StorageReadControl,
234    ) -> StorageBackendResult<Option<Arc<dyn VectorIndex>>> {
235        control.check()?;
236        if let Some(original) = &self.2 {
237            original.check()?;
238        }
239        Ok(self
240            .0
241            .snapshot_with_diskann_changes(changes, self.2.as_ref().unwrap_or(control))?
242            .map(|index| {
243                Arc::new(ReadOnlySnapshot(
244                    index,
245                    self.1.clone(),
246                    self.2.clone(),
247                    None,
248                )) as Arc<dyn VectorIndex>
249            }))
250    }
251
252    fn snapshot_with_control(
253        &self,
254        control: &crate::read_control::StorageReadControl,
255    ) -> StorageBackendResult<Arc<dyn VectorIndex>> {
256        control.check()?;
257        if self.1.is_some() || self.2.is_some() || self.3.is_some() {
258            return self.clone().with_vector_read_control(control)?.snapshot();
259        }
260        Ok(Arc::new(ReadOnlySnapshot::new(
261            self.0.snapshot_with_control(control)?,
262        )))
263    }
264}
265
266impl<T: crate::vector_index::VectorRead + ?Sized> crate::vector_index::VectorRead
267    for ReadOnlySnapshot<T>
268{
269    fn corpus_fingerprint(
270        &self,
271        control: &StorageReadControl,
272    ) -> StorageBackendResult<Option<[u8; 32]>> {
273        self.check_control(control)?;
274        self.0.corpus_fingerprint(control)
275    }
276    fn check_control(&self, control: &StorageReadControl) -> StorageBackendResult<()> {
277        if let Some(original) = &self.2 {
278            original.check()?;
279        }
280        self.0.check_control(control)
281    }
282    fn dimensions(&self) -> u32 {
283        self.0.dimensions()
284    }
285    fn next_document_after(
286        &self,
287        after: Option<DocId>,
288        control: &StorageReadControl,
289    ) -> StorageBackendResult<Option<DocId>> {
290        self.check_control(control)?;
291        self.0.next_document_after(after, control)
292    }
293    fn document_vector_count(
294        &self,
295        document: DocId,
296        control: &StorageReadControl,
297    ) -> StorageBackendResult<u64> {
298        self.check_control(control)?;
299        self.0.document_vector_count(document, control)
300    }
301    fn read_vector(
302        &self,
303        document: DocId,
304        ordinal: u32,
305        control: &StorageReadControl,
306    ) -> StorageBackendResult<Option<uqa_core::memory::BudgetedVec<f32>>> {
307        self.check_control(control)?;
308        let value = self.0.read_vector(document, ordinal, control)?;
309        self.check_control(control)?;
310        Ok(value)
311    }
312}