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