uqa_storage/ivf_index/spill/
canonical.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! Canonical creation preserves complete-document training boundaries without collecting tensors.
8
9use super::{corrupt, key, IVFMetadataSnapshot, IVFPreparedMetadata, IVFRestoreBuilder, IVFState};
10use crate::{
11    read_control::StorageReadControl, vector_index::VectorRead, IVFIndexParams,
12    StorageBackendResult,
13};
14use uqa_core::{
15    memory::{BudgetedVec, MemoryError},
16    DocId,
17};
18
19pub(crate) struct IVFCanonicalBuilder {
20    candidate: IVFPreparedMetadata,
21    previous: Option<(DocId, u32)>,
22}
23
24impl IVFCanonicalBuilder {
25    pub(crate) fn new(
26        dimensions: u32,
27        params: IVFIndexParams,
28        control: &StorageReadControl,
29    ) -> StorageBackendResult<Self> {
30        let header = IVFMetadataSnapshot {
31            state: IVFState::Untrained,
32            centroids: Vec::new(),
33            assignments: Vec::new(),
34            trained_size: 0,
35            deletes_since_train: 0,
36            vector_count: 0,
37        };
38        Ok(Self {
39            candidate: IVFRestoreBuilder::new(dimensions, params, header, control)?.finish()?,
40            previous: None,
41        })
42    }
43
44    pub(crate) fn vector(
45        &mut self,
46        document: DocId,
47        ordinal: u32,
48        raw: &[f32],
49    ) -> StorageBackendResult<()> {
50        self.candidate.control.check()?;
51        let expected = match self.previous {
52            Some((prior, order)) if prior == document => u64::from(order) + 1,
53            Some((prior, _)) if prior > document => {
54                return Err(corrupt("unordered canonical vectors"))
55            }
56            _ => {
57                self.complete_document()?;
58                0
59            }
60        };
61        if u64::from(ordinal) != expected {
62            return Err(corrupt("canonical tensor has non-contiguous ordinals"));
63        }
64        crate::vector_index::validate_vector_values(self.candidate.dimensions, raw)?;
65        self.candidate.insert_assigned(document, ordinal, raw)?;
66        self.previous = Some((document, ordinal));
67        Ok(())
68    }
69
70    fn complete_document(&mut self) -> StorageBackendResult<()> {
71        self.candidate.snapshot.vector_count = self.candidate.vectors.len();
72        if self.candidate.snapshot.state == IVFState::Untrained
73            && self.candidate.vectors.len() >= self.candidate.params.train_threshold
74        {
75            self.candidate.train()?;
76        }
77        Ok(())
78    }
79
80    pub(crate) fn finish(mut self) -> StorageBackendResult<IVFPreparedMetadata> {
81        self.complete_document()?;
82        self.candidate.control.check()?;
83        Ok(self.candidate)
84    }
85}
86
87impl IVFPreparedMetadata {
88    /// Clone an immutable generation under the requested control, sharing roots only when the memory allowance is unchanged.
89    pub fn clone_controlled(&self, control: &StorageReadControl) -> StorageBackendResult<Self> {
90        control.check()?;
91        let bytes = self.snapshot.centroids.iter().try_fold(
92            self.snapshot
93                .centroids
94                .len()
95                .checked_mul(size_of::<Vec<f32>>())
96                .ok_or(MemoryError::SizeOverflow)?,
97            |bytes, vector| {
98                vector
99                    .len()
100                    .checked_mul(4)
101                    .and_then(|count| bytes.checked_add(count))
102                    .ok_or(MemoryError::SizeOverflow)
103            },
104        )?;
105        let memory = control.memory().reserve(bytes)?;
106        let mut snapshot = self.snapshot.clone();
107        snapshot.assignments = Vec::new();
108        let vectors = if self.control.memory().shares_allowance(control.memory()) {
109            self.vectors.clone()
110        } else {
111            let mut copied = super::Map::new(control.memory(), control.memory().limit() / 3);
112            let mut after = None;
113            while let Some((key, vector)) =
114                self.vectors.next_with_memory(after, control.memory())?
115            {
116                control.check()?;
117                let _memory = control
118                    .memory()
119                    .reserve(crate::spill_map::Record::memory_bytes(&*vector)?)?;
120                copied.insert(key, (*vector).clone(), Some(control))?;
121                after = Some(key);
122            }
123            copied
124        };
125        Ok(Self {
126            snapshot,
127            vectors,
128            params: self.params,
129            dimensions: self.dimensions,
130            control: control.clone(),
131            memory,
132        })
133    }
134}
135
136impl VectorRead for IVFPreparedMetadata {
137    fn check_control(&self, control: &StorageReadControl) -> StorageBackendResult<()> {
138        control.check()
139    }
140    fn dimensions(&self) -> u32 {
141        self.dimensions
142    }
143    fn next_document_after(
144        &self,
145        after: Option<DocId>,
146        control: &StorageReadControl,
147    ) -> StorageBackendResult<Option<DocId>> {
148        control.check()?;
149        Ok(self
150            .vectors
151            .next_key(
152                after.map(|document| key(document, u32::MAX)),
153                control.memory(),
154            )?
155            .map(|key| (key >> 32) as DocId))
156    }
157    fn document_vector_count(
158        &self,
159        document: DocId,
160        control: &StorageReadControl,
161    ) -> StorageBackendResult<u64> {
162        let mut after = key(document, 0).checked_sub(1);
163        let mut count = 0_u64;
164        loop {
165            control.check()?;
166            let Some(found) = self.vectors.next_key(after, control.memory())? else {
167                break;
168            };
169            if found >> 32 != u128::from(document) {
170                break;
171            }
172            if u64::from(found as u32) != count {
173                return Err(corrupt("canonical tensor has non-contiguous ordinals"));
174            }
175            count += 1;
176            after = Some(found);
177        }
178        Ok(count)
179    }
180    fn read_vector(
181        &self,
182        document: DocId,
183        ordinal: u32,
184        control: &StorageReadControl,
185    ) -> StorageBackendResult<Option<BudgetedVec<f32>>> {
186        control.check()?;
187        self.vectors
188            .get_with_memory(key(document, ordinal), control.memory())?
189            .map(|vector| crate::vector_index::copy_vector(&vector.raw_vector, control))
190            .transpose()
191    }
192}
uqa_storage/ivf_index/spill/canonical.rs

uqa_storage/ivf_index/spill/
canonical.rs
