[ENH] Deprecate FetchSegmentOperator

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 CollectionSegments {
+    pub collection: Collection,
+    pub metadata_segment: Segment,
+    pub record_segment: Segment,
+    pub vector_segment: Segment,
+}
+
+impl TryFrom<chroma_proto::ScanOperator> for CollectionSegments {
+    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

@@ -17,7 +17,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 = ();
 

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

@@ -42,7 +42,7 @@ pub struct ProjectionOperator {
     pub metadata: bool,
 }
 
-#[derive(Debug)]
+#[derive(Clone, Debug)]
 pub struct ProjectionInput {
     pub logs: Chunk<LogRecord>,
     pub blockfile_provider: BlockfileProvider,