Fix tests and benches

Author: sicheng

rust/worker/benches/filter.rs

@@ -4,13 +4,14 @@ use chroma_benchmark::benchmark::{bench_run, tokio_multi_thread};
 use chroma_log::test::{upsert_generator, LoadFromGenerator};
 use chroma_segment::test::TestDistributedSegment;
 use chroma_system::Operator;
+use chroma_types::operator::Filter;
 use chroma_types::{
     BooleanOperator, Chunk, CompositeExpression, MetadataComparison, MetadataExpression,
     MetadataValue, PrimitiveOperator, Where,
 };
 use criterion::Criterion;
 use criterion::{criterion_group, criterion_main};
-use worker::execution::operators::filter::{FilterInput, FilterOperator};
+use worker::execution::operators::filter::FilterInput;
 
 fn baseline_where_clauses() -> Vec<(&'static str, Option<Where>)> {
     use BooleanOperator::*;
@@ -89,12 +90,12 @@ fn bench_filter(criterion: &mut Criterion) {
         };
 
         for (op, where_clause) in baseline_where_clauses() {
-            let filter_operator = FilterOperator {
+            let filter_operator = Filter {
                 query_ids: None,
                 where_clause: where_clause.clone(),
             };
 
-            let routine = |(op, input): (FilterOperator, FilterInput)| async move {
+            let routine = |(op, input): (Filter, FilterInput)| async move {
                 op.run(&input)
                     .await
                     .expect("FilterOperator should not fail");

rust/worker/benches/get.rs

@@ -111,7 +111,7 @@ async fn bench_routine(input: (System, GetOrchestrator, Vec<String>)) {
         .expect("Orchestrator should not fail");
     assert_eq!(
         output
-            .0
+            .result
             .records
             .into_iter()
             .map(|record| record.id)

rust/worker/benches/limit.rs

@@ -2,10 +2,11 @@ use chroma_benchmark::benchmark::{bench_run, tokio_multi_thread};
 use chroma_log::test::{upsert_generator, LoadFromGenerator};
 use chroma_segment::test::TestDistributedSegment;
 use chroma_system::Operator;
+use chroma_types::operator::Limit;
 use chroma_types::{Chunk, SignedRoaringBitmap};
 use criterion::Criterion;
 use criterion::{criterion_group, criterion_main};
-use worker::execution::operators::limit::{LimitInput, LimitOperator};
+use worker::execution::operators::limit::LimitInput;
 
 const FETCH: usize = 100;
 
@@ -30,12 +31,12 @@ fn bench_limit(criterion: &mut Criterion) {
         };
 
         for offset in [0, record_count / 2, record_count - FETCH] {
-            let limit_operator = LimitOperator {
+            let limit_operator = Limit {
                 skip: offset as u32,
                 fetch: Some(FETCH as u32),
             };
 
-            let routine = |(op, input): (LimitOperator, LimitInput)| async move {
+            let routine = |(op, input): (Limit, LimitInput)| async move {
                 op.run(&input).await.expect("LimitOperator should not fail");
             };
 

rust/worker/benches/load.rs

@@ -2,14 +2,12 @@ use chroma_benchmark::datasets::sift::Sift1MData;
 use chroma_log::{in_memory_log::InMemoryLog, test::modulo_metadata, Log};
 use chroma_segment::test::TestDistributedSegment;
 use chroma_types::{
+    operator::{Filter, Limit, Projection},
     Chunk, CollectionUuid, LogRecord, MetadataComparison, MetadataExpression, MetadataSetValue,
     Operation, OperationRecord, SetOperator, Where,
 };
 use indicatif::ProgressIterator;
-use worker::execution::operators::{
-    fetch_log::FetchLogOperator, filter::FilterOperator, limit::LimitOperator,
-    projection::ProjectionOperator,
-};
+use worker::execution::operators::fetch_log::FetchLogOperator;
 
 const DATA_CHUNK_SIZE: usize = 10000;
 
@@ -62,15 +60,15 @@ pub fn empty_fetch_log(collection_uuid: CollectionUuid) -> FetchLogOperator {
     }
 }
 
-pub fn trivial_filter() -> FilterOperator {
-    FilterOperator {
+pub fn trivial_filter() -> Filter {
+    Filter {
         query_ids: None,
         where_clause: None,
     }
 }
 
-pub fn always_false_filter_for_modulo_metadata() -> FilterOperator {
-    FilterOperator {
+pub fn always_false_filter_for_modulo_metadata() -> Filter {
+    Filter {
         query_ids: None,
         where_clause: Some(Where::disjunction(vec![
             Where::Metadata(MetadataExpression {
@@ -91,8 +89,8 @@ pub fn always_false_filter_for_modulo_metadata() -> FilterOperator {
     }
 }
 
-pub fn always_true_filter_for_modulo_metadata() -> FilterOperator {
-    FilterOperator {
+pub fn always_true_filter_for_modulo_metadata() -> Filter {
+    Filter {
         query_ids: None,
         where_clause: Some(Where::conjunction(vec![
             Where::Metadata(MetadataExpression {
@@ -113,30 +111,30 @@ pub fn always_true_filter_for_modulo_metadata() -> FilterOperator {
     }
 }
 
-pub fn trivial_limit() -> LimitOperator {
-    LimitOperator {
+pub fn trivial_limit() -> Limit {
+    Limit {
         skip: 0,
         fetch: Some(100),
     }
 }
 
-pub fn offset_limit() -> LimitOperator {
-    LimitOperator {
+pub fn offset_limit() -> Limit {
+    Limit {
         skip: 100,
         fetch: Some(100),
     }
 }
 
-pub fn trivial_projection() -> ProjectionOperator {
-    ProjectionOperator {
+pub fn trivial_projection() -> Projection {
+    Projection {
         document: false,
         embedding: false,
         metadata: false,
     }
 }
 
-pub fn all_projection() -> ProjectionOperator {
-    ProjectionOperator {
+pub fn all_projection() -> Projection {
+    Projection {
         document: true,
         embedding: true,
         metadata: true,

rust/worker/benches/query.rs

@@ -8,6 +8,7 @@ use chroma_benchmark::{
 use chroma_config::{registry::Registry, Configurable};
 use chroma_segment::test::TestDistributedSegment;
 use chroma_system::{ComponentHandle, Dispatcher, Orchestrator, System};
+use chroma_types::operator::{Knn, KnnProjection};
 use criterion::{criterion_group, criterion_main, Criterion};
 use futures::{stream, StreamExt, TryStreamExt};
 use load::{
@@ -17,12 +18,9 @@ use load::{
 use rand::{seq::SliceRandom, thread_rng};
 use worker::{
     config::RootConfig,
-    execution::{
-        operators::{knn::KnnOperator, knn_projection::KnnProjectionOperator},
-        orchestration::{
-            knn::KnnOrchestrator,
-            knn_filter::{KnnFilterOrchestrator, KnnFilterOutput},
-        },
+    execution::orchestration::{
+        knn::KnnOrchestrator,
+        knn_filter::{KnnFilterOrchestrator, KnnFilterOutput},
     },
 };
 
@@ -91,11 +89,11 @@ fn knn(
         dispatcher_handle.clone(),
         1000,
         knn_filter_output.clone(),
-        KnnOperator {
+        Knn {
             embedding: query,
             fetch: Sift1MData::k() as u32,
         },
-        KnnProjectionOperator {
+        KnnProjection {
             projection: all_projection(),
             distance: true,
         },

rust/worker/benches/regex.rs

@@ -7,6 +7,7 @@ use chroma_benchmark::datasets::wikipedia::WikipediaDataset;
 use chroma_log::test::{int_as_id, random_embedding};
 use chroma_segment::test::TestDistributedSegment;
 use chroma_system::Operator;
+use chroma_types::operator::Filter;
 use chroma_types::{
     Chunk, DocumentExpression, DocumentOperator, LogRecord, Operation, OperationRecord,
     ScalarEncoding, SignedRoaringBitmap, Where,
@@ -18,7 +19,7 @@ use indicatif::ProgressIterator;
 use regex::Regex;
 use roaring::RoaringBitmap;
 use tokio::time::Instant;
-use worker::execution::operators::filter::{FilterInput, FilterOperator};
+use worker::execution::operators::filter::FilterInput;
 
 const LOG_CHUNK_SIZE: usize = 10000;
 const DOCUMENT_SIZE: usize = 100000;
@@ -110,7 +111,7 @@ fn bench_regex(criterion: &mut Criterion) {
     };
 
     for pattern in REGEX_PATTERNS {
-        let filter_operator = FilterOperator {
+        let filter_operator = Filter {
             query_ids: None,
             where_clause: Some(Where::Document(DocumentExpression {
                 operator: DocumentOperator::Regex,
@@ -119,7 +120,7 @@ fn bench_regex(criterion: &mut Criterion) {
         };
 
         let routine = |(op, input, expected): (
-            FilterOperator,
+            Filter,
             FilterInput,
             HashMap<String, RoaringBitmap>,
         )| async move {

rust/worker/benches/spann.rs

@@ -20,14 +20,14 @@ use chroma_index::{
 };
 use chroma_storage::{local::LocalStorage, Storage};
 use chroma_system::Operator;
-use chroma_types::{CollectionUuid, InternalSpannConfiguration};
+use chroma_types::{operator::KnnMerge, CollectionUuid, InternalSpannConfiguration};
 use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion};
 use futures::StreamExt;
 use rand::seq::SliceRandom;
 use roaring::RoaringBitmap;
 use worker::execution::operators::{
+    knn_merge::KnnMergeInput,
     spann_bf_pl::{SpannBfPlInput, SpannBfPlOperator},
-    spann_knn_merge::{SpannKnnMergeInput, SpannKnnMergeOperator},
 };
 
 fn get_records(runtime: &tokio::runtime::Runtime) -> Vec<(u32, Vec<f32>)> {
@@ -199,10 +199,10 @@ fn calculate_recall<'a>(
                 merge_list.push(bf_output.records);
             }
             // Now merge.
-            let knn_input = SpannKnnMergeInput {
-                records: merge_list,
+            let knn_input = KnnMergeInput {
+                batch_distances: merge_list,
             };
-            let knn_operator = SpannKnnMergeOperator { k: k as u32 };
+            let knn_operator = KnnMerge { fetch: k as u32 };
             let knn_output = knn_operator
                 .run(&knn_input)
                 .await
@@ -233,7 +233,7 @@ fn calculate_recall<'a>(
                 .expect("Error running operator");
             let mut recall = 0;
             for bf_record in bf_output.records.iter() {
-                for spann_record in knn_output.merged_records.iter() {
+                for spann_record in knn_output.distances.iter() {
                     if bf_record.offset_id == spann_record.offset_id {
                         recall += 1;
                     }