uqa_storage/vector_index/retained/
builder.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! Evaluated vectors move into one charged immutable corpus before publication.
8
9use super::{RetainedVectorIndex, VectorEntries};
10use crate::{
11    read_control::StorageReadControl, vector_index::validate_vector_values, StorageBackendError,
12    StorageBackendResult,
13};
14use uqa_core::{
15    memory::{Budgeted, BudgetedVec, MemoryError, MemoryReservation},
16    ordering::sort_by_with_control,
17    DocId,
18};
19
20/// Append selected document tensors without copying their buffers or reserving their payload twice. The caller converts inputs under this builder's allowance and transfers their reservations. Finalization orders identities, validates unique tensor ordinals and retains the same allowance for readers.
21pub struct RetainedVectorIndexBuilder {
22    entries: BudgetedVec<(DocId, u32, Vec<f32>)>,
23    payload: MemoryReservation,
24    dimensions: u32,
25    control: StorageReadControl,
26}
27
28impl RetainedVectorIndexBuilder {
29    pub fn new(dimensions: u32, control: &StorageReadControl) -> Self {
30        Self {
31            entries: BudgetedVec::new(control.memory()),
32            payload: control.memory().empty_reservation(),
33            dimensions,
34            control: control.clone(),
35        }
36    }
37
38    pub fn dimensions(&self) -> u32 {
39        self.dimensions
40    }
41
42    pub fn add_document(
43        &mut self,
44        id: DocId,
45        vectors: Budgeted<Vec<Vec<f32>>>,
46    ) -> StorageBackendResult<()> {
47        self.control.check()?;
48        u32::try_from(vectors.len().saturating_sub(1)).map_err(|_| {
49            StorageBackendError::Other("vector ordinal exceeds u32 index format".into())
50        })?;
51        // Drop unadopted vector buffers before releasing their reservation on every error.
52        let mut pending = vectors.into_parts();
53        if !pending.1.budget().shares_allowance(self.control.memory()) {
54            return Err(StorageBackendError::Other(
55                "retained vectors must use the original index allowance".into(),
56            ));
57        }
58        let mut bytes = 0_usize;
59        for vector in &pending.0 {
60            self.control.check()?;
61            bytes = vector
62                .capacity()
63                .checked_mul(size_of::<f32>())
64                .and_then(|payload| bytes.checked_add(payload))
65                .ok_or(MemoryError::SizeOverflow)?;
66        }
67        let required = pending
68            .0
69            .capacity()
70            .checked_mul(size_of::<Vec<f32>>())
71            .and_then(|headers| headers.checked_add(bytes))
72            .ok_or(MemoryError::SizeOverflow)?;
73        if pending.1.bytes() < required {
74            return Err(StorageBackendError::Other(
75                "retained vector reservation does not cover its buffers".into(),
76            ));
77        }
78        for vector in &pending.0 {
79            self.control.check()?;
80            validate_vector_values(self.dimensions, vector)?;
81        }
82        self.entries.reserve(pending.0.len())?;
83        let original = self.entries.len();
84        let append = (|| -> StorageBackendResult<()> {
85            for (ordinal, vector) in pending.0.drain(..).enumerate() {
86                self.control.check()?;
87                let ordinal = u32::try_from(ordinal).expect("validated vector ordinal count");
88                self.entries.push((id, ordinal, vector))?;
89            }
90            self.control.check()
91        })();
92        if let Err(error) = append {
93            self.entries.truncate(original);
94            return Err(error);
95        }
96        self.payload.absorb(pending.1.split(bytes));
97        // The now-empty outer buffer drops with the rest of its reservation.
98        Ok(())
99    }
100
101    pub fn finish(mut self) -> StorageBackendResult<RetainedVectorIndex> {
102        sort_by_with_control(
103            &mut self.entries,
104            &mut || self.control.check(),
105            |left, right, _| Ok((left.0, left.1).cmp(&(right.0, right.1))),
106        )?;
107        let mut previous: Option<(DocId, u32)> = None;
108        for (id, ordinal, _) in self.entries.iter() {
109            self.control.check()?;
110            let expected = match previous {
111                Some((previous_id, previous_ordinal)) if previous_id == *id => {
112                    previous_ordinal.checked_add(1)
113                }
114                _ => Some(0),
115            };
116            if expected != Some(*ordinal) {
117                return Err(StorageBackendError::Other(
118                    "retained vector input repeats a document or tensor ordinal".into(),
119                ));
120            }
121            previous = Some((*id, *ordinal));
122        }
123        let (entries, mut memory): (VectorEntries, _) = self.entries.into_parts();
124        memory.absorb(self.payload);
125        RetainedVectorIndex::from_entries(
126            Budgeted::new(entries, memory),
127            self.dimensions,
128            "memory-bruteforce",
129            &self.control,
130        )
131    }
132}
uqa_storage/vector_index/retained/builder.rs

uqa_storage/vector_index/retained/
builder.rs
