[BUG]: v0 in version file for forked collection missing segment file paths from source collection (#4550)

## Description of changes

Prior to this fix, `v0` in the version file for forked collections was
missing the segment file paths from the source collection.

There was a similar bug where if segment file paths were provided during
the creation of a new collection, those file paths would also be absent
from `v0` in the version file (we currently never provide segment file
paths upon creation, but seemed good to fix as well).

## Test plan

_How are these changes tested?_

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

Updated existing Go integration tests to check the created version file.

## 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)?_

n/a

Author: codetheweb

go/pkg/sysdb/coordinator/coordinator.go

@@ -97,7 +97,7 @@ func (s *Coordinator) GetTenant(ctx context.Context, getTenant *model.GetTenant)
 	return tenant, nil
 }
 
-func (s *Coordinator) CreateCollectionAndSegments(ctx context.Context, createCollection *model.CreateCollection, createSegments []*model.CreateSegment) (*model.Collection, bool, error) {
+func (s *Coordinator) CreateCollectionAndSegments(ctx context.Context, createCollection *model.CreateCollection, createSegments []*model.Segment) (*model.Collection, bool, error) {
 	collection, created, err := s.catalog.CreateCollectionAndSegments(ctx, createCollection, createSegments, createCollection.Ts)
 	if err != nil {
 		return nil, false, err
@@ -165,7 +165,7 @@ func (s *Coordinator) CountForks(ctx context.Context, sourceCollectionID types.U
 	return s.catalog.CountForks(ctx, sourceCollectionID)
 }
 
-func (s *Coordinator) CreateSegment(ctx context.Context, segment *model.CreateSegment) error {
+func (s *Coordinator) CreateSegment(ctx context.Context, segment *model.Segment) error {
 	if err := verifyCreateSegment(segment); err != nil {
 		return err
 	}
@@ -212,7 +212,7 @@ func verifyCollectionMetadata(metadata *model.CollectionMetadata[model.Collectio
 	return nil
 }
 
-func verifyCreateSegment(segment *model.CreateSegment) error {
+func verifyCreateSegment(segment *model.Segment) error {
 	if err := verifySegmentMetadata(segment.Metadata); err != nil {
 		return err
 	}

go/pkg/sysdb/coordinator/coordinator_test.go

@@ -9,6 +9,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/chroma-core/chroma/go/pkg/proto/coordinatorpb"
 	"github.com/chroma-core/chroma/go/pkg/sysdb/metastore/db/dao"
 	s3metastore "github.com/chroma-core/chroma/go/pkg/sysdb/metastore/s3"
 	"github.com/pingcap/log"
@@ -179,8 +180,8 @@ func testSegment(t *rapid.T) {
 
 	t.Repeat(map[string]func(*rapid.T){
 		"create_segment": func(t *rapid.T) {
-			segment := rapid.Custom[*model.CreateSegment](func(t *rapid.T) *model.CreateSegment {
-				return &model.CreateSegment{
+			segment := rapid.Custom[*model.Segment](func(t *rapid.T) *model.Segment {
+				return &model.Segment{
 					ID:           types.MustParse(rapid.StringMatching(`[0-9a-f]{8}-[0-9a-f]{4}-4[0-9a-f]{3}-[89ab][0-9a-f]{3}-[0-9a-f]{12}`).Draw(t, "segment_id")),
 					Type:         "test-segment-type",
 					Scope:        "test-segment-scope",
@@ -294,7 +295,7 @@ func (suite *APIsTestSuite) TestCreateCollectionAndSegments() {
 		DatabaseName: suite.databaseName,
 	}
 
-	segments := []*model.CreateSegment{
+	segments := []*model.Segment{
 		{
 			ID:           types.NewUniqueID(),
 			Type:         "test_type",
@@ -306,6 +307,9 @@ func (suite *APIsTestSuite) TestCreateCollectionAndSegments() {
 			Type:         "test_type",
 			Scope:        "VECTOR",
 			CollectionID: newCollection.ID,
+			FilePaths: map[string][]string{
+				"test_path": {"test_file1"},
+			},
 		},
 	}
 
@@ -347,6 +351,22 @@ func (suite *APIsTestSuite) TestCreateCollectionAndSegments() {
 	}
 	suite.ElementsMatch(expected_ids, actual_ids)
 
+	// Validate version file
+	versionFilePathPrefix := suite.s3MetaStore.GetVersionFilePath(collection.TenantID, suite.databaseId, newCollection.ID.String(), "0")
+	versionFile, err := suite.s3MetaStore.GetVersionFile(versionFilePathPrefix)
+	suite.NoError(err)
+	suite.NotNil(versionFile)
+	v0 := versionFile.VersionHistory.Versions[0]
+	suite.NotNil(v0)
+
+	// Validate file paths of segments
+	suite.NotNil(v0.SegmentInfo)
+	suite.NotNil(v0.SegmentInfo.SegmentCompactionInfo)
+	suite.Equal(len(v0.SegmentInfo.SegmentCompactionInfo), 2)
+	for _, segment := range segments {
+		assertExpectedSegmentInfoExist(suite, segment, v0.SegmentInfo.SegmentCompactionInfo)
+	}
+
 	// Attempt to create a duplicate collection (should fail)
 	_, _, err = suite.coordinator.CreateCollectionAndSegments(ctx, newCollection, segments)
 	suite.Error(err)
@@ -486,7 +506,7 @@ func (suite *APIsTestSuite) TestCreateGetDeleteCollections() {
 	suite.Equal(createCollection.Metadata, results[0].Metadata)
 
 	// Create segments associated with collection
-	segment := &model.CreateSegment{
+	segment := &model.Segment{
 		ID:           types.MustParse("00000000-0000-0000-0000-000000000001"),
 		CollectionID: createCollection.ID,
 		Type:         "test_segment",
@@ -993,7 +1013,7 @@ func (suite *APIsTestSuite) TestCreateGetDeleteSegments() {
 
 	sampleSegments := SampleSegments(suite.sampleCollections)
 	for _, segment := range sampleSegments {
-		errSegmentCreation := c.CreateSegment(ctx, &model.CreateSegment{
+		errSegmentCreation := c.CreateSegment(ctx, &model.Segment{
 			ID:           segment.ID,
 			Type:         segment.Type,
 			Scope:        segment.Scope,
@@ -1003,7 +1023,7 @@ func (suite *APIsTestSuite) TestCreateGetDeleteSegments() {
 		suite.NoError(errSegmentCreation)
 
 		// Create segment with empty collection id fails
-		errSegmentCreation = c.CreateSegment(ctx, &model.CreateSegment{
+		errSegmentCreation = c.CreateSegment(ctx, &model.Segment{
 			ID:           segment.ID,
 			Type:         segment.Type,
 			Scope:        segment.Scope,
@@ -1014,7 +1034,7 @@ func (suite *APIsTestSuite) TestCreateGetDeleteSegments() {
 
 		// Create segment to test unique constraint violation on segment.id.
 		// This should fail because the id is already taken.
-		errSegmentCreation = c.CreateSegment(ctx, &model.CreateSegment{
+		errSegmentCreation = c.CreateSegment(ctx, &model.Segment{
 			ID:           segment.ID,
 			Type:         segment.Type,
 			Scope:        segment.Scope,
@@ -1037,7 +1057,7 @@ func (suite *APIsTestSuite) TestCreateGetDeleteSegments() {
 	suite.Equal(sampleSegments, results)
 
 	// Duplicate create fails
-	err := c.CreateSegment(ctx, &model.CreateSegment{
+	err := c.CreateSegment(ctx, &model.Segment{
 		ID:           sampleSegments[0].ID,
 		Type:         sampleSegments[0].Type,
 		Scope:        sampleSegments[0].Scope,
@@ -1108,7 +1128,7 @@ func (suite *APIsTestSuite) TestUpdateSegment() {
 	}
 
 	ctx := context.Background()
-	errSegmentCreation := suite.coordinator.CreateSegment(ctx, &model.CreateSegment{
+	errSegmentCreation := suite.coordinator.CreateSegment(ctx, &model.Segment{
 		ID:           segment.ID,
 		Type:         segment.Type,
 		Scope:        segment.Scope,
@@ -1313,7 +1333,7 @@ func (suite *APIsTestSuite) TestCollectionVersioningWithMinio() {
 		DatabaseName: suite.databaseName,
 	}
 
-	segments := []*model.CreateSegment{
+	segments := []*model.Segment{
 		{
 			ID:           types.NewUniqueID(),
 			Type:         "test_type_a",
@@ -1353,6 +1373,33 @@ func (suite *APIsTestSuite) TestCollectionVersioningWithMinio() {
 	// suite.True(exists, "Version file should exist in S3")
 }
 
+func findSegmentInfo(segmentID types.UniqueID, segmentInfos []*coordinatorpb.FlushSegmentCompactionInfo) *coordinatorpb.FlushSegmentCompactionInfo {
+	for _, segmentInfo := range segmentInfos {
+		if segmentInfo.SegmentId == segmentID.String() {
+			return segmentInfo
+		}
+	}
+	return nil
+}
+
+func assertExpectedSegmentInfoExist(suite *APIsTestSuite, expectedSegment *model.Segment, segmentInfos []*coordinatorpb.FlushSegmentCompactionInfo) {
+	segmentInfo := findSegmentInfo(expectedSegment.ID, segmentInfos)
+	suite.NotNil(segmentInfo)
+
+	if expectedSegment.FilePaths == nil {
+		suite.Nil(segmentInfo.FilePaths)
+		return
+	}
+
+	suite.NotNil(segmentInfo.FilePaths)
+
+	filePaths := map[string][]string{}
+	for key, filePath := range segmentInfo.FilePaths {
+		filePaths[key] = filePath.Paths
+	}
+	suite.Equal(filePaths, expectedSegment.FilePaths)
+}
+
 func (suite *APIsTestSuite) TestForkCollection() {
 	ctx := context.Background()
 
@@ -1363,28 +1410,28 @@ func (suite *APIsTestSuite) TestForkCollection() {
 		DatabaseName: suite.databaseName,
 	}
 
-	sourceCreateMetadataSegment := &model.CreateSegment{
+	sourceCreateMetadataSegment := &model.Segment{
 		ID:           types.NewUniqueID(),
 		Type:         "test_blockfile",
 		Scope:        "METADATA",
 		CollectionID: sourceCreateCollection.ID,
 	}
 
-	sourceCreateRecordSegment := &model.CreateSegment{
+	sourceCreateRecordSegment := &model.Segment{
 		ID:           types.NewUniqueID(),
 		Type:         "test_blockfile",
 		Scope:        "RECORD",
 		CollectionID: sourceCreateCollection.ID,
 	}
 
-	sourceCreateVectorSegment := &model.CreateSegment{
+	sourceCreateVectorSegment := &model.Segment{
 		ID:           types.NewUniqueID(),
 		Type:         "test_hnsw",
 		Scope:        "VECTOR",
 		CollectionID: sourceCreateCollection.ID,
 	}
 
-	segments := []*model.CreateSegment{
+	segments := []*model.Segment{
 		sourceCreateMetadataSegment,
 		sourceCreateRecordSegment,
 		sourceCreateVectorSegment,
@@ -1467,6 +1514,22 @@ func (suite *APIsTestSuite) TestForkCollection() {
 		}
 	}
 
+	// Check version file of forked collection
+	versionFilePathPrefix := suite.s3MetaStore.GetVersionFilePath(collection.TenantID, suite.databaseId, forkCollection.TargetCollectionID.String(), "0")
+	versionFile, err := suite.s3MetaStore.GetVersionFile(versionFilePathPrefix)
+	suite.NoError(err)
+	suite.NotNil(versionFile)
+	v0 := versionFile.VersionHistory.Versions[0]
+	suite.NotNil(v0)
+	// Validate file paths of segments
+	suite.NotNil(v0.SegmentInfo)
+	suite.NotNil(v0.SegmentInfo.SegmentCompactionInfo)
+	suite.Equal(len(v0.SegmentInfo.SegmentCompactionInfo), 3)
+
+	for _, segment := range collection_segments {
+		assertExpectedSegmentInfoExist(suite, segment, v0.SegmentInfo.SegmentCompactionInfo)
+	}
+
 	// Attempt to fork a collcetion with same name (should fail)
 	forkCollectionWithSameName := &model.ForkCollection{
 		SourceCollectionID:                   sourceCreateCollection.ID,
@@ -1505,28 +1568,28 @@ func (suite *APIsTestSuite) TestCountForks() {
 		DatabaseName: suite.databaseName,
 	}
 
-	sourceCreateMetadataSegment := &model.CreateSegment{
+	sourceCreateMetadataSegment := &model.Segment{
 		ID:           types.NewUniqueID(),
 		Type:         "test_blockfile",
 		Scope:        "METADATA",
 		CollectionID: sourceCreateCollection.ID,
 	}
 
-	sourceCreateRecordSegment := &model.CreateSegment{
+	sourceCreateRecordSegment := &model.Segment{
 		ID:           types.NewUniqueID(),
 		Type:         "test_blockfile",
 		Scope:        "RECORD",
 		CollectionID: sourceCreateCollection.ID,
 	}
 
-	sourceCreateVectorSegment := &model.CreateSegment{
+	sourceCreateVectorSegment := &model.Segment{
 		ID:           types.NewUniqueID(),
 		Type:         "test_hnsw",
 		Scope:        "VECTOR",
 		CollectionID: sourceCreateCollection.ID,
 	}
 
-	segments := []*model.CreateSegment{
+	segments := []*model.Segment{
 		sourceCreateMetadataSegment,
 		sourceCreateRecordSegment,
 		sourceCreateVectorSegment,

go/pkg/sysdb/coordinator/model/segment.go

@@ -14,15 +14,6 @@ type Segment struct {
 	FilePaths    map[string][]string
 }
 
-type CreateSegment struct {
-	ID           types.UniqueID
-	Type         string
-	Scope        string
-	CollectionID types.UniqueID
-	Metadata     *SegmentMetadata[SegmentMetadataValueType]
-	Ts           types.Timestamp
-}
-
 type UpdateSegment struct {
 	ID            types.UniqueID
 	ResetTopic    bool

go/pkg/sysdb/coordinator/table_catalog.go

@@ -932,17 +932,18 @@ func (tc *Catalog) ForkCollection(ctx context.Context, forkCollection *model.For
 			LastCompactionTimeSecs:     sourceCollection.LastCompactionTimeSecs,
 		}
 
-		createSegments := []*model.CreateSegment{}
+		createSegments := []*model.Segment{}
 		flushFilePaths := []*model.FlushSegmentCompaction{}
 		for _, segment := range sourceSegments {
 			newSegmentID := types.NewUniqueID()
-			createSegment := &model.CreateSegment{
+			createSegment := &model.Segment{
 				ID:           newSegmentID,
 				Type:         segment.Type,
 				Scope:        segment.Scope,
 				CollectionID: forkCollection.TargetCollectionID,
 				Metadata:     segment.Metadata,
 				Ts:           ts.Unix(),
+				FilePaths:    segment.FilePaths,
 			}
 			createSegments = append(createSegments, createSegment)
 			flushFilePath := &model.FlushSegmentCompaction{
@@ -1034,7 +1035,7 @@ func (tc *Catalog) CountForks(ctx context.Context, sourceCollectionID types.Uniq
 	return uint64(len(lineageFile.Dependencies)), nil
 }
 
-func (tc *Catalog) CreateSegment(ctx context.Context, createSegment *model.CreateSegment, ts types.Timestamp) (*model.Segment, error) {
+func (tc *Catalog) CreateSegment(ctx context.Context, createSegment *model.Segment, ts types.Timestamp) (*model.Segment, error) {
 	var result *model.Segment
 
 	err := tc.txImpl.Transaction(ctx, func(txCtx context.Context) error {
@@ -1050,7 +1051,7 @@ func (tc *Catalog) CreateSegment(ctx context.Context, createSegment *model.Creat
 	return result, nil
 }
 
-func (tc *Catalog) createSegmentImpl(txCtx context.Context, createSegment *model.CreateSegment, ts types.Timestamp) (*model.Segment, error) {
+func (tc *Catalog) createSegmentImpl(txCtx context.Context, createSegment *model.Segment, ts types.Timestamp) (*model.Segment, error) {
 	var result *model.Segment
 
 	// insert segment
@@ -1090,7 +1091,21 @@ func (tc *Catalog) createSegmentImpl(txCtx context.Context, createSegment *model
 	return result, nil
 }
 
-func (tc *Catalog) createFirstVersionFile(ctx context.Context, databaseID string, createCollection *model.CreateCollection, createSegments []*model.CreateSegment, ts types.Timestamp) (string, error) {
+func (tc *Catalog) createFirstVersionFile(ctx context.Context, databaseID string, createCollection *model.CreateCollection, createSegments []*model.Segment, ts types.Timestamp) (string, error) {
+	segmentCompactionInfos := make([]*coordinatorpb.FlushSegmentCompactionInfo, 0, len(createSegments))
+	for _, segment := range createSegments {
+		convertedPaths := make(map[string]*coordinatorpb.FilePaths)
+		for k, v := range segment.FilePaths {
+			convertedPaths[k] = &coordinatorpb.FilePaths{Paths: v}
+		}
+
+		info := &coordinatorpb.FlushSegmentCompactionInfo{
+			SegmentId: segment.ID.String(),
+			FilePaths: convertedPaths,
+		}
+		segmentCompactionInfos = append(segmentCompactionInfos, info)
+	}
+
 	collectionVersionFilePb := &coordinatorpb.CollectionVersionFile{
 		CollectionInfoImmutable: &coordinatorpb.CollectionInfoImmutable{
 			TenantId:               createCollection.TenantID,
@@ -1104,6 +1119,9 @@ func (tc *Catalog) createFirstVersionFile(ctx context.Context, databaseID string
 				{
 					Version:       0,
 					CreatedAtSecs: int64(ts),
+					SegmentInfo: &coordinatorpb.CollectionSegmentInfo{
+						SegmentCompactionInfo: segmentCompactionInfos,
+					},
 				},
 			},
 		},
@@ -1117,7 +1135,7 @@ func (tc *Catalog) createFirstVersionFile(ctx context.Context, databaseID string
 	return fullFilePath, nil
 }
 
-func (tc *Catalog) CreateCollectionAndSegments(ctx context.Context, createCollection *model.CreateCollection, createSegments []*model.CreateSegment, ts types.Timestamp) (*model.Collection, bool, error) {
+func (tc *Catalog) CreateCollectionAndSegments(ctx context.Context, createCollection *model.CreateCollection, createSegments []*model.Segment, ts types.Timestamp) (*model.Collection, bool, error) {
 	var resultCollection *model.Collection
 	created := false
 

go/pkg/sysdb/grpc/collection_service.go

@@ -77,9 +77,9 @@ func (s *Server) CreateCollection(ctx context.Context, req *coordinatorpb.Create
 	}
 
 	// Convert the request segments to create segment models
-	createSegments := []*model.CreateSegment{}
+	createSegments := []*model.Segment{}
 	for _, segment := range req.Segments {
-		createSegment, err := convertSegmentToModel(segment)
+		createSegment, err := convertProtoSegment(segment)
 		if err != nil {
 			log.Error("Error in creating segments for the collection", zap.Error(err))
 			res.Collection = nil // We don't need to set the collection in case of error
@@ -90,6 +90,12 @@ func (s *Server) CreateCollection(ctx context.Context, req *coordinatorpb.Create
 			}
 			return res, grpcutils.BuildInternalGrpcError(err.Error())
 		}
+		filePaths := make(map[string][]string)
+		for key, filePath := range segment.FilePaths {
+			filePaths[key] = filePath.Paths
+		}
+		createSegment.FilePaths = filePaths
+
 		createSegments = append(createSegments, createSegment)
 	}