Fix test and benches

Author: sicheng

rust/worker/benches/get.rs

@@ -5,110 +5,52 @@ use chroma_benchmark::benchmark::{bench_run, tokio_multi_thread};
 use chroma_config::{registry::Registry, Configurable};
 use chroma_segment::test::TestDistributedSegment;
 use chroma_system::{ComponentHandle, Dispatcher, Orchestrator, System};
+use chroma_types::operator::{Filter, Limit, Projection};
 use criterion::{criterion_group, criterion_main, Criterion};
 use load::{
     all_projection, always_false_filter_for_modulo_metadata,
     always_true_filter_for_modulo_metadata, empty_fetch_log, offset_limit, sift1m_segments,
     trivial_filter, trivial_limit, trivial_projection,
 };
-use worker::{config::RootConfig, execution::orchestration::get::GetOrchestrator};
-
-fn trivial_get(
-    test_segments: &TestDistributedSegment,
-    dispatcher_handle: ComponentHandle<Dispatcher>,
-) -> GetOrchestrator {
-    let blockfile_provider = test_segments.blockfile_provider.clone();
-    let collection_uuid = test_segments.collection.collection_id;
-    GetOrchestrator::new(
-        blockfile_provider,
-        dispatcher_handle,
-        1000,
-        test_segments.into(),
-        empty_fetch_log(collection_uuid),
-        trivial_filter(),
-        trivial_limit(),
-        trivial_projection(),
-    )
-}
-
-fn get_false_filter(
-    test_segments: &TestDistributedSegment,
-    dispatcher_handle: ComponentHandle<Dispatcher>,
-) -> GetOrchestrator {
-    let blockfile_provider = test_segments.blockfile_provider.clone();
-    let collection_uuid = test_segments.collection.collection_id;
-    GetOrchestrator::new(
-        blockfile_provider,
-        dispatcher_handle,
-        1000,
-        test_segments.into(),
-        empty_fetch_log(collection_uuid),
-        always_false_filter_for_modulo_metadata(),
-        trivial_limit(),
-        trivial_projection(),
-    )
-}
-
-fn get_true_filter(
-    test_segments: &TestDistributedSegment,
-    dispatcher_handle: ComponentHandle<Dispatcher>,
-) -> GetOrchestrator {
-    let blockfile_provider = test_segments.blockfile_provider.clone();
-    let collection_uuid = test_segments.collection.collection_id;
-    GetOrchestrator::new(
-        blockfile_provider,
-        dispatcher_handle,
-        1000,
-        test_segments.into(),
-        empty_fetch_log(collection_uuid),
-        always_true_filter_for_modulo_metadata(),
-        trivial_limit(),
-        trivial_projection(),
-    )
-}
+use worker::{
+    config::RootConfig,
+    execution::orchestration::{filter::FilterOrchestrator, get::GetOrchestrator},
+};
 
-fn get_true_filter_limit(
-    test_segments: &TestDistributedSegment,
-    dispatcher_handle: ComponentHandle<Dispatcher>,
-) -> GetOrchestrator {
-    let blockfile_provider = test_segments.blockfile_provider.clone();
-    let collection_uuid = test_segments.collection.collection_id;
-    GetOrchestrator::new(
-        blockfile_provider,
-        dispatcher_handle,
+async fn bench_routine(
+    (system, dispatcher, test_segments, filter, limit, projection, expected_ids): (
+        System,
+        ComponentHandle<Dispatcher>,
+        &TestDistributedSegment,
+        Filter,
+        Limit,
+        Projection,
+        Vec<String>,
+    ),
+) {
+    let matching_records = FilterOrchestrator::new(
+        test_segments.blockfile_provider.clone(),
+        dispatcher.clone(),
+        test_segments.hnsw_provider.clone(),
         1000,
         test_segments.into(),
-        empty_fetch_log(collection_uuid),
-        always_true_filter_for_modulo_metadata(),
-        offset_limit(),
-        trivial_projection(),
+        empty_fetch_log(test_segments.collection.collection_id),
+        filter,
     )
-}
-
-fn get_true_filter_limit_projection(
-    test_segments: &TestDistributedSegment,
-    dispatcher_handle: ComponentHandle<Dispatcher>,
-) -> GetOrchestrator {
-    let blockfile_provider = test_segments.blockfile_provider.clone();
-    let collection_uuid = test_segments.collection.collection_id;
-    GetOrchestrator::new(
-        blockfile_provider,
-        dispatcher_handle,
+    .run(system.clone())
+    .await
+    .expect("Filter orchestrator should not fail");
+    let output = GetOrchestrator::new(
+        test_segments.blockfile_provider.clone(),
+        dispatcher,
         1000,
-        test_segments.into(),
-        empty_fetch_log(collection_uuid),
-        always_true_filter_for_modulo_metadata(),
-        offset_limit(),
-        all_projection(),
+        matching_records,
+        limit,
+        projection,
     )
-}
-
-async fn bench_routine(input: (System, GetOrchestrator, Vec<String>)) {
-    let (system, orchestrator, expected_ids) = input;
-    let output = orchestrator
-        .run(system)
-        .await
-        .expect("Orchestrator should not fail");
+    .run(system)
+    .await
+    .expect("Get orchestrator should not fail");
     assert_eq!(
         output
             .result
@@ -138,35 +80,55 @@ fn bench_get(criterion: &mut Criterion) {
     let trivial_get_setup = || {
         (
             system.clone(),
-            trivial_get(&test_segments, dispatcher_handle.clone()),
+            dispatcher_handle.clone(),
+            &test_segments,
+            trivial_filter(),
+            trivial_limit(),
+            trivial_projection(),
             (0..100).map(|id| id.to_string()).collect(),
         )
     };
     let get_false_filter_setup = || {
         (
             system.clone(),
-            get_false_filter(&test_segments, dispatcher_handle.clone()),
+            dispatcher_handle.clone(),
+            &test_segments,
+            always_false_filter_for_modulo_metadata(),
+            trivial_limit(),
+            trivial_projection(),
             Vec::new(),
         )
     };
     let get_true_filter_setup = || {
         (
             system.clone(),
-            get_true_filter(&test_segments, dispatcher_handle.clone()),
+            dispatcher_handle.clone(),
+            &test_segments,
+            always_true_filter_for_modulo_metadata(),
+            trivial_limit(),
+            trivial_projection(),
             (0..100).map(|id| id.to_string()).collect(),
         )
     };
     let get_true_filter_limit_setup = || {
         (
             system.clone(),
-            get_true_filter_limit(&test_segments, dispatcher_handle.clone()),
+            dispatcher_handle.clone(),
+            &test_segments,
+            always_true_filter_for_modulo_metadata(),
+            offset_limit(),
+            trivial_projection(),
             (100..200).map(|id| id.to_string()).collect(),
         )
     };
     let get_true_filter_limit_projection_setup = || {
         (
             system.clone(),
-            get_true_filter_limit_projection(&test_segments, dispatcher_handle.clone()),
+            dispatcher_handle.clone(),
+            &test_segments,
+            always_true_filter_for_modulo_metadata(),
+            offset_limit(),
+            all_projection(),
             (100..200).map(|id| id.to_string()).collect(),
         )
     };

rust/worker/src/execution/orchestration/compact.rs

@@ -1058,7 +1058,10 @@ mod tests {
 
     use crate::{
         config::RootConfig,
-        execution::{operators::fetch_log::FetchLogOperator, orchestration::get::GetOrchestrator},
+        execution::{
+            operators::fetch_log::FetchLogOperator,
+            orchestration::{filter::FilterOrchestrator, get::GetOrchestrator},
+        },
     };
 
     use super::CompactOrchestrator;
@@ -1156,6 +1159,19 @@ mod tests {
                 }),
             ])),
         };
+        let filter_orchestrator = FilterOrchestrator::new(
+            test_segments.blockfile_provider.clone(),
+            dispatcher_handle.clone(),
+            test_segments.hnsw_provider.clone(),
+            1000,
+            old_cas.clone(),
+            fetch_log.clone(),
+            filter.clone(),
+        );
+        let matching_records = filter_orchestrator
+            .run(system.clone())
+            .await
+            .expect("Filter orchestrator should not fail");
         let limit = Limit {
             skip: 0,
             fetch: None,
@@ -1169,9 +1185,7 @@ mod tests {
             test_segments.blockfile_provider.clone(),
             dispatcher_handle.clone(),
             1000,
-            old_cas.clone(),
-            fetch_log.clone(),
-            filter.clone(),
+            matching_records,
             limit.clone(),
             project.clone(),
         );
@@ -1231,13 +1245,25 @@ mod tests {
             old_cas.vector_segment.file_path
         );
 
-        let get_orchestrator = GetOrchestrator::new(
+        let filter_orchestrator = FilterOrchestrator::new(
             test_segments.blockfile_provider.clone(),
-            dispatcher_handle,
+            dispatcher_handle.clone(),
+            test_segments.hnsw_provider.clone(),
             1000,
             new_cas,
             fetch_log,
             filter,
+        );
+        let matching_records = filter_orchestrator
+            .run(system.clone())
+            .await
+            .expect("Filter orchestrator should not fail");
+
+        let get_orchestrator = GetOrchestrator::new(
+            test_segments.blockfile_provider.clone(),
+            dispatcher_handle,
+            1000,
+            matching_records,
             limit,
             project,
         );