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    pub(crate) fn clone_controlled(
89        &self,
90        control: &StorageReadControl,
91    ) -> StorageBackendResult<Self> {
92        control.check()?;
93        let bytes = self.snapshot.centroids.iter().try_fold(
94            self.snapshot
95                .centroids
96                .len()
97                .checked_mul(size_of::<Vec<f32>>())
98                .ok_or(MemoryError::SizeOverflow)?,
99            |bytes, vector| {
100                vector
101                    .len()
102                    .checked_mul(4)
103                    .and_then(|count| bytes.checked_add(count))
104                    .ok_or(MemoryError::SizeOverflow)
105            },
106        )?;
107        let memory = control.memory().reserve(bytes)?;
108        let mut snapshot = self.snapshot.clone();
109        snapshot.assignments = Vec::new();
110        let vectors = if self.control.memory().shares_allowance(control.memory()) {
111            self.vectors.clone()
112        } else {
113            let mut copied = super::Map::new(control.memory(), control.memory().limit() / 3);
114            let mut after = None;
115            while let Some((key, vector)) =
116                self.vectors.next_with_memory(after, control.memory())?
117            {
118                control.check()?;
119                let _memory = control
120                    .memory()
121                    .reserve(crate::spill_map::Record::memory_bytes(&*vector)?)?;
122                copied.insert(key, (*vector).clone(), Some(control))?;
123                after = Some(key);
124            }
125            copied
126        };
127        Ok(Self {
128            snapshot,
129            vectors,
130            params: self.params,
131            dimensions: self.dimensions,
132            control: control.clone(),
133            memory,
134        })
135    }
136}
137
138impl VectorRead for IVFPreparedMetadata {
139    fn check_control(&self, control: &StorageReadControl) -> StorageBackendResult<()> {
140        control.check()
141    }
142    fn dimensions(&self) -> u32 {
143        self.dimensions
144    }
145    fn next_document_after(
146        &self,
147        after: Option<DocId>,
148        control: &StorageReadControl,
149    ) -> StorageBackendResult<Option<DocId>> {
150        control.check()?;
151        Ok(self
152            .vectors
153            .next_key(
154                after.map(|document| key(document, u32::MAX)),
155                control.memory(),
156            )?
157            .map(|key| (key >> 32) as DocId))
158    }
159    fn document_vector_count(
160        &self,
161        document: DocId,
162        control: &StorageReadControl,
163    ) -> StorageBackendResult<u64> {
164        let mut after = key(document, 0).checked_sub(1);
165        let mut count = 0_u64;
166        loop {
167            control.check()?;
168            let Some(found) = self.vectors.next_key(after, control.memory())? else {
169                break;
170            };
171            if found >> 32 != u128::from(document) {
172                break;
173            }
174            if u64::from(found as u32) != count {
175                return Err(corrupt("canonical tensor has non-contiguous ordinals"));
176            }
177            count += 1;
178            after = Some(found);
179        }
180        Ok(count)
181    }
182    fn read_vector(
183        &self,
184        document: DocId,
185        ordinal: u32,
186        control: &StorageReadControl,
187    ) -> StorageBackendResult<Option<BudgetedVec<f32>>> {
188        control.check()?;
189        self.vectors
190            .get_with_memory(key(document, ordinal), control.memory())?
191            .map(|vector| crate::vector_index::copy_vector(&vector.raw_vector, control))
192            .transpose()
193    }
194}
uqa_storage/ivf_index/spill/canonical.rs

uqa_storage/ivf_index/spill/
canonical.rs
