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::{HNSWIndex, HNSWMutation, HNSWPersistenceDelta};
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 vectors = self.raw.load_all_from(read)?;
141            let delta = HNSWIndex::prepare_canonical(
142                self.dimensions,
143                self.params,
144                &vectors,
145                read.control(),
146            )?;
147            self.stage_delta(batch, &delta, next_revision(revision)?, false)
148        })
149    }
150
151    fn build_from_canonical(
152        &self,
153        read: &dyn KeyValueRead,
154    ) -> StorageBackendResult<Budgeted<HNSWIndex>> {
155        let entries = self.raw.load_all_from(read)?;
156        HNSWIndex::from_canonical_controlled(self.dimensions, self.params, &entries, read.control())
157    }
158
159    fn stage_delta(
160        &self,
161        batch: &mut dyn KeyValueBatch,
162        delta: &HNSWPersistenceDelta,
163        revision: u64,
164        preview: bool,
165    ) -> StorageBackendResult<()> {
166        let metadata = hnsw_metadata_key(&self.table, &self.field)?;
167        let value = encode_value(&metadata_from_graph(
168            self.dimensions,
169            self.params,
170            delta.meta,
171            revision,
172        )?)?;
173        if preview {
174            batch.preview_hnsw_record(&metadata, Some(&value))?;
175        } else {
176            batch.fence_hnsw_prefix(&metadata)?;
177            batch.put(&metadata, &value)?;
178        }
179        if delta.full_rewrite {
180            let prefix = hnsw_node_prefix(&self.table, &self.field)?;
181            if preview {
182                batch.preview_hnsw_prefix(&prefix)?;
183            } else {
184                batch.delete_prefix(&prefix)?;
185            }
186        }
187        for node in &delta.nodes {
188            let key = hnsw_node_key(&self.table, &self.field, node.node_id)?;
189            let value = encode_value(&PersistedHNSWNode::try_from(node)?)?;
190            if preview {
191                batch.preview_hnsw_record(&key, Some(&value))?;
192            } else {
193                batch.put(&key, &value)?;
194            }
195        }
196        Ok(())
197    }
198}
199
200impl VectorIndex for KeyValueHNSWIndex {
201    fn contains_document(&self, doc_id: DocId) -> StorageBackendResult<bool> {
202        self.raw.contains_document(doc_id)
203    }
204
205    fn dimensions(&self) -> u32 {
206        self.dimensions
207    }
208    fn index_kind(&self) -> &'static str {
209        "hnsw"
210    }
211    fn add(&mut self, doc_id: DocId, vector: Vec<f32>) -> StorageBackendResult<()> {
212        self.add_many(doc_id, vec![vector])
213    }
214    fn add_many(&mut self, doc_id: DocId, vectors: Vec<Vec<f32>>) -> StorageBackendResult<()> {
215        self.mutate_graph(
216            HNSWMutation::Replace {
217                document: doc_id,
218                vectors: &vectors,
219            },
220            |batch| self.raw.stage_replace(batch, doc_id, &vectors),
221        )
222    }
223    fn delete(&mut self, doc_id: DocId) -> StorageBackendResult<()> {
224        self.mutate_graph(HNSWMutation::Delete(doc_id), |batch| {
225            self.raw.stage_replace(batch, doc_id, &[])
226        })
227    }
228    fn clear(&mut self) -> StorageBackendResult<()> {
229        self.mutate_graph(HNSWMutation::Clear, |batch| self.raw.stage_clear(batch))
230    }
231    fn search_knn(&self, query: &[f32], k: usize) -> StorageBackendResult<PostingList> {
232        self.read_graph()?.value.search_knn(query, k)
233    }
234    fn search_threshold(&self, query: &[f32], threshold: f32) -> StorageBackendResult<PostingList> {
235        self.read_graph()?.value.search_threshold(query, threshold)
236    }
237    fn count(&self) -> StorageBackendResult<usize> {
238        self.read_graph()?.value.count()
239    }
240    fn initialize(&mut self) -> StorageBackendResult<()> {
241        self.rebuild_graph()
242    }
243    fn snapshot(&self) -> StorageBackendResult<Arc<dyn VectorIndex>> {
244        Ok(self.read_graph()?.snapshot)
245    }
246    fn vector_read_snapshot(
247        &self,
248        control: &crate::read_control::StorageReadControl,
249    ) -> StorageBackendResult<Option<crate::vector_index::VectorReadSnapshot>> {
250        self.raw.vector_read_snapshot(control)
251    }
252}
253
254fn next_revision(revision: Option<u64>) -> StorageBackendResult<u64> {
255    revision
256        .unwrap_or(0)
257        .checked_add(1)
258        .ok_or_else(|| StorageBackendError::Other("HNSW metadata revision space exhausted".into()))
259}
uqa_storage/key_value/hnsw_index.rs

uqa_storage/key_value/
hnsw_index.rs
