uqa_storage/vector_index/
collection.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! Live vector registrations and immutable retained collections preserve their distinct allocation owners.
8
9use std::{borrow::Borrow, collections::BTreeMap, sync::Arc};
10use uqa_core::{
11    memory::{Budgeted, BudgetedMap, BudgetedMapIter, BudgetedString, MemoryReservation},
12    FieldName,
13};
14
15use super::VectorIndex;
16use crate::{
17    read_control::StorageReadControl, ReadOnlySnapshot, StorageBackendError, StorageBackendResult,
18};
19
20type LiveIndexes = BTreeMap<FieldName, Box<dyn VectorIndex>>;
21type RetainedIndexes = Arc<Budgeted<BudgetedMap<FieldName, Box<dyn VectorIndex>>>>;
22
23/// A table's vector registrations. Retained collections own controlled ordered-map nodes and share their names, index handles and original allowance with nested captures. Live mutation requires an explicit fallible borrow; a retained collection never converts itself to an uncharged writable map.
24pub struct VectorIndexes(Contents);
25
26enum Contents {
27    Live(LiveIndexes),
28    Retained {
29        indexes: Option<RetainedIndexes>,
30        control: StorageReadControl,
31    },
32}
33
34impl Default for VectorIndexes {
35    fn default() -> Self {
36        Self(Contents::Live(BTreeMap::new()))
37    }
38}
39
40impl From<LiveIndexes> for VectorIndexes {
41    fn from(indexes: LiveIndexes) -> Self {
42        Self(Contents::Live(indexes))
43    }
44}
45
46impl VectorIndexes {
47    pub fn live_mut(&mut self) -> StorageBackendResult<&mut LiveIndexes> {
48        match &mut self.0 {
49            Contents::Live(indexes) => Ok(indexes),
50            Contents::Retained { control, .. } => {
51                control.check()?;
52                Err(StorageBackendError::Other(
53                    "cannot mutate retained vector registrations".into(),
54                ))
55            }
56        }
57    }
58
59    pub fn get<Q: Ord + ?Sized>(&self, field: &Q) -> Option<&(dyn VectorIndex + 'static)>
60    where
61        FieldName: Borrow<Q>,
62    {
63        let index = match &self.0 {
64            Contents::Live(indexes) => indexes.get(field),
65            Contents::Retained { indexes, .. } => {
66                indexes.as_ref().and_then(|indexes| indexes.get(field))
67            }
68        };
69        index.map(Box::as_ref)
70    }
71
72    pub fn contains_key<Q: Ord + ?Sized>(&self, field: &Q) -> bool
73    where
74        FieldName: Borrow<Q>,
75    {
76        self.get(field).is_some()
77    }
78
79    pub fn len(&self) -> usize {
80        match &self.0 {
81            Contents::Live(indexes) => indexes.len(),
82            Contents::Retained { indexes, .. } => {
83                indexes.as_ref().map_or(0, |indexes| indexes.len())
84            }
85        }
86    }
87
88    pub fn is_empty(&self) -> bool {
89        self.len() == 0
90    }
91
92    pub fn iter(&self) -> VectorIndexesIter<'_> {
93        VectorIndexesIter(match &self.0 {
94            Contents::Live(indexes) => Entries::Live(indexes.iter()),
95            Contents::Retained { indexes, .. } => indexes
96                .as_ref()
97                .map_or(Entries::Empty, |indexes| Entries::Retained(indexes.iter())),
98        })
99    }
100
101    pub fn keys(&self) -> impl ExactSizeIterator<Item = &FieldName> {
102        self.iter().map(|(field, _)| field)
103    }
104
105    pub fn values(&self) -> impl ExactSizeIterator<Item = &(dyn VectorIndex + 'static)> {
106        self.iter().map(|(_, index)| index)
107    }
108
109    /// Snapshot live providers under a caller allowance, or share an already retained collection without copying names, nodes or handles. Nested captures preserve the original provider allowance and cancellation boundary.
110    pub fn capture(
111        source: &dyn VectorIndexSource,
112        control: &StorageReadControl,
113    ) -> StorageBackendResult<Self> {
114        control.check()?;
115        if let Some(retained) = source.retained_collection()? {
116            match &retained.0 {
117                Contents::Retained { control, .. } => control.check()?,
118                Contents::Live(_) => {
119                    return Err(StorageBackendError::Other(
120                        "vector provider returned live registrations as a retained collection"
121                            .into(),
122                    ))
123                }
124            }
125            return Ok(retained);
126        }
127        let mut retained = RetainedVectorIndexesBuilder::new(control);
128        source.visit(&mut |field, index| {
129            let (field, memory) = copy_field(field, control)?;
130            let index = index.snapshot_with_control(control)?;
131            retained.insert_admitted(field, ReadOnlySnapshot::new(index), memory)
132        })?;
133        retained.finish()
134    }
135}
136
137impl<Q: Ord + ?Sized> std::ops::Index<&Q> for VectorIndexes
138where
139    FieldName: Borrow<Q>,
140{
141    type Output = dyn VectorIndex;
142    fn index(&self, field: &Q) -> &Self::Output {
143        self.get(field).expect("missing vector index field")
144    }
145}
146
147pub struct VectorIndexesIter<'a>(Entries<'a>);
148
149#[expect(
150    clippy::large_enum_variant,
151    reason = "fixed iterator stack avoids unadmitted heap scratch"
152)]
153enum Entries<'a> {
154    Empty,
155    Live(std::collections::btree_map::Iter<'a, FieldName, Box<dyn VectorIndex>>),
156    Retained(BudgetedMapIter<'a, FieldName, Box<dyn VectorIndex>>),
157}
158
159impl<'a> Iterator for VectorIndexesIter<'a> {
160    type Item = (&'a FieldName, &'a (dyn VectorIndex + 'static));
161
162    fn next(&mut self) -> Option<Self::Item> {
163        let next = match &mut self.0 {
164            Entries::Empty => None,
165            Entries::Live(iter) => iter.next(),
166            Entries::Retained(iter) => iter.next(),
167        };
168        next.map(|(field, index)| (field, index.as_ref()))
169    }
170
171    fn size_hint(&self) -> (usize, Option<usize>) {
172        match &self.0 {
173            Entries::Empty => (0, Some(0)),
174            Entries::Live(iter) => iter.size_hint(),
175            Entries::Retained(iter) => iter.size_hint(),
176        }
177    }
178}
179
180impl ExactSizeIterator for VectorIndexesIter<'_> {}
181impl std::iter::FusedIterator for VectorIndexesIter<'_> {}
182
183impl<'a> IntoIterator for &'a VectorIndexes {
184    type Item = (&'a FieldName, &'a (dyn VectorIndex + 'static));
185    type IntoIter = VectorIndexesIter<'a>;
186    fn into_iter(self) -> Self::IntoIter {
187        self.iter()
188    }
189}
190
191/// Borrow physical vector registrations without requiring an owned metadata copy.
192pub trait VectorIndexSource {
193    fn visit(
194        &self,
195        visitor: &mut dyn FnMut(&str, &dyn VectorIndex) -> StorageBackendResult<()>,
196    ) -> StorageBackendResult<()>;
197
198    /// An immutable collection may share its original allocation owner. Implementations must never capture or clone a live collection here.
199    fn retained_collection(&self) -> StorageBackendResult<Option<VectorIndexes>> {
200        Ok(None)
201    }
202}
203
204impl VectorIndexSource for LiveIndexes {
205    fn visit(
206        &self,
207        visitor: &mut dyn FnMut(&str, &dyn VectorIndex) -> StorageBackendResult<()>,
208    ) -> StorageBackendResult<()> {
209        for (field, index) in self {
210            visitor(field, index.as_ref())?;
211        }
212        Ok(())
213    }
214}
215
216impl VectorIndexSource for VectorIndexes {
217    fn visit(
218        &self,
219        visitor: &mut dyn FnMut(&str, &dyn VectorIndex) -> StorageBackendResult<()>,
220    ) -> StorageBackendResult<()> {
221        for (field, index) in self {
222            visitor(field, index)?;
223        }
224        Ok(())
225    }
226
227    fn retained_collection(&self) -> StorageBackendResult<Option<VectorIndexes>> {
228        match &self.0 {
229            Contents::Live(_) => Ok(None),
230            Contents::Retained { indexes, control } => {
231                control.check()?;
232                Ok(Some(Self(Contents::Retained {
233                    indexes: indexes.clone(),
234                    control: control.clone(),
235                })))
236            }
237        }
238    }
239}
240
241/// Assemble reconstructed vector indexes without opaque standard-map nodes. Every node, copied name, boxed adapter and shared owner is reserved before allocation.
242pub struct RetainedVectorIndexesBuilder {
243    indexes: BudgetedMap<FieldName, Box<dyn VectorIndex>>,
244    control: StorageReadControl,
245}
246
247impl RetainedVectorIndexesBuilder {
248    pub fn new(control: &StorageReadControl) -> Self {
249        Self {
250            indexes: BudgetedMap::new(control.memory()),
251            control: control.clone(),
252        }
253    }
254
255    /// Admit and copy a field name before a reconstruction producer retains it in intermediate or final metadata.
256    pub fn copy_field(
257        field: &str,
258        control: &StorageReadControl,
259    ) -> StorageBackendResult<(FieldName, MemoryReservation)> {
260        copy_field(field, control)
261    }
262
263    /// Adopt an already admitted field name and a provider-owned immutable index. The name lease must cover its actual capacity and belong to this collection's allowance.
264    pub fn insert_admitted<T: VectorIndex + 'static>(
265        &mut self,
266        field: FieldName,
267        index: T,
268        mut memory: MemoryReservation,
269    ) -> StorageBackendResult<()> {
270        self.control.check()?;
271        if !memory.budget().shares_allowance(self.control.memory())
272            || memory.bytes() < field.capacity()
273        {
274            return Err(StorageBackendError::Other(
275                "retained vector name lacks its original allowance".into(),
276            ));
277        }
278        if self.indexes.contains_key(&field) {
279            return Err(StorageBackendError::Other(
280                "retained vector field is registered twice".into(),
281            ));
282        }
283        memory.grow(size_of::<ReadOnlySnapshot<T>>())?;
284        let index = ReadOnlySnapshot::from_budgeted(Budgeted::new(index, memory))?;
285        let index: Box<dyn VectorIndex> = Box::new(index);
286        self.control.check()?;
287        self.indexes.insert(field, index)?;
288        Ok(())
289    }
290
291    pub fn finish(self) -> StorageBackendResult<VectorIndexes> {
292        self.control.check()?;
293        let indexes = if self.indexes.is_empty() {
294            None
295        } else {
296            Some(
297                Budgeted::new(self.indexes, self.control.memory().empty_reservation())
298                    .into_shared()?,
299            )
300        };
301        Ok(VectorIndexes(Contents::Retained {
302            indexes,
303            control: self.control,
304        }))
305    }
306}
307
308fn copy_field(
309    field: &str,
310    control: &StorageReadControl,
311) -> StorageBackendResult<(FieldName, MemoryReservation)> {
312    control.check()?;
313    let mut name = BudgetedString::new(control.memory());
314    name.reserve(field.len())?;
315    for (offset, character) in field.chars().enumerate() {
316        if offset % 1024 == 0 {
317            control.check()?;
318        }
319        name.push(character)?;
320    }
321    Ok(name.into_parts())
322}
323
324#[cfg(test)]
325mod tests;
uqa_storage/vector_index/collection.rs

uqa_storage/vector_index/
collection.rs
