uqa_storage/vector_index/canonical/
selection.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! Evaluated document selections compose raw sources without transferring physical provenance.
8
9use super::{VectorRead, VectorReadSnapshot};
10use crate::{mvcc::VersionError, read_control::StorageReadControl, StorageBackendResult};
11use uqa_core::{
12    memory::{Budgeted, BudgetedVec},
13    DocId,
14};
15
16pub struct SelectedVectorRead {
17    base: Option<VectorReadSnapshot>,
18    selected: BudgetedVec<(DocId, Option<VectorReadSnapshot>)>,
19    dimensions: u32,
20    control: StorageReadControl,
21}
22
23impl SelectedVectorRead {
24    /// Retain already evaluated replacements/deletions in strictly increasing document order. None for the base means the supplied selection is the entire visible field. Capture inspects only controls and dimensions.
25    pub fn capture(
26        base: Option<VectorReadSnapshot>,
27        dimensions: u32,
28        selected: impl IntoIterator<Item = (DocId, Option<VectorReadSnapshot>)>,
29        control: &StorageReadControl,
30    ) -> StorageBackendResult<VectorReadSnapshot> {
31        control.check()?;
32        let check = |source: &VectorReadSnapshot| {
33            source.check_control(control)?;
34            if source.dimensions() != dimensions {
35                return Err(
36                    VersionError::InvalidEncoding("selected vector source width mismatch")
37                        .into_storage_error(),
38                );
39            }
40            Ok(())
41        };
42        if let Some(base) = &base {
43            check(base)?;
44        }
45        let mut sources = BudgetedVec::new(control.memory());
46        let mut after = None;
47        for (document, source) in selected {
48            control.check()?;
49            if after.is_some_and(|after| document <= after) {
50                return Err(VersionError::InvalidEncoding(
51                    "selected vector documents are not strictly increasing",
52                )
53                .into_storage_error());
54            }
55            if let Some(source) = &source {
56                check(source)?;
57            }
58            sources.push((document, source))?;
59            after = Some(document);
60        }
61        let source = Self {
62            base,
63            selected: sources,
64            dimensions,
65            control: control.clone(),
66        };
67        Ok(Budgeted::new(source, control.memory().empty_reservation()).into_shared()?)
68    }
69
70    fn source(&self, document: DocId) -> Option<&VectorReadSnapshot> {
71        match self
72            .selected
73            .binary_search_by_key(&document, |entry| entry.0)
74        {
75            Ok(position) => self.selected[position].1.as_ref(),
76            Err(_) => self.base.as_ref(),
77        }
78    }
79}
80
81impl VectorRead for SelectedVectorRead {
82    fn corpus_fingerprint(
83        &self,
84        control: &StorageReadControl,
85    ) -> StorageBackendResult<Option<[u8; 32]>> {
86        self.check_control(control)?;
87        super::selected_fingerprint(
88            self.dimensions,
89            self.base.as_deref(),
90            self.selected
91                .iter()
92                .map(|(document, source)| (*document, source.as_deref())),
93            control,
94        )
95    }
96    fn check_control(&self, control: &StorageReadControl) -> StorageBackendResult<()> {
97        self.control.check()?;
98        control.check()?;
99        if let Some(base) = &self.base {
100            base.check_control(control)?;
101        }
102        Ok(())
103    }
104    fn dimensions(&self) -> u32 {
105        self.dimensions
106    }
107    fn next_document_after(
108        &self,
109        mut after: Option<DocId>,
110        control: &StorageReadControl,
111    ) -> StorageBackendResult<Option<DocId>> {
112        loop {
113            self.check_control(control)?;
114            let base = self
115                .base
116                .as_ref()
117                .map(|base| base.next_document_after(after, control))
118                .transpose()?
119                .flatten();
120            let position = self
121                .selected
122                .partition_point(|entry| after.is_some_and(|after| entry.0 <= after));
123            let selected = self.selected.get(position).map(|entry| entry.0);
124            let Some(document) = base.into_iter().chain(selected).min() else {
125                return Ok(None);
126            };
127            if after.is_some_and(|after| document <= after) {
128                return Err(VersionError::InvalidEncoding(
129                    "selected vector cursor did not advance",
130                )
131                .into_storage_error());
132            }
133            if self.document_vector_count(document, control)? != 0 {
134                return Ok(Some(document));
135            }
136            after = Some(document);
137        }
138    }
139    fn document_vector_count(
140        &self,
141        document: DocId,
142        control: &StorageReadControl,
143    ) -> StorageBackendResult<u64> {
144        self.check_control(control)?;
145        self.source(document).map_or(Ok(0), |source| {
146            source.document_vector_count(document, control)
147        })
148    }
149    fn read_vector(
150        &self,
151        document: DocId,
152        ordinal: u32,
153        control: &StorageReadControl,
154    ) -> StorageBackendResult<Option<BudgetedVec<f32>>> {
155        self.check_control(control)?;
156        self.source(document).map_or(Ok(None), |source| {
157            source.read_vector(document, ordinal, control)
158        })
159    }
160}
uqa_storage/vector_index/canonical/selection.rs

uqa_storage/vector_index/canonical/
selection.rs
