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