Skip to main content

nodedb_types/
hnsw.rs

1// SPDX-License-Identifier: Apache-2.0
2
3//! Shared HNSW types used by both Origin and Lite vector engines.
4
5use crate::vector_distance::DistanceMetric;
6use crate::vector_dtype::VectorStorageDtype;
7use serde::{Deserialize, Serialize};
8
9/// HNSW index parameters shared between Origin and Lite.
10#[derive(Debug, Clone, Serialize, Deserialize)]
11pub struct HnswParams {
12    /// Max bidirectional connections per node at layers > 0.
13    pub m: usize,
14    /// Max connections at layer 0 (typically 2*M for denser base layer).
15    pub m0: usize,
16    /// Dynamic candidate list size during construction.
17    pub ef_construction: usize,
18    /// Distance metric for similarity computation.
19    pub metric: DistanceMetric,
20    /// On-disk + in-memory vector storage dtype. Defaults to F32 for
21    /// backward compatibility with indexes created before this field existed.
22    #[serde(default)]
23    pub dtype: VectorStorageDtype,
24}
25
26impl Default for HnswParams {
27    fn default() -> Self {
28        Self {
29            m: 16,
30            m0: 32,
31            ef_construction: 200,
32            metric: DistanceMetric::Cosine,
33            dtype: VectorStorageDtype::F32,
34        }
35    }
36}
37
38/// HNSW node snapshot for checkpoint serialization.
39///
40/// Shared format between Origin and Lite — both serialize nodes
41/// identically via MessagePack, enabling cross-deployment checkpoint
42/// compatibility.
43#[derive(
44    Debug, Clone, Serialize, Deserialize, zerompk::ToMessagePack, zerompk::FromMessagePack,
45)]
46pub struct HnswNodeSnapshot {
47    pub vector: Vec<f32>,
48    pub neighbors: Vec<Vec<u32>>,
49    pub deleted: bool,
50}
51
52/// HNSW checkpoint snapshot — shared serialization format.
53#[derive(
54    Debug, Clone, Serialize, Deserialize, zerompk::ToMessagePack, zerompk::FromMessagePack,
55)]
56pub struct HnswCheckpoint {
57    pub dim: usize,
58    pub m: usize,
59    pub m0: usize,
60    pub ef_construction: usize,
61    pub metric: u8,
62    pub entry_point: Option<u32>,
63    pub max_layer: usize,
64    pub rng_state: u64,
65    pub nodes: Vec<HnswNodeSnapshot>,
66}
67
68#[cfg(test)]
69mod tests {
70    use super::*;
71
72    #[test]
73    fn default_params() {
74        let p = HnswParams::default();
75        assert_eq!(p.m, 16);
76        assert_eq!(p.m0, 32);
77        assert_eq!(p.ef_construction, 200);
78    }
79
80    #[test]
81    fn checkpoint_serde_roundtrip() {
82        let snap = HnswCheckpoint {
83            dim: 128,
84            m: 16,
85            m0: 32,
86            ef_construction: 200,
87            metric: 1,
88            entry_point: Some(0),
89            max_layer: 3,
90            rng_state: 42,
91            nodes: vec![HnswNodeSnapshot {
92                vector: vec![0.1, 0.2, 0.3],
93                neighbors: vec![vec![1, 2], vec![3]],
94                deleted: false,
95            }],
96        };
97        let bytes = zerompk::to_msgpack_vec(&snap).unwrap();
98        let restored: HnswCheckpoint = zerompk::from_msgpack(&bytes).unwrap();
99        assert_eq!(restored.dim, 128);
100        assert_eq!(restored.nodes.len(), 1);
101        assert_eq!(restored.nodes[0].vector.len(), 3);
102    }
103}