nodedb_lite/engine/sparse_vector/
checkpoint.rs

1// SPDX-License-Identifier: Apache-2.0
2
3//! Checkpoint serialization and restoration for [`SparseVectorManager`].
4//!
5//! Persists every sparse inverted index so a cold open loads them directly —
6//! no rebuild pass over source documents is required.
7//!
8//! ## Key layout under `Namespace::Vector`
9//!
10//! `nodedb_types::Namespace` has no sparse-specific variant, so the sparse
11//! index shares the vector engine's namespace behind a distinct `sparse:`
12//! key prefix that cannot collide with the dense HNSW keys.
13//!
14//! | Key                        | Value                                                 |
15//! |----------------------------|-------------------------------------------------------|
16//! | `sparse:_indices`          | MessagePack `Vec<String>` — index key list            |
17//! | `sparse:{index_key}:docs`  | MessagePack `Vec<(String, Vec<(u32, f32)>)>`          |
18//!
19//! Documents are stored as their own `(dimension, weight)` entries rather than
20//! as posting lists: the inverted index is a pure function of the document
21//! vectors, so storing the source form keeps the blob minimal and makes the
22//! restored index structurally identical to a freshly built one.
23//!
24//! [`SparseVectorManager`]: super::manager::SparseVectorManager
25
26use std::collections::HashMap;
27
28use nodedb_types::Namespace;
29use nodedb_types::error::{NodeDbError, NodeDbResult};
30
31use crate::storage::engine::{StorageEngine, WriteOp};
32
33use super::index::SparseInvertedIndex;
34
35/// Key holding the MessagePack list of live sparse index keys.
36const INDEX_LIST_KEY: &[u8] = b"sparse:_indices";
37
38/// A single document's sparse vector as persisted.
39type SerDocument = (String, Vec<(u32, f32)>);
40
41/// Per-index documents key: `sparse:{index_key}:docs`.
42fn documents_key(index_key: &str) -> String {
43    format!("sparse:{index_key}:docs")
44}
45
46/// Serialize sparse index state into write ops.
47///
48/// Pure and synchronous, so it is safe to call while holding the manager's
49/// mutex guard; the caller performs the I/O after releasing the lock.
50///
51/// Empty indexes are still listed so that a checkpoint taken after every
52/// document was deleted restores as empty rather than resurrecting the
53/// previous checkpoint's contents.
54pub(crate) fn serialize_sparse(
55    indices: &HashMap<String, SparseInvertedIndex>,
56) -> NodeDbResult<Vec<WriteOp>> {
57    let mut ops: Vec<WriteOp> = Vec::with_capacity(indices.len() + 1);
58
59    let mut index_keys: Vec<String> = indices.keys().cloned().collect();
60    index_keys.sort();
61    let keys_bytes = zerompk::to_msgpack_vec(&index_keys)
62        .map_err(|e| NodeDbError::serialization("msgpack", e))?;
63    ops.push(WriteOp::Put {
64        ns: Namespace::Vector,
65        key: INDEX_LIST_KEY.to_vec(),
66        value: keys_bytes,
67    });
68
69    for (index_key, index) in indices {
70        let documents: Vec<SerDocument> = index.documents();
71        let bytes = zerompk::to_msgpack_vec(&documents)
72            .map_err(|e| NodeDbError::serialization("msgpack", e))?;
73        ops.push(WriteOp::Put {
74            ns: Namespace::Vector,
75            key: documents_key(index_key).into_bytes(),
76            value: bytes,
77        });
78    }
79
80    Ok(ops)
81}
82
83/// Write pre-serialized sparse index state to storage.
84pub(crate) async fn write_serialized_sparse<S>(storage: &S, ops: Vec<WriteOp>) -> NodeDbResult<()>
85where
86    S: StorageEngine,
87{
88    if ops.is_empty() {
89        return Ok(());
90    }
91    storage
92        .batch_write(&ops)
93        .await
94        .map_err(|e| NodeDbError::storage(format!("sparse checkpoint batch_write: {e}")))?;
95    Ok(())
96}
97
98/// Restore sparse index state from storage on cold open.
99///
100/// Returns `None` when no checkpoint has ever been written, which the caller
101/// distinguishes from a checkpoint that legitimately holds no documents — only
102/// the former warrants a rebuild from source documents. A corrupt blob for one
103/// index is logged and that index restored empty rather than failing the whole
104/// open; the remaining indexes stay usable.
105pub(crate) async fn restore_sparse<S>(
106    storage: &S,
107) -> NodeDbResult<Option<HashMap<String, SparseInvertedIndex>>>
108where
109    S: StorageEngine,
110{
111    let Some(keys_bytes) = storage.get(Namespace::Vector, INDEX_LIST_KEY).await? else {
112        return Ok(None);
113    };
114    let Ok(index_keys) = zerompk::from_msgpack::<Vec<String>>(&keys_bytes) else {
115        tracing::warn!("sparse checkpoint: index list undecodable — starting fresh");
116        return Ok(None);
117    };
118
119    let mut indices: HashMap<String, SparseInvertedIndex> =
120        HashMap::with_capacity(index_keys.len());
121
122    for index_key in &index_keys {
123        let key = documents_key(index_key);
124        let Some(bytes) = storage.get(Namespace::Vector, key.as_bytes()).await? else {
125            indices.insert(index_key.clone(), SparseInvertedIndex::new());
126            continue;
127        };
128        match zerompk::from_msgpack::<Vec<SerDocument>>(&bytes) {
129            Ok(documents) => {
130                indices.insert(
131                    index_key.clone(),
132                    SparseInvertedIndex::from_documents(documents),
133                );
134            }
135            Err(e) => {
136                tracing::warn!(
137                    index_key,
138                    error = %e,
139                    "sparse checkpoint: document blob undecodable — index restored empty"
140                );
141                indices.insert(index_key.clone(), SparseInvertedIndex::new());
142            }
143        }
144    }
145
146    tracing::debug!(index_count = indices.len(), "sparse checkpoint restored");
147    Ok(Some(indices))
148}
149
150#[cfg(test)]
151mod tests {
152    use nodedb_types::SparseVector;
153
154    use super::*;
155    use crate::PagedbStorageMem;
156
157    fn sv(entries: &[(u32, f32)]) -> SparseVector {
158        SparseVector::from_entries(entries.to_vec()).expect("valid sparse vector")
159    }
160
161    #[tokio::test]
162    async fn round_trip_through_storage() {
163        let storage = PagedbStorageMem::open_in_memory()
164            .await
165            .expect("in-memory pagedb");
166
167        let mut index = SparseInvertedIndex::new();
168        index.insert("d1", &sv(&[(1, 0.5), (9, 0.25)]));
169        index.insert("d2", &sv(&[(9, 1.0)]));
170
171        let mut indices = HashMap::new();
172        indices.insert("docs:emb".to_string(), index);
173
174        let ops = serialize_sparse(&indices).expect("serialize");
175        write_serialized_sparse(&storage, ops).await.expect("write");
176
177        let restored = restore_sparse(&storage)
178            .await
179            .expect("restore")
180            .expect("checkpoint present");
181        let restored_index = restored.get("docs:emb").expect("index restored");
182        assert_eq!(restored_index.doc_count(), 2);
183
184        let hits = restored_index.search(&sv(&[(9, 1.0)]), 10);
185        assert_eq!(hits.len(), 2);
186        assert_eq!(hits[0].doc_id, "d2");
187    }
188
189    #[tokio::test]
190    async fn absent_checkpoint_reports_none() {
191        let storage = PagedbStorageMem::open_in_memory()
192            .await
193            .expect("in-memory pagedb");
194        let restored = restore_sparse(&storage).await.expect("restore");
195        assert!(restored.is_none());
196    }
197}
nodedb_lite/engine/sparse_vector/checkpoint.rs

nodedb_lite/engine/sparse_vector/
checkpoint.rs
