uqa_storage/diskann_index/
scoring.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! Canonical tensor scores precede document selection and posting construction.
8
9use super::{format::DiskANNVectorVersion, DiskANNCanonicalRead};
10use crate::{
11    mvcc::VersionError,
12    read_control::StorageReadControl,
13    vector_index::{
14        cosine_similarity_controlled,
15        query::{postings_from_unique_scores, VectorQueryBuffer},
16        validate_threshold, validate_vector_values_controlled,
17    },
18    StorageBackendResult,
19};
20use uqa_core::{DocId, PostingList};
21
22pub(super) mod selection;
23mod statistics;
24pub use statistics::DiskANNScoringStats;
25pub(crate) use statistics::ScoringStatsExt;
26#[cfg(test)]
27mod tests;
28
29/// One vector-bearing document scored from every ordinal on a fixed canonical source. This is raw cosine, not navigation distance or probability.
30#[derive(Debug, Clone, Copy)]
31pub struct DiskANNDocumentScore {
32    document: DocId,
33    version: DiskANNVectorVersion,
34    vectors: u64,
35    score: f32,
36}
37
38impl DiskANNDocumentScore {
39    /// Canonical document identity before posting construction.
40    pub fn document(&self) -> DocId {
41        self.document
42    }
43    /// Original mutation version shared by every scored ordinal.
44    pub fn version(&self) -> DiskANNVectorVersion {
45        self.version
46    }
47    /// Number of visible tensor ordinals, independent of the document candidate quota.
48    pub fn vector_count(&self) -> u64 {
49        self.vectors
50    }
51    /// Canonical sequential-f32 cosine maximum, including its existing numeric edge behavior.
52    pub fn raw_cosine(&self) -> f32 {
53        self.score
54    }
55}
56
57/// Borrows a validated query and one retained canonical boundary. Exact scans serve explicit exact/threshold and numeric-edge routes; they do not replace missing ANN coverage or recover from graph errors.
58pub struct DiskANNCanonicalScorer<'a> {
59    source: &'a dyn DiskANNCanonicalRead,
60    query: &'a [f32],
61    control: &'a StorageReadControl,
62}
63
64impl<'a> DiskANNCanonicalScorer<'a> {
65    pub fn new(
66        source: &'a dyn DiskANNCanonicalRead,
67        query: &'a [f32],
68        control: &'a StorageReadControl,
69    ) -> StorageBackendResult<Self> {
70        source.check_control(control)?;
71        if source.dimensions() == 0 {
72            return Err(invalid("canonical scoring requires positive dimensions"));
73        }
74        validate_vector_values_controlled(source.dimensions(), query, Some(control))?;
75        source.check_control(control)?;
76        Ok(Self {
77            source,
78            query,
79            control,
80        })
81    }
82
83    /// Score the complete tensor using the canonical total-order maximum. Absent documents and explicit empty tensors have no score.
84    pub fn score_document(
85        &self,
86        document: DocId,
87    ) -> StorageBackendResult<Option<DiskANNDocumentScore>> {
88        self.document(document).map(|(_, score)| score)
89    }
90
91    /// Validate a physical candidate's original version and ordinal before admitting its complete document score. Stale versions are masked; a matching origin with an impossible ordinal is corrupt.
92    pub fn score_candidate(
93        &self,
94        document: DocId,
95        ordinal: u32,
96        version: DiskANNVectorVersion,
97    ) -> StorageBackendResult<Option<DiskANNDocumentScore>> {
98        self.source.check_control(self.control)?;
99        if self.source.origin(document, self.control)? != Some(version) {
100            self.source.check_control(self.control)?;
101            return Ok(None);
102        }
103        let (actual, score) = self.document(document)?;
104        if actual != Some(version) {
105            return Err(invalid("canonical origin changed on a retained source"));
106        }
107        let score = score
108            .filter(|score| u64::from(ordinal) < score.vectors)
109            .ok_or_else(|| invalid("candidate ordinal is outside its canonical tensor"))?;
110        Ok(Some(score))
111    }
112
113    /// Stream one score per vector-bearing document in ascending document order. On any error, the caller must discard all output already produced by its callback.
114    pub fn visit_scores(
115        &self,
116        visit: &mut dyn FnMut(DiskANNDocumentScore) -> StorageBackendResult<()>,
117    ) -> StorageBackendResult<()> {
118        self.source.check_control(self.control)?;
119        let mut after = None;
120        while let Some(document) = self.source.next_document_after(after, self.control)? {
121            if after.is_some_and(|previous| document <= previous) {
122                return Err(invalid("canonical scoring cursor did not advance"));
123            }
124            let (version, score) = self.document(document)?;
125            if version.is_none() {
126                return Err(invalid("enumerated canonical document has no origin"));
127            }
128            if let Some(score) = score {
129                visit(score)?;
130            }
131            after = Some(document);
132        }
133        self.source.check_control(self.control)
134    }
135
136    /// Exact document top-k with workspace proportional to selected documents, not corpus size. Scores descend with document-ID tie breaking; returned postings retain their ordinary document order and caller-owned result boundary.
137    pub fn search_exact_knn(&self, k: usize) -> StorageBackendResult<PostingList> {
138        self.search_exact_knn_with_stats(k)
139            .map(|(postings, _)| postings)
140    }
141
142    /// Return the unchanged exact result with successful tensor-scoring work from this invocation.
143    pub fn search_exact_knn_with_stats(
144        &self,
145        k: usize,
146    ) -> StorageBackendResult<(PostingList, DiskANNScoringStats)> {
147        self.source.check_control(self.control)?;
148        let mut stats = DiskANNScoringStats::default();
149        if k == 0 {
150            return Ok((PostingList::new(), stats));
151        }
152        let mut selected = selection::TopK::new(k, self.control);
153        self.visit_scores(&mut |score| {
154            stats.record(score.vector_count())?;
155            selected.offer(score)
156        })?;
157        let result = selected.finish(self.control)?;
158        self.source.check_control(self.control)?;
159        Ok((result, stats))
160    }
161
162    /// Scan all canonical documents and apply the existing raw-cosine threshold after full tensor reduction. NaN comparison and finite-threshold validation match the existing exact index.
163    pub fn search_threshold(&self, threshold: f32) -> StorageBackendResult<PostingList> {
164        self.search_threshold_with_stats(threshold)
165            .map(|(postings, _)| postings)
166    }
167
168    /// Count every scored tensor before threshold filtering, including tensors that produce no result row.
169    pub fn search_threshold_with_stats(
170        &self,
171        threshold: f32,
172    ) -> StorageBackendResult<(PostingList, DiskANNScoringStats)> {
173        self.source.check_control(self.control)?;
174        validate_threshold(threshold)?;
175        let mut scores = VectorQueryBuffer::new(Some(self.control));
176        let mut stats = DiskANNScoringStats::default();
177        self.visit_scores(&mut |score| {
178            stats.record(score.vector_count())?;
179            if score.score >= threshold {
180                scores.push((score.document, score.score))?;
181            }
182            Ok(())
183        })?;
184        let result = postings_from_unique_scores(scores, None, Some(self.control))?;
185        self.source.check_control(self.control)?;
186        Ok((result, stats))
187    }
188
189    fn document(
190        &self,
191        document: DocId,
192    ) -> StorageBackendResult<(Option<DiskANNVectorVersion>, Option<DiskANNDocumentScore>)> {
193        self.source.check_control(self.control)?;
194        let mut version = None;
195        let mut best: Option<f32> = None;
196        let mut count = 0_u64;
197        let mut failure = None;
198        let returned =
199            self.source
200                .visit_document(document, self.control, &mut |ordinal, current, raw| {
201                    if failure.is_none() {
202                        let result = (|| {
203                            self.source.check_control(self.control)?;
204                            if u64::from(ordinal) != count
205                                || version.is_some_and(|version| version != current)
206                            {
207                                return Err(invalid(
208                                    "canonical tensor identity changed while scoring",
209                                ));
210                            }
211                            validate_vector_values_controlled(
212                                self.source.dimensions(),
213                                raw,
214                                Some(self.control),
215                            )?;
216                            let score = cosine_similarity_controlled(self.query, raw, || {
217                                self.source.check_control(self.control)
218                            })?;
219                            if best.is_none_or(|best| score.total_cmp(&best).is_gt()) {
220                                best = Some(score);
221                            }
222                            version = Some(current);
223                            count += 1;
224                            Ok(())
225                        })();
226                        if let Err(error) = result {
227                            failure = Some(error);
228                        }
229                    }
230                    if failure.is_some() {
231                        Err(invalid("canonical score consumer rejected an ordinal"))
232                    } else {
233                        Ok(())
234                    }
235                });
236        if let Some(error) = failure {
237            return Err(error);
238        }
239        let returned = returned?;
240        self.source.check_control(self.control)?;
241        if count != 0 && returned != version {
242            return Err(invalid("canonical tensor returned a different origin"));
243        }
244        let score = best.map(|score| DiskANNDocumentScore {
245            document,
246            version: version.expect("a scored ordinal has an origin"),
247            vectors: count,
248            score,
249        });
250        Ok((returned, score))
251    }
252}
253
254fn invalid(message: &'static str) -> crate::StorageBackendError {
255    VersionError::InvalidEncoding(message).into_storage_error()
256}
uqa_storage/diskann_index/scoring.rs

uqa_storage/diskann_index/
scoring.rs
