uqa_storage/key_value/
ivf_index.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! IVF state and canonical values share one logical Key/Value visibility boundary.
8
9use std::sync::Arc;
10
11use uqa_core::{memory::Budgeted, DocId, PostingList};
12
13use super::codec::{other_error, vector_field_prefix};
14use super::index_keys::{
15    hnsw_metadata_key, hnsw_node_prefix, ivf_assignment_prefix, ivf_centroid_prefix,
16    ivf_metadata_key,
17};
18use super::index_view::{read_view, IndexState, IndexView};
19use super::ivf_persistence;
20use super::{KeyValueBatch, KeyValueRead, KeyValueStore, KeyValueVectorIndex};
21use crate::ivf_index::{
22    IVFCanonicalBuilder, IVFMutation, IVFPreparedMetadata, IVFReadIndex, IVFState,
23};
24use crate::vector_index::{IVFIndexParams, VectorIndex};
25use crate::{ReadOnlySnapshot, StorageBackendError, StorageBackendResult};
26
27#[cfg(test)]
28mod tests;
29
30pub struct KeyValueIVFIndex {
31    store: Arc<dyn KeyValueStore>,
32    raw: KeyValueVectorIndex,
33    table: String,
34    field: String,
35    dimensions: u32,
36    params: IVFIndexParams,
37    require_persisted: bool,
38    view: IndexView<IVFReadIndex>,
39}
40
41impl KeyValueIVFIndex {
42    pub fn create(
43        store: Arc<dyn KeyValueStore>,
44        table: impl Into<String>,
45        field: impl Into<String>,
46        dimensions: u32,
47        params: IVFIndexParams,
48    ) -> StorageBackendResult<Self> {
49        let index = Self::new(store, table.into(), field.into(), dimensions, params, false)?;
50        index.read_index()?;
51        Ok(index)
52    }
53
54    pub fn restore(
55        store: Arc<dyn KeyValueStore>,
56        table: impl Into<String>,
57        field: impl Into<String>,
58        dimensions: u32,
59        params: IVFIndexParams,
60    ) -> StorageBackendResult<Self> {
61        let index = Self::new(store, table.into(), field.into(), dimensions, params, true)?;
62        index.read_index()?;
63        Ok(index)
64    }
65
66    fn new(
67        store: Arc<dyn KeyValueStore>,
68        table: String,
69        field: String,
70        dimensions: u32,
71        params: IVFIndexParams,
72        require_persisted: bool,
73    ) -> StorageBackendResult<Self> {
74        Ok(Self {
75            raw: KeyValueVectorIndex::new(store.clone(), &table, &field, dimensions),
76            store,
77            table,
78            field,
79            dimensions,
80            params: params.validate()?,
81            require_persisted,
82            view: IndexView::new(!require_persisted),
83        })
84    }
85
86    pub(super) fn drop_metadata(
87        store: &dyn KeyValueStore,
88        table: &str,
89        field: &str,
90    ) -> StorageBackendResult<()> {
91        let mut batch = store.batch();
92        batch.fence_ivf_prefix(&ivf_metadata_key(table, field)?)?;
93        batch.fence_hnsw_prefix(&hnsw_metadata_key(table, field)?)?;
94        batch.delete(&ivf_metadata_key(table, field)?)?;
95        batch.delete_prefix(&ivf_centroid_prefix(table, field)?)?;
96        batch.delete_prefix(&ivf_assignment_prefix(table, field)?)?;
97        batch.delete(&hnsw_metadata_key(table, field)?)?;
98        batch.delete_prefix(&hnsw_node_prefix(table, field)?)?;
99        batch.commit()
100    }
101
102    fn index_at(&self, read: &dyn KeyValueRead) -> StorageBackendResult<IndexState<IVFReadIndex>> {
103        self.view.load(
104            read,
105            &[
106                &ivf_metadata_key(&self.table, &self.field)?,
107                &ivf_centroid_prefix(&self.table, &self.field)?,
108                &ivf_assignment_prefix(&self.table, &self.field)?,
109                &vector_field_prefix(&self.table, &self.field)?,
110            ],
111            |creating| {
112                let revision = ivf_persistence::load_revision(read, &self.table, &self.field)?;
113                let index = if creating || (revision.is_none() && !self.require_persisted) {
114                    self.build_from_canonical(read)?
115                } else {
116                    ivf_persistence::restore_state(
117                        read,
118                        &self.raw,
119                        &self.table,
120                        &self.field,
121                        self.dimensions,
122                        self.params,
123                    )?
124                    .0
125                };
126                Ok((index, revision))
127            },
128        )
129    }
130
131    fn read_index(&self) -> StorageBackendResult<IndexState<IVFReadIndex>> {
132        read_view(self.store.as_ref(), |read| self.index_at(read))
133    }
134
135    fn mutate(
136        &self,
137        mutation: IVFMutation<'_>,
138        canonical: impl FnOnce(&mut dyn KeyValueBatch) -> StorageBackendResult<()>,
139    ) -> StorageBackendResult<()> {
140        let metadata = ivf_metadata_key(&self.table, &self.field)?;
141        let centroids = ivf_centroid_prefix(&self.table, &self.field)?;
142        let assignments = ivf_assignment_prefix(&self.table, &self.field)?;
143        let vectors = vector_field_prefix(&self.table, &self.field)?;
144        self.view.evaluate_candidate(
145            self.store.as_ref(),
146            &[&metadata, &centroids, &assignments, &vectors],
147            |read, batch| {
148                let cached = self.index_at(read)?;
149                let after = cached
150                    .value
151                    .clone_controlled(read.control())?
152                    .prepare(&[mutation])?;
153                let changed_doc = match mutation {
154                    IVFMutation::Replace { document, .. } | IVFMutation::Delete(document) => {
155                        Some(document)
156                    }
157                    IVFMutation::Clear | IVFMutation::Train => None,
158                };
159                let full_rewrite = changed_doc.is_none()
160                    || cached.definition_candidate
161                    || cached.revision.is_none()
162                    || cached.value.header().centroids != after.header().centroids;
163                canonical(batch)?;
164                let preview = !cached.definition_candidate
165                    && cached.revision.is_some()
166                    && changed_doc.is_some();
167                if preview {
168                    batch.ivf_mutation(&ivf_metadata_key(&self.table, &self.field)?, mutation)?;
169                }
170                let revision = next_revision(cached.revision)?;
171                self.stage_snapshot(batch, &after, revision, full_rewrite, changed_doc, preview)?;
172                drop(cached);
173                Ok((IVFReadIndex::new(after)?, revision))
174            },
175        )
176    }
177
178    fn build_from_canonical(
179        &self,
180        read: &dyn KeyValueRead,
181    ) -> StorageBackendResult<Budgeted<IVFReadIndex>> {
182        IVFReadIndex::new(self.canonical_candidate(read)?)
183    }
184
185    fn canonical_candidate(
186        &self,
187        read: &dyn KeyValueRead,
188    ) -> StorageBackendResult<IVFPreparedMetadata> {
189        let mut builder = IVFCanonicalBuilder::new(self.dimensions, self.params, read.control())?;
190        self.raw
191            .visit_canonical_from(read, |document, ordinal, vector| {
192                builder.vector(document, ordinal, vector)
193            })?;
194        builder.finish()
195    }
196
197    fn rebuild(&self) -> StorageBackendResult<()> {
198        self.view.evaluate(self.store.as_ref(), |read, batch| {
199            let revision = ivf_persistence::load_revision(read, &self.table, &self.field)?;
200            if revision.is_none() && self.require_persisted {
201                return Err(other_error(format!(
202                    "missing persisted IVF metadata for {}.{}",
203                    self.table, self.field
204                )));
205            }
206            let candidate = self
207                .canonical_candidate(read)?
208                .prepare(&[IVFMutation::Train])?;
209            self.stage_snapshot(
210                batch,
211                &candidate,
212                next_revision(revision)?,
213                true,
214                None,
215                false,
216            )
217        })
218    }
219
220    fn stage_snapshot(
221        &self,
222        batch: &mut dyn KeyValueBatch,
223        snapshot: &IVFPreparedMetadata,
224        revision: u64,
225        full_rewrite: bool,
226        changed_doc: Option<DocId>,
227        preview: bool,
228    ) -> StorageBackendResult<()> {
229        ivf_persistence::stage_snapshot(
230            batch,
231            &self.table,
232            &self.field,
233            self.dimensions,
234            self.params,
235            snapshot,
236            revision,
237            full_rewrite,
238            changed_doc,
239            preview,
240        )
241    }
242}
243
244impl VectorIndex for KeyValueIVFIndex {
245    fn contains_document(&self, doc_id: DocId) -> StorageBackendResult<bool> {
246        self.raw.contains_document(doc_id)
247    }
248
249    fn dimensions(&self) -> u32 {
250        self.dimensions
251    }
252    fn index_kind(&self) -> &'static str {
253        "ivf"
254    }
255    fn add(&mut self, doc_id: DocId, vector: Vec<f32>) -> StorageBackendResult<()> {
256        self.add_many(doc_id, vec![vector])
257    }
258    fn add_many(&mut self, doc_id: DocId, vectors: Vec<Vec<f32>>) -> StorageBackendResult<()> {
259        self.mutate(
260            IVFMutation::Replace {
261                document: doc_id,
262                vectors: &vectors,
263            },
264            |batch| self.raw.stage_replace(batch, doc_id, &vectors),
265        )
266    }
267    fn delete(&mut self, doc_id: DocId) -> StorageBackendResult<()> {
268        self.mutate(IVFMutation::Delete(doc_id), |batch| {
269            self.raw.stage_replace(batch, doc_id, &[])
270        })
271    }
272    fn clear(&mut self) -> StorageBackendResult<()> {
273        self.mutate(IVFMutation::Clear, |batch| self.raw.stage_clear(batch))
274    }
275    fn search_knn(&self, query: &[f32], k: usize) -> StorageBackendResult<PostingList> {
276        self.snapshot()?.search_knn(query, k)
277    }
278    fn search_threshold(&self, query: &[f32], threshold: f32) -> StorageBackendResult<PostingList> {
279        self.read_index()?.value.search_threshold(query, threshold)
280    }
281    fn count(&self) -> StorageBackendResult<usize> {
282        self.read_index()?.value.count()
283    }
284    fn initialize(&mut self) -> StorageBackendResult<()> {
285        self.rebuild()
286    }
287    fn snapshot(&self) -> StorageBackendResult<Arc<dyn VectorIndex>> {
288        read_view(self.store.as_ref(), |read| {
289            let cached = self.index_at(read)?;
290            if cached.value.header().state != IVFState::Stale {
291                return Ok(cached.snapshot);
292            }
293            let candidate = IVFReadIndex::new(
294                cached
295                    .value
296                    .clone_controlled(read.control())?
297                    .prepare(&[IVFMutation::Train])?,
298            )?;
299            ReadOnlySnapshot::from_budgeted(candidate)?
300                .with_canonical_vectors(Some(read.control()))?
301                .with_vector_read_control(read.control())?
302                .snapshot()
303        })
304    }
305    fn vector_read_snapshot(
306        &self,
307        control: &crate::read_control::StorageReadControl,
308    ) -> StorageBackendResult<Option<crate::vector_index::VectorReadSnapshot>> {
309        self.raw.vector_read_snapshot(control)
310    }
311}
312
313fn next_revision(revision: Option<u64>) -> StorageBackendResult<u64> {
314    revision
315        .unwrap_or(0)
316        .checked_add(1)
317        .ok_or_else(|| StorageBackendError::Other("IVF metadata revision space exhausted".into()))
318}
uqa_storage/key_value/ivf_index.rs

uqa_storage/key_value/
ivf_index.rs
