Repository: chroma-core/chroma
Pull request #4484: [ENH] Add RPC on SysDB to get fork count for a collection
Author: login: drewkim
Created: 2025-05-07T21:29:38Z | Merged: 2025-05-08T01:01:14Z

## Description of changes

_Summarize the changes made by this PR._

- Improvements & Bug fixes
  - N/A
- New functionality
  - Adds a new RPC to the SysDB, `CountForks`, that returns the number of forks for a given collection.

## 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 section](https://github.com/chroma-core/chroma/tree/main/docs/docs.trychroma.com)?_



Diff:
diff --git a/go/pkg/sysdb/coordinator/coordinator.go b/go/pkg/sysdb/coordinator/coordinator.go
index c6aa01b2a64..7290cc40c3c 100644
--- a/go/pkg/sysdb/coordinator/coordinator.go
+++ b/go/pkg/sysdb/coordinator/coordinator.go
@@ -161,6 +161,10 @@ func (s *Coordinator) ForkCollection(ctx context.Context, forkCollection *model.
 	return s.catalog.ForkCollection(ctx, forkCollection)
 }
 
+func (s *Coordinator) CountForks(ctx context.Context, sourceCollectionID types.UniqueID) (uint64, error) {
+	return s.catalog.CountForks(ctx, sourceCollectionID)
+}
+
 func (s *Coordinator) CreateSegment(ctx context.Context, segment *model.CreateSegment) error {
 	if err := verifyCreateSegment(segment); err != nil {
 		return err
diff --git a/go/pkg/sysdb/coordinator/coordinator_test.go b/go/pkg/sysdb/coordinator/coordinator_test.go
index 7c93cf8d648..755a034b6ec 100644
--- a/go/pkg/sysdb/coordinator/coordinator_test.go
+++ b/go/pkg/sysdb/coordinator/coordinator_test.go
@@ -1484,6 +1484,127 @@ func (suite *APIsTestSuite) TestForkCollection() {
 	suite.Empty(collections)
 }
 
+func (suite *APIsTestSuite) TestCountForks() {
+	ctx := context.Background()
+
+	sourceCreateCollection := &model.CreateCollection{
+		ID:           types.NewUniqueID(),
+		Name:         "test_fork_collection_source",
+		TenantID:     suite.tenantName,
+		DatabaseName: suite.databaseName,
+	}
+
+	sourceCreateMetadataSegment := &model.CreateSegment{
+		ID:           types.NewUniqueID(),
+		Type:         "test_blockfile",
+		Scope:        "METADATA",
+		CollectionID: sourceCreateCollection.ID,
+	}
+
+	sourceCreateRecordSegment := &model.CreateSegment{
+		ID:           types.NewUniqueID(),
+		Type:         "test_blockfile",
+		Scope:        "RECORD",
+		CollectionID: sourceCreateCollection.ID,
+	}
+
+	sourceCreateVectorSegment := &model.CreateSegment{
+		ID:           types.NewUniqueID(),
+		Type:         "test_hnsw",
+		Scope:        "VECTOR",
+		CollectionID: sourceCreateCollection.ID,
+	}
+
+	segments := []*model.CreateSegment{
+		sourceCreateMetadataSegment,
+		sourceCreateRecordSegment,
+		sourceCreateVectorSegment,
+	}
+
+	_, _, err := suite.coordinator.CreateCollectionAndSegments(ctx, sourceCreateCollection, segments)
+	suite.NoError(err)
+
+	sourceFlushMetadataSegment := &model.FlushSegmentCompaction{
+		ID: sourceCreateMetadataSegment.ID,
+		FilePaths: map[string][]string{
+			"fts_index": {"metadata_sparse_index_file"},
+		},
+	}
+
+	sourceFlushRecordSegment := &model.FlushSegmentCompaction{
+		ID: sourceCreateRecordSegment.ID,
+		FilePaths: map[string][]string{
+			"data_record": {"record_sparse_index_file"},
+		},
+	}
+
+	sourceFlushVectorSegment := &model.FlushSegmentCompaction{
+		ID: sourceCreateVectorSegment.ID,
+		FilePaths: map[string][]string{
+			"hnsw_index": {"hnsw_source_layer_file"},
+		},
+	}
+
+	sourceFlushCollectionCompaction := &model.FlushCollectionCompaction{
+		ID:                       sourceCreateCollection.ID,
+		TenantID:                 sourceCreateCollection.TenantID,
+		LogPosition:              1000,
+		CurrentCollectionVersion: 0,
+		FlushSegmentCompactions: []*model.FlushSegmentCompaction{
+			sourceFlushMetadataSegment,
+			sourceFlushRecordSegment,
+			sourceFlushVectorSegment,
+		},
+		TotalRecordsPostCompaction: 1000,
+		SizeBytesPostCompaction:    65536,
+	}
+
+	// Flush some data to sourceo collection
+	_, err = suite.coordinator.FlushCollectionCompaction(ctx, sourceFlushCollectionCompaction)
+	suite.NoError(err)
+
+	var forkedCollectionIDs []types.UniqueID
+
+	// Create 5 forks from the source collection
+	for i := 0; i < 5; i++ {
+		forkCollection := &model.ForkCollection{
+			SourceCollectionID:                   sourceCreateCollection.ID,
+			SourceCollectionLogCompactionOffset:  800,
+			SourceCollectionLogEnumerationOffset: 1200,
+			TargetCollectionID:                   types.NewUniqueID(),
+			TargetCollectionName:                 fmt.Sprintf("test_fork_collection_fork_source%d", i),
+		}
+		forkedCollection, _, err := suite.coordinator.ForkCollection(ctx, forkCollection)
+		suite.NoError(err)
+		forkedCollectionIDs = append(forkedCollectionIDs, forkedCollection.ID)
+	}
+
+	// Create 5 forks from one of the forked collections
+	for i := 0; i < 5; i++ {
+		forkCollection := &model.ForkCollection{
+			SourceCollectionID:                   forkedCollectionIDs[0],
+			SourceCollectionLogCompactionOffset:  800,
+			SourceCollectionLogEnumerationOffset: 1200,
+			TargetCollectionID:                   types.NewUniqueID(),
+			TargetCollectionName:                 fmt.Sprintf("test_fork_collection_fork_forked%d", i),
+		}
+		forkedCollection, _, err := suite.coordinator.ForkCollection(ctx, forkCollection)
+		suite.NoError(err)
+		forkedCollectionIDs = append(forkedCollectionIDs, forkedCollection.ID)
+	}
+
+	count, err := suite.coordinator.CountForks(ctx, sourceCreateCollection.ID)
+	suite.NoError(err)
+	suite.Equal(uint64(10), count)
+
+	// Check that each forked collection has 10 forks as well
+	for _, forkedCollectionID := range forkedCollectionIDs {
+		count, err := suite.coordinator.CountForks(ctx, forkedCollectionID)
+		suite.NoError(err)
+		suite.Equal(uint64(10), count)
+	}
+}
+
 func TestAPIsTestSuite(t *testing.T) {
 	testSuite := new(APIsTestSuite)
 	suite.Run(t, testSuite)
diff --git a/go/pkg/sysdb/coordinator/table_catalog.go b/go/pkg/sysdb/coordinator/table_catalog.go
index f6a7666abf5..ad98d10ddd6 100644
--- a/go/pkg/sysdb/coordinator/table_catalog.go
+++ b/go/pkg/sysdb/coordinator/table_catalog.go
@@ -993,6 +993,48 @@ func (tc *Catalog) ForkCollection(ctx context.Context, forkCollection *model.For
 	return tc.GetCollectionWithSegments(ctx, forkCollection.TargetCollectionID)
 }
 
+func (tc *Catalog) CountForks(ctx context.Context, sourceCollectionID types.UniqueID) (uint64, error) {
+	var rootCollectionID types.UniqueID
+
+	sourceCollectionIDStr := sourceCollectionID.String()
+	sourceCollectionDb, err := tc.metaDomain.CollectionDb(ctx).GetCollectionEntry(&sourceCollectionIDStr, nil)
+	if err != nil {
+		return 0, err
+	}
+	if sourceCollectionDb == nil {
+		return 0, common.ErrCollectionNotFound
+	}
+
+	if len(sourceCollectionDb.RootCollectionId) > 0 {
+		rootCollectionID, err = types.Parse(sourceCollectionDb.RootCollectionId)
+		if err != nil {
+			return 0, err
+		}
+	} else {
+		rootCollectionID = sourceCollectionID
+	}
+
+	limit := int32(1)
+	collections, err := tc.GetCollections(ctx, rootCollectionID, nil, "", "", &limit, nil)
+	if err != nil {
+		return 0, err
+	}
+	if len(collections) == 0 {
+		return 0, common.ErrCollectionNotFound
+	}
+	rootCollection := collections[0]
+
+	lineageFile, err := tc.getLineageFile(ctx, rootCollection)
+	if err != nil {
+		return 0, err
+	}
+
+	if lineageFile == nil || lineageFile.Dependencies == nil {
+		return 0, nil
+	}
+	return uint64(len(lineageFile.Dependencies)), nil
+}
+
 func (tc *Catalog) CreateSegment(ctx context.Context, createSegment *model.CreateSegment, ts types.Timestamp) (*model.Segment, error) {
 	var result *model.Segment
 
diff --git a/go/pkg/sysdb/grpc/collection_service.go b/go/pkg/sysdb/grpc/collection_service.go
index 2198bdd40ac..2e2f1777927 100644
--- a/go/pkg/sysdb/grpc/collection_service.go
+++ b/go/pkg/sysdb/grpc/collection_service.go
@@ -393,6 +393,27 @@ func (s *Server) ForkCollection(ctx context.Context, req *coordinatorpb.ForkColl
 	return res, nil
 }
 
+func (s *Server) CountForks(ctx context.Context, req *coordinatorpb.CountForksRequest) (*coordinatorpb.CountForksResponse, error) {
+	res := &coordinatorpb.CountForksResponse{}
+
+	sourceCollectionID := req.SourceCollectionId
+	parsedSourceCollectionID, err := types.ToUniqueID(&sourceCollectionID)
+	if err != nil {
+		log.Error("CountForks failed. Failed to parse source collection id", zap.Error(err), zap.String("collection_id", sourceCollectionID))
+		return res, grpcutils.BuildInternalGrpcError(err.Error())
+	}
+
+	count, err := s.coordinator.CountForks(ctx, parsedSourceCollectionID)
+	if err != nil {
+		if err == common.ErrCollectionNotFound {
+			return res, grpcutils.BuildNotFoundGrpcError(err.Error())
+		}
+		return res, grpcutils.BuildInternalGrpcError(err.Error())
+	}
+	res.Count = count
+	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 {
diff --git a/go/pkg/sysdb/grpc/collection_service_test.go b/go/pkg/sysdb/grpc/collection_service_test.go
index 8633df0a94a..bfbe6bbe297 100644
--- a/go/pkg/sysdb/grpc/collection_service_test.go
+++ b/go/pkg/sysdb/grpc/collection_service_test.go
@@ -2,6 +2,7 @@ package grpc
 
 import (
 	"context"
+	"fmt"
 	"strconv"
 	"testing"
 	"time"
@@ -572,6 +573,51 @@ func (suite *CollectionServiceTestSuite) TestGetCollectionSize() {
 	suite.NoError(err)
 }
 
+func (suite *CollectionServiceTestSuite) TestCountForks() {
+	collectionName := "collection_service_test_count_forks"
+	collectionID, err := dao.CreateTestCollection(suite.db, collectionName, 128, suite.databaseId)
+	suite.NoError(err)
+
+	req := coordinatorpb.CountForksRequest{
+		SourceCollectionId: collectionID,
+	}
+	res, err := suite.s.CountForks(context.Background(), &req)
+	suite.NoError(err)
+	suite.Equal(uint64(0), res.Count)
+
+	var forkedCollectionIDs []string
+
+	// Create 5 forks
+	for i := 0; i < 5; i++ {
+		forkCollectionReq := &coordinatorpb.ForkCollectionRequest{
+			SourceCollectionId:                   collectionID,
+			SourceCollectionLogCompactionOffset:  0,
+			SourceCollectionLogEnumerationOffset: 0,
+			TargetCollectionId:                   types.NewUniqueID().String(),
+			TargetCollectionName:                 fmt.Sprintf("test_fork_collection_fork_%d", i),
+		}
+		forkedCollection, err := suite.s.ForkCollection(context.Background(), forkCollectionReq)
+		suite.NoError(err)
+		forkedCollectionIDs = append(forkedCollectionIDs, forkedCollection.Collection.Id)
+	}
+
+	res, err = suite.s.CountForks(context.Background(), &req)
+	suite.NoError(err)
+	suite.Equal(uint64(5), res.Count)
+
+	// Check that each forked collection has 5 forks as well
+	for _, forkedCollectionID := range forkedCollectionIDs {
+		res, err = suite.s.CountForks(context.Background(), &coordinatorpb.CountForksRequest{
+			SourceCollectionId: forkedCollectionID,
+		})
+		suite.NoError(err)
+		suite.Equal(uint64(5), res.Count)
+	}
+
+	err = dao.CleanUpTestCollection(suite.db, collectionID)
+	suite.NoError(err)
+}
+
 func TestCollectionServiceTestSuite(t *testing.T) {
 	testSuite := new(CollectionServiceTestSuite)
 	suite.Run(t, testSuite)
diff --git a/idl/chromadb/proto/coordinator.proto b/idl/chromadb/proto/coordinator.proto
index 6e061dabf84..aa5d8fd83ea 100644
--- a/idl/chromadb/proto/coordinator.proto
+++ b/idl/chromadb/proto/coordinator.proto
@@ -229,6 +229,14 @@ message ForkCollectionResponse {
   repeated Segment segments = 2;
 }
 
+message CountForksRequest {
+  string source_collection_id = 1;
+}
+
+message CountForksResponse {
+  uint64 count = 1;
+}
+
 message ResetStateResponse {
   reserved 1;
   reserved "status";
@@ -476,6 +484,7 @@ service SysDB {
   rpc CheckCollections(CheckCollectionsRequest) returns (CheckCollectionsResponse) {}
   rpc UpdateCollection(UpdateCollectionRequest) returns (UpdateCollectionResponse) {}
   rpc ForkCollection(ForkCollectionRequest) returns (ForkCollectionResponse) {}
+  rpc CountForks(CountForksRequest) returns (CountForksResponse) {}
   rpc ResetState(google.protobuf.Empty) returns (ResetStateResponse) {}
   rpc GetLastCompactionTimeForTenant(GetLastCompactionTimeForTenantRequest) returns (GetLastCompactionTimeForTenantResponse) {}
   rpc SetLastCompactionTimeForTenant(SetLastCompactionTimeForTenantRequest) returns (google.protobuf.Empty) {}
diff --git a/rust/sysdb/src/sysdb.rs b/rust/sysdb/src/sysdb.rs
index b210c326537..5f4237cba42 100644
--- a/rust/sysdb/src/sysdb.rs
+++ b/rust/sysdb/src/sysdb.rs
@@ -19,9 +19,9 @@ use chroma_types::{
     UpdateCollectionConfiguration, UpdateCollectionError, VectorIndexConfiguration,
 };
 use chroma_types::{
-    Collection, CollectionConversionError, CollectionUuid, FlushCompactionResponse,
-    FlushCompactionResponseConversionError, ForkCollectionError, Segment, SegmentConversionError,
-    SegmentScope, Tenant,
+    Collection, CollectionConversionError, CollectionUuid, CountForksError,
+    FlushCompactionResponse, FlushCompactionResponseConversionError, ForkCollectionError, Segment,
+    SegmentConversionError, SegmentScope, Tenant,
 };
 use std::collections::HashMap;
 use std::fmt::Debug;
@@ -340,6 +340,17 @@ impl SysDb {
         }
     }
 
+    pub async fn count_forks(
+        &mut self,
+        source_collection_id: CollectionUuid,
+    ) -> Result<u64, CountForksError> {
+        match self {
+            SysDb::Grpc(grpc) => grpc.count_forks(source_collection_id).await,
+            SysDb::Sqlite(_) => Err(CountForksError::Local),
+            SysDb::Test(_) => Err(CountForksError::Local),
+        }
+    }
+
     pub async fn get_collections_to_gc(
         &mut self,
         cutoff_time: Option<SystemTime>,
@@ -975,6 +986,25 @@ impl GrpcSysDb {
         })
     }
 
+    pub async fn count_forks(
+        &mut self,
+        source_collection_id: CollectionUuid,
+    ) -> Result<u64, CountForksError> {
+        let res = self
+            .client
+            .count_forks(chroma_proto::CountForksRequest {
+                source_collection_id: source_collection_id.0.to_string(),
+            })
+            .await
+            .map_err(|err| match err.code() {
+                Code::NotFound => CountForksError::NotFound(source_collection_id.0.to_string()),
+                _ => CountForksError::Internal(err.into()),
+            })?
+            .into_inner();
+
+        Ok(res.count)
+    }
+
     pub async fn get_collections_to_gc(
         &mut self,
         cutoff_time: Option<SystemTime>,
diff --git a/rust/types/src/api_types.rs b/rust/types/src/api_types.rs
index 4f12b3f22e6..364bc2a645b 100644
--- a/rust/types/src/api_types.rs
+++ b/rust/types/src/api_types.rs
@@ -841,6 +841,26 @@ impl ChromaError for ForkCollectionError {
     }
 }
 
+#[derive(Debug, Error)]
+pub enum CountForksError {
+    #[error("Collection [{0}] does not exist")]
+    NotFound(String),
+    #[error(transparent)]
+    Internal(#[from] Box<dyn ChromaError>),
+    #[error("Count forks is unsupported for local chroma")]
+    Local,
+}
+
+impl ChromaError for CountForksError {
+    fn code(&self) -> ErrorCodes {
+        match self {
+            CountForksError::NotFound(_) => ErrorCodes::NotFound,
+            CountForksError::Internal(chroma_error) => chroma_error.code(),
+            CountForksError::Local => ErrorCodes::Unimplemented,
+        }
+    }
+}
+
 #[derive(Debug, Error)]
 pub enum GetCollectionSizeError {
     #[error(transparent)]


Reviews:
user: login: propel-code-bot[bot]
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-07T21:32:34Z
user: login: Sicheng-Pan
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-07T22:00:04Z
user: login: propel-code-bot[bot]
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-07T22:10:14Z
user: login: propel-code-bot[bot]
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-07T22:14:28Z
user: login: propel-code-bot[bot]
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-07T22:14:29Z
user: login: propel-code-bot[bot]
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-07T22:14:30Z
user: login: propel-code-bot[bot]
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-07T22:21:59Z
user: login: propel-code-bot[bot]
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-07T22:22:00Z
user: login: codetheweb
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-07T22:58:50Z
user: login: codetheweb
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-05-07T22:58:54Z
user: login: Sicheng-Pan
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-05-07T23:02:32Z

Review comments:
diff_hunk: @@ -993,6 +993,53 @@ func (tc *Catalog) ForkCollection(ctx context.Context, forkCollection *model.For
 	return tc.GetCollectionWithSegments(ctx, forkCollection.TargetCollectionID)
 }
 
+func (tc *Catalog) CountForks(ctx context.Context, sourceCollectionID types.UniqueID) (uint64, error) {
+	var rootCollection *model.Collection
+	var rootCollectionID types.UniqueID
+	var sourceCollection *model.Collection
+
+	sourceCollectionIDStr := sourceCollectionID.String()
+
+	sourceCollectionDb, err := tc.metaDomain.CollectionDb(ctx).GetCollectionEntry(&sourceCollectionIDStr, nil)
+	if err != nil {
+		return 0, err
+	}
+
+	if len(sourceCollectionDb.RootCollectionId) > 0 {
+		rootCollectionID, err = types.Parse(sourceCollectionDb.RootCollectionId)
+		if err != nil {
+			return 0, err
+		}
+	} else {
+		rootCollectionID = sourceCollectionID
+	}
+
+	sourceCollection, _, err = tc.GetCollectionWithSegments(ctx, sourceCollectionID)
+	if err != nil {
+		return 0, err
+	}
+	if rootCollectionID != sourceCollectionID {
+		limit := int32(1)
+		collections, err := tc.GetCollections(ctx, rootCollectionID, nil, "", "", &limit, nil)
+		if err != nil {
+			return 0, err
+		}
+		if len(collections) == 0 {
+			return 0, common.ErrCollectionNotFound
+		}
+		rootCollection = collections[0]
+	} else {
+		rootCollection = sourceCollection
+	}
+
+	lineageFile, err := tc.getLineageFile(ctx, rootCollection)
+	if err != nil {
+		return 0, err
+	}
+
+	return uint64(len(lineageFile.Dependencies)), nil
+}
path: go/pkg/sysdb/coordinator/table_catalog.go
user: login: propel-code-bot[bot]
body: [**BestPractice**]

In the implementation of `CountForks` in `table_catalog.go`, error handling for `GetCollectionEntry` and `GetCollections` is missing specificity. It's best practice to check for 'not found' errors and return a more specific error or 0, rather than always returning a generic error to the user. Consider refining error handling to distinguish between an actual error and a non-existent collection.

```suggestion
func (tc *Catalog) CountForks(ctx context.Context, sourceCollectionID types.UniqueID) (uint64, error) {
	var rootCollection *model.Collection
	var rootCollectionID types.UniqueID
	var sourceCollection *model.Collection

	sourceCollectionIDStr := sourceCollectionID.String()

	sourceCollectionDb, err := tc.metaDomain.CollectionDb(ctx).GetCollectionEntry(&sourceCollectionIDStr, nil)
	if err != nil {
		if errors.Is(err, common.ErrCollectionNotFound) {
			return 0, nil
		}
		return 0, err
	}

	if len(sourceCollectionDb.RootCollectionId) > 0 {
		rootCollectionID, err = types.Parse(sourceCollectionDb.RootCollectionId)
		if err != nil {
			return 0, err
		}
	} else {
		rootCollectionID = sourceCollectionID
	}

	sourceCollection, _, err = tc.GetCollectionWithSegments(ctx, sourceCollectionID)
	if err != nil {
		if errors.Is(err, common.ErrCollectionNotFound) {
			return 0, nil
		}
		return 0, err
	}
	if rootCollectionID != sourceCollectionID {
		limit := int32(1)
		collections, err := tc.GetCollections(ctx, rootCollectionID, nil, "", "", &limit, nil)
		if err != nil {
			return 0, err
		}
		if len(collections) == 0 {
			return 0, nil
		}
		rootCollection = collections[0]
	} else {
		rootCollection = sourceCollection
	}

	lineageFile, err := tc.getLineageFile(ctx, rootCollection)
	if err != nil {
		return 0, err
	}

	return uint64(len(lineageFile.Dependencies)), nil
}
```

⚡ **Committable suggestion**

Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation.
created_at: 2025-05-07T21:32:34Z
author_association: CONTRIBUTOR
diff_hunk: @@ -393,6 +393,24 @@ func (s *Server) ForkCollection(ctx context.Context, req *coordinatorpb.ForkColl
 	return res, nil
 }
 
+func (s *Server) CountForks(ctx context.Context, req *coordinatorpb.CountForksRequest) (*coordinatorpb.CountForksResponse, error) {
+	res := &coordinatorpb.CountForksResponse{}
+
+	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))
path: go/pkg/sysdb/grpc/collection_service.go
user: login: Sicheng-Pan
body: nit: update the error message
created_at: 2025-05-07T22:00:04Z
author_association: CONTRIBUTOR
diff_hunk: @@ -161,6 +161,10 @@ func (s *Coordinator) ForkCollection(ctx context.Context, forkCollection *model.
 	return s.catalog.ForkCollection(ctx, forkCollection)
 }
 
+func (s *Coordinator) CountForks(ctx context.Context, sourceCollectionID types.UniqueID) (uint64, error) {
path: go/pkg/sysdb/coordinator/coordinator.go
user: login: propel-code-bot[bot]
body: [**Documentation**]

The new exported method `CountForks` on the Coordinator struct in `coordinator.go` needs to be documented (at least with a Go doc comment) since it's public API. This helps clarify the method's contract for downstream users.
created_at: 2025-05-07T22:10:13Z
author_association: CONTRIBUTOR
diff_hunk: @@ -1484,6 +1484,117 @@ func (suite *APIsTestSuite) TestForkCollection() {
 	suite.Empty(collections)
 }
 
+func (suite *APIsTestSuite) TestCountForks() {
+	ctx := context.Background()
+
+	sourceCreateCollection := &model.CreateCollection{
+		ID:           types.NewUniqueID(),
+		Name:         "test_fork_collection_source",
+		TenantID:     suite.tenantName,
+		DatabaseName: suite.databaseName,
+	}
+
+	sourceCreateMetadataSegment := &model.CreateSegment{
+		ID:           types.NewUniqueID(),
+		Type:         "test_blockfile",
+		Scope:        "METADATA",
+		CollectionID: sourceCreateCollection.ID,
+	}
+
+	sourceCreateRecordSegment := &model.CreateSegment{
+		ID:           types.NewUniqueID(),
+		Type:         "test_blockfile",
+		Scope:        "RECORD",
+		CollectionID: sourceCreateCollection.ID,
+	}
+
+	sourceCreateVectorSegment := &model.CreateSegment{
+		ID:           types.NewUniqueID(),
+		Type:         "test_hnsw",
+		Scope:        "VECTOR",
+		CollectionID: sourceCreateCollection.ID,
+	}
+
+	segments := []*model.CreateSegment{
+		sourceCreateMetadataSegment,
+		sourceCreateRecordSegment,
+		sourceCreateVectorSegment,
+	}
+
+	_, _, err := suite.coordinator.CreateCollectionAndSegments(ctx, sourceCreateCollection, segments)
+	suite.NoError(err)
+
+	sourceFlushMetadataSegment := &model.FlushSegmentCompaction{
+		ID: sourceCreateMetadataSegment.ID,
+		FilePaths: map[string][]string{
+			"fts_index": {"metadata_sparse_index_file"},
+		},
+	}
+
+	sourceFlushRecordSegment := &model.FlushSegmentCompaction{
+		ID: sourceCreateRecordSegment.ID,
+		FilePaths: map[string][]string{
+			"data_record": {"record_sparse_index_file"},
+		},
+	}
+
+	sourceFlushVectorSegment := &model.FlushSegmentCompaction{
+		ID: sourceCreateVectorSegment.ID,
+		FilePaths: map[string][]string{
+			"hnsw_index": {"hnsw_source_layer_file"},
+		},
+	}
+
+	sourceFlushCollectionCompaction := &model.FlushCollectionCompaction{
+		ID:                       sourceCreateCollection.ID,
+		TenantID:                 sourceCreateCollection.TenantID,
+		LogPosition:              1000,
+		CurrentCollectionVersion: 0,
+		FlushSegmentCompactions: []*model.FlushSegmentCompaction{
+			sourceFlushMetadataSegment,
+			sourceFlushRecordSegment,
+			sourceFlushVectorSegment,
+		},
+		TotalRecordsPostCompaction: 1000,
+		SizeBytesPostCompaction:    65536,
+	}
+
+	// Flush some data to sourceo collection
+	_, err = suite.coordinator.FlushCollectionCompaction(ctx, sourceFlushCollectionCompaction)
+	suite.NoError(err)
+
+	var forkedCollectionId types.UniqueID
+
+	// Create 5 forks from the source collection
+	for i := 0; i < 5; i++ {
+		forkCollection := &model.ForkCollection{
+			SourceCollectionID:                   sourceCreateCollection.ID,
+			SourceCollectionLogCompactionOffset:  800,
+			SourceCollectionLogEnumerationOffset: 1200,
+			TargetCollectionID:                   types.NewUniqueID(),
+			TargetCollectionName:                 fmt.Sprintf("test_fork_collection_fork_source%d", i),
+		}
+		forkedCollection, _, _ := suite.coordinator.ForkCollection(ctx, forkCollection)
+		forkedCollectionId = forkedCollection.ID
path: go/pkg/sysdb/coordinator/coordinator_test.go
user: login: propel-code-bot[bot]
body: [**BestPractice**]

The `TestCountForks` test ignores the error return from `ForkCollection` when creating forks from the forked collection. Consider checking these errors to ensure proper test behavior.

```suggestion
		forkedCollection, _, err := suite.coordinator.ForkCollection(ctx, forkCollection)
		suite.NoError(err)
		forkedCollectionId = forkedCollection.ID
```

⚡ **Committable suggestion**

Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation.
created_at: 2025-05-07T22:14:28Z
author_association: CONTRIBUTOR
diff_hunk: @@ -1484,6 +1484,117 @@ func (suite *APIsTestSuite) TestForkCollection() {
 	suite.Empty(collections)
 }
 
+func (suite *APIsTestSuite) TestCountForks() {
+	ctx := context.Background()
+
+	sourceCreateCollection := &model.CreateCollection{
+		ID:           types.NewUniqueID(),
+		Name:         "test_fork_collection_source",
+		TenantID:     suite.tenantName,
+		DatabaseName: suite.databaseName,
+	}
+
+	sourceCreateMetadataSegment := &model.CreateSegment{
+		ID:           types.NewUniqueID(),
+		Type:         "test_blockfile",
+		Scope:        "METADATA",
+		CollectionID: sourceCreateCollection.ID,
+	}
+
+	sourceCreateRecordSegment := &model.CreateSegment{
+		ID:           types.NewUniqueID(),
+		Type:         "test_blockfile",
+		Scope:        "RECORD",
+		CollectionID: sourceCreateCollection.ID,
+	}
+
+	sourceCreateVectorSegment := &model.CreateSegment{
+		ID:           types.NewUniqueID(),
+		Type:         "test_hnsw",
+		Scope:        "VECTOR",
+		CollectionID: sourceCreateCollection.ID,
+	}
+
+	segments := []*model.CreateSegment{
+		sourceCreateMetadataSegment,
+		sourceCreateRecordSegment,
+		sourceCreateVectorSegment,
+	}
+
+	_, _, err := suite.coordinator.CreateCollectionAndSegments(ctx, sourceCreateCollection, segments)
+	suite.NoError(err)
+
+	sourceFlushMetadataSegment := &model.FlushSegmentCompaction{
+		ID: sourceCreateMetadataSegment.ID,
+		FilePaths: map[string][]string{
+			"fts_index": {"metadata_sparse_index_file"},
+		},
+	}
+
+	sourceFlushRecordSegment := &model.FlushSegmentCompaction{
+		ID: sourceCreateRecordSegment.ID,
+		FilePaths: map[string][]string{
+			"data_record": {"record_sparse_index_file"},
+		},
+	}
+
+	sourceFlushVectorSegment := &model.FlushSegmentCompaction{
+		ID: sourceCreateVectorSegment.ID,
+		FilePaths: map[string][]string{
+			"hnsw_index": {"hnsw_source_layer_file"},
+		},
+	}
+
+	sourceFlushCollectionCompaction := &model.FlushCollectionCompaction{
+		ID:                       sourceCreateCollection.ID,
+		TenantID:                 sourceCreateCollection.TenantID,
+		LogPosition:              1000,
+		CurrentCollectionVersion: 0,
+		FlushSegmentCompactions: []*model.FlushSegmentCompaction{
+			sourceFlushMetadataSegment,
+			sourceFlushRecordSegment,
+			sourceFlushVectorSegment,
+		},
+		TotalRecordsPostCompaction: 1000,
+		SizeBytesPostCompaction:    65536,
+	}
+
+	// Flush some data to sourceo collection
+	_, err = suite.coordinator.FlushCollectionCompaction(ctx, sourceFlushCollectionCompaction)
+	suite.NoError(err)
+
+	var forkedCollectionId types.UniqueID
+
+	// Create 5 forks from the source collection
+	for i := 0; i < 5; i++ {
+		forkCollection := &model.ForkCollection{
+			SourceCollectionID:                   sourceCreateCollection.ID,
+			SourceCollectionLogCompactionOffset:  800,
+			SourceCollectionLogEnumerationOffset: 1200,
+			TargetCollectionID:                   types.NewUniqueID(),
+			TargetCollectionName:                 fmt.Sprintf("test_fork_collection_fork_source%d", i),
+		}
+		forkedCollection, _, _ := suite.coordinator.ForkCollection(ctx, forkCollection)
+		forkedCollectionId = forkedCollection.ID
+	}
+
+	// Create 5 forks from one of the forked collections
+	for i := 0; i < 5; i++ {
+		forkCollection := &model.ForkCollection{
+			SourceCollectionID:                   forkedCollectionId,
+			SourceCollectionLogCompactionOffset:  800,
+			SourceCollectionLogEnumerationOffset: 1200,
+			TargetCollectionID:                   types.NewUniqueID(),
+			TargetCollectionName:                 fmt.Sprintf("test_fork_collection_fork_forked%d", i),
+		}
+		_, _, _ = suite.coordinator.ForkCollection(ctx, forkCollection)
path: go/pkg/sysdb/coordinator/coordinator_test.go
user: login: propel-code-bot[bot]
body: [**BestPractice**]

The test is not checking errors when creating forks from forked collections. Always check errors in tests to ensure proper test behavior.

```suggestion
		_, _, err := suite.coordinator.ForkCollection(ctx, forkCollection)
		suite.NoError(err)
```

⚡ **Committable suggestion**

Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation.
created_at: 2025-05-07T22:14:29Z
author_association: CONTRIBUTOR
diff_hunk: @@ -993,6 +993,43 @@ func (tc *Catalog) ForkCollection(ctx context.Context, forkCollection *model.For
 	return tc.GetCollectionWithSegments(ctx, forkCollection.TargetCollectionID)
 }
 
+func (tc *Catalog) CountForks(ctx context.Context, sourceCollectionID types.UniqueID) (uint64, error) {
+	var rootCollection *model.Collection
+	var rootCollectionID types.UniqueID
+
+	sourceCollectionIDStr := sourceCollectionID.String()
+	sourceCollectionDb, err := tc.metaDomain.CollectionDb(ctx).GetCollectionEntry(&sourceCollectionIDStr, nil)
+	if err != nil {
+		return 0, err
+	}
+
+	if len(sourceCollectionDb.RootCollectionId) > 0 {
+		rootCollectionID, err = types.Parse(sourceCollectionDb.RootCollectionId)
+		if err != nil {
+			return 0, err
+		}
+	} else {
+		rootCollectionID = sourceCollectionID
+	}
+
+	limit := int32(1)
+	collections, err := tc.GetCollections(ctx, rootCollectionID, nil, "", "", &limit, nil)
+	if err != nil {
+		return 0, err
+	}
+	if len(collections) == 0 {
+		return 0, common.ErrCollectionNotFound
+	}
+	rootCollection = collections[0]
+
+	lineageFile, err := tc.getLineageFile(ctx, rootCollection)
+	if err != nil {
+		return 0, err
+	}
+
+	return uint64(len(lineageFile.Dependencies)), nil
path: go/pkg/sysdb/coordinator/table_catalog.go
user: login: propel-code-bot[bot]
body: [**BestPractice**]

The implementation of CountForks returns a count based on the number of lineage file dependencies, which depends on lineage file correctness and availability. To make this robust, consider validating that `lineageFile.Dependencies` is not nil before taking `len(lineageFile.Dependencies)`, to prevent possible nil pointer panics.

```suggestion
	if lineageFile == nil || lineageFile.Dependencies == nil {
		return 0, nil
	}
	return uint64(len(lineageFile.Dependencies)), nil
```

⚡ **Committable suggestion**

Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation.
created_at: 2025-05-07T22:14:30Z
author_association: CONTRIBUTOR
diff_hunk: @@ -476,6 +484,7 @@ service SysDB {
   rpc CheckCollections(CheckCollectionsRequest) returns (CheckCollectionsResponse) {}
   rpc UpdateCollection(UpdateCollectionRequest) returns (UpdateCollectionResponse) {}
   rpc ForkCollection(ForkCollectionRequest) returns (ForkCollectionResponse) {}
+  rpc CountForks(CountForksRequest) returns (CountForksResponse) {}
   rpc ResetState(google.protobuf.Empty) returns (ResetStateResponse) {}
   rpc GetLastCompactionTimeForTenant(GetLastCompactionTimeForTenantRequest) returns (GetLastCompactionTimeForTenantResponse) {}
   rpc SetLastCompactionTimeForTenant(SetLastCompactionTimeForTenantRequest) returns (google.protobuf.Empty) {}
path: idl/chromadb/proto/coordinator.proto
user: login: propel-code-bot[bot]
body: [**BestPractice**]

You need to add `CountForks` to the proto and wire it through the gRPC service, coordinator, catalog, and tests, as you are doing. However, ensure that after adding this new RPC, your client and server code (in Go) actually use the generated Go code (from `protoc`). This includes:

- Confirm the proto-go code is re-generated for all changes (committed to repo if applicable, or documented in repo's README/how-to-build instructions).
- Confirm all new request/response Go structs are available and the gRPC interface implements the new call. 

If there are `*.pb.go` or `*_grpc.pb.go` files, keep them in sync. See company/OSS repo conventions.
created_at: 2025-05-07T22:21:59Z
author_association: CONTRIBUTOR
diff_hunk: @@ -993,6 +993,45 @@ func (tc *Catalog) ForkCollection(ctx context.Context, forkCollection *model.For
 	return tc.GetCollectionWithSegments(ctx, forkCollection.TargetCollectionID)
 }
 
+func (tc *Catalog) CountForks(ctx context.Context, sourceCollectionID types.UniqueID) (uint64, error) {
+	var rootCollectionID types.UniqueID
+
+	sourceCollectionIDStr := sourceCollectionID.String()
path: go/pkg/sysdb/coordinator/table_catalog.go
user: login: propel-code-bot[bot]
body: [**DataTypeCheck**]

In the `CountForks` implementation, be explicit about error handling for nil or missing collection entries; e.g. check for nil before accessing `sourceCollectionDb.RootCollectionId` to avoid panic if the collection does not exist. Consider:

```go
if sourceCollectionDb == nil {
    return 0, common.ErrCollectionNotFound
}
```
created_at: 2025-05-07T22:21:59Z
author_association: CONTRIBUTOR
diff_hunk: @@ -572,6 +573,39 @@ func (suite *CollectionServiceTestSuite) TestGetCollectionSize() {
 	suite.NoError(err)
 }
 
+func (suite *CollectionServiceTestSuite) TestCountForks() {
+	collectionName := "collection_service_test_count_forks"
+	collectionID, err := dao.CreateTestCollection(suite.db, collectionName, 128, suite.databaseId)
+	suite.NoError(err)
+
+	req := coordinatorpb.CountForksRequest{
+		SourceCollectionId: collectionID,
+	}
+	res, err := suite.s.CountForks(context.Background(), &req)
+	suite.NoError(err)
+	suite.Equal(uint64(0), res.Count)
+
+	// Create 5 forks
+	for i := 0; i < 5; i++ {
+		forkCollectionReq := &coordinatorpb.ForkCollectionRequest{
+			SourceCollectionId:                   collectionID,
+			SourceCollectionLogCompactionOffset:  0,
+			SourceCollectionLogEnumerationOffset: 0,
+			TargetCollectionId:                   types.NewUniqueID().String(),
+			TargetCollectionName:                 fmt.Sprintf("test_fork_collection_fork_%d", i),
+		}
+		_, err = suite.s.ForkCollection(context.Background(), forkCollectionReq)
+		suite.NoError(err)
+	}
+
+	res, err = suite.s.CountForks(context.Background(), &req)
path: go/pkg/sysdb/grpc/collection_service_test.go
user: login: codetheweb
body: wondering if this should be called `CalculateForkTreeSize` or similar
also might be good to test here that the result is 5 for `CountForks` on every collection in the set
created_at: 2025-05-07T22:58:50Z
author_association: CONTRIBUTOR

PR comments:
user: login: github-actions[bot]
created_at: 2025-05-07T21:29:48Z
body: # Reviewer Checklist
  Please leverage this checklist to ensure your code review is thorough before approving
  ## Testing, Bugs, Errors, Logs, Documentation
  - [ ] Can you think of any use case in which the code does not behave as intended? Have they been tested?
  - [ ] Can you think of any inputs or external events that could break the code? Is user input validated and safe? Have they been tested?
  - [ ] If appropriate, are there adequate property based tests?
  - [ ] If appropriate, are there adequate unit tests?
  - [ ] Should any logging, debugging, tracing information be added or removed?
  - [ ] Are error messages user-friendly?
  - [ ] Have all documentation changes needed been made?
  - [ ] Have all non-obvious changes been commented?
  ## System Compatibility
  - [ ] Are there any potential impacts on other parts of the system or backward compatibility?
  - [ ] Does this change intersect with any items on our roadmap, and if so, is there a plan for fitting them together?
  ## Quality
  - [ ] Is this code of a unexpectedly high quality (Readability, Modularity, Intuitiveness)
author_association: NONE
user: login: drewkim
created_at: 2025-05-07T21:29:52Z
body: * **#4484** <a href="https://app.graphite.dev/github/pr/chroma-core/chroma/4484?utm_source=stack-comment-icon" target="_blank"><img src="https://static.graphite.dev/graphite-32x32-black.png" alt="Graphite" width="10px" height="10px"/></a> 👈 <a href="https://app.graphite.dev/github/pr/chroma-core/chroma/4484?utm_source=stack-comment-view-in-graphite" target="_blank">(View in Graphite)</a>
* `main`


<h2></h2>

This stack of pull requests is managed by <a href="https://graphite.dev?utm-source=stack-comment"><b>Graphite</b></a>. Learn more about <a href="https://stacking.dev/?utm_source=stack-comment">stacking</a>.
<!-- Current dependencies on/for this PR: -->
author_association: CONTRIBUTOR
user: login: propel-code-bot[bot]
created_at: 2025-05-07T21:31:47Z
body: <!-- Summary by @propel-code-bot -->

**Add RPC on SysDB to get fork count for a collection**

This PR adds a new RPC method 'CountForks' to the SysDB service that allows clients to count the number of forks derived from a collection. The implementation traverses the collection lineage tree to determine the total number of related collections, providing this functionality across both Go and Rust interfaces with appropriate error handling and test coverage.

**Key Changes:**
• Added new CountForks RPC method to the SysDB proto definition
• Implemented CountForks functionality in table_catalog.go
• Added server implementation in collection_service.go
• Exposed API through Rust client interface
• Added comprehensive tests in Go and integrated with Rust error handling


**Affected Areas:**
• SysDB gRPC service
• Coordinator layer
• Rust client API
• Protocol definitions


*This summary was automatically generated by @propel-code-bot*
author_association: CONTRIBUTOR