Implement Flat Index

Author: 4adex

crates/core/src/error.rs

@@ -4,4 +4,5 @@ pub enum DbError {
     StorageError(String),
     SerializationError(String),
     DeserializationError,
+    IndexError(String),
 }

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 {

crates/index/src/flat.rs

@@ -1,11 +1,73 @@
-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<_>>())
     }
-}
+}

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 { 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 {