Add check in api crate for dimension mismatch while insertion; Add tests for the same

Author: TanmayArya-1p

crates/api/src/lib.rs

@@ -19,15 +19,27 @@ fn generate_point_id() -> u64 {
 pub struct VectorDb {
     storage: Arc<dyn StorageEngine>,
     index: Arc<RwLock<dyn VectorIndex>>, // Using a RwLock instead of Mutex to improve concurrency
+    dimension: usize,
 }
 
 impl VectorDb {
-    fn _new(storage: Arc<dyn StorageEngine>, index: Arc<RwLock<dyn VectorIndex>>) -> Self {
-        Self { storage, index }
+    fn _new(
+        storage: Arc<dyn StorageEngine>,
+        index: Arc<RwLock<dyn VectorIndex>>,
+        dimension: usize,
+    ) -> Self {
+        Self {
+            storage,
+            index,
+            dimension,
+        }
     }
 
     //TODO: Make this an atomic operation
     pub fn insert(&self, vector: DenseVector, payload: Payload) -> Result<PointId, DbError> {
+        if vector.len() != self.dimension {
+            return Err(DbError::DimensionMismatch);
+        }
         // Generate a new point id
         let point_id = generate_point_id();
         self.storage
@@ -105,7 +117,7 @@ pub fn init_api(config: DbConfig) -> Result<VectorDb, DbError> {
     };
 
     // Init the db
-    let db = VectorDb::_new(storage, index);
+    let db = VectorDb::_new(storage, index, config.dimension);
 
     Ok(db)
 }
@@ -147,6 +159,22 @@ mod tests {
         assert_eq!(point.payload.as_ref().unwrap(), &payload);
     }
 
+    #[test]
+    fn test_dimension_mismatch() {
+        let db = create_test_db();
+        let v1 = vec![1.0, 2.0, 3.0];
+        let v2 = vec![1.0, 2.0];
+        let payload = Payload {};
+
+        let res1 = db.insert(v1, payload);
+        assert!(res1.is_ok());
+
+        // Insert vector of dimension 2 != 3
+        let res2 = db.insert(v2, payload);
+        assert!(res2.is_err());
+        assert_eq!(res2.unwrap_err(), DbError::DimensionMismatch);
+    }
+
     #[test]
     fn test_delete() {
         let db = create_test_db();

crates/defs/src/error.rs

@@ -1,9 +1,10 @@
-#[derive(Debug)]
+#[derive(Debug, PartialEq, Eq)]
 pub enum DbError {
     ParseError,
     StorageError(String),
     SerializationError(String),
     DeserializationError,
     IndexError(String),
     LockError,
+    DimensionMismatch,
 }

crates/grpc_server/.sample.env

@@ -1,9 +1,9 @@
-GRPC_SERVER_ROOT_PASSWORD=123 // required
-GRPC_SERVER_DIMENSION=3 // required
+GRPC_SERVER_ROOT_PASSWORD=123 # required
+GRPC_SERVER_DIMENSION=3 # required
 
-GRPC_SERVER_HOST=localhost // defaults to 127.0.0.1 aka localhost
-GRPC_SERVER_PORT=8080 // defaults to 8080
-GRPC_SERVER_STORAGE_TYPE=inmemory // (inmemory/rocksdb) defaults to 'inmemory'
-GRPC_SERVER_INDEX_TYPE=flat // defaults to flat
-GRPC_SERVER_DATA_PATH=data // defaults to a temporary directory
-GRPC_SERVER_LOGGING=true // defaults to true
+GRPC_SERVER_HOST=localhost # defaults to 127.0.0.1 aka localhost
+GRPC_SERVER_PORT=8080 # defaults to 8080
+GRPC_SERVER_STORAGE_TYPE=inmemory # (inmemory/rocksdb) defaults to 'inmemory'
+GRPC_SERVER_INDEX_TYPE=flat # defaults to flat
+GRPC_SERVER_DATA_PATH=data # defaults to a temporary directory
+GRPC_SERVER_LOGGING=true # defaults to true

crates/grpc_server/src/tests.rs

@@ -0,0 +1 @@
+