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

uqa_storage/key_value/
hnsw_index.rs
