uqa_storage/key_value/
hnsw_index.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! HNSW readers and evaluated mutations use one logical Key/Value visibility boundary.
8
9use std::sync::Arc;
10use uqa_core::{memory::Budgeted, DocId, PostingList};
11
12use super::codec::encode_value;
13use super::codec::{other_error, vector_field_prefix};
14use super::hnsw_persistence::{self, metadata_from_graph, PersistedHNSWNode};
15use super::index_keys::{hnsw_metadata_key, hnsw_node_key, hnsw_node_prefix};
16use super::{KeyValueBatch, KeyValueRead, KeyValueStore, KeyValueVectorIndex};
17use crate::hnsw_index::{HNSWCanonicalBuilder, HNSWGraphDelta, HNSWIndex, HNSWMutation};
18use crate::vector_index::{HNSWIndexParams, VectorIndex};
19use crate::{StorageBackendError, StorageBackendResult};
20
21use super::index_view::{read_view, IndexState, IndexView};
22
23pub struct KeyValueHNSWIndex {
24    store: Arc<dyn KeyValueStore>,
25    raw: KeyValueVectorIndex,
26    table: String,
27    field: String,
28    dimensions: u32,
29    params: HNSWIndexParams,
30    require_persisted: bool,
31    view: IndexView<HNSWIndex>,
32}
33
34impl KeyValueHNSWIndex {
35    pub fn create(
36        store: Arc<dyn KeyValueStore>,
37        table: impl Into<String>,
38        field: impl Into<String>,
39        dimensions: u32,
40        params: HNSWIndexParams,
41    ) -> StorageBackendResult<Self> {
42        let index = Self::new(store, table.into(), field.into(), dimensions, params, false)?;
43        index.read_graph()?;
44        Ok(index)
45    }
46
47    pub fn restore(
48        store: Arc<dyn KeyValueStore>,
49        table: impl Into<String>,
50        field: impl Into<String>,
51        dimensions: u32,
52        params: HNSWIndexParams,
53    ) -> StorageBackendResult<Self> {
54        let index = Self::new(store, table.into(), field.into(), dimensions, params, true)?;
55        index.read_graph()?;
56        Ok(index)
57    }
58
59    fn new(
60        store: Arc<dyn KeyValueStore>,
61        table: String,
62        field: String,
63        dimensions: u32,
64        params: HNSWIndexParams,
65        require_persisted: bool,
66    ) -> StorageBackendResult<Self> {
67        Ok(Self {
68            raw: KeyValueVectorIndex::new(Arc::clone(&store), &table, &field, dimensions),
69            store,
70            table,
71            field,
72            dimensions,
73            params: params.validate()?,
74            require_persisted,
75            view: IndexView::new(!require_persisted),
76        })
77    }
78
79    fn graph_at(&self, read: &dyn KeyValueRead) -> StorageBackendResult<IndexState<HNSWIndex>> {
80        self.view.load(
81            read,
82            &[
83                &hnsw_metadata_key(&self.table, &self.field)?,
84                &hnsw_node_prefix(&self.table, &self.field)?,
85                &vector_field_prefix(&self.table, &self.field)?,
86            ],
87            |creating| {
88                let revision = hnsw_persistence::load_revision(read, &self.table, &self.field)?;
89                let graph = if creating || (revision.is_none() && !self.require_persisted) {
90                    self.build_from_canonical(read)?
91                } else {
92                    hnsw_persistence::restore_graph(
93                        read,
94                        &self.raw,
95                        &self.table,
96                        &self.field,
97                        self.dimensions,
98                        self.params,
99                    )?
100                    .0
101                };
102                Ok((graph, revision))
103            },
104        )
105    }
106
107    fn read_graph(&self) -> StorageBackendResult<IndexState<HNSWIndex>> {
108        read_view(self.store.as_ref(), |read| self.graph_at(read))
109    }
110
111    fn mutate_graph(
112        &self,
113        mutation: HNSWMutation<'_>,
114        canonical: impl FnOnce(&mut dyn KeyValueBatch) -> StorageBackendResult<()>,
115    ) -> StorageBackendResult<()> {
116        self.view.evaluate(self.store.as_ref(), |read, batch| {
117            let cached = self.graph_at(read)?;
118            let preview = !cached.definition_candidate
119                && cached.revision.is_some()
120                && !matches!(mutation, HNSWMutation::Clear);
121            if preview {
122                // Admit the retained input before constructing its derived graph.
123                batch.hnsw_mutation(&hnsw_metadata_key(&self.table, &self.field)?, mutation)?;
124            }
125            let delta = cached.value.prepare_delta(mutation, read.control())?;
126            canonical(batch)?;
127            // Only a later reader publishes the graph with its actual committed/private identity.
128            self.stage_delta(batch, &delta, next_revision(cached.revision)?, preview)
129        })
130    }
131
132    fn rebuild_graph(&self) -> StorageBackendResult<()> {
133        self.view.evaluate(self.store.as_ref(), |read, batch| {
134            let revision = hnsw_persistence::load_revision(read, &self.table, &self.field)?;
135            if revision.is_none() && self.require_persisted {
136                return Err(other_error(format!(
137                    "missing persisted HNSW metadata for {}.{}",
138                    self.table, self.field
139                )));
140            }
141            let delta = self.canonical_builder(read)?.finish_delta()?;
142            self.stage_delta(batch, &delta, next_revision(revision)?, false)
143        })
144    }
145
146    fn build_from_canonical(
147        &self,
148        read: &dyn KeyValueRead,
149    ) -> StorageBackendResult<Budgeted<HNSWIndex>> {
150        self.canonical_builder(read)?.finish()
151    }
152
153    fn canonical_builder(
154        &self,
155        read: &dyn KeyValueRead,
156    ) -> StorageBackendResult<HNSWCanonicalBuilder> {
157        let mut builder = HNSWCanonicalBuilder::new(self.dimensions, self.params, read.control())?;
158        self.raw
159            .visit_canonical_from(read, |document, ordinal, vector| {
160                builder.push(document, ordinal, vector)
161            })?;
162        Ok(builder)
163    }
164
165    fn stage_delta(
166        &self,
167        batch: &mut dyn KeyValueBatch,
168        delta: &HNSWGraphDelta,
169        revision: u64,
170        preview: bool,
171    ) -> StorageBackendResult<()> {
172        let metadata = hnsw_metadata_key(&self.table, &self.field)?;
173        let value = encode_value(&metadata_from_graph(
174            self.dimensions,
175            self.params,
176            delta.meta,
177            revision,
178        )?)?;
179        if preview {
180            batch.preview_hnsw_record(&metadata, Some(&value))?;
181        } else {
182            batch.fence_hnsw_prefix(&metadata)?;
183            batch.put(&metadata, &value)?;
184        }
185        if delta.full_rewrite {
186            let prefix = hnsw_node_prefix(&self.table, &self.field)?;
187            if preview {
188                batch.preview_hnsw_prefix(&prefix)?;
189            } else {
190                batch.delete_prefix(&prefix)?;
191            }
192        }
193        for node in delta.nodes() {
194            let node = node?;
195            let key = hnsw_node_key(&self.table, &self.field, node.node_id)?;
196            let value = encode_value(&PersistedHNSWNode::try_from(&*node)?)?;
197            if preview {
198                batch.preview_hnsw_record(&key, Some(&value))?;
199            } else {
200                batch.put(&key, &value)?;
201            }
202        }
203        Ok(())
204    }
205}
206
207impl VectorIndex for KeyValueHNSWIndex {
208    fn contains_document(&self, doc_id: DocId) -> StorageBackendResult<bool> {
209        self.raw.contains_document(doc_id)
210    }
211
212    fn dimensions(&self) -> u32 {
213        self.dimensions
214    }
215    fn index_kind(&self) -> &'static str {
216        "hnsw"
217    }
218    fn add(&mut self, doc_id: DocId, vector: Vec<f32>) -> StorageBackendResult<()> {
219        self.add_many(doc_id, vec![vector])
220    }
221    fn add_many(&mut self, doc_id: DocId, vectors: Vec<Vec<f32>>) -> StorageBackendResult<()> {
222        self.mutate_graph(
223            HNSWMutation::Replace {
224                document: doc_id,
225                vectors: &vectors,
226            },
227            |batch| self.raw.stage_replace(batch, doc_id, &vectors),
228        )
229    }
230    fn delete(&mut self, doc_id: DocId) -> StorageBackendResult<()> {
231        self.mutate_graph(HNSWMutation::Delete(doc_id), |batch| {
232            self.raw.stage_replace(batch, doc_id, &[])
233        })
234    }
235    fn clear(&mut self) -> StorageBackendResult<()> {
236        self.mutate_graph(HNSWMutation::Clear, |batch| self.raw.stage_clear(batch))
237    }
238    fn search_knn(&self, query: &[f32], k: usize) -> StorageBackendResult<PostingList> {
239        self.read_graph()?.value.search_knn(query, k)
240    }
241    fn search_threshold(&self, query: &[f32], threshold: f32) -> StorageBackendResult<PostingList> {
242        self.read_graph()?.value.search_threshold(query, threshold)
243    }
244    fn count(&self) -> StorageBackendResult<usize> {
245        self.read_graph()?.value.count()
246    }
247    fn initialize(&mut self) -> StorageBackendResult<()> {
248        self.rebuild_graph()
249    }
250    fn snapshot(&self) -> StorageBackendResult<Arc<dyn VectorIndex>> {
251        Ok(self.read_graph()?.snapshot)
252    }
253    fn vector_read_snapshot(
254        &self,
255        control: &crate::read_control::StorageReadControl,
256    ) -> StorageBackendResult<Option<crate::vector_index::VectorReadSnapshot>> {
257        self.raw.vector_read_snapshot(control)
258    }
259}
260
261fn next_revision(revision: Option<u64>) -> StorageBackendResult<u64> {
262    revision
263        .unwrap_or(0)
264        .checked_add(1)
265        .ok_or_else(|| StorageBackendError::Other("HNSW metadata revision space exhausted".into()))
266}
uqa_storage/key_value/hnsw_index.rs

uqa_storage/key_value/
hnsw_index.rs
