[ENH] Propagate segment information from frontend to query node

chromadb/proto/convert.py

@@ -573,9 +573,9 @@ 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,9 +6,9 @@ import "chromadb/proto/chroma.proto";
 
 message ScanOperator {
     Collection collection = 1;
-    string knn_id = 2;
-    string metadata_id = 3;
-    string record_id = 4;
+    Segment knn = 2;
+    Segment metadata = 3;
+    Segment record = 4;
 }
 
 message FilterOperator {

rust/types/src/segment.rs

@@ -6,6 +6,7 @@ use crate::chroma_proto;
 use chroma_error::{ChromaError, ErrorCodes};
 use std::{collections::HashMap, str::FromStr};
 use thiserror::Error;
+use tonic::Status;
 use uuid::Uuid;
 
 /// SegmentUuid is a wrapper around Uuid to provide a type for the segment id.
@@ -106,6 +107,12 @@ impl ChromaError for SegmentConversionError {
     }
 }
 
+impl From<SegmentConversionError> for Status {
+    fn from(value: SegmentConversionError) -> Self {
+        Status::invalid_argument(value.to_string())
+    }
+}
+
 impl TryFrom<chroma_proto::Segment> for Segment {
     type Error = SegmentConversionError;