uqa_storage/diskann_index/query/
retained.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! Owned vector-index snapshots keep canonical visibility and prepared physical readers together.
8
9use super::selection::{DiskANNReadChanges, DiskANNReadSnapshot, SelectedCanonical};
10use super::{invalid, DiskANNQuery};
11use crate::diskann_index::{
12    format::{DiskANNGeneration, DiskANNManifest},
13    pages::{DiskANNOriginReader, DiskANNPageSource, DiskANNReadLimits, DiskANNReader},
14    DiskANNCanonicalCounts, DiskANNCanonicalRead, DiskANNQueryRead, ExactVectorReason,
15};
16use crate::{
17    read_control::StorageReadControl,
18    vector_index::{DiskANNExecutionRoute, DiskANNIndexParams, VectorQueryResult},
19    StorageBackendError, StorageBackendResult, VectorIndex,
20};
21use std::sync::{Arc, OnceLock};
22use uqa_core::{memory::Budgeted, DocId, PostingList};
23
24mod values;
25
26struct Retained<S> {
27    canonical: S,
28    reader: DiskANNReader,
29    origins: DiskANNOriginReader,
30    control: StorageReadControl,
31    observed_counts: OnceLock<DiskANNCanonicalCounts>,
32}
33
34/// Read-only `VectorIndex` over an owned fixed canonical source and its prepared generation. Cloning and nested snapshots share the same leases, resident data and original query allowance, without reopening pages or copying the canonical corpus.
35pub struct RetainedDiskANNIndex<S> {
36    retained: Arc<Budgeted<Retained<S>>>,
37}
38
39impl<S> Clone for RetainedDiskANNIndex<S> {
40    fn clone(&self) -> Self {
41        Self {
42            retained: Arc::clone(&self.retained),
43        }
44    }
45}
46
47impl<S: DiskANNQueryRead + Send + Sync + 'static> RetainedDiskANNIndex<S> {
48    /// Prepare the source already selected on the supplied canonical/catalog view, then retain both owners. Provider adapters establish this actual association; a matching generation label by itself does not establish visibility.
49    pub fn open(
50        canonical: S,
51        source: Arc<dyn DiskANNPageSource>,
52        parameters: DiskANNIndexParams,
53        limits: DiskANNReadLimits,
54        control: &StorageReadControl,
55    ) -> StorageBackendResult<Self> {
56        let DiskANNQuery {
57            reader,
58            origins,
59            control: original,
60            ..
61        } = DiskANNQuery::open(&canonical, source, parameters, limits, control)?;
62        let retained = Retained {
63            canonical,
64            reader,
65            origins,
66            control: original,
67            observed_counts: OnceLock::new(),
68        };
69        let retained =
70            Budgeted::new(retained, control.memory().empty_reservation()).into_shared()?;
71        control.check()?;
72        Ok(Self { retained })
73    }
74
75    /// Immutable physical identity and effective configuration from the actual selected generation. Reading metadata does not perform a logical vector query.
76    pub fn manifest(&self) -> &DiskANNManifest {
77        self.retained.reader.manifest()
78    }
79
80    pub(in crate::diskann_index) fn canonical(&self) -> &S {
81        &self.retained.canonical
82    }
83
84    /// Advance canonical visibility within the same storage-owned lineage while sharing its prepared physical generation. Only the lifecycle owner may establish this association.
85    pub(in crate::diskann_index) fn with_canonical(
86        &self,
87        canonical: S,
88    ) -> StorageBackendResult<Self> {
89        self.check()?;
90        canonical.check_control(&self.retained.control)?;
91        if canonical.dimensions() != self.dimensions() {
92            return Err(invalid("replacement canonical dimensions differ"));
93        }
94        let retained = Retained {
95            canonical,
96            reader: self.retained.reader.clone(),
97            origins: self.retained.origins.clone(),
98            control: self.retained.control.clone(),
99            observed_counts: OnceLock::new(),
100        };
101        let retained = Budgeted::new(retained, self.retained.control.memory().empty_reservation())
102            .into_shared()?;
103        self.check()?;
104        Ok(Self { retained })
105    }
106
107    fn check(&self) -> StorageBackendResult<()> {
108        self.retained.control.check()?;
109        self.retained
110            .canonical
111            .check_control(&self.retained.control)
112    }
113
114    fn query(&self, invoking: &StorageReadControl) -> DiskANNQuery<'_> {
115        // The search uses the original memory allowance. This additional control is checked throughout scoring/traversal so an independent invocation can cancel without replacing retained ownership.
116        DiskANNQuery {
117            canonical: &self.retained.canonical,
118            reader: self.retained.reader.clone(),
119            origins: self.retained.origins.clone(),
120            control: invoking.clone(),
121        }
122    }
123}
124
125impl<S: DiskANNQueryRead + Send + Sync + 'static> VectorIndex for RetainedDiskANNIndex<S> {
126    fn dimensions(&self) -> u32 {
127        self.retained.canonical.dimensions()
128    }
129    fn index_kind(&self) -> &'static str {
130        "diskann"
131    }
132    fn diskann_query_metadata(
133        &self,
134        control: &StorageReadControl,
135    ) -> StorageBackendResult<Option<crate::diskann_index::DiskANNQueryMetadata>> {
136        self.check()?;
137        control.check()?;
138        let corpus_fingerprint = self.retained.canonical.corpus_fingerprint(control)?;
139        let canonical_counts = self
140            .retained
141            .population_counts(self.manifest().input().generation, control)?;
142        self.check()?;
143        Ok(Some(crate::diskann_index::DiskANNQueryMetadata {
144            manifest: *self.manifest(),
145            corpus_fingerprint,
146            canonical_counts,
147            read_limits: self.retained.reader.limits(),
148            read_capabilities: self.retained.reader.capabilities(),
149        }))
150    }
151    fn add(&mut self, _: DocId, _: Vec<f32>) -> StorageBackendResult<()> {
152        Err(read_only())
153    }
154    fn add_many(&mut self, _: DocId, _: Vec<Vec<f32>>) -> StorageBackendResult<()> {
155        Err(read_only())
156    }
157    fn delete(&mut self, _: DocId) -> StorageBackendResult<()> {
158        Err(read_only())
159    }
160    fn clear(&mut self) -> StorageBackendResult<()> {
161        Err(read_only())
162    }
163    fn initialize(&mut self) -> StorageBackendResult<()> {
164        Err(read_only())
165    }
166
167    fn search_knn(&self, query: &[f32], k: usize) -> StorageBackendResult<PostingList> {
168        self.search_knn_with_control(query, k, &self.retained.control)
169    }
170    fn search_threshold(&self, query: &[f32], threshold: f32) -> StorageBackendResult<PostingList> {
171        self.search_threshold_with_control(query, threshold, &self.retained.control)
172    }
173    fn search_knn_with_control(
174        &self,
175        query: &[f32],
176        k: usize,
177        control: &StorageReadControl,
178    ) -> StorageBackendResult<PostingList> {
179        self.search_knn_with_statistics(query, k, Some(control))
180            .map(|result| result.postings)
181    }
182    fn search_threshold_with_control(
183        &self,
184        query: &[f32],
185        threshold: f32,
186        control: &StorageReadControl,
187    ) -> StorageBackendResult<PostingList> {
188        self.search_threshold_with_statistics(query, threshold, Some(control))
189            .map(|result| result.postings)
190    }
191
192    fn search_knn_with_statistics(
193        &self,
194        query: &[f32],
195        k: usize,
196        control: Option<&StorageReadControl>,
197    ) -> StorageBackendResult<VectorQueryResult> {
198        let result = self
199            .query(control.unwrap_or(&self.retained.control))
200            .search_knn(query, k, &self.retained.control)?;
201        if let Some(counts) = result.populations {
202            let previous = self.retained.observed_counts.get_or_init(|| counts);
203            if *previous != counts {
204                return Err(invalid(
205                    "fixed canonical population changed between queries",
206                ));
207            }
208        }
209        let route = if k == 0 {
210            DiskANNExecutionRoute::EmptyK
211        } else {
212            match result.exact_reason {
213                None => DiskANNExecutionRoute::Approximate,
214                Some(ExactVectorReason::ZeroNorm) => DiskANNExecutionRoute::ExactZeroNorm,
215                Some(ExactVectorReason::NonFiniteNorm) => DiskANNExecutionRoute::ExactNonFiniteNorm,
216            }
217        };
218        Ok(VectorQueryResult::measured(
219            result,
220            self.manifest().input().generation,
221            route,
222        ))
223    }
224
225    fn search_threshold_with_statistics(
226        &self,
227        query: &[f32],
228        threshold: f32,
229        control: Option<&StorageReadControl>,
230    ) -> StorageBackendResult<VectorQueryResult> {
231        let result = self
232            .query(control.unwrap_or(&self.retained.control))
233            .search_threshold_with_stats(query, threshold, &self.retained.control)?;
234        Ok(VectorQueryResult::measured(
235            result,
236            self.manifest().input().generation,
237            DiskANNExecutionRoute::ExactThreshold,
238        ))
239    }
240
241    fn count(&self) -> StorageBackendResult<usize> {
242        self.check()?;
243        if let Some(source) = self.retained.canonical.unversioned_vectors() {
244            return super::values::count(&**source, &self.retained.control);
245        }
246        let count = self
247            .retained
248            .canonical
249            .vector_count(&self.retained.control)?;
250        self.check()?;
251        Ok(count)
252    }
253
254    fn contains_document(&self, document: DocId) -> StorageBackendResult<bool> {
255        self.check()?;
256        if let Some(source) = self.retained.canonical.unversioned_vectors() {
257            return source
258                .document_vector_count(document, &self.retained.control)
259                .map(|count| count != 0);
260        }
261        let contains = self
262            .retained
263            .canonical
264            .contains_vectors(document, &self.retained.control)?;
265        self.check()?;
266        Ok(contains)
267    }
268
269    fn snapshot(&self) -> StorageBackendResult<Arc<dyn VectorIndex>> {
270        self.check()?;
271        Ok(Arc::new(self.clone()))
272    }
273    fn snapshot_with_control(
274        &self,
275        control: &StorageReadControl,
276    ) -> StorageBackendResult<Arc<dyn VectorIndex>> {
277        control.check()?;
278        self.snapshot()
279    }
280
281    fn diskann_read_snapshot(
282        &self,
283        control: &StorageReadControl,
284    ) -> StorageBackendResult<Option<DiskANNReadSnapshot>> {
285        control.check()?;
286        self.check()?;
287        if self.retained.canonical.unversioned_vectors().is_some() {
288            return Ok(None);
289        }
290        Ok(Some(DiskANNReadSnapshot::new(
291            self.retained.clone(),
292            self.manifest().input().generation,
293            &self.retained.control,
294        )))
295    }
296
297    fn vector_read_snapshot(
298        &self,
299        control: &StorageReadControl,
300    ) -> StorageBackendResult<Option<crate::vector_index::VectorReadSnapshot>> {
301        control.check()?;
302        self.check()?;
303        Ok(Some(self.retained.clone()))
304    }
305
306    fn snapshot_with_vector_read(
307        &self,
308        source: crate::vector_index::VectorReadSnapshot,
309        control: &StorageReadControl,
310    ) -> StorageBackendResult<Option<Arc<dyn VectorIndex>>> {
311        self.check()?;
312        let actual = DiskANNReadSnapshot::new(
313            self.retained.clone(),
314            self.manifest().input().generation,
315            &self.retained.control,
316        );
317        let canonical = super::matching::MatchingCanonical::new(actual, source, control)?;
318        let retained = Retained {
319            canonical,
320            reader: self.retained.reader.clone(),
321            origins: self.retained.origins.clone(),
322            control: self.retained.control.clone(),
323            observed_counts: OnceLock::new(),
324        };
325        let retained = Budgeted::new(retained, self.retained.control.memory().empty_reservation())
326            .into_shared()?;
327        self.check()?;
328        control.check()?;
329        Ok(Some(Arc::new(RetainedDiskANNIndex { retained })))
330    }
331
332    fn snapshot_with_diskann_changes(
333        &self,
334        changes: &DiskANNReadChanges,
335        control: &StorageReadControl,
336    ) -> StorageBackendResult<Option<Arc<dyn VectorIndex>>> {
337        control.check()?;
338        self.check()?;
339        if self.retained.canonical.unversioned_vectors().is_some() {
340            return Ok(None);
341        }
342        let base = DiskANNReadSnapshot::new(
343            self.retained.clone(),
344            self.manifest().input().generation,
345            &self.retained.control,
346        );
347        let population_control =
348            StorageReadControl::new(self.retained.control.memory(), control.cancellation());
349        let canonical =
350            SelectedCanonical::new(base, changes, &self.retained.origins, &population_control)?;
351        let retained = Retained {
352            canonical,
353            reader: self.retained.reader.clone(),
354            origins: self.retained.origins.clone(),
355            control: self.retained.control.clone(),
356            observed_counts: OnceLock::new(),
357        };
358        let retained = Budgeted::new(retained, self.retained.control.memory().empty_reservation())
359            .into_shared()?;
360        self.check()?;
361        control.check()?;
362        Ok(Some(Arc::new(RetainedDiskANNIndex { retained })))
363    }
364}
365
366impl<S: DiskANNQueryRead> crate::diskann_index::DiskANNCanonicalRead for Budgeted<Retained<S>> {
367    fn population_counts(
368        &self,
369        generation: DiskANNGeneration,
370        control: &StorageReadControl,
371    ) -> StorageBackendResult<Option<DiskANNCanonicalCounts>> {
372        self.control.check()?;
373        let counts = self.canonical.population_counts(generation, control)?;
374        self.canonical.check_control(control)?;
375        Ok(counts.or_else(|| {
376            (generation == self.reader.manifest().input().generation)
377                .then(|| self.observed_counts.get().copied())
378                .flatten()
379        }))
380    }
381
382    fn corpus_fingerprint(
383        &self,
384        control: &StorageReadControl,
385    ) -> StorageBackendResult<Option<[u8; 32]>> {
386        self.control.check()?;
387        self.canonical.corpus_fingerprint(control)
388    }
389    fn check_control(&self, control: &StorageReadControl) -> StorageBackendResult<()> {
390        self.control.check()?;
391        self.canonical.check_control(control)
392    }
393    fn dimensions(&self) -> u32 {
394        self.canonical.dimensions()
395    }
396    fn next_document_after(
397        &self,
398        after: Option<DocId>,
399        control: &StorageReadControl,
400    ) -> StorageBackendResult<Option<DocId>> {
401        self.canonical.next_document_after(after, control)
402    }
403    fn origin(
404        &self,
405        document: DocId,
406        control: &StorageReadControl,
407    ) -> StorageBackendResult<Option<crate::diskann_index::format::DiskANNVectorVersion>> {
408        self.canonical.origin(document, control)
409    }
410    fn visit_document(
411        &self,
412        document: DocId,
413        control: &StorageReadControl,
414        visit: &mut crate::diskann_index::DiskANNCanonicalVectorVisitor<'_>,
415    ) -> StorageBackendResult<Option<crate::diskann_index::format::DiskANNVectorVersion>> {
416        self.canonical.visit_document(document, control, visit)
417    }
418
419    fn read_vector(
420        &self,
421        document: DocId,
422        ordinal: u32,
423        control: &StorageReadControl,
424    ) -> StorageBackendResult<Option<uqa_core::memory::BudgetedVec<f32>>> {
425        self.control.check()?;
426        self.canonical.read_vector(document, ordinal, control)
427    }
428}
429
430impl<S: DiskANNQueryRead> DiskANNQueryRead for Budgeted<Retained<S>> {
431    fn document_origin(
432        &self,
433        document: DocId,
434        control: &StorageReadControl,
435    ) -> StorageBackendResult<Option<crate::diskann_index::format::DiskANNCanonicalOrigin>> {
436        self.canonical.document_origin(document, control)
437    }
438    fn next_change_after(
439        &self,
440        after: Option<DocId>,
441        control: &StorageReadControl,
442    ) -> StorageBackendResult<Option<crate::diskann_index::format::DiskANNChangeIdentity>> {
443        self.canonical.next_change_after(after, control)
444    }
445}
446
447fn read_only() -> StorageBackendError {
448    StorageBackendError::Other("cannot write a retained DiskANN snapshot".into())
449}
uqa_storage/diskann_index/query/retained.rs

uqa_storage/diskann_index/query/
retained.rs
