uqa_storage/key_value/vector_index/
origin.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! `DiskANN` origins accompany the existing canonical vector keys, without a duplicate corpus.
8
9pub(in crate::key_value) mod journal;
10mod lifecycle;
11mod live;
12mod maintenance;
13pub(in crate::key_value) mod populations;
14mod retained;
15#[cfg(test)]
16mod tests;
17pub use live::KeyValueDiskANNHandle;
18pub use populations::KeyValueDiskANNPopulationRecords;
19pub use retained::{DiskANNCanonicalVectorVisitor, RetainedDiskANNCanonical};
20
21use std::sync::Arc;
22
23use super::KeyValueVectorIndex;
24use crate::diskann_index::format::{
25    DiskANNCanonicalOrigin as Record, DiskANNChangeIdentity, DiskANNVectorVersion,
26    CANONICAL_ORIGIN_BYTES as BYTES,
27};
28use crate::key_value::{codec, KeyValueRead};
29use crate::mvcc::VersionError;
30use crate::read_control::StorageReadControl;
31use crate::{KeyValueStore, RelationIdentity, StorageBackendResult};
32use uqa_core::DocId;
33
34pub(in crate::key_value) const ROOT: &[u8] = b"\0uqa-diskann-canonical-v1\0";
35
36/// Canonical tensor mutation owner for the Key/Value layout. Public `DiskANN` catalog routing remains unavailable until publication and recovery are integrated.
37pub struct KeyValueDiskANNCanonical {
38    index: KeyValueVectorIndex,
39}
40
41impl KeyValueDiskANNCanonical {
42    pub fn new(
43        store: Arc<dyn KeyValueStore>,
44        table: impl Into<String>,
45        field: impl Into<String>,
46        dimensions: u32,
47    ) -> StorageBackendResult<Self> {
48        if dimensions == 0 || !store.transaction_model().is_versioned() {
49            return Err(invalid(
50                "canonical origins require dimensions and a versioned store",
51            ));
52        }
53        Ok(Self {
54            index: KeyValueVectorIndex::new(store, table, field, dimensions),
55        })
56    }
57
58    /// Atomically replace every ordinal and its origin; an empty tensor persists an explicit zero-count replacement. The publishing transaction is also used for private origins, without assigning commit order.
59    pub fn replace(
60        &self,
61        document: DocId,
62        vectors: &[Vec<f32>],
63        control: &StorageReadControl,
64    ) -> StorageBackendResult<DiskANNVectorVersion> {
65        self.replace_guarded(document, vectors, control, &mut |_, _| Ok(()))
66    }
67
68    fn replace_guarded(
69        &self,
70        document: DocId,
71        vectors: &[Vec<f32>],
72        control: &StorageReadControl,
73        guard: &mut crate::key_value::KeyValueMutation<'_>,
74    ) -> StorageBackendResult<DiskANNVectorVersion> {
75        control.check()?;
76        for vector in vectors {
77            crate::vector_index::validate_vector_values_controlled(
78                self.index.dimensions,
79                vector,
80                Some(control),
81            )?;
82        }
83        let count = codec::usize_to_u64(vectors.len(), "canonical vector count")?;
84        codec::validate_vector_ordinal_count(count)?;
85        let _workspace = control.memory().reserve(self.workspace_bytes(true)?)?;
86        let key = key(&self.index.table, &self.index.field, document)?;
87        let mut version = None;
88        self.index
89            .store
90            .with_versioned_mutation(&mut |origin, read, batch| {
91                control.check()?;
92                guard(read, batch)?;
93                let current = DiskANNVectorVersion::new(origin.transaction(), origin.revision())?;
94                let record = Record::new(current, self.index.dimensions, count)?;
95                self.index.stage_replace_values(batch, document, vectors)?;
96                batch.replace_diskann_origin(&key, &record.encode())?;
97                batch.put(
98                    &journal::key(
99                        &self.index.table,
100                        &self.index.field,
101                        DiskANNChangeIdentity::new(document, current),
102                    )?,
103                    &record.encode(),
104                )?;
105                control.check()?;
106                version = Some(current);
107                Ok(())
108            })?;
109        version.ok_or_else(|| invalid("canonical mutation did not execute"))
110    }
111
112    /// Retain the exact private/committed boundary without enumerating canonical values.
113    pub fn retain(
114        &self,
115        control: &StorageReadControl,
116    ) -> StorageBackendResult<RetainedDiskANNCanonical> {
117        self.retain_selected(None, control)
118    }
119
120    /// Capture the real table/index definitions with canonical input on one fixed view. Ordinary unbound sources cannot authorize a later catalog publication.
121    pub fn retain_for_index(
122        &self,
123        index: &RelationIdentity,
124        control: &StorageReadControl,
125    ) -> StorageBackendResult<RetainedDiskANNCanonical> {
126        self.retain_selected(Some(index), control)
127    }
128
129    fn retain_selected(
130        &self,
131        index: Option<&RelationIdentity>,
132        control: &StorageReadControl,
133    ) -> StorageBackendResult<RetainedDiskANNCanonical> {
134        control.check()?;
135        let _workspace = control.memory().reserve(self.workspace_bytes(false)?)?;
136        let vectors = codec::vector_field_prefix(&self.index.table, &self.index.field)?;
137        let origins = prefix(&self.index.table, &self.index.field)?;
138        let changes = journal::prefix(&self.index.table, &self.index.field)?;
139        let populations = populations::field_prefixes(&self.index.table, &self.index.field)?;
140        let marker = super::guards::reference_key(&vectors, control)?;
141        let mut selected = None;
142        self.index
143            .store
144            .with_read_view(&mut |read: &dyn KeyValueRead| {
145                let binding = index
146                    .map(|index| {
147                        crate::key_value::catalog::diskann::Binding::capture(
148                            read,
149                            &self.index.table,
150                            &self.index.field,
151                            self.index.dimensions,
152                            index,
153                            control,
154                        )
155                    })
156                    .transpose()?;
157                let mut prefixes = uqa_core::memory::BudgetedVec::new(control.memory());
158                prefixes.extend_from_slice(&[&*vectors, &*origins, &*changes, &*marker])?;
159                prefixes.extend_from_slice(&[&populations[0], &populations[1]])?;
160                if let Some(binding) = &binding {
161                    prefixes.extend_from_slice(&binding.prefixes())?;
162                    prefixes.push(crate::key_value::diskann::READ_PREFIX)?;
163                }
164                let mut source = RetainedDiskANNCanonical::new(
165                    read.retain(&prefixes)?,
166                    &vectors,
167                    &origins,
168                    &changes,
169                    self.index.dimensions,
170                    control,
171                )?;
172                drop(prefixes);
173                source.binding = binding;
174                selected = Some(source);
175                Ok(())
176            })?;
177        selected.ok_or_else(|| invalid("canonical read did not execute"))
178    }
179
180    fn workspace_bytes(&self, writing: bool) -> StorageBackendResult<usize> {
181        // Existing key/blob encoders own ordinary Vec buffers; reserve their simultaneous capacity before invoking them.
182        self.index
183            .table
184            .len()
185            .checked_add(self.index.field.len())
186            .and_then(|bytes| bytes.checked_add(ROOT.len() + 128))
187            .and_then(|bytes| bytes.checked_mul(8))
188            .and_then(|bytes| {
189                bytes.checked_add(
190                    usize::try_from(if writing { self.index.dimensions } else { 0 })
191                        .ok()?
192                        .checked_mul(4)?,
193                )
194            })
195            .ok_or_else(|| uqa_core::memory::MemoryError::SizeOverflow.into())
196    }
197}
198
199pub(in crate::key_value) fn table_prefix(table: &str) -> StorageBackendResult<Vec<u8>> {
200    Ok(namespaced(&codec::vector_key_prefix(table)?))
201}
202
203pub(in crate::key_value) fn prefix(table: &str, field: &str) -> StorageBackendResult<Vec<u8>> {
204    Ok(namespaced(&codec::vector_field_prefix(table, field)?))
205}
206
207fn namespaced(canonical: &[u8]) -> Vec<u8> {
208    let mut prefix = Vec::with_capacity(ROOT.len() + canonical.len());
209    prefix.extend_from_slice(ROOT);
210    prefix.extend_from_slice(canonical);
211    prefix
212}
213
214pub(super) fn key(table: &str, field: &str, document: DocId) -> StorageBackendResult<Vec<u8>> {
215    let mut key = prefix(table, field)?;
216    key.extend_from_slice(&document.to_be_bytes());
217    Ok(key)
218}
219
220fn invalid(message: &'static str) -> crate::StorageBackendError {
221    VersionError::InvalidEncoding(message).into_storage_error()
222}
uqa_storage/key_value/vector_index/origin.rs

uqa_storage/key_value/vector_index/
origin.rs
