feat: support sparse vector

crates/service/src/algorithms/clustering/elkan_k_means.rs

@@ -7,18 +7,18 @@ use std::ops::{Index, IndexMut};
 pub struct ElkanKMeans<S: G> {
     dims: u16,
     c: usize,
-    pub centroids: Vec2<S>,
+    pub centroids: Vec2<S::Scalar>,
     lowerbound: Square,
     upperbound: Vec<F32>,
     assign: Vec<usize>,
     rand: StdRng,
-    samples: Vec2<S>,
+    samples: Vec2<S::Scalar>,
 }
 
 const DELTA: f32 = 1.0 / 1024.0;
 
 impl<S: G> ElkanKMeans<S> {
-    pub fn new(c: usize, samples: Vec2<S>) -> Self {
+    pub fn new(c: usize, samples: Vec2<S::Scalar>) -> Self {
         let n = samples.len();
         let dims = samples.dims();
 
@@ -249,7 +249,7 @@ impl<S: G> ElkanKMeans<S> {
         change == 0
     }
 
-    pub fn finish(self) -> Vec2<S> {
+    pub fn finish(self) -> Vec2<S::Scalar> {
         self.centroids
     }
 }

crates/service/src/algorithms/flat.rs

@@ -8,7 +8,7 @@ use crate::utils::dir_ops::sync_dir;
 use std::cmp::Reverse;
 use std::collections::BinaryHeap;
 use std::fs::create_dir;
-use std::path::PathBuf;
+use std::path::Path;
 use std::sync::Arc;
 
 pub struct Flat<S: G> {
@@ -17,37 +17,26 @@ pub struct Flat<S: G> {
 
 impl<S: G> Flat<S> {
     pub fn create(
-        path: PathBuf,
+        path: &Path,
         options: IndexOptions,
         sealed: Vec<Arc<SealedSegment<S>>>,
         growing: Vec<Arc<GrowingSegment<S>>>,
     ) -> Self {
-        create_dir(&path).unwrap();
-        let ram = make(path.clone(), sealed, growing, options.clone());
-        let mmap = save(ram, path.clone());
-        sync_dir(&path);
+        create_dir(path).unwrap();
+        let ram = make(path, sealed, growing, options);
+        let mmap = save(path, ram);
+        sync_dir(path);
         Self { mmap }
     }
-    pub fn open(path: PathBuf, options: IndexOptions) -> Self {
-        let mmap = load(path, options.clone());
-        Self { mmap }
-    }
-
-    pub fn len(&self) -> u32 {
-        self.mmap.raw.len()
-    }
 
-    pub fn vector(&self, i: u32) -> &[S::Scalar] {
-        self.mmap.raw.vector(i)
-    }
-
-    pub fn payload(&self, i: u32) -> Payload {
-        self.mmap.raw.payload(i)
+    pub fn open(path: &Path, options: IndexOptions) -> Self {
+        let mmap = open(path, options);
+        Self { mmap }
     }
 
     pub fn basic(
         &self,
-        vector: &[S::Scalar],
+        vector: S::VectorRef<'_>,
         _opts: &SearchOptions,
         filter: impl Filter,
     ) -> BinaryHeap<Reverse<Element>> {
@@ -56,12 +45,24 @@ impl<S: G> Flat<S> {
 
     pub fn vbase<'a>(
         &'a self,
-        vector: &'a [S::Scalar],
+        vector: S::VectorRef<'a>,
         _opts: &'a SearchOptions,
         filter: impl Filter + 'a,
     ) -> (Vec<Element>, Box<(dyn Iterator<Item = Element> + 'a)>) {
         vbase(&self.mmap, vector, filter)
     }
+
+    pub fn len(&self) -> u32 {
+        self.mmap.raw.len()
+    }
+
+    pub fn vector(&self, i: u32) -> S::VectorRef<'_> {
+        self.mmap.raw.vector(i)
+    }
+
+    pub fn payload(&self, i: u32) -> Payload {
+        self.mmap.raw.payload(i)
+    }
 }
 
 unsafe impl<S: G> Send for Flat<S> {}
@@ -81,39 +82,39 @@ unsafe impl<S: G> Send for FlatMmap<S> {}
 unsafe impl<S: G> Sync for FlatMmap<S> {}
 
 pub fn make<S: G>(
-    path: PathBuf,
+    path: &Path,
     sealed: Vec<Arc<SealedSegment<S>>>,
     growing: Vec<Arc<GrowingSegment<S>>>,
     options: IndexOptions,
 ) -> FlatRam<S> {
     let idx_opts = options.indexing.clone().unwrap_flat();
     let raw = Arc::new(Raw::create(
-        path.join("raw"),
+        &path.join("raw"),
         options.clone(),
         sealed,
         growing,
     ));
     let quantization = Quantization::create(
-        path.join("quantization"),
+        &path.join("quantization"),
         options.clone(),
         idx_opts.quantization,
         &raw,
     );
     FlatRam { raw, quantization }
 }
 
-pub fn save<S: G>(ram: FlatRam<S>, _: PathBuf) -> FlatMmap<S> {
+pub fn save<S: G>(_: &Path, ram: FlatRam<S>) -> FlatMmap<S> {
     FlatMmap {
         raw: ram.raw,
         quantization: ram.quantization,
     }
 }
 
-pub fn load<S: G>(path: PathBuf, options: IndexOptions) -> FlatMmap<S> {
+pub fn open<S: G>(path: &Path, options: IndexOptions) -> FlatMmap<S> {
     let idx_opts = options.indexing.clone().unwrap_flat();
-    let raw = Arc::new(Raw::open(path.join("raw"), options.clone()));
+    let raw = Arc::new(Raw::open(&path.join("raw"), options.clone()));
     let quantization = Quantization::open(
-        path.join("quantization"),
+        &path.join("quantization"),
         options.clone(),
         idx_opts.quantization,
         &raw,
@@ -123,7 +124,7 @@ pub fn load<S: G>(path: PathBuf, options: IndexOptions) -> FlatMmap<S> {
 
 pub fn basic<S: G>(
     mmap: &FlatMmap<S>,
-    vector: &[S::Scalar],
+    vector: S::VectorRef<'_>,
     mut filter: impl Filter,
 ) -> BinaryHeap<Reverse<Element>> {
     let mut result = BinaryHeap::new();
@@ -139,7 +140,7 @@ pub fn basic<S: G>(
 
 pub fn vbase<'a, S: G>(
     mmap: &'a FlatMmap<S>,
-    vector: &'a [S::Scalar],
+    vector: S::VectorRef<'a>,
     mut filter: impl Filter + 'a,
 ) -> (Vec<Element>, Box<dyn Iterator<Item = Element> + 'a>) {
     let mut result = Vec::new();