[ENH]: Execute task with no backfill or incremental

Author: tanujnay112

go/pkg/sysdb/coordinator/coordinator.go

@@ -13,6 +13,7 @@ import (
 	"github.com/chroma-core/chroma/go/pkg/sysdb/metastore/db/dbmodel"
 	s3metastore "github.com/chroma-core/chroma/go/pkg/sysdb/metastore/s3"
 	"github.com/chroma-core/chroma/go/pkg/types"
+	"github.com/google/uuid"
 	"github.com/pingcap/log"
 	"go.uber.org/zap"
 )
@@ -286,6 +287,24 @@ func (s *Coordinator) FlushCollectionCompaction(ctx context.Context, flushCollec
 	return s.catalog.FlushCollectionCompaction(ctx, flushCollectionCompaction)
 }
 
+func (s *Coordinator) FlushCollectionCompactionAndAttachedFunction(
+	ctx context.Context,
+	flushCollectionCompaction *model.FlushCollectionCompaction,
+	attachedFunctionID uuid.UUID,
+	completionOffset int64,
+) (*model.FlushCollectionInfo, error) {
+	return s.catalog.FlushCollectionCompactionAndAttachedFunction(ctx, flushCollectionCompaction, attachedFunctionID, completionOffset)
+}
+
+func (s *Coordinator) FlushCollectionCompactionAndAttachedFunctionExtended(
+	ctx context.Context,
+	collectionCompactions []*model.FlushCollectionCompaction,
+	attachedFunctionID uuid.UUID,
+	completionOffset int64,
+) (*model.ExtendedFlushCollectionInfo, error) {
+	return s.catalog.FlushCollectionCompactionAndAttachedFunctionExtended(ctx, collectionCompactions, attachedFunctionID, completionOffset)
+}
+
 func (s *Coordinator) ListCollectionsToGc(ctx context.Context, cutoffTimeSecs *uint64, limit *uint64, tenantID *string, minVersionsIfAlive *uint64) ([]*model.CollectionToGc, error) {
 	return s.catalog.ListCollectionsToGc(ctx, cutoffTimeSecs, limit, tenantID, minVersionsIfAlive)
 }

go/pkg/sysdb/coordinator/create_task_test.go

@@ -176,7 +176,8 @@ func (suite *HeapClientIntegrationTestSuite) TestAttachFunctionPushesScheduleToH
 	})
 	suite.NoError(err, "Should attached function successfully")
 	suite.NotNil(response)
-	suite.NotEmpty(response.Id, "Attached function ID should be returned")
+	suite.NotNil(response.AttachedFunction)
+	suite.NotEmpty(response.AttachedFunction.Id, "Attached function ID should be returned")
 
 	// Get updated heap summary
 	updatedSummary, err := suite.heapClient.Summary(ctx, &coordinatorpb.HeapSummaryRequest{})
@@ -263,9 +264,15 @@ func (suite *HeapClientIntegrationTestSuite) TestPartialTaskRecovery_HybridAppro
 		return
 	}
 	suite.NotNil(taskResp)
-	// TODO: Uncomment after proto regeneration
-	// originalTaskID := taskResp.Id
-	suite.T().Skip("Test requires proto regeneration for AttachFunctionResponse.Id field")
+	suite.NotNil(taskResp.AttachedFunction)
+	originalTaskID := taskResp.AttachedFunction.Id
+	suite.T().Logf("Created fully initialized task: %s", originalTaskID)
+
+	// STEP 2: Directly UPDATE database to make task partial (simulate Phase 3 failure)
+	// Set lowest_live_nonce = NULL to simulate the task being stuck
+	_, err = db.Exec(`UPDATE public.tasks SET lowest_live_nonce = NULL WHERE task_id = $1`, originalTaskID)
+	suite.NoError(err, "Should be able to corrupt task in database")
+	suite.T().Logf("Made task partial by setting lowest_live_nonce = NULL")
 
 	// STEP 3: Try to create task with same name but DIFFERENT parameters → should fail
 	_, err = suite.sysdbClient.AttachFunction(ctx, &coordinatorpb.AttachFunctionRequest{
@@ -358,8 +365,8 @@ func (suite *HeapClientIntegrationTestSuite) TestPartialTaskCleanup_ThenRecreate
 		return
 	}
 	suite.NotNil(taskResp)
-	// TODO: Uncomment after proto regeneration
-	suite.T().Skip("Test requires proto regeneration for AttachFunctionResponse.Id field")
+	suite.NotNil(taskResp.AttachedFunction)
+	suite.T().Logf("Created task: %s", taskResp.AttachedFunction.Id)
 
 	// STEP 2: Call CleanupExpiredPartialAttachedFunctions (with short timeout to test it doesn't affect complete tasks)
 	cleanupResp, err := suite.sysdbClient.CleanupExpiredPartialAttachedFunctions(ctx, &coordinatorpb.CleanupExpiredPartialAttachedFunctionsRequest{
@@ -377,12 +384,12 @@ func (suite *HeapClientIntegrationTestSuite) TestPartialTaskCleanup_ThenRecreate
 	})
 	suite.NoError(err, "Task should still exist after cleanup")
 	suite.NotNil(getResp)
-	suite.Equal(taskResp.Id, getResp.AttachedFunction.Id)
+	suite.Equal(taskResp.AttachedFunction.Id, getResp.AttachedFunction.Id)
 	suite.T().Logf("Task still exists after cleanup: %s", getResp.AttachedFunction.Id)
 
 	// STEP 4: Delete the task
 	_, err = suite.sysdbClient.DetachFunction(ctx, &coordinatorpb.DetachFunctionRequest{
-		AttachedFunctionId: taskResp.Id,
+		AttachedFunctionId: taskResp.AttachedFunction.Id,
 		DeleteOutput:       true,
 	})
 	suite.NoError(err, "Should delete task")
@@ -399,8 +406,9 @@ func (suite *HeapClientIntegrationTestSuite) TestPartialTaskCleanup_ThenRecreate
 	})
 	suite.NoError(err, "Should be able to recreate task after deletion")
 	suite.NotNil(taskResp2)
-	suite.NotEqual(taskResp.Id, taskResp2.Id, "New task should have different ID")
-	suite.T().Logf("Successfully recreated task: %s", taskResp2.Id)
+	suite.NotNil(taskResp2.AttachedFunction)
+	suite.NotEqual(taskResp.AttachedFunction.Id, taskResp2.AttachedFunction.Id, "New task should have different ID")
+	suite.T().Logf("Successfully recreated task: %s", taskResp2.AttachedFunction.Id)
 }
 
 func TestHeapClientIntegrationSuite(t *testing.T) {

go/pkg/sysdb/coordinator/heap_client_integration_test.go

@@ -102,6 +102,10 @@ type FlushCollectionInfo struct {
 	AttachedFunctionCompletionOffset *int64
 }
 
+type ExtendedFlushCollectionInfo struct {
+	Collections []*FlushCollectionInfo
+}
+
 func FilterCollection(collection *Collection, collectionID types.UniqueID, collectionName *string) bool {
 	if collectionID != types.NilUniqueID() && collectionID != collection.ID {
 		return false

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

@@ -1736,6 +1736,119 @@ func (tc *Catalog) FlushCollectionCompaction(ctx context.Context, flushCollectio
 	return flushCollectionInfo, nil
 }
 
+// FlushCollectionCompactionAndAttachedFunction atomically updates collection compaction data and attached function completion offset.
+// TODO(tanujnay112): Deprecate this
+func (tc *Catalog) FlushCollectionCompactionAndAttachedFunction(
+	ctx context.Context,
+	flushCollectionCompaction *model.FlushCollectionCompaction,
+	attachedFunctionID uuid.UUID,
+	completionOffset int64,
+) (*model.FlushCollectionInfo, error) {
+	if !tc.versionFileEnabled {
+		// Attached-function-based compactions are only supported with versioned collections
+		log.Error("FlushCollectionCompactionAndAttachedFunction is only supported for versioned collections")
+		return nil, errors.New("attached-function-based compaction requires versioned collections")
+	}
+
+	var flushCollectionInfo *model.FlushCollectionInfo
+
+	err := tc.txImpl.Transaction(ctx, func(txCtx context.Context) error {
+		var err error
+		// Get the transaction from context to pass to FlushCollectionCompactionForVersionedCollection
+		tx := dbcore.GetDB(txCtx)
+		flushCollectionInfo, err = tc.FlushCollectionCompactionForVersionedCollection(txCtx, flushCollectionCompaction, tx)
+		if err != nil {
+			return err
+		}
+
+		// Update ONLY completion_offset
+		err = tc.metaDomain.AttachedFunctionDb(txCtx).UpdateCompletionOffset(attachedFunctionID, completionOffset)
+		if err != nil {
+			return err
+		}
+
+		return nil
+	})
+
+	if err != nil {
+		return nil, err
+	}
+
+	// Populate attached function fields with authoritative values from database
+	flushCollectionInfo.AttachedFunctionCompletionOffset = &completionOffset
+
+	log.Info("FlushCollectionCompactionAndAttachedFunction",
+		zap.String("collection_id", flushCollectionCompaction.ID.String()),
+		zap.String("attached_function_id", attachedFunctionID.String()),
+		zap.Int64("completion_offset", completionOffset))
+
+	return flushCollectionInfo, nil
+}
+
+// FlushCollectionCompactionAndAttachedFunctionExtended atomically updates multiple collection compaction data
+// and attached function completion offset in a single transaction.
+func (tc *Catalog) FlushCollectionCompactionAndAttachedFunctionExtended(
+	ctx context.Context,
+	collectionCompactions []*model.FlushCollectionCompaction,
+	attachedFunctionID uuid.UUID,
+	completionOffset int64,
+) (*model.ExtendedFlushCollectionInfo, error) {
+	if !tc.versionFileEnabled {
+		// Attached-function-based compactions are only supported with versioned collections
+		log.Error("FlushCollectionCompactionAndAttachedFunctionExtended is only supported for versioned collections")
+		return nil, errors.New("attached-function-based compaction requires versioned collections")
+	}
+
+	if len(collectionCompactions) == 0 {
+		return nil, errors.New("at least one collection compaction is required")
+	}
+
+	flushInfos := make([]*model.FlushCollectionInfo, 0, len(collectionCompactions))
+
+	err := tc.txImpl.Transaction(ctx, func(txCtx context.Context) error {
+		var err error
+		// Get the transaction from context to pass to FlushCollectionCompactionForVersionedCollection
+		tx := dbcore.GetDB(txCtx)
+
+		// Handle all collection compactions
+		for _, collectionCompaction := range collectionCompactions {
+			log.Info("FlushCollectionCompactionAndAttachedFunctionExtended", zap.String("collection_id", collectionCompaction.ID.String()))
+			flushInfo, err := tc.FlushCollectionCompactionForVersionedCollection(txCtx, collectionCompaction, tx)
+			if err != nil {
+				return err
+			}
+			flushInfos = append(flushInfos, flushInfo)
+		}
+
+		err = tc.metaDomain.AttachedFunctionDb(txCtx).UpdateCompletionOffset(attachedFunctionID, completionOffset)
+		if err != nil {
+			return err
+		}
+
+		return nil
+	})
+
+	if err != nil {
+		return nil, err
+	}
+
+	// Populate attached function fields with authoritative values from database
+	for _, flushInfo := range flushInfos {
+		flushInfo.AttachedFunctionCompletionOffset = &completionOffset
+	}
+
+	// Log with first collection ID (typically the output collection)
+	log.Info("FlushCollectionCompactionAndAttachedFunctionExtended",
+		zap.String("first_collection_id", collectionCompactions[0].ID.String()),
+		zap.Int("collection_count", len(collectionCompactions)),
+		zap.String("attached_function_id", attachedFunctionID.String()),
+		zap.Int64("completion_offset", completionOffset))
+
+	return &model.ExtendedFlushCollectionInfo{
+		Collections: flushInfos,
+	}, nil
+}
+
 func (tc *Catalog) validateVersionFile(versionFile *coordinatorpb.CollectionVersionFile, collectionID string, version int64) error {
 	if versionFile.GetCollectionInfoImmutable().GetCollectionId() != collectionID {
 		log.Error("collection id mismatch", zap.String("collection_id", collectionID), zap.String("version_file_collection_id", versionFile.GetCollectionInfoImmutable().GetCollectionId()))

go/pkg/sysdb/coordinator/table_catalog.go

@@ -76,7 +76,7 @@ func (s *Coordinator) validateAttachedFunctionMatchesRequest(ctx context.Context
 	return nil
 }
 
-// AttachFunction creates a new attached function in the database
+// AttachFunction creates an output collection and attached function in a single transaction
 func (s *Coordinator) AttachFunction(ctx context.Context, req *coordinatorpb.AttachFunctionRequest) (*coordinatorpb.AttachFunctionResponse, error) {
 	log := log.With(zap.String("method", "AttachFunction"))
 
@@ -108,8 +108,6 @@ func (s *Coordinator) AttachFunction(ctx context.Context, req *coordinatorpb.Att
 
 			// Validation passed, reuse the concurrent attached function ID (idempotent)
 			attachedFunctionID = concurrentAttachedFunction.ID
-			// Already created, skip Phase 2
-			skipPhase2 = true
 			return nil
 		}
 
@@ -145,16 +143,59 @@ func (s *Coordinator) AttachFunction(ctx context.Context, req *coordinatorpb.Att
 			return common.ErrCollectionNotFound
 		}
 
-		// Check if output collection already exists
-		outputCollectionName := req.OutputCollectionName
-		existingOutputCollections, err := s.catalog.metaDomain.CollectionDb(txCtx).GetCollections(nil, &outputCollectionName, req.TenantId, req.Database, nil, nil, false)
+		// Create output collection with segments
+		outputCollectionID := types.NewUniqueID()
+
+		// Set a default dimension to ensure segment writers can be initialized
+		dimension := int32(1) // Default dimension for attached function output collections
+
+		createCollection := &model.CreateCollection{
+			ID:                   outputCollectionID,
+			Name:                 req.OutputCollectionName,
+			ConfigurationJsonStr: `{"hnsw": {"space": "cosine", "M": 16, "ef_construction": 64}}`,
+			DatabaseName:         req.Database,
+			TenantID:             req.TenantId,
+			GetOrCreate:          false, // We want to fail if it already exists
+			Dimension:            &dimension,
+			Metadata:             nil,
+		}
+
+		// Create segments for the collection (distributed setup)
+		segments := []*model.Segment{
+			{
+				ID:           types.NewUniqueID(),
+				Type:         "urn:chroma:segment/vector/hnsw-distributed",
+				Scope:        "VECTOR",
+				CollectionID: outputCollectionID,
+			},
+			{
+				ID:           types.NewUniqueID(),
+				Type:         "urn:chroma:segment/metadata/blockfile",
+				Scope:        "METADATA",
+				CollectionID: outputCollectionID,
+			},
+			{
+				ID:           types.NewUniqueID(),
+				Type:         "urn:chroma:segment/record/blockfile",
+				Scope:        "RECORD",
+				CollectionID: outputCollectionID,
+			},
+		}
+
+		// Create output collection and segments directly to avoid nested transaction
+		outputCollection, _, err := s.catalog.createCollectionImpl(txCtx, createCollection, "", 0)
 		if err != nil {
-			log.Error("AttachFunction: failed to check output collection", zap.Error(err))
+			log.Error("AttachFunction: failed to create output collection", zap.Error(err))
 			return err
 		}
-		if len(existingOutputCollections) > 0 {
-			log.Error("AttachFunction: output collection already exists")
-			return common.ErrCollectionUniqueConstraintViolation
+
+		// Create segments for the collection
+		for _, segment := range segments {
+			_, err := s.catalog.createSegmentImpl(txCtx, segment, 0)
+			if err != nil {
+				log.Error("AttachFunction: failed to create segment", zap.Error(err))
+				return err
+			}
 		}
 
 		// Serialize params
@@ -170,6 +211,7 @@ func (s *Coordinator) AttachFunction(ctx context.Context, req *coordinatorpb.Att
 			paramsJSON = "{}"
 		}
 
+		// Create attached function
 		now := time.Now()
 		attachedFunction := &dbmodel.AttachedFunction{
 			ID:                      attachedFunctionID,
@@ -178,6 +220,7 @@ func (s *Coordinator) AttachFunction(ctx context.Context, req *coordinatorpb.Att
 			DatabaseID:              databases[0].ID,
 			InputCollectionID:       req.InputCollectionId,
 			OutputCollectionName:    req.OutputCollectionName,
+			OutputCollectionID:      &[]string{outputCollection.ID.String()}[0],
 			FunctionID:              function.ID,
 			FunctionParams:          paramsJSON,
 			CompletionOffset:        0,
@@ -198,6 +241,7 @@ func (s *Coordinator) AttachFunction(ctx context.Context, req *coordinatorpb.Att
 
 		log.Debug("AttachFunction: attached function created with is_ready=false",
 			zap.String("attached_function_id", attachedFunctionID.String()),
+			zap.String("output_collection_name", req.OutputCollectionName),
 			zap.String("name", req.Name))
 		return nil
 	})
@@ -207,7 +251,9 @@ func (s *Coordinator) AttachFunction(ctx context.Context, req *coordinatorpb.Att
 	}
 
 	return &coordinatorpb.AttachFunctionResponse{
-		Id: attachedFunctionID.String(),
+		AttachedFunction: &coordinatorpb.AttachedFunction{
+			Id: attachedFunctionID.String(),
+		},
 	}, nil
 }