uqa_storage/key_value/
ivf_index.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! Transactional IVF physical state over the logical Key/Value backend.
8
9use std::sync::Arc;
10
11use parking_lot::Mutex;
12use uqa_core::{DocId, PostingList};
13
14use super::codec::other_error;
15use super::index_keys::{
16    hnsw_metadata_key, hnsw_node_prefix, ivf_assignment_prefix, ivf_centroid_prefix,
17    ivf_metadata_key,
18};
19use super::ivf_persistence;
20use super::{KeyValueStore, KeyValueVectorIndex};
21use crate::ivf_index::{IVFIndex, IVFMetadataSnapshot, IVFState};
22use crate::vector_index::{IVFIndexParams, VectorIndex};
23use crate::{StorageBackendError, StorageBackendResult};
24
25struct CachedIVF {
26    index: IVFIndex,
27    revision: Option<u64>,
28}
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    cached: Mutex<CachedIVF>,
38}
39
40impl KeyValueIVFIndex {
41    pub fn create(
42        store: Arc<dyn KeyValueStore>,
43        table: impl Into<String>,
44        field: impl Into<String>,
45        dimensions: u32,
46        params: IVFIndexParams,
47    ) -> StorageBackendResult<Self> {
48        let params = params.validate()?;
49        let table = table.into();
50        let field = field.into();
51        let raw = KeyValueVectorIndex::new(Arc::clone(&store), &table, &field, dimensions);
52        let index = build_from_canonical(&raw, dimensions, params)?;
53        Ok(Self {
54            store,
55            raw,
56            table,
57            field,
58            dimensions,
59            params,
60            cached: Mutex::new(CachedIVF {
61                index,
62                revision: None,
63            }),
64        })
65    }
66
67    pub fn restore(
68        store: Arc<dyn KeyValueStore>,
69        table: impl Into<String>,
70        field: impl Into<String>,
71        dimensions: u32,
72        params: IVFIndexParams,
73    ) -> StorageBackendResult<Self> {
74        let params = params.validate()?;
75        let table = table.into();
76        let field = field.into();
77        let raw = KeyValueVectorIndex::new(Arc::clone(&store), &table, &field, dimensions);
78        let (index, revision) = ivf_persistence::restore_state(
79            store.as_ref(),
80            &raw,
81            &table,
82            &field,
83            dimensions,
84            params,
85        )?;
86        Ok(Self {
87            store,
88            raw,
89            table,
90            field,
91            dimensions,
92            params,
93            cached: Mutex::new(CachedIVF {
94                index,
95                revision: Some(revision),
96            }),
97        })
98    }
99
100    pub(super) fn drop_metadata(
101        store: &dyn KeyValueStore,
102        table: &str,
103        field: &str,
104    ) -> StorageBackendResult<()> {
105        let mut batch = store.batch();
106        batch.delete(&ivf_metadata_key(table, field)?)?;
107        batch.delete_prefix(&ivf_centroid_prefix(table, field)?)?;
108        batch.delete_prefix(&ivf_assignment_prefix(table, field)?)?;
109        batch.delete(&hnsw_metadata_key(table, field)?)?;
110        batch.delete_prefix(&hnsw_node_prefix(table, field)?)?;
111        batch.commit()
112    }
113
114    fn replace_document(&self, doc_id: DocId, vectors: &[Vec<f32>]) -> StorageBackendResult<()> {
115        let mut cached = self.cached.lock();
116        self.verify_revision(cached.revision)?;
117        let before = cached.index.metadata_snapshot();
118        let mut staged = cached.index.detached_clone();
119        staged.add_many(doc_id, vectors.to_vec())?;
120        train_if_stale(&staged)?;
121        let after = staged.metadata_snapshot();
122        let full_rewrite = cached.revision.is_none() || before.centroids != after.centroids;
123        let revision = next_revision(cached.revision)?;
124        let mut batch = self.store.batch();
125        self.raw.stage_replace(batch.as_mut(), doc_id, vectors)?;
126        self.stage_snapshot(batch.as_mut(), &after, revision, full_rewrite, Some(doc_id))?;
127        batch.commit()?;
128        *cached = CachedIVF {
129            index: staged,
130            revision: Some(revision),
131        };
132        Ok(())
133    }
134
135    fn delete_document(&self, doc_id: DocId) -> StorageBackendResult<()> {
136        let mut cached = self.cached.lock();
137        self.verify_revision(cached.revision)?;
138        let before = cached.index.metadata_snapshot();
139        let mut staged = cached.index.detached_clone();
140        staged.delete(doc_id)?;
141        train_if_stale(&staged)?;
142        let after = staged.metadata_snapshot();
143        let full_rewrite = cached.revision.is_none() || before.centroids != after.centroids;
144        let revision = next_revision(cached.revision)?;
145        let mut batch = self.store.batch();
146        self.raw.stage_replace(batch.as_mut(), doc_id, &[])?;
147        self.stage_snapshot(batch.as_mut(), &after, revision, full_rewrite, Some(doc_id))?;
148        batch.commit()?;
149        *cached = CachedIVF {
150            index: staged,
151            revision: Some(revision),
152        };
153        Ok(())
154    }
155
156    fn replace_all(&self, clear_vectors: bool, train: bool) -> StorageBackendResult<()> {
157        let mut cached = self.cached.lock();
158        self.verify_revision(cached.revision)?;
159        let mut staged = cached.index.detached_clone();
160        if clear_vectors {
161            staged.clear()?;
162        } else if train {
163            staged.initialize()?;
164        }
165        let snapshot = staged.metadata_snapshot();
166        let revision = next_revision(cached.revision)?;
167        let mut batch = self.store.batch();
168        if clear_vectors {
169            self.raw.stage_clear(batch.as_mut())?;
170        }
171        self.stage_snapshot(batch.as_mut(), &snapshot, revision, true, None)?;
172        batch.commit()?;
173        *cached = CachedIVF {
174            index: staged,
175            revision: Some(revision),
176        };
177        Ok(())
178    }
179
180    fn verify_revision(&self, expected: Option<u64>) -> StorageBackendResult<()> {
181        let Some(expected) = expected else {
182            return Ok(());
183        };
184        let actual = ivf_persistence::load_revision(self.store.as_ref(), &self.table, &self.field)?
185            .ok_or_else(|| {
186                other_error(format!(
187                    "missing persisted IVF metadata for {}.{}",
188                    self.table, self.field
189                ))
190            })?;
191        if actual != expected {
192            return Err(other_error(format!(
193                "concurrent IVF metadata change for {}.{}: expected revision {expected}, found {actual}",
194                self.table, self.field
195            )));
196        }
197        Ok(())
198    }
199
200    fn stage_snapshot(
201        &self,
202        batch: &mut dyn super::KeyValueBatch,
203        snapshot: &IVFMetadataSnapshot,
204        revision: u64,
205        full_rewrite: bool,
206        changed_doc: Option<DocId>,
207    ) -> StorageBackendResult<()> {
208        ivf_persistence::stage_snapshot(
209            batch,
210            &self.table,
211            &self.field,
212            self.dimensions,
213            self.params,
214            snapshot,
215            revision,
216            full_rewrite,
217            changed_doc,
218        )
219    }
220}
221
222impl VectorIndex for KeyValueIVFIndex {
223    fn dimensions(&self) -> u32 {
224        self.dimensions
225    }
226
227    fn index_kind(&self) -> &'static str {
228        "ivf"
229    }
230
231    fn add(&mut self, doc_id: DocId, vector: Vec<f32>) -> StorageBackendResult<()> {
232        self.replace_document(doc_id, &[vector])
233    }
234
235    fn add_many(&mut self, doc_id: DocId, vectors: Vec<Vec<f32>>) -> StorageBackendResult<()> {
236        self.replace_document(doc_id, &vectors)
237    }
238
239    fn delete(&mut self, doc_id: DocId) -> StorageBackendResult<()> {
240        self.delete_document(doc_id)
241    }
242
243    fn clear(&mut self) -> StorageBackendResult<()> {
244        self.replace_all(true, false)
245    }
246
247    fn search_knn(&self, query: &[f32], k: usize) -> StorageBackendResult<PostingList> {
248        let cached = self.cached.lock();
249        if cached.index.state() != IVFState::Stale {
250            return cached.index.search_knn(query, k);
251        }
252        let staged = cached.index.detached_clone();
253        drop(cached);
254        staged.train()?;
255        staged.search_knn(query, k)
256    }
257
258    fn search_threshold(&self, query: &[f32], threshold: f32) -> StorageBackendResult<PostingList> {
259        self.cached.lock().index.search_threshold(query, threshold)
260    }
261
262    fn count(&self) -> StorageBackendResult<usize> {
263        self.cached.lock().index.count()
264    }
265
266    fn initialize(&mut self) -> StorageBackendResult<()> {
267        self.replace_all(false, true)
268    }
269
270    fn snapshot(&self) -> StorageBackendResult<Arc<dyn VectorIndex>> {
271        let staged = self.cached.lock().index.detached_clone();
272        train_if_stale(&staged)?;
273        Ok(Arc::new(staged))
274    }
275}
276
277fn build_from_canonical(
278    raw: &KeyValueVectorIndex,
279    dimensions: u32,
280    params: IVFIndexParams,
281) -> StorageBackendResult<IVFIndex> {
282    let mut index = IVFIndex::with_params(
283        dimensions,
284        params.nlist,
285        params.nprobe,
286        params.train_threshold,
287    );
288    for (doc_id, vectors) in raw.load_by_document()? {
289        index.add_many(doc_id, vectors)?;
290    }
291    Ok(index)
292}
293
294fn next_revision(revision: Option<u64>) -> StorageBackendResult<u64> {
295    revision
296        .unwrap_or(0)
297        .checked_add(1)
298        .ok_or_else(|| StorageBackendError::Other("IVF metadata revision space exhausted".into()))
299}
300
301fn train_if_stale(index: &IVFIndex) -> StorageBackendResult<()> {
302    if index.state() == IVFState::Stale {
303        index.train()?;
304    }
305    Ok(())
306}
uqa_storage/key_value/ivf_index.rs

uqa_storage/key_value/
ivf_index.rs
