uqa_storage/diskann_index/query/selection/
source.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! Opaque canonical capabilities retain actual provider ownership and every enclosing read guard.
8
9use crate::diskann_index::{
10    format::{
11        DiskANNCanonicalOrigin, DiskANNChangeIdentity, DiskANNGeneration, DiskANNVectorVersion,
12    },
13    DiskANNCanonicalCounts, DiskANNCanonicalRead, DiskANNCanonicalVectorVisitor, DiskANNQueryRead,
14};
15use crate::{read_control::StorageReadControl, StorageBackendResult};
16use std::sync::Arc;
17use uqa_core::{
18    memory::{Budgeted, MemoryReservation},
19    DocId,
20};
21
22/// An actual selected `DiskANN` canonical source. Only storage snapshot owners construct this capability; callers cannot substitute a generation label or fabricate vector origins. Clones share canonical rows, prepared ownership and read controls.
23#[derive(Clone)]
24pub struct DiskANNReadSnapshot {
25    canonical: Arc<dyn DiskANNQueryRead + Send + Sync>,
26    generation: DiskANNGeneration,
27    control: StorageReadControl,
28    guard: Option<Arc<Budgeted<Guard>>>,
29}
30
31struct Guard {
32    previous: Option<Arc<Budgeted<Guard>>>,
33    control: Option<StorageReadControl>,
34    _memory: Option<Arc<MemoryReservation>>,
35}
36
37impl DiskANNReadSnapshot {
38    pub(in crate::diskann_index::query) fn new(
39        canonical: Arc<dyn DiskANNQueryRead + Send + Sync>,
40        generation: DiskANNGeneration,
41        control: &StorageReadControl,
42    ) -> Self {
43        Self {
44            canonical,
45            generation,
46            control: control.clone(),
47            guard: None,
48        }
49    }
50
51    /// Whether both actual sources belong to the same database, table and index incarnation. This metadata comparison does not establish equal visibility, generation, configuration or canonical width.
52    pub fn same_lineage(&self, other: &Self) -> bool {
53        self.generation.database() == other.generation.database()
54            && self.generation.table() == other.generation.table()
55            && self.generation.index() == other.generation.index()
56    }
57
58    pub(crate) fn with_guard(
59        mut self,
60        memory: Option<Arc<MemoryReservation>>,
61        control: Option<StorageReadControl>,
62    ) -> StorageBackendResult<Self> {
63        self.check_control(&self.control)?;
64        if let Some(control) = &control {
65            control.check()?;
66        }
67        if memory.is_some() || control.is_some() {
68            let guard = Guard {
69                previous: self.guard.take(),
70                control,
71                _memory: memory,
72            };
73            self.guard = Some(
74                Budgeted::new(guard, self.control.memory().empty_reservation()).into_shared()?,
75            );
76        }
77        Ok(self)
78    }
79
80    fn read<T>(
81        &self,
82        control: &StorageReadControl,
83        operation: impl FnOnce() -> StorageBackendResult<T>,
84    ) -> StorageBackendResult<T> {
85        self.check_control(control)?;
86        let result = operation()?;
87        self.check_control(control)?;
88        Ok(result)
89    }
90}
91
92impl DiskANNCanonicalRead for DiskANNReadSnapshot {
93    fn population_counts(
94        &self,
95        generation: DiskANNGeneration,
96        control: &StorageReadControl,
97    ) -> StorageBackendResult<Option<DiskANNCanonicalCounts>> {
98        self.read(control, || {
99            self.canonical.population_counts(generation, control)
100        })
101    }
102
103    fn corpus_fingerprint(
104        &self,
105        control: &StorageReadControl,
106    ) -> StorageBackendResult<Option<[u8; 32]>> {
107        self.read(control, || self.canonical.corpus_fingerprint(control))
108    }
109    fn check_control(&self, control: &StorageReadControl) -> StorageBackendResult<()> {
110        self.control.check()?;
111        control.check()?;
112        let mut guard = self.guard.as_ref();
113        while let Some(current) = guard {
114            if let Some(control) = &current.control {
115                control.check()?;
116            }
117            guard = current.previous.as_ref();
118        }
119        self.canonical.check_control(control)
120    }
121
122    fn dimensions(&self) -> u32 {
123        self.canonical.dimensions()
124    }
125
126    fn next_document_after(
127        &self,
128        after: Option<DocId>,
129        control: &StorageReadControl,
130    ) -> StorageBackendResult<Option<DocId>> {
131        self.read(control, || {
132            self.canonical.next_document_after(after, control)
133        })
134    }
135
136    fn origin(
137        &self,
138        document: DocId,
139        control: &StorageReadControl,
140    ) -> StorageBackendResult<Option<DiskANNVectorVersion>> {
141        self.read(control, || self.canonical.origin(document, control))
142    }
143
144    fn visit_document(
145        &self,
146        document: DocId,
147        control: &StorageReadControl,
148        visit: &mut DiskANNCanonicalVectorVisitor<'_>,
149    ) -> StorageBackendResult<Option<DiskANNVectorVersion>> {
150        self.read(control, || {
151            self.canonical.visit_document(document, control, visit)
152        })
153    }
154
155    fn read_vector(
156        &self,
157        document: DocId,
158        ordinal: u32,
159        control: &StorageReadControl,
160    ) -> StorageBackendResult<Option<uqa_core::memory::BudgetedVec<f32>>> {
161        self.read(control, || {
162            self.canonical.read_vector(document, ordinal, control)
163        })
164    }
165}
166
167impl crate::vector_index::VectorRead for DiskANNReadSnapshot {
168    fn corpus_fingerprint(
169        &self,
170        control: &StorageReadControl,
171    ) -> StorageBackendResult<Option<[u8; 32]>> {
172        DiskANNCanonicalRead::corpus_fingerprint(self, control)
173    }
174    fn check_control(&self, control: &StorageReadControl) -> StorageBackendResult<()> {
175        DiskANNCanonicalRead::check_control(self, control)
176    }
177    fn dimensions(&self) -> u32 {
178        DiskANNCanonicalRead::dimensions(self)
179    }
180    fn next_document_after(
181        &self,
182        after: Option<DocId>,
183        control: &StorageReadControl,
184    ) -> StorageBackendResult<Option<DocId>> {
185        DiskANNCanonicalRead::next_document_after(self, after, control)
186    }
187    fn document_vector_count(
188        &self,
189        document: DocId,
190        control: &StorageReadControl,
191    ) -> StorageBackendResult<u64> {
192        Ok(self
193            .document_origin(document, control)?
194            .map_or(0, DiskANNCanonicalOrigin::count))
195    }
196    fn read_vector(
197        &self,
198        document: DocId,
199        ordinal: u32,
200        control: &StorageReadControl,
201    ) -> StorageBackendResult<Option<uqa_core::memory::BudgetedVec<f32>>> {
202        DiskANNCanonicalRead::read_vector(self, document, ordinal, control)
203    }
204}
205
206impl DiskANNQueryRead for DiskANNReadSnapshot {
207    fn document_origin(
208        &self,
209        document: DocId,
210        control: &StorageReadControl,
211    ) -> StorageBackendResult<Option<DiskANNCanonicalOrigin>> {
212        self.read(control, || {
213            self.canonical.document_origin(document, control)
214        })
215    }
216
217    fn next_change_after(
218        &self,
219        after: Option<DocId>,
220        control: &StorageReadControl,
221    ) -> StorageBackendResult<Option<DiskANNChangeIdentity>> {
222        self.read(control, || self.canonical.next_change_after(after, control))
223    }
224}
uqa_storage/diskann_index/query/selection/source.rs

uqa_storage/diskann_index/query/selection/
source.rs
