[ENH] Propagate segment information from frontend to query node (#3255)

## Description of changes

*Summarize the changes made by this PR.*
 - Improvements & Bug fixes
   - N/A
 - New functionality
   - Propagate full segment information from frontend to query node

## Test plan
*How are these changes tested?*

- [x] Tests pass locally with `pytest` for python, `yarn test` for js, `cargo test` for rust

## Documentation Changes
*Are all docstrings for user-facing APIs updated if required? Do we need to make documentation changes in the [docs repository](https://github.com/chroma-core/docs)?*
N/A

chromadb/proto/convert.py

@@ -576,6 +576,9 @@ def to_proto_scan(scan: Scan) -> query_pb.ScanOperator:
         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,15 @@ 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;
+    Segment knn = 5;
+    Segment metadata = 6;
+    Segment record = 7;
 }
 
 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;