Merge pull request #29 from sdslabs/adesh-refactor

h4shk4t · web-flow · commit dedb7df9a708 · 2025-07-02T23:02:43.000+05:30
diff --git a/crates/core/src/error.rs b/crates/core/src/error.rs
@@ -4,4 +4,5 @@ pub enum DbError {
     StorageError(String),
     SerializationError(String),
     DeserializationError,
+    IndexError(String),
 }
diff --git a/crates/core/src/types.rs b/crates/core/src/types.rs
@@ -27,6 +27,21 @@ pub struct Point {
     pub payload: Option<Payload>,
 }
 
+/// Struct which will be stored in the vector index
+#[derive(Serialize, Deserialize, Clone, Debug, PartialEq)]
+pub struct IndexedVector {
+    pub id: PointId,
+    pub vector: DenseVector,
+}
+
+#[derive(Copy, Clone)]
+pub enum Similarity {
+    Euclidean,
+    Manhattan,
+    Hamming,
+    Cosine,
+}
+
 // Query Vector. Basically the type of query results that can be generated. Not implementing this but referencing here for furture reference
 // #[derive(Debug, Clone)]
 // pub enum QueryVector {
diff --git a/crates/index/src/flat.rs b/crates/index/src/flat.rs
@@ -1,11 +1,257 @@
-use core::DenseVector;
+use core::{DbError, DenseVector, IndexedVector, PointId, Similarity};
 
-struct FlatIndex {
-    index: Vec<DenseVector>
+use crate::{distance, VectorIndex};
+
+pub struct FlatIndex {
+    index: Vec<IndexedVector>,
+}
+
+impl FlatIndex {
+    pub fn new() -> Self {
+        Self { index: Vec::new() }
+    }
+
+    pub fn build(vectors: Vec<IndexedVector>) -> Self {
+        FlatIndex { index: vectors }
+    }
+}
+
+impl Default for FlatIndex {
+    fn default() -> Self {
+        Self::new()
+    }
 }
 
 impl VectorIndex for FlatIndex {
-    fn insert() {
-        
+    fn insert(&mut self, vector: IndexedVector) -> Result<(), DbError> {
+        self.index.push(vector);
+        Ok(())
     }
-}
+
+    fn delete(&mut self, point_id: PointId) -> Result<bool, DbError> {
+        if let Some(pos) = self.index.iter().position(|vector| vector.id == point_id) {
+            self.index.remove(pos);
+            Ok(true)
+        } else {
+            Ok(false)
+        }
+    }
+
+    fn search(
+        &self,
+        query_vector: DenseVector,
+        similarity: Similarity,
+        k: usize,
+    ) -> Result<Vec<PointId>, DbError> {
+        let mut scores = self
+            .index
+            .iter()
+            .map(|point| {
+                (
+                    point.id,
+                    distance(point.vector.clone(), query_vector.clone(), similarity),
+                )
+            })
+            .collect::<Vec<_>>();
+
+        // Sorting logic according to type of metric used
+        match similarity {
+            Similarity::Euclidean | Similarity::Manhattan | Similarity::Hamming => {
+                scores.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap());
+            }
+            Similarity::Cosine => {
+                scores.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap());
+            }
+        }
+
+        Ok(scores
+            .into_iter()
+            .take(k)
+            .map(|(id, _)| id)
+            .collect::<Vec<_>>())
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_flat_index_new() {
+        let index = FlatIndex::new();
+        assert_eq!(index.index.len(), 0);
+    }
+
+    #[test]
+    fn test_flat_index_build() {
+        let vectors = vec![
+            IndexedVector {
+                id: 1,
+                vector: vec![1.0, 2.0, 3.0],
+            },
+            IndexedVector {
+                id: 2,
+                vector: vec![4.0, 5.0, 6.0],
+            },
+        ];
+        let index = FlatIndex::build(vectors.clone());
+        assert_eq!(index.index, vectors);
+    }
+
+    #[test]
+    fn test_insert() {
+        let mut index = FlatIndex::new();
+        let vector = IndexedVector {
+            id: 1,
+            vector: vec![1.0, 2.0, 3.0],
+        };
+
+        assert!(index.insert(vector.clone()).is_ok());
+        assert_eq!(index.index.len(), 1);
+        assert_eq!(index.index[0], vector);
+    }
+
+    #[test]
+    fn test_delete_existing() {
+        let mut index = FlatIndex::new();
+        let vector = IndexedVector {
+            id: 1,
+            vector: vec![1.0, 2.0, 3.0],
+        };
+        index.insert(vector).unwrap();
+
+        let result = index.delete(1).unwrap();
+        assert!(result);
+        assert_eq!(index.index.len(), 0);
+    }
+
+    #[test]
+    fn test_delete_non_existing() {
+        let mut index = FlatIndex::new();
+        let vector = IndexedVector {
+            id: 1,
+            vector: vec![1.0, 2.0, 3.0],
+        };
+        index.insert(vector).unwrap();
+
+        let result = index.delete(999).unwrap();
+        assert!(!result);
+        assert_eq!(index.index.len(), 1);
+    }
+
+    #[test]
+    fn test_search_euclidean() {
+        let mut index = FlatIndex::new();
+        index
+            .insert(IndexedVector {
+                id: 1,
+                vector: vec![1.0, 1.0],
+            })
+            .unwrap();
+        index
+            .insert(IndexedVector {
+                id: 2,
+                vector: vec![2.0, 2.0],
+            })
+            .unwrap();
+        index
+            .insert(IndexedVector {
+                id: 3,
+                vector: vec![10.0, 10.0],
+            })
+            .unwrap();
+
+        let results = index
+            .search(vec![0.0, 0.0], Similarity::Euclidean, 2)
+            .unwrap();
+        assert_eq!(results, vec![1, 2]);
+    }
+
+    #[test]
+    fn test_search_cosine() {
+        let mut index = FlatIndex::new();
+        index
+            .insert(IndexedVector {
+                id: 1,
+                vector: vec![1.0, 0.0],
+            })
+            .unwrap();
+        index
+            .insert(IndexedVector {
+                id: 2,
+                vector: vec![0.5, 0.5],
+            })
+            .unwrap();
+        index
+            .insert(IndexedVector {
+                id: 3,
+                vector: vec![0.0, 1.0],
+            })
+            .unwrap();
+
+        let results = index.search(vec![1.0, 1.0], Similarity::Cosine, 2).unwrap();
+        assert_eq!(results, vec![2, 1]);
+    }
+
+    #[test]
+    fn test_search_manhattan() {
+        let mut index = FlatIndex::new();
+        index
+            .insert(IndexedVector {
+                id: 1,
+                vector: vec![1.0, 1.0],
+            })
+            .unwrap();
+        index
+            .insert(IndexedVector {
+                id: 2,
+                vector: vec![2.0, 2.0],
+            })
+            .unwrap();
+        index
+            .insert(IndexedVector {
+                id: 3,
+                vector: vec![5.0, 5.0],
+            })
+            .unwrap();
+
+        let results = index
+            .search(vec![0.0, 0.0], Similarity::Manhattan, 2)
+            .unwrap();
+        assert_eq!(results, vec![1, 2]);
+    }
+
+    #[test]
+    fn test_search_hamming() {
+        let mut index = FlatIndex::new();
+        index
+            .insert(IndexedVector {
+                id: 1,
+                vector: vec![1.0, 0.0, 1.0, 0.0],
+            })
+            .unwrap();
+        index
+            .insert(IndexedVector {
+                id: 2,
+                vector: vec![1.0, 0.0, 0.0, 0.0],
+            })
+            .unwrap();
+        index
+            .insert(IndexedVector {
+                id: 3,
+                vector: vec![0.0, 0.0, 0.0, 0.0],
+            })
+            .unwrap();
+
+        let results = index
+            .search(vec![1.0, 0.0, 0.0, 0.0], Similarity::Hamming, 2)
+            .unwrap();
+        assert_eq!(results, vec![2, 3]);
+    }
+
+    #[test]
+    fn test_default() {
+        let index = FlatIndex::default();
+        assert_eq!(index.index.len(), 0);
+    }
+}
diff --git a/crates/index/src/lib.rs b/crates/index/src/lib.rs
@@ -1,11 +1,61 @@
-use core::{DenseVector, PointId};
-use std::fmt::Error;
+use core::{DbError, DenseVector, IndexedVector, PointId, Similarity};
+
+pub mod flat;
 
 pub trait VectorIndex {
-    fn insert(&self, vector: DenseVector) -> Result<(), Error>;
-    fn delete(&self, point_id: PointId) -> Result<(), Error>;
-    fn search(&self, query_vector: DenseVector) -> Result<DenseVector, Error>;
-    // fn build() -> Result<(), Error>; move this to impl for dyn compatibility
+    fn insert(&mut self, vector: IndexedVector) -> Result<(), DbError>;
+
+    // Returns true if point id existed and is deleted, else returns false
+    fn delete(&mut self, point_id: PointId) -> Result<bool, DbError>;
+
+    fn search(
+        &self,
+        query_vector: DenseVector,
+        similarity: Similarity,
+        k: usize,
+    ) -> Result<Vec<PointId>, DbError>; // Return a Vec of ids of closest vectors (length max k)
+
+    // fn build() -> Result<(), DbError>; move this to impl for dyn compatibility
+}
+
+/// Distance function to get the distance between two vectors (taken from old version)
+pub fn distance(a: DenseVector, b: DenseVector, dist_type: Similarity) -> f32 {
+    assert_eq!(a.len(), b.len());
+    match dist_type {
+        Similarity::Euclidean => {
+            let score: Vec<f32> = a
+                .iter()
+                .zip(b.iter())
+                .map(|(&x, &y)| (x - y) * (x - y))
+                .collect();
+            score.iter().sum::<f32>().sqrt()
+        }
+        Similarity::Manhattan => {
+            let score: Vec<f32> = a
+                .iter()
+                .zip(b.iter())
+                .map(|(&x, &y)| (x - y).abs())
+                .collect();
+            score.iter().sum::<f32>()
+        }
+        Similarity::Hamming => {
+            let score: Vec<f32> = a
+                .iter()
+                .zip(b.iter())
+                .map(|(&x, &y)| (if (x - y) > 1e-8 { 1f32 } else { 0f32 }))
+                .collect();
+            score.iter().sum::<f32>()
+        }
+        Similarity::Cosine => {
+            let p_score: Vec<f32> = a.iter().zip(b.iter()).map(|(&x, &y)| x * y).collect();
+            let p = p_score.iter().sum::<f32>();
+            let q_score: Vec<f32> = a.iter().map(|&n| n * n).collect();
+            let q = q_score.iter().sum::<f32>().sqrt();
+            let r_score: Vec<f32> = b.iter().map(|&n| n * n).collect();
+            let r = r_score.iter().sum::<f32>().sqrt();
+            p / (q * r)
+        }
+    }
 }
 
 pub enum IndexType {

Original file line number	Diff line number	Diff line change
`@@ -4,4 +4,5 @@ pub enum DbError {`
`4`	`4`	`StorageError(String),`
`5`	`5`	`SerializationError(String),`
`6`	`6`	`DeserializationError,`
	`7`	`+ IndexError(String),`
`7`	`8`	`}`