[ENH] Wire up proto defs for sysdb fork endpoint (#4299)

## Description of changes

*Summarize the changes made by this PR.*
 - Improvements & Bug fixes
   - N/A
 - New functionality
   - Defines the SysDB fork service request and response
   - Wires up the request upto `table_catalog.go`

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

- [ ] 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)?*

Author: Sicheng-Pan

go/pkg/sysdb/coordinator/coordinator.go

@@ -157,6 +157,10 @@ func (s *Coordinator) UpdateCollection(ctx context.Context, collection *model.Up
 	return s.catalog.UpdateCollection(ctx, collection, collection.Ts)
 }
 
+func (s *Coordinator) ForkCollection(ctx context.Context, forkCollection *model.ForkCollection) (*model.Collection, []*model.Segment, error) {
+	return s.catalog.ForkCollection(ctx, forkCollection)
+}
+
 func (s *Coordinator) CreateSegment(ctx context.Context, segment *model.CreateSegment) error {
 	if err := verifyCreateSegment(segment); err != nil {
 		return err

go/pkg/sysdb/coordinator/model/collection.go

@@ -60,6 +60,14 @@ type UpdateCollection struct {
 	Ts                      types.Timestamp
 }
 
+type ForkCollection struct {
+	SourceCollectionID                   types.UniqueID
+	SourceCollectionLogCompactionOffset  uint64
+	SourceCollectionLogEnumerationOffset uint64
+	TargetCollectionID                   types.UniqueID
+	TargetCollectionName                 string
+}
+
 type FlushCollectionCompaction struct {
 	ID                         types.UniqueID
 	TenantID                   string

go/pkg/sysdb/coordinator/table_catalog.go

@@ -813,6 +813,25 @@ func (tc *Catalog) UpdateCollection(ctx context.Context, updateCollection *model
 	return result, nil
 }
 
+func (tc *Catalog) ForkCollection(ctx context.Context, forkCollection *model.ForkCollection) (*model.Collection, []*model.Segment, error) {
+	log.Info("Forking collection", zap.String("sourceCollectionId", forkCollection.SourceCollectionID.String()), zap.String("targetCollectionName", forkCollection.TargetCollectionName))
+
+	var source_collection *model.Collection
+	var source_segments []*model.Segment
+
+	err := tc.txImpl.Transaction(ctx, func(txCtx context.Context) error {
+		var err error
+		source_collection, source_segments, err = tc.GetCollectionWithSegments(ctx, forkCollection.SourceCollectionID)
+		return err
+	})
+	if err != nil {
+		return nil, nil, err
+	}
+
+	// TODO: Implement forking logic
+	return source_collection, source_segments, nil
+}
+
 func (tc *Catalog) CreateSegment(ctx context.Context, createSegment *model.CreateSegment, ts types.Timestamp) (*model.Segment, error) {
 	var result *model.Segment
 

go/pkg/sysdb/grpc/collection_service.go

@@ -3,7 +3,6 @@ package grpc
 import (
 	"context"
 	"encoding/json"
-	"fmt"
 
 	"github.com/chroma-core/chroma/go/pkg/grpcutils"
 
@@ -250,31 +249,15 @@ func (s *Server) GetCollectionWithSegments(ctx context.Context, req *coordinator
 		}
 		return res, grpcutils.BuildInternalGrpcError(err.Error())
 	}
-
 	res.Collection = convertCollectionToProto(collection)
+
 	segmentpbList := make([]*coordinatorpb.Segment, 0, len(segments))
-	scopeToSegmentMap := map[coordinatorpb.SegmentScope]*coordinatorpb.Segment{}
 	for _, segment := range segments {
 		segmentpb := convertSegmentToProto(segment)
-		scopeToSegmentMap[segmentpb.GetScope()] = segmentpb
 		segmentpbList = append(segmentpbList, segmentpb)
 	}
-
-	if len(segmentpbList) != 3 {
-		log.Error("GetCollectionWithSegments failed. Unexpected number of collection segments", zap.String("collection_id", collectionID))
-		return res, grpcutils.BuildInternalGrpcError(fmt.Sprintf("Unexpected number of segments for collection %s: %d", collectionID, len(segmentpbList)))
-	}
-
-	scopes := []coordinatorpb.SegmentScope{coordinatorpb.SegmentScope_METADATA, coordinatorpb.SegmentScope_RECORD, coordinatorpb.SegmentScope_VECTOR}
-
-	for _, scope := range scopes {
-		if _, exists := scopeToSegmentMap[scope]; !exists {
-			log.Error("GetCollectionWithSegments failed. Collection segment scope not found", zap.String("collection_id", collectionID), zap.String("missing_scope", scope.String()))
-			return res, grpcutils.BuildInternalGrpcError(fmt.Sprintf("Missing segment scope for collection %s: %s", collectionID, scope.String()))
-		}
-	}
-
 	res.Segments = segmentpbList
+
 	return res, nil
 }
 
@@ -360,6 +343,53 @@ func (s *Server) UpdateCollection(ctx context.Context, req *coordinatorpb.Update
 	return res, nil
 }
 
+func (s *Server) ForkCollection(ctx context.Context, req *coordinatorpb.ForkCollectionRequest) (*coordinatorpb.ForkCollectionResponse, error) {
+	res := &coordinatorpb.ForkCollectionResponse{}
+
+	sourceCollectionID := req.SourceCollectionId
+	parsedSourceCollectionID, err := types.ToUniqueID(&sourceCollectionID)
+	if err != nil {
+		log.Error("ForkCollection failed. Failed to parse source collection id", zap.Error(err), zap.String("collection_id", sourceCollectionID))
+		return res, grpcutils.BuildInternalGrpcError(err.Error())
+	}
+
+	targetCollectionID := req.TargetCollectionId
+	parsedTargetCollectionID, err := types.ToUniqueID(&targetCollectionID)
+	if err != nil {
+		log.Error("ForkCollection failed. Failed to parse target collection id", zap.Error(err), zap.String("collection_id", targetCollectionID))
+		return res, grpcutils.BuildInternalGrpcError(err.Error())
+	}
+
+	forkCollection := &model.ForkCollection{
+		SourceCollectionID:                   parsedSourceCollectionID,
+		SourceCollectionLogCompactionOffset:  req.SourceCollectionLogCompactionOffset,
+		SourceCollectionLogEnumerationOffset: req.SourceCollectionLogEnumerationOffset,
+		TargetCollectionID:                   parsedTargetCollectionID,
+		TargetCollectionName:                 req.TargetCollectionName,
+	}
+	collection, segments, err := s.coordinator.ForkCollection(ctx, forkCollection)
+	if err != nil {
+		log.Error("ForkCollection failed. ", zap.Error(err), zap.String("collection_id", sourceCollectionID))
+		if err == common.ErrCollectionNotFound {
+			return res, grpcutils.BuildNotFoundGrpcError(err.Error())
+		}
+		if err == common.ErrCollectionUniqueConstraintViolation {
+			return res, grpcutils.BuildAlreadyExistsGrpcError(err.Error())
+		}
+		return res, grpcutils.BuildInternalGrpcError(err.Error())
+	}
+	res.Collection = convertCollectionToProto(collection)
+
+	segmentpbList := make([]*coordinatorpb.Segment, 0, len(segments))
+	for _, segment := range segments {
+		segmentpb := convertSegmentToProto(segment)
+		segmentpbList = append(segmentpbList, segmentpb)
+	}
+	res.Segments = segmentpbList
+
+	return res, nil
+}
+
 func (s *Server) ListCollectionVersions(ctx context.Context, req *coordinatorpb.ListCollectionVersionsRequest) (*coordinatorpb.ListCollectionVersionsResponse, error) {
 	collectionID, err := types.ToUniqueID(&req.CollectionId)
 	if err != nil {

idl/chromadb/proto/coordinator.proto

@@ -216,6 +216,19 @@ message UpdateCollectionResponse {
   reserved "status";
 }
 
+message ForkCollectionRequest {
+  string source_collection_id = 1;
+  uint64 source_collection_log_compaction_offset = 2;
+  uint64 source_collection_log_enumeration_offset = 3;
+  string target_collection_id = 4;
+  string target_collection_name = 5;
+}
+
+message ForkCollectionResponse {
+  Collection collection = 1;
+  repeated Segment segments = 2;
+}
+
 message ResetStateResponse {
   reserved 1;
   reserved "status";
@@ -333,34 +346,13 @@ message CollectionSegmentInfo {
     // GC's {collection,version} selection policy.
 }
 
-// Tuple of collection ID, tenant ID, and version.
-message CollectionVersionTuple {
-    string collection_id = 1;
-    string tenant_id = 2;
-    int64 version = 3;
-}
-
 message VersionListForCollection {
   string tenant_id = 1;
   string database_id = 2;
   string collection_id = 3;
   repeated int64 versions = 4;
 }
 
-// Contains information about the lineage of a collection.
-message CollectionLineageInfo {
-    // ID of the collection.
-    string collection_id = 1;
-    // ID of the tenant.
-    string tenant_id = 2;
-    // Whether the collection is a root collection.
-    bool is_root_collection = 3;
-    // An ordered list of descendant collections.
-    // The first element is the root {collection, version}, and the last
-    // element is the direct parent of the current collection.
-    repeated CollectionVersionTuple parent_collections = 4;
-}
-
 // Request to list versions of a collection.
 message ListCollectionVersionsRequest {
   string collection_id = 1;
@@ -382,6 +374,16 @@ message ListCollectionVersionsResponse {
   bool list_is_truncated = 2;
 }
 
+message CollectionLineageFile {
+	repeated CollectionVersionDependency dependencies = 1;
+}
+
+message CollectionVersionDependency {
+	string source_collection_id = 1; // The forked collection
+	uint64 source_collection_version = 2; // The forked collection version
+	string target_collection_id = 3; // The forking collection
+}
+
 // Request to restore a collection.
 message RestoreCollectionRequest {
   string collection_id = 1;
@@ -473,6 +475,7 @@ service SysDB {
   rpc GetCollectionWithSegments(GetCollectionWithSegmentsRequest) returns (GetCollectionWithSegmentsResponse) {}
   rpc CheckCollections(CheckCollectionsRequest) returns (CheckCollectionsResponse) {}
   rpc UpdateCollection(UpdateCollectionRequest) returns (UpdateCollectionResponse) {}
+  rpc ForkCollection(ForkCollectionRequest) returns (ForkCollectionResponse) {}
   rpc ResetState(google.protobuf.Empty) returns (ResetStateResponse) {}
   rpc GetLastCompactionTimeForTenant(GetLastCompactionTimeForTenantRequest) returns (GetLastCompactionTimeForTenantResponse) {}
   rpc SetLastCompactionTimeForTenant(SetLastCompactionTimeForTenantRequest) returns (google.protobuf.Empty) {}