In progress

Author: Sicheng-Pan

Cargo.lock

@@ -22,6 +22,7 @@ validator = { workspace = true }
 regex = { workspace = true }
 regex-syntax = { workspace = true }
 utoipa = { workspace = true }
+sprs = { workspace = true }
 
 # (Cross-crate testing dependencies)
 proptest = { workspace = true, optional = true }

rust/types/Cargo.toml

@@ -1,6 +1,7 @@
 use chroma_error::{ChromaError, ErrorCodes};
 use serde::{Deserialize, Serialize};
 use serde_json::{Number, Value};
+use sprs::CsVec;
 use std::{
     cmp::Ordering,
     collections::{HashMap, HashSet},
@@ -84,6 +85,23 @@ impl From<SparseVector> for chroma_proto::SparseVector {
     }
 }
 
+/// Convert SparseVector to sprs::CsVec for efficient sparse operations
+impl From<&SparseVector> for CsVec<f32> {
+    fn from(sparse: &SparseVector) -> Self {
+        let (indices, values) = sparse
+            .iter()
+            .map(|(index, value)| (index as usize, value))
+            .unzip();
+        CsVec::new(u32::MAX as usize, indices, values)
+    }
+}
+
+impl From<SparseVector> for CsVec<f32> {
+    fn from(sparse: SparseVector) -> Self {
+        (&sparse).into()
+    }
+}
+
 #[cfg(feature = "pyo3")]
 impl<'py> pyo3::IntoPyObject<'py> for SparseVector {
     type Target = pyo3::PyAny;

rust/types/src/metadata.rs

@@ -184,12 +184,12 @@ impl Operator<ScoreInput, ScoreOutput> for Score {
             ranks: &input.ranks,
         };
         let score_domain = score_provider.eval(self.clone());
-        Ok(ScoreOutput {
-            scores: score_domain
-                .support
-                .into_iter()
-                .map(|(offset_id, measure)| RecordDistance { offset_id, measure })
-                .collect(),
-        })
+        let mut scores = score_domain
+            .support
+            .into_iter()
+            .map(|(offset_id, measure)| RecordDistance { offset_id, measure })
+            .collect::<Vec<_>>();
+        scores.sort_unstable();
+        Ok(ScoreOutput { scores })
     }
 }

rust/worker/src/execution/operators/score.rs

@@ -1,7 +1,4 @@
-use std::{
-    cmp::Reverse,
-    collections::{BinaryHeap, HashMap},
-};
+use std::{cmp::Reverse, collections::BinaryHeap};
 
 use async_trait::async_trait;
 use chroma_blockstore::provider::BlockfileProvider;
@@ -16,7 +13,7 @@ use chroma_segment::{
 use chroma_system::Operator;
 use chroma_types::{
     operator::{Rank, RecordDistance},
-    MaterializedLogOperation, MetadataValue, Segment,
+    MaterializedLogOperation, MetadataValue, Segment, SparseVector,
 };
 use sprs::CsVec;
 use thiserror::Error;
@@ -64,7 +61,7 @@ impl ChromaError for SparseKnnError {
 
 #[derive(Clone, Debug)]
 pub struct SparseKnn {
-    pub embedding: HashMap<u32, f32>,
+    pub embedding: SparseVector,
     pub key: String,
     pub limit: u32,
 }
@@ -74,17 +71,8 @@ impl Operator<SparseKnnInput, SparseKnnOutput> for SparseKnn {
     type Error = SparseKnnError;
 
     async fn run(&self, input: &SparseKnnInput) -> Result<SparseKnnOutput, SparseKnnError> {
-        let mut query_raw_vector = self
-            .embedding
-            .iter()
-            .map(|(index, val)| (*index as usize, *val))
-            .collect::<Vec<_>>();
-        query_raw_vector.sort_unstable_by_key(|(index, _)| *index);
-        let (query_raw_indexes, query_raw_values) = query_raw_vector
-            .into_iter()
-            .unzip::<usize, f32, Vec<_>, Vec<_>>();
-        let query_sparse_verctor =
-            CsVec::new(u32::MAX as usize, query_raw_indexes, query_raw_values);
+        // Convert SparseVector to sprs::CsVec
+        let query_sparse_verctor: CsVec<f32> = (&self.embedding).into();
         let record_segment_reader = match RecordSegmentReader::from_segment(
             &input.record_segment,
             &input.blockfile_provider,
@@ -118,16 +106,8 @@ impl Operator<SparseKnnInput, SparseKnnOutput> for SparseKnn {
                 else {
                     continue;
                 };
-                let mut log_raw_vector = sparse_vector
-                    .iter()
-                    .map(|(index, val)| (*index as usize, *val))
-                    .collect::<Vec<_>>();
-                log_raw_vector.sort_unstable_by_key(|(index, _)| *index);
-                let (log_raw_indexes, log_raw_values) = log_raw_vector
-                    .into_iter()
-                    .unzip::<usize, f32, Vec<_>, Vec<_>>();
-                let log_sparse_verctor =
-                    CsVec::new(u32::MAX as usize, log_raw_indexes, log_raw_values);
+                // Convert SparseVector to sprs::CsVec
+                let log_sparse_verctor: CsVec<f32> = sparse_vector.into();
                 let score = query_sparse_verctor.dot(&log_sparse_verctor);
                 if (min_heap.len() as u32) < self.limit {
                     min_heap.push(Reverse(RecordDistance {
@@ -170,7 +150,11 @@ impl Operator<SparseKnnInput, SparseKnnOutput> for SparseKnn {
 
         let sorted_compact_records = sparse_reader
             .wand(
-                self.embedding.clone(),
+                self.embedding
+                    .indices
+                    .iter()
+                    .copied()
+                    .zip(self.embedding.values.iter().copied()),
                 self.limit,
                 input.mask.compact_offset_ids.clone(),
             )

rust/worker/src/execution/operators/sparse_knn.rs

@@ -7,4 +7,4 @@ mod compact;
 pub mod get;
 pub mod knn;
 pub mod knn_filter;
-pub mod retrieve;
+pub mod search;