uqa_storage/read_only_snapshot/
vectors.rs1use 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 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}