1use crate::vector_distance::DistanceMetric;
6use crate::vector_dtype::VectorStorageDtype;
7use serde::{Deserialize, Serialize};
8
9#[derive(Debug, Clone, Serialize, Deserialize)]
11pub struct HnswParams {
12 pub m: usize,
14 pub m0: usize,
16 pub ef_construction: usize,
18 pub metric: DistanceMetric,
20 #[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#[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#[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}