add(snapshots): implement snapshots for hnsw index

Author: TanmayArya-1p

crates/api/src/lib.rs

@@ -233,7 +233,7 @@ mod tests {
         let temp_dir = tempdir().unwrap();
         let config = DbConfig {
             storage_type: StorageType::RocksDb,
-            index_type: IndexType::Flat,
+            index_type: IndexType::HNSW,
             data_path: temp_dir.path().to_path_buf(),
             dimension: 3,
             similarity: Similarity::Cosine,

crates/defs/src/types.rs

@@ -47,7 +47,7 @@ pub struct IndexedVector {
     pub vector: DenseVector,
 }
 
-#[derive(Debug, Deserialize, Copy, Clone)]
+#[derive(Debug, Serialize , Deserialize, Copy, Clone)]
 pub enum Similarity {
     Euclidean,
     Manhattan,

crates/index/src/hnsw/index.rs

@@ -18,7 +18,7 @@ pub struct HnswIndex {
     // Default query beam width (ef); recommended ef ≥ k at query time
     pub ef: usize,
     // In-memory vector cache owned by the index
-    cache: HashMap<PointId, DenseVector>,
+    pub cache: HashMap<PointId, DenseVector>,
     // Fixed metric for this index; used consistently in insert and search
     pub similarity: Similarity,
 }

crates/index/src/hnsw/mod.rs

@@ -5,8 +5,10 @@ pub mod index;
 pub mod search;
 pub mod serialize;
 pub mod types;
-
+use defs::Magic;
 pub use index::HnswIndex;
 
+pub const HNSW_MAGIC_BYTES: Magic = [0x02, 0x01, 0x03, 0x00];
+
 #[cfg(test)]
 mod tests;

crates/index/src/hnsw/serialize.rs

@@ -1,21 +1,146 @@
-use defs::DbError;
+use std::collections::HashMap;
+
+use defs::{DbError, Dimension, PointId, Similarity};
+use serde::{Deserialize, Serialize};
 use storage::StorageEngine;
 
-use crate::{IndexSnapshot, SerializableIndex, hnsw::HnswIndex};
+use crate::{IndexSnapshot, IndexType, SerializableIndex, hnsw::{HNSW_MAGIC_BYTES, HnswIndex, types::{LevelGenerator, Node, PointIndexation}}};
+
+#[repr(packed)]
+#[derive(Serialize, Deserialize)]
+pub struct HnswMetadataPack{
+    pub ef_construction: usize,
+    pub data_dimension: Dimension,
+    pub ef: usize,
+    pub similarity: Similarity,
+}
+
+#[derive(Serialize, Deserialize)]
+pub struct HnswIndexPack {
+    pub max_connections: usize,
+    pub max_connections_0: usize,
+    pub max_layer: usize,
+    pub points_by_layer: Vec<Vec<PointId>>,
+    pub nodes: Vec<Node>,
+    pub entry_point: Option<PointId>,
+    pub level_scale: f64,
+}
+
 
 impl SerializableIndex for HnswIndex {
     fn serialize_topology(&self) -> Result<Vec<u8>, DbError> {
-        return Err(DbError::SerializationError("not implemented".to_string()));
+        let mut buffer = Vec::new();
+
+        let nodes: Vec<Node> = self.index.nodes.values().cloned().collect();
+        let index_pack = HnswIndexPack {
+            max_connections: self.index.max_connections,
+            max_connections_0: self.index.max_connections_0,
+            max_layer: self.index.max_layer,
+            points_by_layer: self.index.points_by_layer.clone(),
+            nodes,
+            entry_point: self.index.entry_point,
+            level_scale: self.index.level_generator.level_scale,
+        };
+
+        let index_bytes = bincode::serialize(&index_pack).map_err(|e| DbError::SerializationError(e.to_string()))?;
+        buffer.extend(index_bytes);
+
+        return Ok(buffer);
     }
+
     fn serialize_metadata(&self) -> Result<Vec<u8>, DbError> {
-        return Err(DbError::SerializationError("not implemented".to_string()));
+        let mut buffer = Vec::new();
+        let index_pack = HnswMetadataPack {
+            ef_construction: self.ef_construction,
+            data_dimension: self.data_dimension,
+            ef: self.ef,
+            similarity: self.similarity
+        };
+
+        let metadata_bytes = bincode::serialize(&index_pack).map_err(|e| DbError::SerializationError(e.to_string()))?;
+        buffer.extend(metadata_bytes);
+        return Ok(buffer);
     }
 
     fn snapshot(&self) -> Result<IndexSnapshot, DbError> {
-        return Err(DbError::SerializationError("not implemented".to_string()));
+        let topology_bytes = self.serialize_topology()?;
+        let metadata_bytes = self.serialize_metadata()?;
+        Ok(IndexSnapshot {
+            index_type: crate::IndexType::HNSW,
+            magic: HNSW_MAGIC_BYTES,
+            topology_b: topology_bytes,
+            metadata_b: metadata_bytes,
+        })
+    }
+
+    fn populate_vectors(&mut self, storage: &dyn StorageEngine) -> Result<(), DbError> {
+        // assumes index topology is restored
+        for id in self.index.nodes.keys() {
+            let vec = storage.get_vector(*id)?.ok_or(DbError::SerializationError(format!("Failed to locate vector for id: {} in storage", id)))?;
+            self.cache.insert(*id, vec);
+        }
+        Ok(())
     }
+}
+
+
+
+impl HnswIndex {
+    pub fn deserialize(
+        IndexSnapshot {
+            index_type,
+            magic,
+            topology_b,
+            metadata_b,
+        }: &IndexSnapshot,
+    ) -> Result<HnswIndex, DbError> {
+        if index_type != &IndexType::HNSW {
+            return Err(DbError::SerializationError(
+                "Invalid index type".to_string(),
+            ));
+        }
+
+        if magic != &HNSW_MAGIC_BYTES {
+            return Err(DbError::SerializationError(
+                "Invalid magic bytes".to_string(),
+            ));
+        }
+
+        let metadata: HnswMetadataPack = bincode::deserialize(metadata_b).map_err(|e| {
+            DbError::SerializationError(format!("Failed to deserialize HNSW Metadata: {}", e))
+        })?;
+
+        let index_pack : HnswIndexPack = bincode::deserialize(topology_b).map_err(|e| {
+            DbError::SerializationError(format!("Failed to deserialize HNSW Index: {}", e))
+        })?;
+
+        let mut hnsw_index_restored = PointIndexation {
+            max_connections: index_pack.max_connections,
+            max_connections_0: index_pack.max_connections_0,
+            max_layer: index_pack.max_layer,
+            points_by_layer: index_pack.points_by_layer,
+            entry_point: index_pack.entry_point,
+            nodes: HashMap::new(),
+            level_generator: LevelGenerator{
+                level_scale: index_pack.level_scale
+            }
+        };
+
+        // restore nodes hashmap
+        for i in index_pack.nodes {
+            hnsw_index_restored.nodes.insert(i.id, i);
+        }
+
+
+        let hnsw = HnswIndex {
+            ef_construction : metadata.ef_construction,
+            data_dimension: metadata.data_dimension,
+            ef: metadata.ef,
+            cache: HashMap::new(),
+            similarity: metadata.similarity,
+            index: hnsw_index_restored
+        };
 
-    fn populate_vectors(&mut self, _storage: &dyn StorageEngine) -> Result<(), DbError> {
-        return Err(DbError::SerializationError("not implemented".to_string()));
+        Ok(hnsw)
     }
 }

crates/index/src/hnsw/types.rs

@@ -2,6 +2,7 @@ use std::collections::HashMap;
 
 use defs::PointId;
 use rand::Rng;
+use serde::{Deserialize, Serialize};
 
 // Compact storage for layered points and adjacency used by `HnswIndex`.
 pub struct PointIndexation {
@@ -22,6 +23,7 @@ pub struct PointIndexation {
 }
 
 // Node with highest level and per-level neighbor lists
+#[derive(Serialize, Deserialize, Clone)]
 pub struct Node {
     pub id: PointId,
     // Highest level (0-based; level 0 is the base layer)

crates/index/src/kd_tree/mod.rs

@@ -8,4 +8,4 @@ pub mod types;
 #[cfg(test)]
 mod tests;
 
-pub const KD_TREE_MAGIC_BYTES: Magic = [0x00, 0x00, 0x00, 0x00];
+pub const KD_TREE_MAGIC_BYTES: Magic = [0x00, 0x01, 0x02, 0x00];

crates/snapshot/src/lib.rs

@@ -15,7 +15,7 @@ use chrono::{DateTime, Local};
 use defs::DbError;
 use flate2::read::GzDecoder;
 use index::{
-    IndexSnapshot, IndexType, VectorIndex, flat::index::FlatIndex, kd_tree::index::KDTree,
+    IndexSnapshot, IndexType, VectorIndex, flat::index::FlatIndex, hnsw::HnswIndex, kd_tree::index::KDTree
 };
 use semver::Version;
 use std::{
@@ -281,6 +281,7 @@ impl Snapshot {
         let vector_index: Arc<RwLock<dyn VectorIndex>> = match manifest.index_type {
             IndexType::Flat => Arc::new(RwLock::new(FlatIndex::deserialize(&index_snapshot)?)),
             IndexType::KDTree => Arc::new(RwLock::new(KDTree::deserialize(&index_snapshot)?)),
+            IndexType::HNSW => Arc::new(RwLock::new(HnswIndex::deserialize(&index_snapshot)?)),
             _ => return Err(DbError::SnapshotError("Unsupported index type".to_string())),
         };