[ENH] Implement retrieve orchestrartor

Author: Sicheng-Pan

Cargo.lock

@@ -69,6 +69,7 @@ humantime = { version = "2.2.0" }
 petgraph = { version = "0.8.1" }
 base64 = "0.22"
 tikv-jemallocator = { version = "0.6.0", features = ["profiling"] }
+sprs = "0.11.3"
 
 chroma-benchmark = { path = "rust/benchmark" }
 chroma-blockstore = { path = "rust/blockstore" }

Cargo.toml

@@ -44,6 +44,7 @@ num_cpus = { workspace = true }
 flatbuffers = { workspace = true }
 tantivy = { workspace = true }
 clap = { workspace = true }
+sprs = { workspace = true }
 
 chroma-blockstore = { workspace = true }
 chroma-cache = { workspace = true }

rust/worker/Cargo.toml

@@ -6,7 +6,10 @@ use chroma_benchmark::{
     datasets::sift::Sift1MData,
 };
 use chroma_config::{registry::Registry, Configurable};
-use chroma_segment::test::TestDistributedSegment;
+use chroma_segment::{
+    distributed_hnsw::{DistributedHNSWSegmentFromSegmentError, DistributedHNSWSegmentReader},
+    test::TestDistributedSegment,
+};
 use chroma_system::{ComponentHandle, Dispatcher, Orchestrator, System};
 use chroma_types::operator::{Knn, KnnProjection};
 use criterion::{criterion_group, criterion_main, Criterion};
@@ -29,12 +32,10 @@ fn trivial_knn_filter(
     dispatcher_handle: ComponentHandle<Dispatcher>,
 ) -> KnnFilterOrchestrator {
     let blockfile_provider = test_segments.blockfile_provider.clone();
-    let hnsw_provider = test_segments.hnsw_provider.clone();
     let collection_uuid = test_segments.collection.collection_id;
     KnnFilterOrchestrator::new(
         blockfile_provider,
         dispatcher_handle,
-        hnsw_provider,
         1000,
         test_segments.into(),
         empty_fetch_log(collection_uuid),
@@ -47,12 +48,10 @@ fn always_true_knn_filter(
     dispatcher_handle: ComponentHandle<Dispatcher>,
 ) -> KnnFilterOrchestrator {
     let blockfile_provider = test_segments.blockfile_provider.clone();
-    let hnsw_provider = test_segments.hnsw_provider.clone();
     let collection_uuid = test_segments.collection.collection_id;
     KnnFilterOrchestrator::new(
         blockfile_provider,
         dispatcher_handle,
-        hnsw_provider,
         1000,
         test_segments.into(),
         empty_fetch_log(collection_uuid),
@@ -65,12 +64,10 @@ fn always_false_knn_filter(
     dispatcher_handle: ComponentHandle<Dispatcher>,
 ) -> KnnFilterOrchestrator {
     let blockfile_provider = test_segments.blockfile_provider.clone();
-    let hnsw_provider = test_segments.hnsw_provider.clone();
     let collection_uuid = test_segments.collection.collection_id;
     KnnFilterOrchestrator::new(
         blockfile_provider,
         dispatcher_handle,
-        hnsw_provider,
         1000,
         test_segments.into(),
         empty_fetch_log(collection_uuid),
@@ -81,13 +78,15 @@ fn always_false_knn_filter(
 fn knn(
     test_segments: &TestDistributedSegment,
     dispatcher_handle: ComponentHandle<Dispatcher>,
+    hnsw_reader: Option<DistributedHNSWSegmentReader>,
     knn_filter_output: KnnFilterOutput,
     query: Vec<f32>,
 ) -> KnnOrchestrator {
     KnnOrchestrator::new(
         test_segments.blockfile_provider.clone(),
         dispatcher_handle.clone(),
         1000,
+        hnsw_reader,
         knn_filter_output.clone(),
         Knn {
             embedding: query,
@@ -126,6 +125,21 @@ fn bench_query(criterion: &mut Criterion) {
     let runtime = tokio_multi_thread();
     let test_segments = runtime.block_on(sift1m_segments());
 
+    let hnsw_reader = match runtime.block_on(DistributedHNSWSegmentReader::from_segment(
+        &test_segments.collection.clone(),
+        &test_segments.vector_segment,
+        test_segments
+            .collection
+            .dimension
+            .expect("Collection dimension should be non-zero") as usize,
+        test_segments.hnsw_provider.clone(),
+    )) {
+        Ok(hnsw_reader) => Some(*hnsw_reader),
+        Err(err) if matches!(*err, DistributedHNSWSegmentFromSegmentError::Uninitialized) => None,
+
+        Err(err) => panic!("{err}"),
+    };
+
     let config = RootConfig::default();
     let system = System::default();
     let registry = Registry::new();
@@ -159,6 +173,7 @@ fn bench_query(criterion: &mut Criterion) {
                             knn(
                                 &test_segments,
                                 dispatcher_handle.clone(),
+                                hnsw_reader.clone(),
                                 knn_filter_output.clone(),
                                 query.clone(),
                             ),
@@ -183,6 +198,7 @@ fn bench_query(criterion: &mut Criterion) {
                             knn(
                                 &test_segments,
                                 dispatcher_handle.clone(),
+                                hnsw_reader.clone(),
                                 knn_filter_output.clone(),
                                 query.clone(),
                             ),
@@ -207,6 +223,7 @@ fn bench_query(criterion: &mut Criterion) {
                             knn(
                                 &test_segments,
                                 dispatcher_handle.clone(),
+                                hnsw_reader.clone(),
                                 knn_filter_output.clone(),
                                 query.clone(),
                             ),

rust/worker/benches/query.rs

@@ -21,4 +21,7 @@ pub mod prefetch_segment;
 pub mod projection;
 pub mod purge_dirty_log;
 pub mod repair_log_offsets;
+pub mod reverse_project;
+pub mod score;
 pub mod source_record_segment;
+pub mod sparse_knn;

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

@@ -0,0 +1,106 @@
+use std::collections::HashMap;
+
+use async_trait::async_trait;
+use chroma_blockstore::provider::BlockfileProvider;
+use chroma_segment::{
+    blockfile_record::{RecordSegmentReader, RecordSegmentReaderCreationError},
+    types::{materialize_logs, LogMaterializerError},
+};
+use chroma_system::Operator;
+use chroma_types::{
+    operator::{KnnProjectionOutput, Rank, RecordDistance},
+    Segment,
+};
+use futures::future::try_join_all;
+use tracing::{Instrument, Span};
+
+use crate::execution::operators::{fetch_log::FetchLogOutput, projection::ProjectionError};
+
+#[derive(Clone, Debug)]
+pub struct ReverseProjectionInput {
+    pub logs: FetchLogOutput,
+    pub blockfile_provider: BlockfileProvider,
+    pub record_segment: Segment,
+    pub projection_outputs: HashMap<Rank, KnnProjectionOutput>,
+}
+
+#[derive(Clone, Debug)]
+pub struct ReverseProjectionOutput {
+    pub rank_records: HashMap<Rank, Vec<RecordDistance>>,
+}
+
+// NOTE: This is a temporary operator that aims to reverse
+// the projection by converting user id to offset id
+#[derive(Clone, Debug)]
+pub struct ReverseProjection {}
+
+#[async_trait]
+impl Operator<ReverseProjectionInput, ReverseProjectionOutput> for ReverseProjection {
+    type Error = ProjectionError;
+
+    async fn run(
+        &self,
+        input: &ReverseProjectionInput,
+    ) -> Result<ReverseProjectionOutput, ProjectionError> {
+        tracing::trace!(
+            "Reversing projection on {} ranks",
+            input.projection_outputs.len()
+        );
+        let record_segment_reader = match RecordSegmentReader::from_segment(
+            &input.record_segment,
+            &input.blockfile_provider,
+        )
+        .await
+        {
+            Ok(reader) => Ok(Some(reader)),
+            Err(e) if matches!(*e, RecordSegmentReaderCreationError::UninitializedSegment) => {
+                Ok(None)
+            }
+            Err(e) => Err(*e),
+        }?;
+        let materialized_logs = materialize_logs(&record_segment_reader, input.logs.clone(), None)
+            .instrument(tracing::trace_span!(parent: Span::current(), "Materialize logs"))
+            .await?;
+        let borrowed_logs = materialized_logs.iter().collect::<Vec<_>>();
+
+        let hydrated_futures = borrowed_logs.iter().map(|log| async {
+            let hydrated_log = log.hydrate(record_segment_reader.as_ref()).await?;
+            <Result<_, LogMaterializerError>>::Ok((
+                hydrated_log.get_user_id(),
+                hydrated_log.get_offset_id(),
+            ))
+        });
+        let log_user_id_to_offset_id = try_join_all(hydrated_futures)
+            .await?
+            .into_iter()
+            .collect::<HashMap<_, _>>();
+        let mut rank_records = HashMap::with_capacity(input.projection_outputs.len());
+
+        for (rank, projection_output) in &input.projection_outputs {
+            let resolve_futures = projection_output.records.iter().map(|record| async {
+                match log_user_id_to_offset_id.get(record.record.id.as_str()) {
+                    Some(&offset_id) => Ok(RecordDistance {
+                        offset_id,
+                        measure: record.distance.unwrap_or_default(),
+                    }),
+                    None => {
+                        if let Some(reader) = &record_segment_reader {
+                            match reader.get_offset_id_for_user_id(&record.record.id).await? {
+                                Some(offset_id) => Ok(RecordDistance {
+                                    offset_id,
+                                    measure: record.distance.unwrap_or_default(),
+                                }),
+                                None => Err(ProjectionError::RecordSegmentPhantomRecord(u32::MAX)),
+                            }
+                        } else {
+                            Err(ProjectionError::RecordSegmentUninitialized)
+                        }
+                    }
+                }
+            });
+            rank_records.insert(rank.clone(), try_join_all(resolve_futures).await?);
+        }
+
+        Ok(ReverseProjectionOutput { rank_records })
+    }
+}