uqa_storage/diskann_index/
query.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! Document selection combines one physical generation with its retained canonical/change view.
8
9use std::sync::Arc;
10use uqa_core::PostingList;
11
12use super::{
13    pages::{DiskANNOriginReader, DiskANNPageSource, DiskANNReadLimits, DiskANNReader},
14    search::{DiskANNTraversal, DiskANNTraversalStats},
15    DiskANNCanonicalCounts, DiskANNCanonicalRead, DiskANNCanonicalScorer,
16    DiskANNCanonicalVectorVisitor, DiskANNQueryRead, ExactVectorReason, NavigationInput,
17};
18use crate::{
19    read_control::StorageReadControl, vector_index::DiskANNIndexParams, StorageBackendResult,
20};
21
22mod candidates;
23mod matching;
24mod retained;
25mod selection;
26mod statistics;
27mod values;
28pub use retained::RetainedDiskANNIndex;
29pub use selection::{DiskANNReadChanges, DiskANNReadSnapshot};
30pub use statistics::DiskANNQueryWork;
31#[cfg(test)]
32mod tests;
33
34/// Raw canonical document scores and the route that produced their candidate support. PQ estimates never enter the posting payloads; these counters describe logical traversal, not physical I/O.
35pub struct DiskANNQueryResult {
36    pub postings: PostingList,
37    pub exact_reason: Option<ExactVectorReason>,
38    pub traversal: DiskANNTraversalStats,
39    pub work: DiskANNQueryWork,
40    /// Complete raw-view populations observed by this successful invocation. Routes that do not verify all raw coverage leave this unknown.
41    pub populations: Option<DiskANNCanonicalCounts>,
42}
43
44/// A reusable prepared generation on one fixed canonical/catalog view. Opening verifies resident PQ and complete origins once; later queries reuse those owners and load graph pages lazily. Execution owns logical read observations before invoking this storage algorithm.
45pub struct DiskANNQuery<'a> {
46    canonical: &'a dyn DiskANNQueryRead,
47    reader: DiskANNReader,
48    origins: DiskANNOriginReader,
49    control: StorageReadControl,
50}
51
52impl<'a> DiskANNQuery<'a> {
53    /// Open the physical source selected by the caller's actual retained catalog/canonical view. Provider query adapters perform that binding; matching a caller-supplied generation label alone does not establish it. All failures propagate without an exact-scan fallback.
54    pub fn open(
55        canonical: &'a dyn DiskANNQueryRead,
56        source: Arc<dyn DiskANNPageSource>,
57        parameters: DiskANNIndexParams,
58        limits: DiskANNReadLimits,
59        control: &StorageReadControl,
60    ) -> StorageBackendResult<Self> {
61        canonical.check_control(control)?;
62        let reader = DiskANNReader::open(
63            source.clone(),
64            canonical.dimensions(),
65            parameters,
66            limits,
67            control,
68        )?;
69        let origins = DiskANNOriginReader::open(source, limits.max_record_bytes, control)?;
70        if reader.manifest() != origins.manifest() {
71            return Err(invalid(
72                "physical manifest changed while preparing the query",
73            ));
74        }
75        canonical.check_control(control)?;
76        Ok(Self {
77            canonical,
78            reader,
79            origins,
80            control: control.clone(),
81        })
82    }
83
84    /// Approximate document top-k after complete tensor reranking, exact uncovered changes and numeric side entries. Exhausted ANN support continues through unexpanded nodes until the document quota is met or the entire base is exhausted.
85    pub fn search_knn(
86        &self,
87        query: &[f32],
88        k: usize,
89        control: &StorageReadControl,
90    ) -> StorageBackendResult<DiskANNQueryResult> {
91        self.check(control)?;
92        let scorer = DiskANNCanonicalScorer::new(self, query, control)?;
93        if k == 0 {
94            self.check(control)?;
95            return Ok(result(
96                PostingList::new(),
97                None,
98                DiskANNTraversalStats::default(),
99                DiskANNQueryWork::default(),
100            ));
101        }
102        let navigation =
103            match NavigationInput::from_raw(self.canonical.dimensions(), query, control)? {
104                NavigationInput::Exact(reason) => {
105                    let (postings, exact) =
106                        if let Some(source) = self.canonical.unversioned_vectors() {
107                            values::exact(&**source, query, k, control)?
108                        } else {
109                            scorer.search_exact_knn_with_stats(k)?
110                        };
111                    self.check(control)?;
112                    return Ok(result(
113                        postings,
114                        Some(reason),
115                        DiskANNTraversalStats::default(),
116                        DiskANNQueryWork {
117                            exact,
118                            ..DiskANNQueryWork::default()
119                        },
120                    ));
121                }
122                NavigationInput::Navigable(navigation) => navigation,
123            };
124        let mut traversal = DiskANNTraversal::new(self.reader.clone(), &navigation, control)?;
125        drop(navigation);
126        let mut candidates = candidates::Candidates::new(self, scorer, k, control);
127        loop {
128            self.check(control)?;
129            let nodes = traversal.next_beam()?;
130            if nodes.is_empty() {
131                break;
132            }
133            candidates.nodes(&nodes)?;
134        }
135        candidates.side()?;
136        candidates.changes()?;
137        let mut raw_vectors = None;
138        if let Some(source) = self.canonical.unversioned_vectors() {
139            let mut current = 0_u64;
140            values::visit(&**source, control, &mut |document| {
141                self.check(control)?;
142                let vectors =
143                    if let Some(origin) = self.canonical.document_origin(document, control)? {
144                        origin.count()
145                    } else {
146                        let (score, vectors) = values::score(&**source, document, query, control)?;
147                        if let Some(score) = score {
148                            candidates.offer_unversioned(document, score, vectors)?;
149                        }
150                        vectors
151                    };
152                current = current
153                    .checked_add(vectors)
154                    .ok_or_else(|| invalid("raw canonical population overflow"))?;
155                Ok(())
156            })?;
157            raw_vectors = Some(current);
158        }
159        while candidates.len() < k {
160            self.check(control)?;
161            let nodes = traversal.complete_next_beam()?;
162            if nodes.is_empty() {
163                break;
164            }
165            candidates.nodes(&nodes)?;
166        }
167        let mut work = candidates.work();
168        work.pages = traversal.page_stats();
169        let postings = candidates.finish()?;
170        self.check(control)?;
171        let mut result = result(postings, None, traversal.stats(), work);
172        if let Some(current) = raw_vectors {
173            let changed = work
174                .changed
175                .vectors
176                .checked_add(work.unversioned.vectors)
177                .ok_or_else(|| invalid("raw uncovered population overflow"))?;
178            result.populations = Some(DiskANNCanonicalCounts::new(current, changed)?);
179        }
180        Ok(result)
181    }
182
183    /// Exact threshold evaluation streams all visible canonical tensors, including changes and empty replacements, independently of ANN membership.
184    pub fn search_threshold(
185        &self,
186        query: &[f32],
187        threshold: f32,
188        control: &StorageReadControl,
189    ) -> StorageBackendResult<PostingList> {
190        self.search_threshold_with_stats(query, threshold, control)
191            .map(|result| result.postings)
192    }
193
194    /// Return threshold results and scoring work from their one exact invocation, without a second scan.
195    pub fn search_threshold_with_stats(
196        &self,
197        query: &[f32],
198        threshold: f32,
199        control: &StorageReadControl,
200    ) -> StorageBackendResult<DiskANNQueryResult> {
201        self.check(control)?;
202        let scorer = DiskANNCanonicalScorer::new(self, query, control)?;
203        let (postings, exact) = if let Some(source) = self.canonical.unversioned_vectors() {
204            values::threshold(&**source, query, threshold, control)?
205        } else {
206            scorer.search_threshold_with_stats(threshold)?
207        };
208        self.check(control)?;
209        Ok(result(
210            postings,
211            None,
212            DiskANNTraversalStats::default(),
213            DiskANNQueryWork {
214                exact,
215                ..DiskANNQueryWork::default()
216            },
217        ))
218    }
219
220    fn check(&self, control: &StorageReadControl) -> StorageBackendResult<()> {
221        self.control.check()?;
222        self.canonical.check_control(control)
223    }
224}
225
226impl DiskANNCanonicalRead for DiskANNQuery<'_> {
227    fn check_control(&self, control: &StorageReadControl) -> StorageBackendResult<()> {
228        self.check(control)
229    }
230    fn dimensions(&self) -> u32 {
231        self.canonical.dimensions()
232    }
233    fn next_document_after(
234        &self,
235        after: Option<uqa_core::DocId>,
236        control: &StorageReadControl,
237    ) -> StorageBackendResult<Option<uqa_core::DocId>> {
238        self.check(control)?;
239        let result = self.canonical.next_document_after(after, control)?;
240        self.check(control)?;
241        Ok(result)
242    }
243    fn origin(
244        &self,
245        document: uqa_core::DocId,
246        control: &StorageReadControl,
247    ) -> StorageBackendResult<Option<super::format::DiskANNVectorVersion>> {
248        self.check(control)?;
249        let result = self.canonical.origin(document, control)?;
250        self.check(control)?;
251        Ok(result)
252    }
253    fn visit_document(
254        &self,
255        document: uqa_core::DocId,
256        control: &StorageReadControl,
257        visit: &mut DiskANNCanonicalVectorVisitor<'_>,
258    ) -> StorageBackendResult<Option<super::format::DiskANNVectorVersion>> {
259        self.check(control)?;
260        let result =
261            self.canonical
262                .visit_document(document, control, &mut |ordinal, version, raw| {
263                    self.check(control)?;
264                    visit(ordinal, version, raw)
265                })?;
266        self.check(control)?;
267        Ok(result)
268    }
269}
270
271fn result(
272    postings: PostingList,
273    exact_reason: Option<ExactVectorReason>,
274    traversal: DiskANNTraversalStats,
275    work: DiskANNQueryWork,
276) -> DiskANNQueryResult {
277    DiskANNQueryResult {
278        postings,
279        exact_reason,
280        traversal,
281        work,
282        populations: None,
283    }
284}
285
286fn invalid(message: &'static str) -> crate::StorageBackendError {
287    crate::mvcc::VersionError::InvalidEncoding(message).into_storage_error()
288}
uqa_storage/diskann_index/query.rs

uqa_storage/diskann_index/
query.rs
