[ENH] Deprecate FetchSegmentOperator

chromadb/proto/convert.py

@@ -573,9 +573,6 @@ def to_proto_where_document(where_document: WhereDocument) -> chroma_pb.WhereDoc
 def to_proto_scan(scan: Scan) -> query_pb.ScanOperator:
     return query_pb.ScanOperator(
         collection=to_proto_collection(scan.collection),
-        knn_id=scan.knn["id"].hex,
-        metadata_id=scan.metadata["id"].hex,
-        record_id=scan.record["id"].hex,
         knn=to_proto_segment(scan.knn),
         metadata=to_proto_segment(scan.metadata),
         record=to_proto_segment(scan.record),

idl/chromadb/proto/query_executor.proto

@@ -6,12 +6,8 @@ import "chromadb/proto/chroma.proto";
 
 message ScanOperator {
     Collection collection = 1;
-    // Deprecated
-    string knn_id = 2;
-    // Deprecated
-    string metadata_id = 3;
-    // Deprecated
-    string record_id = 4;
+    // Reserve for deprecated fields
+    reserved 2, 3, 4;
     Segment knn = 5;
     Segment metadata = 6;
     Segment record = 7;

rust/types/src/collection.rs

@@ -1,5 +1,5 @@
 use super::{Metadata, MetadataValueConversionError};
-use crate::chroma_proto;
+use crate::{chroma_proto, ConversionError, Segment};
 use chroma_error::{ChromaError, ErrorCodes};
 use thiserror::Error;
 use uuid::Uuid;
@@ -89,6 +89,43 @@ impl TryFrom<chroma_proto::Collection> for Collection {
     }
 }
 
+#[derive(Clone, Debug)]
+pub struct CollectionAndSegments {
+    pub collection: Collection,
+    pub metadata_segment: Segment,
+    pub record_segment: Segment,
+    pub vector_segment: Segment,
+}
+
+impl TryFrom<chroma_proto::ScanOperator> for CollectionAndSegments {
+    type Error = ConversionError;
+
+    fn try_from(value: chroma_proto::ScanOperator) -> Result<Self, ConversionError> {
+        Ok(Self {
+            collection: value
+                .collection
+                .ok_or(ConversionError::DecodeError)?
+                .try_into()
+                .map_err(|_| ConversionError::DecodeError)?,
+            metadata_segment: value
+                .metadata
+                .ok_or(ConversionError::DecodeError)?
+                .try_into()
+                .map_err(|_| ConversionError::DecodeError)?,
+            record_segment: value
+                .record
+                .ok_or(ConversionError::DecodeError)?
+                .try_into()
+                .map_err(|_| ConversionError::DecodeError)?,
+            vector_segment: value
+                .knn
+                .ok_or(ConversionError::DecodeError)?
+                .try_into()
+                .map_err(|_| ConversionError::DecodeError)?,
+        })
+    }
+}
+
 #[cfg(test)]
 mod test {
     use super::*;

rust/types/src/types.rs

@@ -28,7 +28,7 @@ pub enum ConversionError {
 impl ChromaError for ConversionError {
     fn code(&self) -> ErrorCodes {
         match self {
-            ConversionError::DecodeError => ErrorCodes::Internal,
+            ConversionError::DecodeError => ErrorCodes::InvalidArgument,
         }
     }
 }

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

@@ -10,7 +10,6 @@ pub(super) mod write_segments;
 
 // Required for benchmark
 pub mod fetch_log;
-pub mod fetch_segment;
 pub mod filter;
 pub mod knn;
 pub mod knn_hnsw;

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

@@ -1,8 +1,6 @@
 use std::collections::HashSet;
 
-use chroma_blockstore::provider::BlockfileProvider;
 use chroma_error::{ChromaError, ErrorCodes};
-use chroma_types::{Chunk, LogRecord, Segment};
 use thiserror::Error;
 use tonic::async_trait;
 use tracing::{trace, Instrument, Span};
@@ -16,16 +14,15 @@ use crate::{
     },
 };
 
+use super::projection::ProjectionInput;
+
 /// The `PrefetchRecordOperator` prefetches the relevant records from the record segments to the cache
 ///
 /// # Parameters
 /// None
 ///
 /// # Inputs
-/// - `logs`: The latest logs of the collection
-/// - `blockfile_provider`: The blockfile provider
-/// - `record_segment`: The record segment information
-/// - `offset_ids`: The offset ids of the records to prefetch
+/// Identical to ProjectionInput
 ///
 /// # Outputs
 /// None
@@ -35,13 +32,7 @@ use crate::{
 #[derive(Debug)]
 pub struct PrefetchRecordOperator {}
 
-#[derive(Debug)]
-pub struct PrefetchRecordInput {
-    pub logs: Chunk<LogRecord>,
-    pub blockfile_provider: BlockfileProvider,
-    pub record_segment: Segment,
-    pub offset_ids: Vec<u32>,
-}
+pub type PrefetchRecordInput = ProjectionInput;
 
 pub type PrefetchRecordOutput = ();