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 delta = cached.value.prepare_delta(mutation, read.control())?;
119            canonical(batch)?;
120            let preview = !cached.definition_candidate
121                && cached.revision.is_some()
122                && !matches!(mutation, HNSWMutation::Clear);
123            if preview {
124                batch.hnsw_mutation(&hnsw_metadata_key(&self.table, &self.field)?, mutation)?;
125            }
126            // Only a later reader publishes the graph with its actual committed/private identity.
127            self.stage_delta(batch, &delta, next_revision(cached.revision)?, preview)
128        })
129    }
130
131    fn rebuild_graph(&self) -> StorageBackendResult<()> {
132        self.view.evaluate(self.store.as_ref(), |read, batch| {
133            let revision = hnsw_persistence::load_revision(read, &self.table, &self.field)?;
134            if revision.is_none() && self.require_persisted {
135                return Err(other_error(format!(
136                    "missing persisted HNSW metadata for {}.{}",
137                    self.table, self.field
138                )));
139            }
140            let delta = self.canonical_builder(read)?.finish_delta()?;
141            self.stage_delta(batch, &delta, next_revision(revision)?, false)
142        })
143    }
144
145    fn build_from_canonical(
146        &self,
147        read: &dyn KeyValueRead,
148    ) -> StorageBackendResult<Budgeted<HNSWIndex>> {
149        self.canonical_builder(read)?.finish()
150    }
151
152    fn canonical_builder(
153        &self,
154        read: &dyn KeyValueRead,
155    ) -> StorageBackendResult<HNSWCanonicalBuilder> {
156        let mut builder = HNSWCanonicalBuilder::new(self.dimensions, self.params, read.control())?;
157        self.raw
158            .visit_canonical_from(read, |document, ordinal, vector| {
159                builder.push(document, ordinal, vector)
160            })?;
161        Ok(builder)
162    }
163
164    fn stage_delta(
165        &self,
166        batch: &mut dyn KeyValueBatch,
167        delta: &HNSWGraphDelta,
168        revision: u64,
169        preview: bool,
170    ) -> StorageBackendResult<()> {
171        let metadata = hnsw_metadata_key(&self.table, &self.field)?;
172        let value = encode_value(&metadata_from_graph(
173            self.dimensions,
174            self.params,
175            delta.meta,
176            revision,
177        )?)?;
178        if preview {
179            batch.preview_hnsw_record(&metadata, Some(&value))?;
180        } else {
181            batch.fence_hnsw_prefix(&metadata)?;
182            batch.put(&metadata, &value)?;
183        }
184        if delta.full_rewrite {
185            let prefix = hnsw_node_prefix(&self.table, &self.field)?;
186            if preview {
187                batch.preview_hnsw_prefix(&prefix)?;
188            } else {
189                batch.delete_prefix(&prefix)?;
190            }
191        }
192        for node in delta.nodes() {
193            let node = node?;
194            let key = hnsw_node_key(&self.table, &self.field, node.node_id)?;
195            let value = encode_value(&PersistedHNSWNode::try_from(&*node)?)?;
196            if preview {
197                batch.preview_hnsw_record(&key, Some(&value))?;
198            } else {
199                batch.put(&key, &value)?;
200            }
201        }
202        Ok(())
203    }
204}
205
206impl VectorIndex for KeyValueHNSWIndex {
207    fn contains_document(&self, doc_id: DocId) -> StorageBackendResult<bool> {
208        self.raw.contains_document(doc_id)
209    }
210
211    fn dimensions(&self) -> u32 {
212        self.dimensions
213    }
214    fn index_kind(&self) -> &'static str {
215        "hnsw"
216    }
217    fn add(&mut self, doc_id: DocId, vector: Vec<f32>) -> StorageBackendResult<()> {
218        self.add_many(doc_id, vec![vector])
219    }
220    fn add_many(&mut self, doc_id: DocId, vectors: Vec<Vec<f32>>) -> StorageBackendResult<()> {
221        self.mutate_graph(
222            HNSWMutation::Replace {
223                document: doc_id,
224                vectors: &vectors,
225            },
226            |batch| self.raw.stage_replace(batch, doc_id, &vectors),
227        )
228    }
229    fn delete(&mut self, doc_id: DocId) -> StorageBackendResult<()> {
230        self.mutate_graph(HNSWMutation::Delete(doc_id), |batch| {
231            self.raw.stage_replace(batch, doc_id, &[])
232        })
233    }
234    fn clear(&mut self) -> StorageBackendResult<()> {
235        self.mutate_graph(HNSWMutation::Clear, |batch| self.raw.stage_clear(batch))
236    }
237    fn search_knn(&self, query: &[f32], k: usize) -> StorageBackendResult<PostingList> {
238        self.read_graph()?.value.search_knn(query, k)
239    }
240    fn search_threshold(&self, query: &[f32], threshold: f32) -> StorageBackendResult<PostingList> {
241        self.read_graph()?.value.search_threshold(query, threshold)
242    }
243    fn count(&self) -> StorageBackendResult<usize> {
244        self.read_graph()?.value.count()
245    }
246    fn initialize(&mut self) -> StorageBackendResult<()> {
247        self.rebuild_graph()
248    }
249    fn snapshot(&self) -> StorageBackendResult<Arc<dyn VectorIndex>> {
250        Ok(self.read_graph()?.snapshot)
251    }
252    fn vector_read_snapshot(
253        &self,
254        control: &crate::read_control::StorageReadControl,
255    ) -> StorageBackendResult<Option<crate::vector_index::VectorReadSnapshot>> {
256        self.raw.vector_read_snapshot(control)
257    }
258}
259
260fn next_revision(revision: Option<u64>) -> StorageBackendResult<u64> {
261    revision
262        .unwrap_or(0)
263        .checked_add(1)
264        .ok_or_else(|| StorageBackendError::Other("HNSW metadata revision space exhausted".into()))
265}
uqa_storage/key_value/hnsw_index.rs

uqa_storage/key_value/
hnsw_index.rs
