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

uqa_storage/key_value/
ivf_index.rs
