uqa_storage/key_value/
vector_index.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! Vector-index adapter over an ordered key/value store.
8
9mod canonical;
10pub(crate) mod guards;
11pub(in crate::key_value) mod origin;
12mod read;
13pub use origin::{
14    DiskANNCanonicalVectorVisitor, KeyValueDiskANNCanonical, KeyValueDiskANNHandle,
15    KeyValueDiskANNPopulationRecords, RetainedDiskANNCanonical,
16};
17
18use crate::vector_index::RetainedVectorIndex;
19
20use super::index_view::read_view;
21
22use super::codec::{
23    other_error, usize_to_u64, validate_vector_ordinal_count, vector_doc_prefix,
24    vector_field_prefix, vector_key, vector_to_blob,
25};
26use super::{
27    validate_vector_values, Arc, DocId, KeyValueBatch, KeyValueStore, PostingList,
28    StorageBackendResult, VectorIndex,
29};
30
31/// Brute-force vector index implemented over [`KeyValueStore`].
32#[derive(Clone)]
33pub struct KeyValueVectorIndex {
34    store: Arc<dyn KeyValueStore>,
35    table: String,
36    field: String,
37    dimensions: u32,
38}
39
40impl KeyValueVectorIndex {
41    pub fn new(
42        store: Arc<dyn KeyValueStore>,
43        table: impl Into<String>,
44        field: impl Into<String>,
45        dimensions: u32,
46    ) -> Self {
47        Self {
48            store,
49            table: table.into(),
50            field: field.into(),
51            dimensions,
52        }
53    }
54
55    fn read_snapshot(&self) -> StorageBackendResult<RetainedVectorIndex> {
56        read_view(self.store.as_ref(), |read| {
57            RetainedVectorIndex::from_entries(
58                self.load_all_from(read)?,
59                self.dimensions,
60                "keyvalue-bruteforce",
61                read.control(),
62            )
63        })
64    }
65
66    pub(super) fn stage_replace(
67        &self,
68        batch: &mut dyn KeyValueBatch,
69        doc_id: DocId,
70        vectors: &[Vec<f32>],
71    ) -> StorageBackendResult<()> {
72        self.stage_replace_values(batch, doc_id, vectors)?;
73        batch.invalidate_diskann_origin(&origin::key(&self.table, &self.field, doc_id)?)
74    }
75
76    /// Keep the origin mutation with its owner so a typed canonical replacement retains its MVCC effect.
77    pub(super) fn stage_replace_values(
78        &self,
79        batch: &mut dyn KeyValueBatch,
80        doc_id: DocId,
81        vectors: &[Vec<f32>],
82    ) -> StorageBackendResult<()> {
83        for vector in vectors {
84            self.validate_dimensions(vector)?;
85        }
86        validate_vector_ordinal_count(usize_to_u64(vectors.len(), "vector count")?)?;
87        self.coordinate_field(batch, false)?;
88        batch.delete_prefix(&vector_doc_prefix(&self.table, &self.field, doc_id)?)?;
89        for (ordinal, vector) in vectors.iter().enumerate() {
90            let ordinal = u32::try_from(ordinal)
91                .map_err(|_| other_error("vector ordinal exceeds u32 index format"))?;
92            batch.put(
93                &vector_key(&self.table, &self.field, doc_id, ordinal)?,
94                &vector_to_blob(vector)?,
95            )?;
96        }
97        Ok(())
98    }
99
100    pub(super) fn stage_clear(&self, batch: &mut dyn KeyValueBatch) -> StorageBackendResult<()> {
101        self.coordinate_field(batch, true)?;
102        batch.delete_prefix(&origin::prefix(&self.table, &self.field)?)?;
103        batch.delete_prefix_allow_absent(&origin::journal::prefix(&self.table, &self.field)?)?;
104        for prefix in origin::populations::field_prefixes(&self.table, &self.field)? {
105            batch.delete_prefix(&prefix)?;
106        }
107        batch.delete_prefix(&vector_field_prefix(&self.table, &self.field)?)
108    }
109
110    fn validate_dimensions(&self, vector: &[f32]) -> StorageBackendResult<()> {
111        validate_vector_values(self.dimensions, vector).map_err(|error| {
112            other_error(format!(
113                "invalid vector for {}.{}: {error}",
114                self.table, self.field
115            ))
116        })
117    }
118}
119
120impl VectorIndex for KeyValueVectorIndex {
121    fn contains_document(&self, doc_id: DocId) -> StorageBackendResult<bool> {
122        read_view(self.store.as_ref(), |read| {
123            read.control().check()?;
124            // The existing key encoder can hold both old and growing buffers; retain their allowance until the key is dropped.
125            let bytes = self
126                .table
127                .len()
128                .checked_add(self.field.len())
129                .and_then(|length| length.checked_add(25))
130                .and_then(|length| length.checked_mul(3))
131                .ok_or(uqa_core::memory::MemoryError::SizeOverflow)?;
132            let _key_memory = read.control().memory().reserve(bytes)?;
133            let key = vector_key(&self.table, &self.field, doc_id, 0)?;
134            let mut found = false;
135            read.visit_keys_after(&key, None, 1, read.control(), &mut |candidate| {
136                found = candidate == key;
137                Ok(())
138            })?;
139            Ok(found)
140        })
141    }
142
143    fn dimensions(&self) -> u32 {
144        self.dimensions
145    }
146
147    fn index_kind(&self) -> &'static str {
148        "keyvalue-bruteforce"
149    }
150
151    fn add(&mut self, doc_id: DocId, vector: Vec<f32>) -> StorageBackendResult<()> {
152        self.add_many(doc_id, vec![vector])
153    }
154
155    fn add_many(&mut self, doc_id: DocId, vectors: Vec<Vec<f32>>) -> StorageBackendResult<()> {
156        let mut batch = self.store.batch();
157        self.stage_replace(batch.as_mut(), doc_id, &vectors)?;
158        batch.commit()
159    }
160
161    fn delete(&mut self, doc_id: DocId) -> StorageBackendResult<()> {
162        let mut batch = self.store.batch();
163        self.stage_replace(batch.as_mut(), doc_id, &[])?;
164        batch.commit()
165    }
166
167    fn clear(&mut self) -> StorageBackendResult<()> {
168        let mut batch = self.store.batch();
169        self.stage_clear(batch.as_mut())?;
170        batch.commit()
171    }
172
173    fn search_knn(&self, query: &[f32], k: usize) -> StorageBackendResult<PostingList> {
174        self.validate_dimensions(query)?;
175        if k == 0 {
176            return Ok(PostingList::new());
177        }
178        self.read_snapshot()?.search_knn(query, k)
179    }
180
181    fn search_threshold(&self, query: &[f32], threshold: f32) -> StorageBackendResult<PostingList> {
182        self.validate_dimensions(query)?;
183        if !threshold.is_finite() {
184            return Err(other_error(format!(
185                "vector similarity threshold must be finite, got {threshold}"
186            )));
187        }
188        self.read_snapshot()?.search_threshold(query, threshold)
189    }
190
191    fn count(&self) -> StorageBackendResult<usize> {
192        self.read_snapshot()?.count()
193    }
194
195    fn snapshot(&self) -> StorageBackendResult<Arc<dyn VectorIndex>> {
196        Ok(Arc::new(self.read_snapshot()?))
197    }
198    fn vector_read_snapshot(
199        &self,
200        control: &crate::read_control::StorageReadControl,
201    ) -> StorageBackendResult<Option<crate::vector_index::VectorReadSnapshot>> {
202        canonical::capture(self, control).map(Some)
203    }
204}
uqa_storage/key_value/vector_index.rs

uqa_storage/key_value/
vector_index.rs
