Feat: implement hnsw index

.env.example

@@ -14,6 +14,7 @@ GRPC_ROOT_PASSWORD=your-secure-password
 STORAGE_TYPE=rocksdb 
 INDEX_TYPE=flat
 DIMENSION=512
+SIMILARITY=cosine
 
 DATA_PATH=./data
 

Cargo.toml

@@ -26,6 +26,7 @@ dotenv = "0.15.0"
 dotenvy = "0.15"
 prost = "0.14.1"
 prost-types = "0.14.1"
+rand = "0.9.2"
 ratatui = "0.26"
 reqwest = { version = "0.12", features = ["json", "blocking", "multipart"] }
 rocksdb = "0.21.0"

crates/api/src/lib.rs

@@ -1,6 +1,7 @@
-use defs::{DbError, IndexedVector, Similarity};
+use defs::{DbError, Dimension, IndexedVector, Similarity};
 
 use defs::{DenseVector, Payload, Point, PointId};
+use index::hnsw::HnswIndex;
 use std::path::PathBuf;
 // use std::sync::atomic::{AtomicU64, Ordering};
 use std::sync::{Arc, RwLock};
@@ -136,7 +137,8 @@ pub struct DbConfig {
     pub storage_type: StorageType,
     pub index_type: IndexType,
     pub data_path: PathBuf,
-    pub dimension: usize,
+    pub dimension: Dimension,
+    pub similarity: Similarity,
 }
 
 pub fn init_api(config: DbConfig) -> Result<VectorDb, DbError> {
@@ -149,6 +151,10 @@ pub fn init_api(config: DbConfig) -> Result<VectorDb, DbError> {
     // Initialize the vector index
     let index: Arc<RwLock<dyn VectorIndex>> = match config.index_type {
         IndexType::Flat => Arc::new(RwLock::new(FlatIndex::new())),
+        IndexType::HNSW => Arc::new(RwLock::new(HnswIndex::new(
+            config.similarity,
+            config.dimension,
+        ))),
         _ => Arc::new(RwLock::new(FlatIndex::new())),
     };
 
@@ -178,6 +184,7 @@ mod tests {
             index_type: IndexType::Flat,
             data_path: temp_dir.path().to_path_buf(),
             dimension: 3,
+            similarity: Similarity::Cosine,
         };
         init_api(config).unwrap()
     }

crates/defs/src/error.rs

@@ -1,4 +1,6 @@
 use std::io;
+
+use crate::Dimension;
 #[derive(Debug, PartialEq, Eq)]
 pub enum DbError {
     ParseError,
@@ -8,6 +10,7 @@ pub enum DbError {
     IndexError(String),
     LockError,
     DimensionMismatch,
+    InvalidDimension { expected: Dimension, got: Dimension },
 }
 
 #[derive(Debug)]

crates/defs/src/types.rs

@@ -9,6 +9,8 @@ pub type Element = f32;
 // pub type ElementHalf = f16; - Unstable https://github.com/rust-lang/rust/issues/116909
 pub type ElementByte = u8;
 
+pub type Dimension = usize;
+
 // Dense Vector and Vector are considered same
 // Sparse vector implementation not supported yet. Refer lib/sparse/src/common/sparse_vector.rs
 pub type DenseVector = Vec<Element>;
@@ -43,7 +45,7 @@ pub struct IndexedVector {
     pub vector: DenseVector,
 }
 
-#[derive(Deserialize, Copy, Clone)]
+#[derive(Debug, Deserialize, Copy, Clone)]
 pub enum Similarity {
     Euclidean,
     Manhattan,
@@ -89,3 +91,29 @@ impl<'q> Eq for DistanceOrderedVector<'q> {}
 //     Discovery(DiscoveryQuery<VectorInternal>),
 //     Context(ContextQuery<VectorInternal>),
 // }
+
+#[derive(Clone, Copy, Debug, PartialEq)]
+pub struct OrdF32(f32);
+
+impl OrdF32 {
+    pub fn new(x: f32) -> Self {
+        Self(x)
+    }
+    pub fn into_inner(self) -> f32 {
+        self.0
+    }
+}
+
+impl Eq for OrdF32 {}
+
+impl Ord for OrdF32 {
+    fn cmp(&self, other: &Self) -> Ordering {
+        self.0.total_cmp(&other.0)
+    }
+}
+
+impl PartialOrd for OrdF32 {
+    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
+        Some(self.cmp(other))
+    }
+}

crates/grpc/src/tests.rs

@@ -4,6 +4,7 @@ use crate::service::vectordb::{DenseVector, InsertVectorRequest, Payload, PointI
 use crate::service::{VectorDBService, run_server};
 use crate::utils::ServerEndpoint;
 use api::DbConfig;
+use defs::Similarity;
 use index::IndexType;
 use std::net::SocketAddr;
 use std::sync::Arc;
@@ -31,6 +32,7 @@ async fn start_test_server() -> Result<SocketAddr, Box<dyn std::error::Error>> {
         index_type: IndexType::Flat,
         data_path: temp_dir.path().to_path_buf(),
         dimension: 3,
+        similarity: Similarity::Cosine,
     };
 
     let vector_db_api = api::init_api(db_config)?;

crates/index/Cargo.toml

@@ -8,4 +8,5 @@ license.workspace = true
 
 [dependencies]
 defs.workspace = true
+rand.workspace = true
 uuid.workspace = true