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

uqa_storage/key_value/
ivf_index.rs
