Skip to content

Commit

Permalink
[ENH] Propagate segment information from frontend to query node
Browse files Browse the repository at this point in the history
  • Loading branch information
Sicheng Pan committed Dec 9, 2024
1 parent 5519de7 commit 56c2d2d
Show file tree
Hide file tree
Showing 6 changed files with 165 additions and 76 deletions.
6 changes: 3 additions & 3 deletions chromadb/proto/convert.py
Original file line number Diff line number Diff line change
Expand Up @@ -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),
)


Expand Down
66 changes: 33 additions & 33 deletions chromadb/proto/query_executor_pb2.py

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

16 changes: 8 additions & 8 deletions chromadb/proto/query_executor_pb2.pyi

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

6 changes: 3 additions & 3 deletions idl/chromadb/proto/query_executor.proto
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down
7 changes: 7 additions & 0 deletions rust/types/src/segment.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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;

Expand Down
Loading

0 comments on commit 56c2d2d

Please sign in to comment.