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 read;
10
11use crate::vector_index::RetainedVectorIndex;
12
13use super::index_view::read_view;
14
15use super::codec::{
16    other_error, usize_to_u64, validate_vector_ordinal_count, vector_doc_prefix,
17    vector_field_prefix, vector_key, vector_to_blob,
18};
19use super::{
20    validate_vector_values, Arc, DocId, KeyValueBatch, KeyValueStore, PostingList,
21    StorageBackendResult, VectorIndex,
22};
23
24/// Brute-force vector index implemented over [`KeyValueStore`].
25#[derive(Clone)]
26pub struct KeyValueVectorIndex {
27    store: Arc<dyn KeyValueStore>,
28    table: String,
29    field: String,
30    dimensions: u32,
31}
32
33impl KeyValueVectorIndex {
34    pub fn new(
35        store: Arc<dyn KeyValueStore>,
36        table: impl Into<String>,
37        field: impl Into<String>,
38        dimensions: u32,
39    ) -> Self {
40        Self {
41            store,
42            table: table.into(),
43            field: field.into(),
44            dimensions,
45        }
46    }
47
48    fn read_snapshot(&self) -> StorageBackendResult<RetainedVectorIndex> {
49        read_view(self.store.as_ref(), |read| {
50            RetainedVectorIndex::from_entries(
51                self.load_all_from(read)?,
52                self.dimensions,
53                "keyvalue-bruteforce",
54                read.control(),
55            )
56        })
57    }
58
59    pub(super) fn stage_replace(
60        &self,
61        batch: &mut dyn KeyValueBatch,
62        doc_id: DocId,
63        vectors: &[Vec<f32>],
64    ) -> StorageBackendResult<()> {
65        for vector in vectors {
66            self.validate_dimensions(vector)?;
67        }
68        validate_vector_ordinal_count(usize_to_u64(vectors.len(), "vector count")?)?;
69        batch.delete_prefix(&vector_doc_prefix(&self.table, &self.field, doc_id)?)?;
70        for (ordinal, vector) in vectors.iter().enumerate() {
71            let ordinal = u32::try_from(ordinal)
72                .map_err(|_| other_error("vector ordinal exceeds u32 index format"))?;
73            batch.put(
74                &vector_key(&self.table, &self.field, doc_id, ordinal)?,
75                &vector_to_blob(vector)?,
76            )?;
77        }
78        Ok(())
79    }
80
81    pub(super) fn stage_clear(&self, batch: &mut dyn KeyValueBatch) -> StorageBackendResult<()> {
82        batch.delete_prefix(&vector_field_prefix(&self.table, &self.field)?)
83    }
84
85    fn validate_dimensions(&self, vector: &[f32]) -> StorageBackendResult<()> {
86        validate_vector_values(self.dimensions, vector).map_err(|error| {
87            other_error(format!(
88                "invalid vector for {}.{}: {error}",
89                self.table, self.field
90            ))
91        })
92    }
93}
94
95impl VectorIndex for KeyValueVectorIndex {
96    fn contains_document(&self, doc_id: DocId) -> StorageBackendResult<bool> {
97        read_view(self.store.as_ref(), |read| {
98            read.control().check()?;
99            // The existing key encoder can hold both old and growing buffers; retain their allowance until the key is dropped.
100            let bytes = self
101                .table
102                .len()
103                .checked_add(self.field.len())
104                .and_then(|length| length.checked_add(25))
105                .and_then(|length| length.checked_mul(3))
106                .ok_or(uqa_core::memory::MemoryError::SizeOverflow)?;
107            let _key_memory = read.control().memory().reserve(bytes)?;
108            let key = vector_key(&self.table, &self.field, doc_id, 0)?;
109            let mut found = false;
110            read.visit_keys_after(&key, None, 1, read.control(), &mut |candidate| {
111                found = candidate == key;
112                Ok(())
113            })?;
114            Ok(found)
115        })
116    }
117
118    fn dimensions(&self) -> u32 {
119        self.dimensions
120    }
121
122    fn index_kind(&self) -> &'static str {
123        "keyvalue-bruteforce"
124    }
125
126    fn add(&mut self, doc_id: DocId, vector: Vec<f32>) -> StorageBackendResult<()> {
127        self.add_many(doc_id, vec![vector])
128    }
129
130    fn add_many(&mut self, doc_id: DocId, vectors: Vec<Vec<f32>>) -> StorageBackendResult<()> {
131        let mut batch = self.store.batch();
132        self.stage_replace(batch.as_mut(), doc_id, &vectors)?;
133        batch.commit()
134    }
135
136    fn delete(&mut self, doc_id: DocId) -> StorageBackendResult<()> {
137        let mut batch = self.store.batch();
138        batch.delete_prefix(&vector_doc_prefix(&self.table, &self.field, doc_id)?)?;
139        batch.commit()
140    }
141
142    fn clear(&mut self) -> StorageBackendResult<()> {
143        let mut batch = self.store.batch();
144        self.stage_clear(batch.as_mut())?;
145        batch.commit()
146    }
147
148    fn search_knn(&self, query: &[f32], k: usize) -> StorageBackendResult<PostingList> {
149        self.validate_dimensions(query)?;
150        if k == 0 {
151            return Ok(PostingList::new());
152        }
153        self.read_snapshot()?.search_knn(query, k)
154    }
155
156    fn search_threshold(&self, query: &[f32], threshold: f32) -> StorageBackendResult<PostingList> {
157        self.validate_dimensions(query)?;
158        if !threshold.is_finite() {
159            return Err(other_error(format!(
160                "vector similarity threshold must be finite, got {threshold}"
161            )));
162        }
163        self.read_snapshot()?.search_threshold(query, threshold)
164    }
165
166    fn count(&self) -> StorageBackendResult<usize> {
167        self.read_snapshot()?.count()
168    }
169
170    fn snapshot(&self) -> StorageBackendResult<Arc<dyn VectorIndex>> {
171        Ok(Arc::new(self.read_snapshot()?))
172    }
173}
uqa_storage/key_value/vector_index.rs

uqa_storage/key_value/
vector_index.rs
