Add tests and fix hamming distance

Author: 4adex

crates/index/src/flat.rs

@@ -71,3 +71,187 @@ impl VectorIndex for FlatIndex {
             .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);
+    }
+}

crates/index/src/lib.rs

@@ -42,7 +42,7 @@ pub fn distance(a: DenseVector, b: DenseVector, dist_type: Similarity) -> f32 {
             let score: Vec<f32> = a
                 .iter()
                 .zip(b.iter())
-                .map(|(&x, &y)| (if x != y { 1f32 } else { 0f32 }))
+                .map(|(&x, &y)| (if (x - y) > 1e-8 { 1f32 } else { 0f32 }))
                 .collect();
             score.iter().sum::<f32>()
         }