uqa_storage/diskann_index/query/
selection.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! Sparse document selections compose actual retained canonical sources without replacing the physical reader.
8
9mod populations;
10mod source;
11pub use source::DiskANNReadSnapshot;
12
13use super::invalid;
14use crate::diskann_index::{
15    format::{
16        DiskANNCanonicalOrigin, DiskANNChangeIdentity, DiskANNGeneration, DiskANNVectorVersion,
17    },
18    pages::DiskANNOriginReader,
19    DiskANNCanonicalCounts, DiskANNCanonicalRead, DiskANNCanonicalVectorVisitor, DiskANNQueryRead,
20};
21use crate::{read_control::StorageReadControl, StorageBackendResult};
22use std::sync::Arc;
23use uqa_core::{
24    memory::{Budgeted, BudgetedVec},
25    DocId,
26};
27
28type SelectedSources = Arc<Budgeted<BudgetedVec<(DocId, Option<DiskANNReadSnapshot>)>>>;
29
30/// Immutable, ordered replacements selected from actual `DiskANN` snapshots. None selects an evaluated deletion independently of retained source membership. A selected document absent from its source is also deleted; a zero-ordinal origin remains an explicit empty replacement. None of these cases falls through to the original tensor. Clones share the sparse directory and its original allowance.
31#[derive(Clone)]
32pub struct DiskANNReadChanges {
33    sources: SelectedSources,
34    control: StorageReadControl,
35    complete: bool,
36}
37
38impl DiskANNReadChanges {
39    /// Capture a strictly increasing document selection. Callers supply the already evaluated private source for each replacement and None for each deletion; unselected documents remain on the original view.
40    pub fn capture(
41        sources: impl IntoIterator<Item = StorageBackendResult<(DocId, Option<DiskANNReadSnapshot>)>>,
42        control: &StorageReadControl,
43    ) -> StorageBackendResult<Self> {
44        control.check()?;
45        let mut selected = BudgetedVec::new(control.memory());
46        let mut after = None;
47        for entry in sources {
48            control.check()?;
49            let (document, source) = entry?;
50            if after.is_some_and(|after| document <= after) {
51                return Err(invalid(
52                    "private document selection is not strictly increasing",
53                ));
54            }
55            if let Some(source) = &source {
56                source.check_control(control)?;
57            }
58            selected.push((document, source))?;
59            after = Some(document);
60        }
61        let sources =
62            Budgeted::new(selected, control.memory().empty_reservation()).into_shared()?;
63        control.check()?;
64        Ok(Self {
65            sources,
66            control: control.clone(),
67            complete: false,
68        })
69    }
70
71    /// Select only these evaluated documents, masking every unselected document in the receiver. The caller must establish that this is the complete visible row selection; source lineage and width are still validated against the retained physical index.
72    pub fn without_unselected_documents(mut self) -> Self {
73        self.complete = true;
74        self
75    }
76
77    fn source(&self, document: DocId) -> Option<&Option<DiskANNReadSnapshot>> {
78        self.sources
79            .binary_search_by_key(&document, |(id, _)| *id)
80            .ok()
81            .map(|position| &self.sources[position].1)
82    }
83
84    fn next(&self, after: Option<DocId>) -> Option<DocId> {
85        let position = self
86            .sources
87            .partition_point(|(document, _)| after.is_some_and(|after| *document <= after));
88        self.sources.get(position).map(|(document, _)| *document)
89    }
90}
91
92pub(super) struct SelectedCanonical {
93    base: DiskANNReadSnapshot,
94    changes: DiskANNReadChanges,
95    generation: DiskANNGeneration,
96    counts: Option<DiskANNCanonicalCounts>,
97}
98
99impl SelectedCanonical {
100    pub(super) fn new(
101        base: DiskANNReadSnapshot,
102        changes: &DiskANNReadChanges,
103        built: &DiskANNOriginReader,
104        control: &StorageReadControl,
105    ) -> StorageBackendResult<Self> {
106        base.check_control(control)?;
107        changes.control.check()?;
108        for source in changes
109            .sources
110            .iter()
111            .filter_map(|(_, source)| source.as_ref())
112        {
113            source.check_control(control)?;
114            if !base.same_lineage(source) || base.dimensions() != source.dimensions() {
115                return Err(invalid(
116                    "private canonical source belongs to a different index lineage or width",
117                ));
118            }
119        }
120        let counts = populations::capture(&base, changes, built, control)?;
121        Ok(Self {
122            base,
123            changes: changes.clone(),
124            generation: built.manifest().input().generation,
125            counts,
126        })
127    }
128
129    fn source(&self, document: DocId) -> Option<&DiskANNReadSnapshot> {
130        self.changes.source(document).map_or_else(
131            || (!self.changes.complete).then_some(&self.base),
132            Option::as_ref,
133        )
134    }
135}
136
137impl DiskANNCanonicalRead for SelectedCanonical {
138    fn population_counts(
139        &self,
140        generation: DiskANNGeneration,
141        control: &StorageReadControl,
142    ) -> StorageBackendResult<Option<DiskANNCanonicalCounts>> {
143        self.check_control(control)?;
144        for source in self
145            .changes
146            .sources
147            .iter()
148            .filter_map(|(_, source)| source.as_ref())
149        {
150            source.check_control(control)?;
151        }
152        Ok((generation == self.generation)
153            .then_some(self.counts)
154            .flatten())
155    }
156
157    fn corpus_fingerprint(
158        &self,
159        control: &StorageReadControl,
160    ) -> StorageBackendResult<Option<[u8; 32]>> {
161        self.check_control(control)?;
162        crate::vector_index::selected_fingerprint(
163            self.dimensions(),
164            (!self.changes.complete).then_some(&self.base as &dyn crate::vector_index::VectorRead),
165            self.changes.sources.iter().map(|(document, source)| {
166                (
167                    *document,
168                    source
169                        .as_ref()
170                        .map(|source| source as &dyn crate::vector_index::VectorRead),
171                )
172            }),
173            control,
174        )
175    }
176    fn check_control(&self, control: &StorageReadControl) -> StorageBackendResult<()> {
177        self.changes.control.check()?;
178        self.base.check_control(control)
179    }
180
181    fn dimensions(&self) -> u32 {
182        self.base.dimensions()
183    }
184
185    fn next_document_after(
186        &self,
187        mut after: Option<DocId>,
188        control: &StorageReadControl,
189    ) -> StorageBackendResult<Option<DocId>> {
190        loop {
191            self.check_control(control)?;
192            let base = if self.changes.complete {
193                None
194            } else {
195                self.base.next_document_after(after, control)?
196            };
197            let selected = self.changes.next(after);
198            let Some(document) = first(base, selected) else {
199                return Ok(None);
200            };
201            if after.is_some_and(|after| document <= after) {
202                return Err(invalid("private canonical document cursor did not advance"));
203            }
204            // Preserve the base cursor's integrity checks, including unstamped values. Only a selected, genuinely absent replacement removes an enumerated identity.
205            if self.changes.source(document).is_none() || self.origin(document, control)?.is_some()
206            {
207                return Ok(Some(document));
208            }
209            after = Some(document);
210        }
211    }
212
213    fn origin(
214        &self,
215        document: DocId,
216        control: &StorageReadControl,
217    ) -> StorageBackendResult<Option<DiskANNVectorVersion>> {
218        self.check_control(control)?;
219        self.source(document)
220            .map(|source| source.origin(document, control))
221            .transpose()
222            .map(Option::flatten)
223    }
224
225    fn visit_document(
226        &self,
227        document: DocId,
228        control: &StorageReadControl,
229        visit: &mut DiskANNCanonicalVectorVisitor<'_>,
230    ) -> StorageBackendResult<Option<DiskANNVectorVersion>> {
231        self.check_control(control)?;
232        self.source(document)
233            .map(|source| source.visit_document(document, control, visit))
234            .transpose()
235            .map(Option::flatten)
236    }
237
238    fn read_vector(
239        &self,
240        document: DocId,
241        ordinal: u32,
242        control: &StorageReadControl,
243    ) -> StorageBackendResult<Option<uqa_core::memory::BudgetedVec<f32>>> {
244        self.check_control(control)?;
245        self.source(document)
246            .map(|source| source.read_vector(document, ordinal, control))
247            .transpose()
248            .map(Option::flatten)
249    }
250}
251
252impl DiskANNQueryRead for SelectedCanonical {
253    fn document_origin(
254        &self,
255        document: DocId,
256        control: &StorageReadControl,
257    ) -> StorageBackendResult<Option<DiskANNCanonicalOrigin>> {
258        self.check_control(control)?;
259        self.source(document)
260            .map(|source| source.document_origin(document, control))
261            .transpose()
262            .map(Option::flatten)
263    }
264
265    fn next_change_after(
266        &self,
267        mut after: Option<DocId>,
268        control: &StorageReadControl,
269    ) -> StorageBackendResult<Option<DiskANNChangeIdentity>> {
270        loop {
271            self.check_control(control)?;
272            let base = if self.changes.complete {
273                None
274            } else {
275                self.base.next_change_after(after, control)?
276            };
277            let selected = self.changes.next(after);
278            let Some(document) = first(base.map(DiskANNChangeIdentity::document), selected) else {
279                return Ok(None);
280            };
281            if after.is_some_and(|after| document <= after) {
282                return Err(invalid("private canonical change cursor did not advance"));
283            }
284            if let Some(source) = self.changes.source(document) {
285                // A private source may already have compacted its own journal. Its actual origin still marks coverage relative to the retained base generation.
286                if let Some(version) = source
287                    .as_ref()
288                    .map(|source| source.origin(document, control))
289                    .transpose()?
290                    .flatten()
291                {
292                    return Ok(Some(DiskANNChangeIdentity::new(document, version)));
293                }
294            } else {
295                return Ok(base);
296            }
297            after = Some(document);
298        }
299    }
300}
301
302fn first(left: Option<DocId>, right: Option<DocId>) -> Option<DocId> {
303    match (left, right) {
304        (Some(left), Some(right)) => Some(left.min(right)),
305        (left, right) => left.or(right),
306    }
307}
uqa_storage/diskann_index/query/selection.rs

uqa_storage/diskann_index/query/
selection.rs
