[ENH] Split knn orchestrators for reuse

idl/chromadb/proto/query_executor.proto

@@ -25,8 +25,8 @@ message KNNOperator {
 }
 
 message LimitOperator {
-    uint32 skip = 1;
-    optional uint32 fetch = 2;
+    uint32 offset = 1;
+    optional uint32 limit = 2;
 }
 
 message ProjectionOperator {

rust/frontend/src/executor/local.rs

@@ -12,7 +12,7 @@ use chroma_system::ComponentHandle;
 use chroma_types::{
     operator::{
         CountResult, Filter, GetResult, KnnBatchResult, KnnProjectionOutput, KnnProjectionRecord,
-        Limit, Projection, ProjectionRecord, RecordDistance, SearchResult,
+        Limit, Projection, ProjectionRecord, RecordMeasure, SearchResult,
     },
     plan::{Count, Get, Knn, Search},
     CollectionAndSegments, CollectionUuid, ExecutorError, HnswSpace, SegmentType,
@@ -162,8 +162,8 @@ impl LocalExecutor {
                     scan: plan.scan.clone(),
                     filter: filter.clone(),
                     limit: Limit {
-                        skip: 0,
-                        fetch: None,
+                        offset: 0,
+                        limit: None,
                     },
                     proj: Default::default(),
                 };
@@ -234,7 +234,7 @@ impl LocalExecutor {
                 .map_err(|err| ExecutorError::Internal(Box::new(err)))?;
 
             let mut records = Vec::new();
-            for RecordDistance { offset_id, measure } in distances {
+            for RecordMeasure { offset_id, measure } in distances {
                 let user_id = hnsw_reader
                     .get_user_id_by_offset_id(offset_id)
                     .await
@@ -268,8 +268,8 @@ impl LocalExecutor {
                     where_clause: None,
                 },
                 limit: Limit {
-                    skip: 0,
-                    fetch: None,
+                    offset: 0,
+                    limit: None,
                 },
                 proj: Projection {
                     document: plan.proj.projection.document,

rust/frontend/src/impls/in_memory_frontend.rs

@@ -572,10 +572,7 @@ impl InMemoryFrontend {
                     },
                 },
                 filter,
-                limit: Limit {
-                    skip: offset,
-                    fetch: limit,
-                },
+                limit: Limit { offset, limit },
                 proj: Projection {
                     document: include.0.contains(&Include::Document),
                     embedding: include.0.contains(&Include::Embedding),

rust/frontend/src/impls/service_based_frontend.rs

@@ -1035,8 +1035,8 @@ impl ServiceBasedFrontend {
                     },
                     filter,
                     limit: Limit {
-                        skip: 0,
-                        fetch: None,
+                        offset: 0,
+                        limit: None,
                     },
                     proj: Projection {
                         document: false,
@@ -1339,10 +1339,7 @@ impl ServiceBasedFrontend {
                     query_ids: ids,
                     where_clause: r#where,
                 },
-                limit: Limit {
-                    skip: offset,
-                    fetch: limit,
-                },
+                limit: Limit { offset, limit },
                 proj: Projection {
                     document: include.0.contains(&Include::Document),
                     embedding: include.0.contains(&Include::Embedding),

rust/garbage_collector/src/helper.rs

@@ -201,8 +201,8 @@ impl ChromaGrpcClients {
                 where_document: None,
             }),
             limit: Some(LimitOperator {
-                skip: 0,
-                fetch: None,
+                offset: 0,
+                limit: None,
             }),
             projection: Some(ProjectionOperator {
                 document: false, // include_documents,

rust/segment/src/local_hnsw.rs

@@ -10,7 +10,7 @@ use chroma_error::{ChromaError, ErrorCodes};
 use chroma_index::{HnswIndex, HnswIndexConfig, Index, IndexConfig, PersistentIndex};
 use chroma_sqlite::{db::SqliteDb, table::MaxSeqId};
 use chroma_types::{
-    operator::RecordDistance, Chunk, Collection, HnswParametersFromSegmentError, LogRecord,
+    operator::RecordMeasure, Chunk, Collection, HnswParametersFromSegmentError, LogRecord,
     Operation, OperationRecord, Segment, SegmentUuid,
 };
 use rayon::iter::{IntoParallelIterator, ParallelIterator};
@@ -277,7 +277,7 @@ impl LocalHnswSegmentReader {
         allowed_offset_ids: &[u32],
         embedding: Vec<f32>,
         k: u32,
-    ) -> Result<Vec<RecordDistance>, LocalHnswSegmentReaderError> {
+    ) -> Result<Vec<RecordMeasure>, LocalHnswSegmentReaderError> {
         let guard = self.index.inner.read().await;
         let len_with_deleted = guard.index.len_with_deleted();
         let actual_len = guard.index.len();
@@ -319,7 +319,7 @@ impl LocalHnswSegmentReader {
                                     .distance_function
                                     .distance(curr_embedding.as_slice(), embedding.as_slice());
                                 if max_heap.len() < k as usize {
-                                    max_heap.push(RecordDistance {
+                                    max_heap.push(RecordMeasure {
                                         offset_id: *curr_id as u32,
                                         measure: curr_distance,
                                     });
@@ -329,7 +329,7 @@ impl LocalHnswSegmentReader {
                                     let top = max_heap.peek().unwrap();
                                     if top.measure > curr_distance {
                                         max_heap.pop();
-                                        max_heap.push(RecordDistance {
+                                        max_heap.push(RecordMeasure {
                                             offset_id: *curr_id as u32,
                                             measure: curr_distance,
                                         });
@@ -358,7 +358,7 @@ impl LocalHnswSegmentReader {
             Ok(offset_ids
                 .into_iter()
                 .zip(distances)
-                .map(|(offset_id, measure)| RecordDistance {
+                .map(|(offset_id, measure)| RecordMeasure {
                     offset_id: offset_id as u32,
                     measure,
                 })

rust/segment/src/sqlite_metadata.rs

@@ -762,7 +762,7 @@ impl SqliteMetadataReader {
                 query_ids,
                 where_clause,
             },
-            limit: Limit { skip, fetch },
+            limit: Limit { offset, limit },
             proj: Projection {
                 document, metadata, ..
             },
@@ -789,8 +789,8 @@ impl SqliteMetadataReader {
 
         filter_limit_query
             .order_by((Embeddings::Table, Embeddings::Id), sea_query::Order::Asc)
-            .offset(skip as u64)
-            .limit(fetch.unwrap_or(u32::MAX) as u64);
+            .offset(offset as u64)
+            .limit(limit.unwrap_or(u32::MAX) as u64);
 
         let alias = Alias::new(SUBQ_ALIAS);
         let mut projection_query = Query::select();
@@ -954,8 +954,8 @@ mod tests {
                     where_clause: Some(where_clause.clause),
                 },
                 limit: Limit {
-                    skip: 3,
-                    fetch: Some(6),
+                    offset: 3,
+                    limit: Some(6),
                 },
                 proj: Projection {
                     document: true,
@@ -1046,8 +1046,8 @@ mod tests {
                 where_clause: Some(where_clause),
             },
             limit: Limit {
-                skip: 0,
-                fetch: None,
+                offset: 0,
+                limit: None,
             },
             proj: Projection {
                 document: false,
@@ -1083,8 +1083,8 @@ mod tests {
                 where_clause: Some(where_clause2),
             },
             limit: Limit {
-                skip: 0,
-                fetch: None,
+                offset: 0,
+                limit: None,
             },
             proj: Projection {
                 document: false,
@@ -1188,8 +1188,8 @@ mod tests {
                 where_clause: Some(fts_where_clause),
             },
             limit: Limit {
-                skip: 0,
-                fetch: None,
+                offset: 0,
+                limit: None,
             },
             proj: Projection {
                 document: true,
@@ -1217,8 +1217,8 @@ mod tests {
                 where_clause: Some(metadata_where_clause),
             },
             limit: Limit {
-                skip: 0,
-                fetch: None,
+                offset: 0,
+                limit: None,
             },
             proj: Projection {
                 document: false,
@@ -1246,8 +1246,8 @@ mod tests {
                 where_clause: Some(hybrid_where_clause),
             },
             limit: Limit {
-                skip: 0,
-                fetch: None,
+                offset: 0,
+                limit: None,
             },
             proj: Projection {
                 document: true,

rust/segment/src/test.rs

@@ -352,8 +352,8 @@ impl TestReferenceSegment {
             result: ProjectionOutput {
                 records: records
                     .into_iter()
-                    .skip(plan.limit.skip as usize)
-                    .take(plan.limit.fetch.unwrap_or(u32::MAX) as usize)
+                    .skip(plan.limit.offset as usize)
+                    .take(plan.limit.limit.unwrap_or(u32::MAX) as usize)
                     .map(|(_, mut rec)| {
                         let Projection {
                             document,

rust/types/Cargo.toml

@@ -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 }