enhance: size compaction result ID preallocation (#49919)

Derive compaction result segment ID reservations from the expected
output segment count instead of fixed ranges.

- Estimate output segment counts from total input size and target
segment size for single, sort, clustering, storage-version-upgrade,
force-merge, and schema-version compactions.
- Allocate result segment IDs and task metadata IDs in one contiguous
RootCoord block, with the result segment range fixed at the front and
metadata IDs consumed from the tail.
- Add CompactionView helpers for total input size and collection TTL so
scheduler submission paths share the same size-based allocation flow.
- Reserve result-segment IDs with
dataCoord.compaction.preAllocateIDExpansionFactor so result segment and
metadata ID preallocation share the existing compaction ID factor.
- Update milvus-proto to e90ee63ad3fc for snapshot export and external
restore APIs, and refresh affected service mocks.
- Validate nil or external force-merge and schema-version collections
before ID allocation, and keep compaction-local oversized allocation
rejection before RootCoord AllocN.
- Keep schema-version compaction on the current BumpSchemaVersion path
and preserve the frozen schema captured by the view.
- Cover allocation sizing, range layout, defensive returns, force-merge
target counts, schema-version submission, and oversized allocation
rejection.

See also: #49955

Signed-off-by: yangxuan <xuan.yang@zilliz.com>

---------

Signed-off-by: yangxuan <xuan.yang@zilliz.com>

internal/allocator/local_allocator.go

@@ -19,9 +19,24 @@ package allocator
 import (
 	"sync"
 
+	"github.com/cockroachdb/errors"
+
 	"github.com/milvus-io/milvus/pkg/v3/util/merr"
 )
 
+var errIDExhausted = errors.New("ID is exhausted")
+
+// NewIDExhaustedError returns a typed error that still carries the allocator
+// exhaustion signal for in-process control flow.
+func NewIDExhaustedError(start, end, count int64) error {
+	return merr.WrapErrServiceInternalErr(errIDExhausted, "ID is exhausted, start=%d, end=%d, count=%d", start, end, count)
+}
+
+// IsIDExhausted reports whether err came from allocator ID exhaustion.
+func IsIDExhausted(err error) bool {
+	return errors.Is(err, errIDExhausted)
+}
+
 // localAllocator implements the Interface.
 // It is constructed from a range of IDs.
 // Once all IDs are allocated, an error will be returned.
@@ -46,7 +61,7 @@ func (a *localAllocator) Alloc(count uint32) (int64, int64, error) {
 	a.mu.Lock()
 	defer a.mu.Unlock()
 	if a.idStart+cnt > a.idEnd {
-		return 0, 0, merr.WrapErrServiceInternalMsg("ID is exhausted, start=%d, end=%d, count=%d", a.idStart, a.idEnd, cnt)
+		return 0, 0, NewIDExhaustedError(a.idStart, a.idEnd, cnt)
 	}
 	start := a.idStart
 	a.idStart += cnt

internal/datacoord/compaction_l0_view.go

@@ -4,6 +4,7 @@ import (
 	"fmt"
 	"math"
 	"sort"
+	"time"
 
 	"github.com/samber/lo"
 
@@ -85,6 +86,17 @@ func (v *LevelZeroCompactionView) GetSegmentsView() []*SegmentView {
 	return v.l0Segments
 }
 
+func (v *LevelZeroCompactionView) GetTotalSize() float64 {
+	if v == nil {
+		return 0
+	}
+	return sumSegmentSize(v.l0Segments)
+}
+
+func (v *LevelZeroCompactionView) GetCollectionTTL() time.Duration {
+	return 0
+}
+
 // ForceTrigger triggers all qualified LevelZeroSegments according to views
 func (v *LevelZeroCompactionView) ForceTrigger() (CompactionView, string) {
 	sort.Slice(v.l0Segments, func(i, j int) bool {

internal/datacoord/compaction_policy_bump_schema_version.go

@@ -19,6 +19,7 @@ package datacoord
 import (
 	"context"
 	"fmt"
+	"time"
 
 	"go.uber.org/zap"
 	"google.golang.org/protobuf/proto"
@@ -170,6 +171,17 @@ func (v *BumpSchemaVersionView) GetSegmentsView() []*SegmentView {
 	return v.segments
 }
 
+func (v *BumpSchemaVersionView) GetTotalSize() float64 {
+	if v == nil {
+		return 0
+	}
+	return sumSegmentSize(v.segments)
+}
+
+func (v *BumpSchemaVersionView) GetCollectionTTL() time.Duration {
+	return 0
+}
+
 func (v *BumpSchemaVersionView) Append(segments ...*SegmentView) {
 	v.segments = append(v.segments, segments...)
 }

internal/datacoord/compaction_policy_clustering.go

@@ -350,6 +350,20 @@ func (v *ClusteringSegmentsView) GetSegmentsView() []*SegmentView {
 	return v.segments
 }
 
+func (v *ClusteringSegmentsView) GetTotalSize() float64 {
+	if v == nil {
+		return 0
+	}
+	return sumSegmentSize(v.segments)
+}
+
+func (v *ClusteringSegmentsView) GetCollectionTTL() time.Duration {
+	if v == nil {
+		return 0
+	}
+	return v.collectionTTL
+}
+
 func (v *ClusteringSegmentsView) Append(segments ...*SegmentView) {
 	if v.segments == nil {
 		v.segments = segments

internal/datacoord/compaction_policy_single.go

@@ -336,6 +336,20 @@ func (v *MixSegmentView) GetSegmentsView() []*SegmentView {
 	return v.segments
 }
 
+func (v *MixSegmentView) GetTotalSize() float64 {
+	if v == nil {
+		return 0
+	}
+	return sumSegmentSize(v.segments)
+}
+
+func (v *MixSegmentView) GetCollectionTTL() time.Duration {
+	if v == nil {
+		return 0
+	}
+	return v.collectionTTL
+}
+
 func (v *MixSegmentView) Append(segments ...*SegmentView) {
 	if v.segments == nil {
 		v.segments = segments

internal/datacoord/compaction_trigger.go

@@ -394,8 +394,14 @@ func (t *compactionTrigger) handleSignal(signal *compactionSignal) error {
 			}
 			totalRows, inputSegmentIDs := plan.A, plan.B
 
-			n := 11 * paramtable.Get().DataCoordCfg.CompactionPreAllocateIDExpansionFactor.GetAsInt64()
-			startID, endID, err := t.allocator.AllocN(n)
+			inputs := typeutil.NewSet[int64](inputSegmentIDs...)
+			totalSize := lo.SumBy(group.segments, func(s *SegmentInfo) int64 {
+				if inputs.Contain(s.GetID()) {
+					return s.getSegmentSize()
+				}
+				return 0
+			})
+			planID, preAllocatedSegmentIDs, err := allocCompactionPlanIDs(t.allocator, float64(totalSize), float64(expectedSize))
 			if err != nil {
 				log.Warn("fail to allocate id", zap.Error(err))
 				return err
@@ -403,24 +409,21 @@ func (t *compactionTrigger) handleSignal(signal *compactionSignal) error {
 			start := time.Now()
 			pts, _ := tsoutil.ParseTS(ct.startTime)
 			task := &datapb.CompactionTask{
-				PlanID:         startID,
-				TriggerID:      signal.id,
-				State:          datapb.CompactionTaskState_pipelining,
-				StartTime:      pts.Unix(),
-				Type:           datapb.CompactionType_MixCompaction,
-				CollectionTtl:  ct.collectionTTL.Nanoseconds(),
-				CollectionID:   group.collectionID,
-				PartitionID:    group.partitionID,
-				Channel:        group.channelName,
-				InputSegments:  inputSegmentIDs,
-				ResultSegments: []int64{},
-				TotalRows:      totalRows,
-				Schema:         coll.Schema,
-				MaxSize:        expectedSize,
-				PreAllocatedSegmentIDs: &datapb.IDRange{
-					Begin: startID + 1,
-					End:   endID,
-				},
+				PlanID:                 planID,
+				TriggerID:              signal.id,
+				State:                  datapb.CompactionTaskState_pipelining,
+				StartTime:              pts.Unix(),
+				Type:                   datapb.CompactionType_MixCompaction,
+				CollectionTtl:          ct.collectionTTL.Nanoseconds(),
+				CollectionID:           group.collectionID,
+				PartitionID:            group.partitionID,
+				Channel:                group.channelName,
+				InputSegments:          inputSegmentIDs,
+				ResultSegments:         []int64{},
+				TotalRows:              totalRows,
+				Schema:                 coll.Schema,
+				MaxSize:                expectedSize,
+				PreAllocatedSegmentIDs: preAllocatedSegmentIDs,
 			}
 			err = t.inspector.enqueueCompaction(task)
 			if err != nil {

internal/datacoord/compaction_trigger_test.go

@@ -528,6 +528,10 @@ func Test_compactionTrigger_force(t *testing.T) {
 	if err != nil {
 		panic(err)
 	}
+	preAllocateIDExpansionFactor := paramtable.Get().DataCoordCfg.CompactionPreAllocateIDExpansionFactor.GetAsInt64()
+	preAllocatedSegmentIDBegin := int64(100)
+	preAllocatedSegmentIDEnd := preAllocatedSegmentIDBegin + preAllocateIDExpansionFactor
+	preAllocatedLogIDEnd := preAllocatedSegmentIDBegin + 4*preAllocateIDExpansionFactor
 
 	tests := []struct {
 		name         string
@@ -576,7 +580,7 @@ func Test_compactionTrigger_force(t *testing.T) {
 			},
 			[]*datapb.CompactionPlan{
 				{
-					PlanID: 100,
+					PlanID: preAllocatedSegmentIDEnd,
 					SegmentBinlogs: []*datapb.CompactionSegmentBinlogs{
 						{
 							SegmentID: 1,
@@ -629,8 +633,8 @@ func Test_compactionTrigger_force(t *testing.T) {
 					Channel:                "ch1",
 					TotalRows:              200,
 					Schema:                 schema,
-					PreAllocatedSegmentIDs: &datapb.IDRange{Begin: 101, End: 200},
-					PreAllocatedLogIDs:     &datapb.IDRange{Begin: 100, End: 200},
+					PreAllocatedSegmentIDs: &datapb.IDRange{Begin: preAllocatedSegmentIDBegin, End: preAllocatedSegmentIDEnd},
+					PreAllocatedLogIDs:     &datapb.IDRange{Begin: preAllocatedSegmentIDBegin, End: preAllocatedLogIDEnd},
 					MaxSize:                1073741824,
 					SlotUsage:              paramtable.Get().DataCoordCfg.MixCompactionSlotUsage.GetAsInt64(),
 					JsonParams:             params,
@@ -2575,6 +2579,55 @@ func (s *CompactionTriggerSuite) TestHandleSignal() {
 			isForce:      true,
 		})
 	})
+
+	s.Run("force compaction preallocates segment IDs from input size", func() {
+		defer s.SetupTest()
+		pt := paramtable.Get()
+		pt.Save(pt.DataCoordCfg.IndexBasedCompaction.Key, "false")
+		defer pt.Reset(pt.DataCoordCfg.IndexBasedCompaction.Key)
+		pt.Save(pt.DataCoordCfg.CompactionPreAllocateIDExpansionFactor.Key, "1")
+		defer pt.Reset(pt.DataCoordCfg.CompactionPreAllocateIDExpansionFactor.Key)
+		pt.Save(pt.DataCoordCfg.SegmentMaxSize.Key, "100")
+		defer pt.Reset(pt.DataCoordCfg.SegmentMaxSize.Key)
+
+		const mb = 1024 * 1024
+		s.meta.segments.segments[1].Binlogs[0].Binlogs[0].MemorySize = 250 * mb
+
+		tr := s.tr
+		handler := NewNMockHandler(s.T())
+		handler.EXPECT().GetCollection(mock.Anything, s.collectionID).
+			Return(&collectionInfo{
+				ID: s.collectionID,
+				Schema: &schemapb.CollectionSchema{
+					Fields: []*schemapb.FieldSchema{
+						{FieldID: s.vecFieldID, DataType: schemapb.DataType_FloatVector},
+					},
+				},
+			}, nil).Once()
+		tr.handler = handler
+
+		const (
+			startID = int64(20000)
+			planID  = int64(20003)
+			endID   = int64(20004)
+		)
+		s.allocator.EXPECT().AllocN(int64(4)).Return(startID, endID, nil).Once()
+		s.inspector.EXPECT().enqueueCompaction(mock.MatchedBy(func(task *datapb.CompactionTask) bool {
+			s.EqualValues(planID, task.GetPlanID())
+			s.ElementsMatch([]int64{1}, task.GetInputSegments())
+			s.EqualValues(startID, task.GetPreAllocatedSegmentIDs().GetBegin())
+			s.EqualValues(planID, task.GetPreAllocatedSegmentIDs().GetEnd())
+			return true
+		})).Return(nil).Once()
+
+		err := tr.handleSignal(NewCompactionSignal().
+			WithCollectionID(s.collectionID).
+			WithPartitionID(s.partitionID).
+			WithChannel(s.channel).
+			WithSegmentIDs(1).
+			WithIsForce(true))
+		s.NoError(err)
+	})
 }
 
 func TestIsCollectionAutoCompactionEnabledExternal(t *testing.T) {

internal/datacoord/compaction_trigger_v2.go

@@ -19,6 +19,7 @@ package datacoord
 import (
 	"context"
 	"fmt"
+	"math"
 	"sync"
 	"time"
 
@@ -443,11 +444,6 @@ func (m *CompactionTriggerManager) SubmitL0ViewToScheduler(ctx context.Context,
 
 func (m *CompactionTriggerManager) SubmitClusteringViewToScheduler(ctx context.Context, view CompactionView) {
 	log := log.Ctx(ctx).With(zap.String("view", view.String()))
-	taskID, _, err := m.allocator.AllocN(2)
-	if err != nil {
-		log.Warn("Failed to submit compaction view to scheduler because allocate id fail", zap.Error(err))
-		return
-	}
 	collection, err := m.handler.GetCollection(ctx, view.GetGroupLabel().CollectionID)
 	if err != nil {
 		log.Warn("Failed to submit compaction view to scheduler because get collection fail", zap.Error(err))
@@ -469,37 +465,35 @@ func (m *CompactionTriggerManager) SubmitClusteringViewToScheduler(ctx context.C
 		return
 	}
 
-	resultSegmentNum := (totalRows/preferSegmentRows + 1) * 2
-	n := resultSegmentNum * paramtable.Get().DataCoordCfg.CompactionPreAllocateIDExpansionFactor.GetAsInt64()
-	start, end, err := m.allocator.AllocN(n)
+	totalSize := view.GetTotalSize()
+	preferSegmentSize := float64(expectedSegmentSize) * paramtable.Get().DataCoordCfg.ClusteringCompactionPreferSegmentSizeRatio.GetAsFloat()
+	planID, analyzeTaskID, preAllocatedSegmentIDs, err := allocClusteringCompactionPlanIDs(m.allocator, totalSize, preferSegmentSize)
 	if err != nil {
-		log.Warn("pre-allocate result segments failed", zap.String("view", view.String()), zap.Error(err))
+		log.Warn("failed to pre-allocate result segment IDs", zap.Error(err))
 		return
 	}
+	now := time.Now().Unix()
 	task := &datapb.CompactionTask{
-		PlanID:             taskID,
-		TriggerID:          view.(*ClusteringSegmentsView).triggerID,
-		State:              datapb.CompactionTaskState_pipelining,
-		StartTime:          time.Now().Unix(),
-		CollectionTtl:      view.(*ClusteringSegmentsView).collectionTTL.Nanoseconds(),
-		Type:               datapb.CompactionType_ClusteringCompaction,
-		CollectionID:       view.GetGroupLabel().CollectionID,
-		PartitionID:        view.GetGroupLabel().PartitionID,
-		Channel:            view.GetGroupLabel().Channel,
-		Schema:             collection.Schema,
-		ClusteringKeyField: view.(*ClusteringSegmentsView).clusteringKeyField,
-		InputSegments:      lo.Map(view.GetSegmentsView(), func(segmentView *SegmentView, _ int) int64 { return segmentView.ID }),
-		ResultSegments:     []int64{},
-		MaxSegmentRows:     maxSegmentRows,
-		PreferSegmentRows:  preferSegmentRows,
-		TotalRows:          totalRows,
-		AnalyzeTaskID:      taskID + 1,
-		LastStateStartTime: time.Now().Unix(),
-		PreAllocatedSegmentIDs: &datapb.IDRange{
-			Begin: start,
-			End:   end,
-		},
-		MaxSize: expectedSegmentSize,
+		PlanID:                 planID,
+		TriggerID:              view.(*ClusteringSegmentsView).triggerID,
+		State:                  datapb.CompactionTaskState_pipelining,
+		StartTime:              now,
+		CollectionTtl:          view.GetCollectionTTL().Nanoseconds(),
+		Type:                   datapb.CompactionType_ClusteringCompaction,
+		CollectionID:           view.GetGroupLabel().CollectionID,
+		PartitionID:            view.GetGroupLabel().PartitionID,
+		Channel:                view.GetGroupLabel().Channel,
+		Schema:                 collection.Schema,
+		ClusteringKeyField:     view.(*ClusteringSegmentsView).clusteringKeyField,
+		InputSegments:          lo.Map(view.GetSegmentsView(), func(segmentView *SegmentView, _ int) int64 { return segmentView.ID }),
+		ResultSegments:         []int64{},
+		MaxSegmentRows:         maxSegmentRows,
+		PreferSegmentRows:      preferSegmentRows,
+		TotalRows:              totalRows,
+		AnalyzeTaskID:          analyzeTaskID,
+		LastStateStartTime:     now,
+		PreAllocatedSegmentIDs: preAllocatedSegmentIDs,
+		MaxSize:                expectedSegmentSize,
 	}
 	err = m.inspector.enqueueCompaction(task)
 	if err != nil {
@@ -517,16 +511,7 @@ func (m *CompactionTriggerManager) SubmitClusteringViewToScheduler(ctx context.C
 }
 
 func (m *CompactionTriggerManager) SubmitSingleViewToScheduler(ctx context.Context, view CompactionView, triggerType CompactionTriggerType) {
-	// single view is definitely one-one mapping
 	log := log.Ctx(ctx).With(zap.String("trigger type", triggerType.String()), zap.String("view", view.String()))
-	// TODO[GOOSE], 11 = 1 planID + 10 segmentID, this is a hack need to be removed.
-	// Any plan that output segment number greater than 10 will be marked as invalid plan for now.
-	n := 11 * paramtable.Get().DataCoordCfg.CompactionPreAllocateIDExpansionFactor.GetAsInt64()
-	startID, endID, err := m.allocator.AllocN(n)
-	if err != nil {
-		log.Warn("Failed to submit compaction view to scheduler because allocate id fail", zap.Error(err))
-		return
-	}
 
 	collection, err := m.handler.GetCollection(ctx, view.GetGroupLabel().CollectionID)
 	if err != nil {
@@ -541,32 +526,38 @@ func (m *CompactionTriggerManager) SubmitSingleViewToScheduler(ctx context.Conte
 		log.Info("skip submitting single compaction for external collection", zap.Int64("collectionID", collection.ID))
 		return
 	}
+
+	expectedSize := getExpectedSegmentSize(m.meta, collection.ID, collection.Schema)
+	totalSize := view.GetTotalSize()
+	planID, preAllocatedSegmentIDs, err := allocCompactionPlanIDs(m.allocator, totalSize, float64(expectedSize))
+	if err != nil {
+		log.Warn("Failed to submit compaction view to scheduler because allocate id fail", zap.Error(err))
+		return
+	}
+
 	var totalRows int64 = 0
 	for _, s := range view.GetSegmentsView() {
 		totalRows += s.NumOfRows
 	}
 
-	expectedSize := getExpectedSegmentSize(m.meta, collection.ID, collection.Schema)
+	now := time.Now().Unix()
 	task := &datapb.CompactionTask{
-		PlanID:             startID,
-		TriggerID:          view.(*MixSegmentView).triggerID,
-		State:              datapb.CompactionTaskState_pipelining,
-		StartTime:          time.Now().Unix(),
-		CollectionTtl:      view.(*MixSegmentView).collectionTTL.Nanoseconds(),
-		Type:               triggerType.GetCompactionType(),
-		CollectionID:       view.GetGroupLabel().CollectionID,
-		PartitionID:        view.GetGroupLabel().PartitionID,
-		Channel:            view.GetGroupLabel().Channel,
-		Schema:             collection.Schema,
-		InputSegments:      lo.Map(view.GetSegmentsView(), func(segmentView *SegmentView, _ int) int64 { return segmentView.ID }),
-		ResultSegments:     []int64{},
-		TotalRows:          totalRows,
-		LastStateStartTime: time.Now().Unix(),
-		MaxSize:            expectedSize,
-		PreAllocatedSegmentIDs: &datapb.IDRange{
-			Begin: startID + 1,
-			End:   endID,
-		},
+		PlanID:                 planID,
+		TriggerID:              view.(*MixSegmentView).triggerID,
+		State:                  datapb.CompactionTaskState_pipelining,
+		StartTime:              now,
+		CollectionTtl:          view.GetCollectionTTL().Nanoseconds(),
+		Type:                   triggerType.GetCompactionType(),
+		CollectionID:           view.GetGroupLabel().CollectionID,
+		PartitionID:            view.GetGroupLabel().PartitionID,
+		Channel:                view.GetGroupLabel().Channel,
+		Schema:                 collection.Schema,
+		InputSegments:          lo.Map(view.GetSegmentsView(), func(segmentView *SegmentView, _ int) int64 { return segmentView.ID }),
+		ResultSegments:         []int64{},
+		TotalRows:              totalRows,
+		LastStateStartTime:     now,
+		MaxSize:                expectedSize,
+		PreAllocatedSegmentIDs: preAllocatedSegmentIDs,
 	}
 	err = m.inspector.enqueueCompaction(task)
 	if err != nil {
@@ -585,51 +576,149 @@ func (m *CompactionTriggerManager) SubmitSingleViewToScheduler(ctx context.Conte
 	)
 }
 
+func allocCompactionPlanIDs(allocator allocator.Allocator, totalSize float64, preferredSize float64) (int64, *datapb.IDRange, error) {
+	segmentIDCount := estimateResultSegmentCount(totalSize, preferredSize)
+	// Reserve one metadata ID because this helper returns one task ID: PlanID.
+	block, err := createCompactionIDBlock(allocator, segmentIDCount, 1)
+	if err != nil {
+		return 0, nil, err
+	}
+	planID, err := block.take()
+	if err != nil {
+		return 0, nil, err
+	}
+	return planID, block.segmentIDRange(), nil
+}
+
+func allocClusteringCompactionPlanIDs(allocator allocator.Allocator, totalSize float64, preferredSize float64) (int64, int64, *datapb.IDRange, error) {
+	segmentIDCount := estimateResultSegmentCount(totalSize, preferredSize)
+	// Reserve two metadata IDs because this helper returns PlanID and AnalyzeTaskID.
+	block, err := createCompactionIDBlock(allocator, segmentIDCount, 2)
+	if err != nil {
+		return 0, 0, nil, err
+	}
+	planID, err := block.take()
+	if err != nil {
+		return 0, 0, nil, err
+	}
+	analyzeTaskID, err := block.take()
+	if err != nil {
+		return 0, 0, nil, err
+	}
+	return planID, analyzeTaskID, block.segmentIDRange(), nil
+}
+
+// createCompactionIDBlock allocates one contiguous block where the
+// first segmentIDCount IDs are result segment IDs and the rest are task metadata IDs.
+func createCompactionIDBlock(allocator allocator.Allocator, segmentIDCount int64, metaIDCount int64) (*compactionIDBlock, error) {
+	segmentIDCount = max(segmentIDCount, 1)
+	expansionFactor := paramtable.Get().DataCoordCfg.CompactionPreAllocateIDExpansionFactor.GetAsInt64()
+	if expansionFactor <= 0 {
+		return nil, merr.WrapErrParameterInvalidMsg("compaction pre-allocate ID expansion factor must be positive: %d", expansionFactor)
+	}
+	maxBatchSize := int64(math.MaxUint32)
+	if metaIDCount < 0 || metaIDCount > maxBatchSize || segmentIDCount > (maxBatchSize-metaIDCount)/expansionFactor {
+		return nil, merr.WrapErrServiceInternal(
+			fmt.Sprintf(
+				"compaction too large to allocate IDs in a single batch: segment ID count %d with expansion factor %d and %d metadata IDs exceeds max batch size %d",
+				segmentIDCount, expansionFactor, metaIDCount, maxBatchSize,
+			),
+		)
+	}
+	segmentIDCount = segmentIDCount * expansionFactor
+	n := metaIDCount + segmentIDCount
+	startID, endID, err := allocator.AllocN(n)
+	if err != nil {
+		return nil, err
+	}
+	if endID-startID != n {
+		return nil, merr.WrapErrServiceInternal(
+			fmt.Sprintf("compaction ID allocation range mismatch: expected %d IDs, got %d IDs", n, endID-startID),
+		)
+	}
+	segmentIDEnd := startID + segmentIDCount
+	return &compactionIDBlock{
+		segments: &datapb.IDRange{Begin: startID, End: segmentIDEnd},
+		next:     segmentIDEnd,
+		end:      endID,
+	}, nil
+}
+
+// compactionIDBlock splits one contiguous RootCoord ID allocation into a fixed
+// result segment ID range followed by metadata IDs consumed by take().
+type compactionIDBlock struct {
+	segments *datapb.IDRange
+	next     int64
+	end      int64
+}
+
+func (b *compactionIDBlock) take() (int64, error) {
+	if b.next >= b.end {
+		return 0, merr.WrapErrServiceInternal("compaction metadata ID range exhausted")
+	}
+	id := b.next
+	b.next++
+	return id, nil
+}
+
+func (b *compactionIDBlock) segmentIDRange() *datapb.IDRange {
+	return b.segments
+}
+
+func estimateResultSegmentCount(totalSize float64, targetSize float64) int64 {
+	if targetSize <= 0 {
+		return 1
+	}
+	if totalSize <= 0 {
+		return 1
+	}
+	return max(int64(math.Ceil(totalSize/targetSize)), 1)
+}
+
 func (m *CompactionTriggerManager) SubmitForceMergeViewToScheduler(ctx context.Context, view CompactionView) {
 	log := log.Ctx(ctx).With(zap.String("view", view.String()))
 
-	taskID, err := m.allocator.AllocID(ctx)
-	if err != nil {
-		log.Warn("Failed to allocate task ID", zap.Error(err))
-		return
-	}
-
 	collection, err := m.handler.GetCollection(ctx, view.GetGroupLabel().CollectionID)
 	if err != nil {
 		log.Warn("Failed to get collection", zap.Error(err))
 		return
 	}
+	if collection == nil {
+		log.Warn("collection not found when submitting force merge compaction", zap.Int64("collectionID", view.GetGroupLabel().CollectionID))
+		return
+	}
+	if collection.IsExternal() {
+		log.Info("skip submitting force merge compaction for external collection", zap.Int64("collectionID", collection.ID))
+		return
+	}
 
 	totalRows := lo.SumBy(view.GetSegmentsView(), func(v *SegmentView) int64 { return v.NumOfRows })
 
-	targetCount := view.(*ForceMergeSegmentView).targetSegmentCount
-	n := targetCount * paramtable.Get().DataCoordCfg.CompactionPreAllocateIDExpansionFactor.GetAsInt64()
-	startID, endID, err := m.allocator.AllocN(n)
+	totalSize := view.GetTotalSize()
+	planID, preAllocatedSegmentIDs, err := allocCompactionPlanIDs(m.allocator, totalSize, view.(*ForceMergeSegmentView).GetTargetSegmentSize())
 	if err != nil {
 		log.Warn("Failed to submit compaction view to scheduler because allocate id fail", zap.Error(err))
 		return
 	}
 
+	now := time.Now().Unix()
 	task := &datapb.CompactionTask{
-		PlanID:             taskID,
-		TriggerID:          view.GetTriggerID(),
-		State:              datapb.CompactionTaskState_pipelining,
-		StartTime:          time.Now().Unix(),
-		CollectionTtl:      view.(*ForceMergeSegmentView).collectionTTL.Nanoseconds(),
-		Type:               datapb.CompactionType_MixCompaction,
-		CollectionID:       view.GetGroupLabel().CollectionID,
-		PartitionID:        view.GetGroupLabel().PartitionID,
-		Channel:            view.GetGroupLabel().Channel,
-		Schema:             collection.Schema,
-		InputSegments:      lo.Map(view.GetSegmentsView(), func(segmentView *SegmentView, _ int) int64 { return segmentView.ID }),
-		ResultSegments:     []int64{},
-		TotalRows:          totalRows,
-		LastStateStartTime: time.Now().Unix(),
-		MaxSize:            int64(view.(*ForceMergeSegmentView).targetSegmentSize),
-		PreAllocatedSegmentIDs: &datapb.IDRange{
-			Begin: startID + 1,
-			End:   endID,
-		},
+		PlanID:                 planID,
+		TriggerID:              view.GetTriggerID(),
+		State:                  datapb.CompactionTaskState_pipelining,
+		StartTime:              now,
+		CollectionTtl:          view.GetCollectionTTL().Nanoseconds(),
+		Type:                   datapb.CompactionType_MixCompaction,
+		CollectionID:           view.GetGroupLabel().CollectionID,
+		PartitionID:            view.GetGroupLabel().PartitionID,
+		Channel:                view.GetGroupLabel().Channel,
+		Schema:                 collection.Schema,
+		InputSegments:          lo.Map(view.GetSegmentsView(), func(segmentView *SegmentView, _ int) int64 { return segmentView.ID }),
+		ResultSegments:         []int64{},
+		TotalRows:              totalRows,
+		LastStateStartTime:     now,
+		MaxSize:                int64(view.(*ForceMergeSegmentView).GetTargetSegmentSize()),
+		PreAllocatedSegmentIDs: preAllocatedSegmentIDs,
 	}
 
 	err = m.inspector.enqueueCompaction(task)
@@ -639,7 +728,7 @@ func (m *CompactionTriggerManager) SubmitForceMergeViewToScheduler(ctx context.C
 	}
 
 	log.Info("Finish to submit force merge task",
-		zap.Int64("planID", taskID),
+		zap.Int64("planID", task.GetPlanID()),
 		zap.Int64("triggerID", task.GetTriggerID()),
 		zap.Int64("collectionID", task.GetCollectionID()),
 		zap.Int64("targetSize", task.GetMaxSize()),
@@ -654,11 +743,6 @@ func (m *CompactionTriggerManager) SubmitBumpSchemaVersionViewToScheduler(ctx co
 			zap.String("actualType", fmt.Sprintf("%T", view)))
 		return
 	}
-	planID, endID, err := m.allocator.AllocN(2)
-	if err != nil {
-		log.Warn("Failed to submit compaction view to scheduler because allocate id fail", zap.Error(err))
-		return
-	}
 	collection, err := m.handler.GetCollection(ctx, view.GetGroupLabel().CollectionID)
 	if err != nil {
 		log.Warn("Failed to submit compaction view to scheduler because get collection fail", zap.Error(err))
@@ -682,26 +766,29 @@ func (m *CompactionTriggerManager) SubmitBumpSchemaVersionViewToScheduler(ctx co
 		totalRows += s.NumOfRows
 	}
 	expectedSize := getExpectedSegmentSize(m.meta, collection.ID, collection.Schema)
+	planID, preAllocatedSegmentIDs, err := allocCompactionPlanIDs(m.allocator, view.GetTotalSize(), float64(expectedSize))
+	if err != nil {
+		log.Warn("Failed to submit compaction view to scheduler because allocate id fail", zap.Error(err))
+		return
+	}
+	now := time.Now().Unix()
 	task := &datapb.CompactionTask{
-		PlanID:             planID,
-		TriggerID:          bumpView.triggerID,
-		State:              datapb.CompactionTaskState_pipelining,
-		StartTime:          time.Now().Unix(),
-		CollectionTtl:      collectionTTL.Nanoseconds(),
-		Type:               datapb.CompactionType_BumpSchemaVersionCompaction,
-		CollectionID:       view.GetGroupLabel().CollectionID,
-		PartitionID:        view.GetGroupLabel().PartitionID,
-		Channel:            view.GetGroupLabel().Channel,
-		Schema:             bumpView.schema,
-		InputSegments:      lo.Map(view.GetSegmentsView(), func(segmentView *SegmentView, _ int) int64 { return segmentView.ID }),
-		ResultSegments:     []int64{},
-		TotalRows:          totalRows,
-		LastStateStartTime: time.Now().Unix(),
-		MaxSize:            expectedSize,
-		PreAllocatedSegmentIDs: &datapb.IDRange{
-			Begin: planID + 1,
-			End:   endID,
-		},
+		PlanID:                 planID,
+		TriggerID:              bumpView.triggerID,
+		State:                  datapb.CompactionTaskState_pipelining,
+		StartTime:              now,
+		CollectionTtl:          collectionTTL.Nanoseconds(),
+		Type:                   datapb.CompactionType_BumpSchemaVersionCompaction,
+		CollectionID:           view.GetGroupLabel().CollectionID,
+		PartitionID:            view.GetGroupLabel().PartitionID,
+		Channel:                view.GetGroupLabel().Channel,
+		Schema:                 bumpView.schema,
+		InputSegments:          lo.Map(view.GetSegmentsView(), func(segmentView *SegmentView, _ int) int64 { return segmentView.ID }),
+		ResultSegments:         []int64{},
+		TotalRows:              totalRows,
+		LastStateStartTime:     now,
+		MaxSize:                expectedSize,
+		PreAllocatedSegmentIDs: preAllocatedSegmentIDs,
 	}
 	err = m.inspector.enqueueCompaction(task)
 	if err != nil {

internal/datacoord/compaction_trigger_v2_test.go

@@ -2,7 +2,9 @@ package datacoord
 
 import (
 	"context"
+	"math"
 	"strconv"
+	"strings"
 	"testing"
 	"time"
 
@@ -28,6 +30,134 @@ func TestCompactionTriggerManagerSuite(t *testing.T) {
 	suite.Run(t, new(CompactionTriggerManagerSuite))
 }
 
+func TestEstimateResultSegmentCount(t *testing.T) {
+	tests := []struct {
+		name       string
+		totalSize  float64
+		targetSize float64
+		want       int64
+	}{
+		{name: "exact multiple", totalSize: 300, targetSize: 100, want: 3},
+		{name: "fractional rounds up", totalSize: 301, targetSize: 100, want: 4},
+		{name: "zero total size", totalSize: 0, targetSize: 100, want: 1},
+		{name: "zero target size", totalSize: 100, targetSize: 0, want: 1},
+		{name: "large ratio", totalSize: 10_000, targetSize: 64, want: 157},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			got := estimateResultSegmentCount(test.totalSize, test.targetSize)
+			if got != test.want {
+				t.Fatalf("estimateResultSegmentCount(%v, %v) = %d, want %d", test.totalSize, test.targetSize, got, test.want)
+			}
+		})
+	}
+}
+
+func TestCompactionIDBlockTakeFromMetadataTail(t *testing.T) {
+	t.Run("returns fixed segment ID range", func(t *testing.T) {
+		block := &compactionIDBlock{
+			segments: &datapb.IDRange{Begin: 100, End: 103},
+			next:     103,
+			end:      104,
+		}
+
+		planID, err := block.take()
+		if err != nil {
+			t.Fatalf("unexpected error: %v", err)
+		}
+		segmentIDRange := block.segmentIDRange()
+
+		if planID != 103 {
+			t.Fatalf("planID = %d, want 103", planID)
+		}
+		if segmentIDRange.GetBegin() != 100 || segmentIDRange.GetEnd() != 103 {
+			t.Fatalf("segment ID range = [%d, %d), want [100, 103)", segmentIDRange.GetBegin(), segmentIDRange.GetEnd())
+		}
+	})
+
+	t.Run("rejects future metadata over-consumption", func(t *testing.T) {
+		block := &compactionIDBlock{
+			segments: &datapb.IDRange{Begin: 100, End: 103},
+			next:     103,
+			end:      104,
+		}
+
+		_, err := block.take()
+		if err != nil {
+			t.Fatalf("unexpected error: %v", err)
+		}
+		_, err = block.take()
+		if err == nil {
+			t.Fatal("expected metadata ID exhaustion error")
+		}
+		segmentIDRange := block.segmentIDRange()
+		if segmentIDRange.GetBegin() != 100 || segmentIDRange.GetEnd() != 103 {
+			t.Fatalf("segment ID range = [%d, %d), want [100, 103)", segmentIDRange.GetBegin(), segmentIDRange.GetEnd())
+		}
+	})
+}
+
+func TestCreateCompactionIDBlockRejectsTooLargeBatch(t *testing.T) {
+	pt := paramtable.Get()
+	pt.Save(pt.DataCoordCfg.CompactionPreAllocateIDExpansionFactor.Key, "10")
+	defer pt.Reset(pt.DataCoordCfg.CompactionPreAllocateIDExpansionFactor.Key)
+
+	mockAlloc := allocator.NewMockAllocator(t)
+	_, err := createCompactionIDBlock(mockAlloc, math.MaxUint32/10+1, 1)
+	if err == nil {
+		t.Fatal("expected too-large allocation error")
+	}
+	if !strings.Contains(err.Error(), "compaction too large to allocate IDs in a single batch") {
+		t.Fatalf("error = %q, want too-large allocation message", err.Error())
+	}
+}
+
+func TestCreateCompactionIDBlockUsesIDExpansionFactor(t *testing.T) {
+	pt := paramtable.Get()
+	pt.Save(pt.DataCoordCfg.CompactionPreAllocateIDExpansionFactor.Key, "2")
+	defer pt.Reset(pt.DataCoordCfg.CompactionPreAllocateIDExpansionFactor.Key)
+
+	mockAlloc := allocator.NewMockAllocator(t)
+	mockAlloc.EXPECT().AllocN(int64(7)).Return(int64(100), int64(107), nil).Once()
+
+	block, err := createCompactionIDBlock(mockAlloc, 3, 1)
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+	segmentIDRange := block.segmentIDRange()
+	if segmentIDRange.GetBegin() != 100 || segmentIDRange.GetEnd() != 106 {
+		t.Fatalf("segment ID range = [%d, %d), want [100, 106)", segmentIDRange.GetBegin(), segmentIDRange.GetEnd())
+	}
+}
+
+func TestCompactionViewsExposeTotalSizeAndCollectionTTL(t *testing.T) {
+	ttl := 3 * time.Hour
+	segments := []*SegmentView{{ID: 1, Size: 10}, {ID: 2, Size: 2.5}}
+	tests := []struct {
+		name    string
+		view    CompactionView
+		wantTTL time.Duration
+	}{
+		{name: "single", view: &MixSegmentView{segments: segments, collectionTTL: ttl}, wantTTL: ttl},
+		{name: "clustering", view: &ClusteringSegmentsView{segments: segments, collectionTTL: ttl}, wantTTL: ttl},
+		{name: "force merge", view: &ForceMergeSegmentView{segments: segments, collectionTTL: ttl}, wantTTL: ttl},
+		{name: "level zero", view: &LevelZeroCompactionView{l0Segments: segments}},
+		{name: "bump schema version", view: &BumpSchemaVersionView{segments: segments}},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			if got := test.view.GetTotalSize(); got != 12.5 {
+				t.Fatalf("GetTotalSize() = %v, want 12.5", got)
+			}
+			if got := test.view.GetCollectionTTL(); got != test.wantTTL {
+				t.Fatalf("GetCollectionTTL() = %v, want %v", got, test.wantTTL)
+			}
+		})
+	}
+}
+
 // testCompactionPolicy is a minimal CompactionPolicy stub for handleTicker tests.
 type testCompactionPolicy struct {
 	enabled       bool
@@ -58,6 +188,8 @@ func (stubDispatchableView) Trigger() (CompactionView, string)           { retur
 func (stubDispatchableView) ForceTrigger() (CompactionView, string)      { return nil, "" }
 func (stubDispatchableView) ForceTriggerAll() ([]CompactionView, string) { return nil, "" }
 func (stubDispatchableView) GetTriggerID() int64                         { return 0 }
+func (stubDispatchableView) GetTotalSize() float64                       { return 0 }
+func (stubDispatchableView) GetCollectionTTL() time.Duration             { return 0 }
 
 type CompactionTriggerManagerSuite struct {
 	suite.Suite
@@ -425,7 +557,380 @@ func (s *CompactionTriggerManagerSuite) TestManualTriggerInvalidParams() {
 	s.Equal(int64(0), triggerID)
 }
 
+func (s *CompactionTriggerManagerSuite) TestSubmitSingleViewToScheduler() {
+	makeMixView := func(segments []*SegmentView) *MixSegmentView {
+		return &MixSegmentView{
+			label:         s.testLabel,
+			segments:      segments,
+			collectionTTL: 100,
+			triggerID:     1001,
+		}
+	}
+
+	s.Run("mix compaction allocates estimated result segments", func() {
+		s.SetupTest()
+		pt := paramtable.Get()
+		pt.Save(pt.DataCoordCfg.CompactionPreAllocateIDExpansionFactor.Key, "1")
+		defer pt.Reset(pt.DataCoordCfg.CompactionPreAllocateIDExpansionFactor.Key)
+		pt.Save(pt.DataCoordCfg.SegmentMaxSize.Key, "100")
+		defer pt.Reset(pt.DataCoordCfg.SegmentMaxSize.Key)
+
+		s.meta.indexMeta = &indexMeta{indexes: make(map[UniqueID]map[UniqueID]*model.Index)}
+		collectionSchema := &schemapb.CollectionSchema{
+			Name: "test_coll",
+			Fields: []*schemapb.FieldSchema{
+				{FieldID: 1, Name: "pk", DataType: schemapb.DataType_Int64, IsPrimaryKey: true},
+				{FieldID: 100, Name: "vec", DataType: schemapb.DataType_FloatVector},
+			},
+		}
+		handler := NewNMockHandler(s.T())
+		handler.EXPECT().GetCollection(mock.Anything, s.testLabel.CollectionID).
+			Return(&collectionInfo{ID: s.testLabel.CollectionID, Schema: collectionSchema}, nil).Once()
+		s.triggerManager.handler = handler
+
+		const (
+			startID = int64(500)
+			endID   = int64(504)
+			planID  = int64(503)
+		)
+		s.mockAlloc.EXPECT().AllocN(int64(4)).Return(startID, endID, nil).Once()
+		s.inspector.EXPECT().enqueueCompaction(mock.Anything).
+			RunAndReturn(func(task *datapb.CompactionTask) error {
+				s.EqualValues(planID, task.GetPlanID())
+				s.EqualValues(1001, task.GetTriggerID())
+				s.Equal(datapb.CompactionType_MixCompaction, task.GetType())
+				s.Equal(s.testLabel.CollectionID, task.GetCollectionID())
+				s.Equal(s.testLabel.PartitionID, task.GetPartitionID())
+				s.Equal(s.testLabel.Channel, task.GetChannel())
+				s.Equal(collectionSchema, task.GetSchema())
+				s.ElementsMatch([]int64{200, 201}, task.GetInputSegments())
+				s.EqualValues(300, task.GetTotalRows())
+				s.Equal(&datapb.IDRange{Begin: startID, End: planID}, task.GetPreAllocatedSegmentIDs())
+				s.Equal(task.GetStartTime(), task.GetLastStateStartTime())
+				return nil
+			}).Return(nil).Once()
+
+		view := makeMixView([]*SegmentView{
+			{ID: 200, label: s.testLabel, NumOfRows: 100, Size: 150 * 1024 * 1024},
+			{ID: 201, label: s.testLabel, NumOfRows: 200, Size: 150 * 1024 * 1024},
+		})
+		s.triggerManager.SubmitSingleViewToScheduler(context.Background(), view, TriggerTypeSingle)
+	})
+
+	s.Run("sort compaction allocates estimated result segments", func() {
+		s.SetupTest()
+		pt := paramtable.Get()
+		pt.Save(pt.DataCoordCfg.CompactionPreAllocateIDExpansionFactor.Key, "1")
+		defer pt.Reset(pt.DataCoordCfg.CompactionPreAllocateIDExpansionFactor.Key)
+		pt.Save(pt.DataCoordCfg.SegmentMaxSize.Key, "100")
+		defer pt.Reset(pt.DataCoordCfg.SegmentMaxSize.Key)
+
+		s.meta.indexMeta = &indexMeta{indexes: make(map[UniqueID]map[UniqueID]*model.Index)}
+		collectionSchema := &schemapb.CollectionSchema{
+			Name: "test_coll",
+			Fields: []*schemapb.FieldSchema{
+				{FieldID: 1, Name: "pk", DataType: schemapb.DataType_Int64, IsPrimaryKey: true},
+				{FieldID: 100, Name: "vec", DataType: schemapb.DataType_FloatVector},
+			},
+		}
+		handler := NewNMockHandler(s.T())
+		handler.EXPECT().GetCollection(mock.Anything, s.testLabel.CollectionID).
+			Return(&collectionInfo{ID: s.testLabel.CollectionID, Schema: collectionSchema}, nil).Once()
+		s.triggerManager.handler = handler
+
+		const (
+			startID = int64(600)
+			endID   = int64(604)
+			planID  = int64(603)
+		)
+		s.mockAlloc.EXPECT().AllocN(int64(4)).Return(startID, endID, nil).Once()
+		s.inspector.EXPECT().enqueueCompaction(mock.Anything).
+			RunAndReturn(func(task *datapb.CompactionTask) error {
+				s.Equal(datapb.CompactionType_SortCompaction, task.GetType())
+				s.Equal(&datapb.IDRange{Begin: startID, End: planID}, task.GetPreAllocatedSegmentIDs())
+				return nil
+			}).Return(nil).Once()
+
+		view := makeMixView([]*SegmentView{
+			{ID: 200, label: s.testLabel, NumOfRows: 100, Size: 300 * 1024 * 1024},
+		})
+		s.triggerManager.SubmitSingleViewToScheduler(context.Background(), view, TriggerTypeSort)
+	})
+
+	s.Run("storage version upgrade keeps size based estimation", func() {
+		s.SetupTest()
+		pt := paramtable.Get()
+		pt.Save(pt.DataCoordCfg.CompactionPreAllocateIDExpansionFactor.Key, "1")
+		defer pt.Reset(pt.DataCoordCfg.CompactionPreAllocateIDExpansionFactor.Key)
+		pt.Save(pt.DataCoordCfg.SegmentMaxSize.Key, "100")
+		defer pt.Reset(pt.DataCoordCfg.SegmentMaxSize.Key)
+
+		s.meta.indexMeta = &indexMeta{indexes: make(map[UniqueID]map[UniqueID]*model.Index)}
+		collectionSchema := &schemapb.CollectionSchema{
+			Name: "test_coll",
+			Fields: []*schemapb.FieldSchema{
+				{FieldID: 1, Name: "pk", DataType: schemapb.DataType_Int64, IsPrimaryKey: true},
+				{FieldID: 100, Name: "vec", DataType: schemapb.DataType_FloatVector},
+			},
+		}
+		handler := NewNMockHandler(s.T())
+		handler.EXPECT().GetCollection(mock.Anything, s.testLabel.CollectionID).
+			Return(&collectionInfo{ID: s.testLabel.CollectionID, Schema: collectionSchema}, nil).Once()
+		s.triggerManager.handler = handler
+
+		const (
+			startID = int64(700)
+			endID   = int64(704)
+			planID  = int64(703)
+		)
+		s.mockAlloc.EXPECT().AllocN(int64(4)).Return(startID, endID, nil).Once()
+		s.inspector.EXPECT().enqueueCompaction(mock.Anything).
+			RunAndReturn(func(task *datapb.CompactionTask) error {
+				s.Equal(datapb.CompactionType_MixCompaction, task.GetType())
+				s.Equal(&datapb.IDRange{Begin: startID, End: planID}, task.GetPreAllocatedSegmentIDs())
+				return nil
+			}).Return(nil).Once()
+
+		view := makeMixView([]*SegmentView{
+			{ID: 200, label: s.testLabel, NumOfRows: 100, Size: 300 * 1024 * 1024},
+		})
+		s.triggerManager.SubmitSingleViewToScheduler(context.Background(), view, TriggerTypeStorageVersionUpgrade)
+	})
+}
+
+func (s *CompactionTriggerManagerSuite) TestSubmitClusteringViewToScheduler() {
+	s.SetupTest()
+	pt := paramtable.Get()
+	pt.Save(pt.DataCoordCfg.CompactionPreAllocateIDExpansionFactor.Key, "2")
+	defer pt.Reset(pt.DataCoordCfg.CompactionPreAllocateIDExpansionFactor.Key)
+	pt.Save(pt.DataCoordCfg.SegmentMaxSize.Key, "100")
+	defer pt.Reset(pt.DataCoordCfg.SegmentMaxSize.Key)
+	pt.Save(pt.DataCoordCfg.ClusteringCompactionPreferSegmentSizeRatio.Key, "1")
+	defer pt.Reset(pt.DataCoordCfg.ClusteringCompactionPreferSegmentSizeRatio.Key)
+
+	s.meta.indexMeta = &indexMeta{indexes: make(map[UniqueID]map[UniqueID]*model.Index)}
+	clusteringKey := &schemapb.FieldSchema{FieldID: 2, Name: "cluster_key", DataType: schemapb.DataType_Int64, IsClusteringKey: true}
+	collectionSchema := &schemapb.CollectionSchema{
+		Name: "test_coll",
+		Fields: []*schemapb.FieldSchema{
+			{FieldID: 1, Name: "pk", DataType: schemapb.DataType_Int64, IsPrimaryKey: true},
+			clusteringKey,
+			{FieldID: 100, Name: "vec", DataType: schemapb.DataType_FloatVector},
+		},
+	}
+	handler := NewNMockHandler(s.T())
+	handler.EXPECT().GetCollection(mock.Anything, s.testLabel.CollectionID).
+		Return(&collectionInfo{ID: s.testLabel.CollectionID, Schema: collectionSchema}, nil).Once()
+	s.triggerManager.handler = handler
+
+	const (
+		startID       = int64(800)
+		segmentIDEnd  = int64(806)
+		planID        = int64(806)
+		analyzeTaskID = int64(807)
+		endID         = int64(808)
+	)
+	s.mockAlloc.EXPECT().AllocN(int64(8)).Return(startID, endID, nil).Once()
+	s.inspector.EXPECT().enqueueCompaction(mock.Anything).
+		RunAndReturn(func(task *datapb.CompactionTask) error {
+			s.Equal(datapb.CompactionType_ClusteringCompaction, task.GetType())
+			s.EqualValues(planID, task.GetPlanID())
+			s.EqualValues(analyzeTaskID, task.GetAnalyzeTaskID())
+			s.Equal(&datapb.IDRange{Begin: startID, End: segmentIDEnd}, task.GetPreAllocatedSegmentIDs())
+			s.Equal(task.GetStartTime(), task.GetLastStateStartTime())
+			return nil
+		}).Return(nil).Once()
+
+	view := &ClusteringSegmentsView{
+		label:              s.testLabel,
+		segments:           []*SegmentView{{ID: 200, label: s.testLabel, NumOfRows: 100, Size: 150 * 1024 * 1024}, {ID: 201, label: s.testLabel, NumOfRows: 200, Size: 150 * 1024 * 1024}},
+		clusteringKeyField: clusteringKey,
+		collectionTTL:      100,
+		triggerID:          1001,
+	}
+	s.triggerManager.SubmitClusteringViewToScheduler(context.Background(), view)
+}
+
+func (s *CompactionTriggerManagerSuite) TestSubmitForceMergeViewToScheduler() {
+	s.SetupTest()
+	pt := paramtable.Get()
+	pt.Save(pt.DataCoordCfg.CompactionPreAllocateIDExpansionFactor.Key, "1")
+	defer pt.Reset(pt.DataCoordCfg.CompactionPreAllocateIDExpansionFactor.Key)
+
+	collectionSchema := &schemapb.CollectionSchema{
+		Name: "test_coll",
+		Fields: []*schemapb.FieldSchema{
+			{FieldID: 1, Name: "pk", DataType: schemapb.DataType_Int64, IsPrimaryKey: true},
+		},
+	}
+	handler := NewNMockHandler(s.T())
+	handler.EXPECT().GetCollection(mock.Anything, s.testLabel.CollectionID).
+		Return(&collectionInfo{ID: s.testLabel.CollectionID, Schema: collectionSchema}, nil).Once()
+	s.triggerManager.handler = handler
+
+	const (
+		startID      = int64(500)
+		segmentIDEnd = int64(502)
+		planID       = int64(502)
+		endID        = int64(503)
+	)
+	s.mockAlloc.EXPECT().AllocN(int64(3)).Return(startID, endID, nil).Once()
+	s.inspector.EXPECT().enqueueCompaction(mock.Anything).
+		RunAndReturn(func(task *datapb.CompactionTask) error {
+			s.EqualValues(planID, task.GetPlanID())
+			s.Equal(datapb.CompactionType_MixCompaction, task.GetType())
+			s.Equal(&datapb.IDRange{Begin: startID, End: segmentIDEnd}, task.GetPreAllocatedSegmentIDs())
+			return nil
+		}).Return(nil).Once()
+
+	view := &ForceMergeSegmentView{
+		label:              s.testLabel,
+		segments:           []*SegmentView{{ID: 200, label: s.testLabel, NumOfRows: 100, Size: 150 * 1024 * 1024}},
+		triggerID:          1001,
+		targetSegmentSize:  100 * 1024 * 1024,
+		targetSegmentCount: 999,
+	}
+	s.triggerManager.SubmitForceMergeViewToScheduler(context.Background(), view)
+}
+
+func (s *CompactionTriggerManagerSuite) TestSubmitViewToSchedulerDefensiveReturns() {
+	makeCollectionSchema := func(external bool) *schemapb.CollectionSchema {
+		field := &schemapb.FieldSchema{FieldID: 1, Name: "pk", DataType: schemapb.DataType_Int64, IsPrimaryKey: true}
+		if external {
+			field.ExternalField = "pk_col"
+		}
+		return &schemapb.CollectionSchema{
+			Name:   "test_coll",
+			Fields: []*schemapb.FieldSchema{field},
+		}
+	}
+
+	makeMixView := func() *MixSegmentView {
+		return &MixSegmentView{
+			label:         s.testLabel,
+			segments:      []*SegmentView{{ID: 200, label: s.testLabel, NumOfRows: 100, Size: 150 * 1024 * 1024}},
+			collectionTTL: 100,
+			triggerID:     1001,
+		}
+	}
+
+	makeClusteringView := func(clusteringKey *schemapb.FieldSchema) *ClusteringSegmentsView {
+		return &ClusteringSegmentsView{
+			label:              s.testLabel,
+			segments:           []*SegmentView{{ID: 200, label: s.testLabel, NumOfRows: 100, Size: 150 * 1024 * 1024}},
+			clusteringKeyField: clusteringKey,
+			collectionTTL:      100,
+			triggerID:          1001,
+		}
+	}
+
+	makeForceMergeView := func() *ForceMergeSegmentView {
+		return &ForceMergeSegmentView{
+			label:              s.testLabel,
+			segments:           []*SegmentView{{ID: 200, label: s.testLabel, NumOfRows: 100, Size: 150 * 1024 * 1024}},
+			triggerID:          1001,
+			targetSegmentSize:  100 * 1024 * 1024,
+			targetSegmentCount: 999,
+		}
+	}
+
+	s.Run("single GetCollection error returns before allocation", func() {
+		s.SetupTest()
+		handler := NewNMockHandler(s.T())
+		handler.EXPECT().GetCollection(mock.Anything, s.testLabel.CollectionID).
+			Return(nil, errors.New("get collection error")).Once()
+		s.triggerManager.handler = handler
+
+		s.triggerManager.SubmitSingleViewToScheduler(context.Background(), makeMixView(), TriggerTypeSingle)
+	})
+
+	s.Run("single AllocN error returns before enqueue", func() {
+		s.SetupTest()
+		s.meta.indexMeta = &indexMeta{indexes: make(map[UniqueID]map[UniqueID]*model.Index)}
+		handler := NewNMockHandler(s.T())
+		handler.EXPECT().GetCollection(mock.Anything, s.testLabel.CollectionID).
+			Return(&collectionInfo{ID: s.testLabel.CollectionID, Schema: makeCollectionSchema(false)}, nil).Once()
+		s.triggerManager.handler = handler
+		s.mockAlloc.EXPECT().AllocN(mock.Anything).Return(int64(0), int64(0), errors.New("alloc error")).Once()
+
+		s.triggerManager.SubmitSingleViewToScheduler(context.Background(), makeMixView(), TriggerTypeSingle)
+	})
+
+	s.Run("clustering GetCollection error returns before allocation", func() {
+		s.SetupTest()
+		clusteringKey := &schemapb.FieldSchema{FieldID: 2, Name: "cluster_key", DataType: schemapb.DataType_Int64, IsClusteringKey: true}
+		handler := NewNMockHandler(s.T())
+		handler.EXPECT().GetCollection(mock.Anything, s.testLabel.CollectionID).
+			Return(nil, errors.New("get collection error")).Once()
+		s.triggerManager.handler = handler
+
+		s.triggerManager.SubmitClusteringViewToScheduler(context.Background(), makeClusteringView(clusteringKey))
+	})
+
+	s.Run("clustering AllocN error returns before enqueue", func() {
+		s.SetupTest()
+		s.meta.indexMeta = &indexMeta{indexes: make(map[UniqueID]map[UniqueID]*model.Index)}
+		clusteringKey := &schemapb.FieldSchema{FieldID: 2, Name: "cluster_key", DataType: schemapb.DataType_Int64, IsClusteringKey: true}
+		collectionSchema := makeCollectionSchema(false)
+		collectionSchema.Fields = append(collectionSchema.Fields, clusteringKey)
+		handler := NewNMockHandler(s.T())
+		handler.EXPECT().GetCollection(mock.Anything, s.testLabel.CollectionID).
+			Return(&collectionInfo{ID: s.testLabel.CollectionID, Schema: collectionSchema}, nil).Once()
+		s.triggerManager.handler = handler
+		s.mockAlloc.EXPECT().AllocN(mock.Anything).Return(int64(0), int64(0), errors.New("alloc error")).Once()
+
+		s.triggerManager.SubmitClusteringViewToScheduler(context.Background(), makeClusteringView(clusteringKey))
+	})
+
+	s.Run("force merge GetCollection error returns before allocation", func() {
+		s.SetupTest()
+		handler := NewNMockHandler(s.T())
+		handler.EXPECT().GetCollection(mock.Anything, s.testLabel.CollectionID).
+			Return(nil, errors.New("get collection error")).Once()
+		s.triggerManager.handler = handler
+
+		s.triggerManager.SubmitForceMergeViewToScheduler(context.Background(), makeForceMergeView())
+	})
+
+	s.Run("force merge nil collection returns before allocation", func() {
+		s.SetupTest()
+		handler := NewNMockHandler(s.T())
+		handler.EXPECT().GetCollection(mock.Anything, s.testLabel.CollectionID).Return(nil, nil).Once()
+		s.triggerManager.handler = handler
+
+		s.triggerManager.SubmitForceMergeViewToScheduler(context.Background(), makeForceMergeView())
+	})
+
+	s.Run("force merge external collection returns before allocation", func() {
+		s.SetupTest()
+		handler := NewNMockHandler(s.T())
+		handler.EXPECT().GetCollection(mock.Anything, s.testLabel.CollectionID).
+			Return(&collectionInfo{ID: s.testLabel.CollectionID, Schema: makeCollectionSchema(true)}, nil).Once()
+		s.triggerManager.handler = handler
+
+		s.triggerManager.SubmitForceMergeViewToScheduler(context.Background(), makeForceMergeView())
+	})
+
+	s.Run("force merge AllocN error returns before enqueue", func() {
+		s.SetupTest()
+		handler := NewNMockHandler(s.T())
+		handler.EXPECT().GetCollection(mock.Anything, s.testLabel.CollectionID).
+			Return(&collectionInfo{ID: s.testLabel.CollectionID, Schema: makeCollectionSchema(false)}, nil).Once()
+		s.triggerManager.handler = handler
+		s.mockAlloc.EXPECT().AllocN(mock.Anything).Return(int64(0), int64(0), errors.New("alloc error")).Once()
+
+		s.triggerManager.SubmitForceMergeViewToScheduler(context.Background(), makeForceMergeView())
+	})
+}
+
 func (s *CompactionTriggerManagerSuite) TestSubmitBumpSchemaVersionViewToScheduler() {
+	Params.Save(Params.DataCoordCfg.CompactionPreAllocateIDExpansionFactor.Key, "1")
+	defer Params.Reset(Params.DataCoordCfg.CompactionPreAllocateIDExpansionFactor.Key)
+	Params.Save(Params.DataCoordCfg.SegmentMaxSize.Key, "100")
+	defer Params.Reset(Params.DataCoordCfg.SegmentMaxSize.Key)
+	Params.Save(Params.DataCoordCfg.DiskSegmentMaxSize.Key, "100")
+	defer Params.Reset(Params.DataCoordCfg.DiskSegmentMaxSize.Key)
+
 	collectionSchema := &schemapb.CollectionSchema{
 		Name: "test_coll",
 		Fields: []*schemapb.FieldSchema{
@@ -440,6 +945,7 @@ func (s *CompactionTriggerManagerSuite) TestSubmitBumpSchemaVersionViewToSchedul
 			ID:        200,
 			label:     s.testLabel,
 			NumOfRows: 1000,
+			Size:      300 * 1024 * 1024,
 		}
 		return &BumpSchemaVersionView{
 			label:     s.testLabel,
@@ -451,16 +957,18 @@ func (s *CompactionTriggerManagerSuite) TestSubmitBumpSchemaVersionViewToSchedul
 
 	s.Run("AllocN fails", func() {
 		s.SetupTest()
-		s.mockAlloc.EXPECT().AllocN(int64(2)).Return(int64(0), int64(0), errors.New("alloc error")).Once()
+		s.meta.indexMeta = &indexMeta{indexes: make(map[UniqueID]map[UniqueID]*model.Index)}
+		handler := NewNMockHandler(s.T())
+		handler.EXPECT().GetCollection(mock.Anything, s.testLabel.CollectionID).
+			Return(&collectionInfo{ID: s.testLabel.CollectionID, Schema: collectionSchema}, nil).Once()
+		s.triggerManager.handler = handler
+		s.mockAlloc.EXPECT().AllocN(int64(4)).Return(int64(0), int64(0), errors.New("alloc error")).Once()
 		view := makeBumpSchemaVersionView(111)
-		// Should return early with no panic — no other mock calls expected.
 		s.triggerManager.SubmitBumpSchemaVersionViewToScheduler(context.Background(), view)
 	})
 
 	s.Run("GetCollection fails", func() {
 		s.SetupTest()
-		const planID = int64(500)
-		s.mockAlloc.EXPECT().AllocN(int64(2)).Return(planID, planID+2, nil).Once()
 		handler := NewNMockHandler(s.T())
 		handler.EXPECT().GetCollection(mock.Anything, s.testLabel.CollectionID).
 			Return(nil, errors.New("get collection error")).Once()
@@ -472,8 +980,6 @@ func (s *CompactionTriggerManagerSuite) TestSubmitBumpSchemaVersionViewToSchedul
 
 	s.Run("collection is nil", func() {
 		s.SetupTest()
-		const planID = int64(501)
-		s.mockAlloc.EXPECT().AllocN(int64(2)).Return(planID, planID+2, nil).Once()
 		handler := NewNMockHandler(s.T())
 		handler.EXPECT().GetCollection(mock.Anything, s.testLabel.CollectionID).
 			Return(nil, nil).Once()
@@ -485,8 +991,6 @@ func (s *CompactionTriggerManagerSuite) TestSubmitBumpSchemaVersionViewToSchedul
 
 	s.Run("collection is external", func() {
 		s.SetupTest()
-		const planID = int64(502)
-		s.mockAlloc.EXPECT().AllocN(int64(2)).Return(planID, planID+2, nil).Once()
 		handler := NewNMockHandler(s.T())
 		handler.EXPECT().GetCollection(mock.Anything, s.testLabel.CollectionID).
 			Return(&collectionInfo{
@@ -519,10 +1023,11 @@ func (s *CompactionTriggerManagerSuite) TestSubmitBumpSchemaVersionViewToSchedul
 		s.SetupTest()
 		s.meta.indexMeta = &indexMeta{indexes: make(map[UniqueID]map[UniqueID]*model.Index)}
 		const (
-			planID    = int64(504)
+			startID   = int64(504)
+			endID     = int64(508)
 			triggerID = int64(999)
 		)
-		s.mockAlloc.EXPECT().AllocN(int64(2)).Return(planID, planID+2, nil).Once()
+		s.mockAlloc.EXPECT().AllocN(int64(4)).Return(startID, endID, nil).Once()
 		handler := NewNMockHandler(s.T())
 		handler.EXPECT().GetCollection(mock.Anything, s.testLabel.CollectionID).
 			Return(&collectionInfo{
@@ -540,10 +1045,12 @@ func (s *CompactionTriggerManagerSuite) TestSubmitBumpSchemaVersionViewToSchedul
 		s.SetupTest()
 		s.meta.indexMeta = &indexMeta{indexes: make(map[UniqueID]map[UniqueID]*model.Index)}
 		const (
-			planID    = int64(600)
+			startID   = int64(600)
+			endID     = int64(604)
+			planID    = int64(603)
 			triggerID = int64(1001)
 		)
-		s.mockAlloc.EXPECT().AllocN(int64(2)).Return(planID, planID+2, nil).Once()
+		s.mockAlloc.EXPECT().AllocN(int64(4)).Return(startID, endID, nil).Once()
 		handler := NewNMockHandler(s.T())
 		handler.EXPECT().GetCollection(mock.Anything, s.testLabel.CollectionID).
 			Return(&collectionInfo{
@@ -566,8 +1073,8 @@ func (s *CompactionTriggerManagerSuite) TestSubmitBumpSchemaVersionViewToSchedul
 				s.NotZero(task.GetStartTime())
 				s.NotZero(task.GetLastStateStartTime())
 				s.NotNil(task.GetPreAllocatedSegmentIDs())
-				s.EqualValues(planID+1, task.GetPreAllocatedSegmentIDs().GetBegin())
-				s.EqualValues(planID+2, task.GetPreAllocatedSegmentIDs().GetEnd())
+				s.EqualValues(startID, task.GetPreAllocatedSegmentIDs().GetBegin())
+				s.EqualValues(planID, task.GetPreAllocatedSegmentIDs().GetEnd())
 				s.EqualValues(time.Hour.Nanoseconds(), task.GetCollectionTtl())
 				return nil
 			}).Once()
@@ -580,7 +1087,9 @@ func (s *CompactionTriggerManagerSuite) TestSubmitBumpSchemaVersionViewToSchedul
 		s.SetupTest()
 		s.meta.indexMeta = &indexMeta{indexes: make(map[UniqueID]map[UniqueID]*model.Index)}
 		const (
-			planID    = int64(601)
+			startID   = int64(601)
+			endID     = int64(605)
+			planID    = int64(604)
 			triggerID = int64(1002)
 		)
 		frozenSchema := &schemapb.CollectionSchema{
@@ -593,7 +1102,7 @@ func (s *CompactionTriggerManagerSuite) TestSubmitBumpSchemaVersionViewToSchedul
 			Version: 3,
 			Fields:  []*schemapb.FieldSchema{{FieldID: 1, Name: "pk", DataType: schemapb.DataType_Int64, IsPrimaryKey: true}},
 		}
-		s.mockAlloc.EXPECT().AllocN(int64(2)).Return(planID, planID+2, nil).Once()
+		s.mockAlloc.EXPECT().AllocN(int64(4)).Return(startID, endID, nil).Once()
 		handler := NewNMockHandler(s.T())
 		handler.EXPECT().GetCollection(mock.Anything, s.testLabel.CollectionID).
 			Return(&collectionInfo{ID: s.testLabel.CollectionID, Schema: liveSchema}, nil).Once()
@@ -604,7 +1113,8 @@ func (s *CompactionTriggerManagerSuite) TestSubmitBumpSchemaVersionViewToSchedul
 				s.EqualValues(2, task.GetSchema().GetVersion())
 				s.Equal(frozenSchema, task.GetSchema())
 				s.NotNil(task.GetPreAllocatedSegmentIDs())
-				s.EqualValues(task.GetPreAllocatedSegmentIDs().GetBegin()+1, task.GetPreAllocatedSegmentIDs().GetEnd())
+				s.EqualValues(startID, task.GetPreAllocatedSegmentIDs().GetBegin())
+				s.EqualValues(planID, task.GetPreAllocatedSegmentIDs().GetEnd())
 				return nil
 			}).Once()
 

internal/datacoord/compaction_view.go

@@ -18,6 +18,7 @@ package datacoord
 
 import (
 	"fmt"
+	"time"
 
 	"github.com/samber/lo"
 
@@ -36,6 +37,12 @@ type CompactionView interface {
 	ForceTrigger() (CompactionView, string)
 	ForceTriggerAll() ([]CompactionView, string)
 	GetTriggerID() int64
+	GetTotalSize() float64
+	GetCollectionTTL() time.Duration
+}
+
+func sumSegmentSize(views []*SegmentView) float64 {
+	return lo.SumBy(views, func(v *SegmentView) float64 { return v.Size })
 }
 
 type FullViews struct {

internal/datacoord/compaction_view_forcemerge.go

@@ -44,7 +44,10 @@ type ForceMergeSegmentView struct {
 
 	topology *CollectionTopology
 
-	targetSegmentSize  float64
+	targetSegmentSize float64
+	// targetSegmentCount records the ForceTriggerAll planning result for logging
+	// and callers of GetTargetSegmentCount. Scheduler ID allocation is derived
+	// from targetSegmentSize so stale counts cannot over-reserve IDs.
 	targetSegmentCount int64
 }
 
@@ -64,6 +67,20 @@ func (v *ForceMergeSegmentView) GetSegmentsView() []*SegmentView {
 	return v.segments
 }
 
+func (v *ForceMergeSegmentView) GetTotalSize() float64 {
+	if v == nil {
+		return 0
+	}
+	return sumSegmentSize(v.segments)
+}
+
+func (v *ForceMergeSegmentView) GetCollectionTTL() time.Duration {
+	if v == nil {
+		return 0
+	}
+	return v.collectionTTL
+}
+
 func (v *ForceMergeSegmentView) Append(segments ...*SegmentView) {
 	v.segments = append(v.segments, segments...)
 }
@@ -108,8 +125,8 @@ func (v *ForceMergeSegmentView) calculateTargetSizeCount() (maxSafeSize float64,
 		}
 	}
 
-	totalSize := sumSegmentSize(v.segments)
-	targetCount = max(1, int64(totalSize/maxSafeSize))
+	totalSize := v.GetTotalSize()
+	targetCount = estimateResultSegmentCount(totalSize, maxSafeSize)
 
 	queryNodeCount := int64(len(v.topology.QueryNodeMemory))
 	numReplicas := int64(v.topology.NumReplicas)
@@ -157,6 +174,13 @@ func (v *ForceMergeSegmentView) ForceTriggerAll() ([]CompactionView, string) {
 
 	results := make([]CompactionView, 0, len(groups))
 	for _, group := range groups {
+		groupTargetCount := targetCount
+		if len(groups) > 1 {
+			// Once adaptive grouping splits the input, each subgroup needs its
+			// own estimate; the topology-adjusted count applies only to the
+			// unsplit single-group case.
+			groupTargetCount = estimateResultSegmentCount(sumSegmentSize(group), targetSizePerSegment)
+		}
 		results = append(results, &ForceMergeSegmentView{
 			label:              v.label,
 			segments:           group,
@@ -165,7 +189,7 @@ func (v *ForceMergeSegmentView) ForceTriggerAll() ([]CompactionView, string) {
 			configMaxSize:      v.configMaxSize,
 			expectedTargetSize: v.expectedTargetSize,
 			targetSegmentSize:  targetSizePerSegment,
-			targetSegmentCount: targetCount,
+			targetSegmentCount: groupTargetCount,
 			topology:           v.topology,
 		})
 	}
@@ -372,7 +396,3 @@ func (v *ForceMergeSegmentView) calculateMaxSafeSize() float64 {
 		zap.Float64("configMaxSize", v.configMaxSize))
 	return maxSafeSize
 }
-
-func sumSegmentSize(views []*SegmentView) float64 {
-	return lo.SumBy(views, func(v *SegmentView) float64 { return v.Size })
-}

internal/datacoord/compaction_view_forcemerge_test.go

@@ -630,6 +630,27 @@ func BenchmarkGroupByPartitionChannel(b *testing.B) {
 }
 
 func TestCalculateTargetSizeCount_QueryNodeParallelism(t *testing.T) {
+	t.Run("fractional target count rounds up", func(t *testing.T) {
+		view := &ForceMergeSegmentView{
+			label: &CompactionGroupLabel{
+				CollectionID: 1,
+				PartitionID:  1,
+				Channel:      "ch1",
+			},
+			segments: []*SegmentView{
+				{ID: 1, Size: 150 * 1024 * 1024},
+			},
+			triggerID:     1,
+			configMaxSize: 100 * 1024 * 1024,
+			topology:      &CollectionTopology{},
+		}
+
+		maxSafeSize, targetCount := view.calculateTargetSizeCount()
+
+		assert.Equal(t, int64(2), targetCount)
+		assert.Equal(t, float64(100*1024*1024), maxSafeSize)
+	})
+
 	t.Run("single QueryNode - no adjustment", func(t *testing.T) {
 		topology := &CollectionTopology{
 			QueryNodeMemory: map[int64]uint64{1: 8 * 1024 * 1024 * 1024},

internal/datacoord/mock_test.go

@@ -129,7 +129,9 @@ func newMock0Allocator(t *testing.T) *allocator.MockAllocator {
 	mock0Allocator := allocator.NewMockAllocator(t)
 	mock0Allocator.EXPECT().AllocID(mock.Anything).Return(100, nil).Maybe()
 	mock0Allocator.EXPECT().AllocTimestamp(mock.Anything).Return(1000, nil).Maybe()
-	mock0Allocator.EXPECT().AllocN(mock.Anything).Return(100, 200, nil).Maybe()
+	mock0Allocator.EXPECT().AllocN(mock.Anything).RunAndReturn(func(i int64) (int64, int64, error) {
+		return 100, 100 + i, nil
+	}).Maybe()
 	return mock0Allocator
 }
 

internal/datanode/compactor/multi_segment_writer_test.go

@@ -25,6 +25,9 @@ import (
 	"github.com/stretchr/testify/mock"
 	"github.com/stretchr/testify/suite"
 	"go.uber.org/atomic"
+	"go.uber.org/zap"
+	"go.uber.org/zap/zapcore"
+	"go.uber.org/zap/zaptest/observer"
 
 	"github.com/milvus-io/milvus-proto/go-api/v3/commonpb"
 	"github.com/milvus-io/milvus-proto/go-api/v3/schemapb"
@@ -33,6 +36,8 @@ import (
 	"github.com/milvus-io/milvus/internal/mocks/flushcommon/mock_util"
 	"github.com/milvus-io/milvus/internal/storage"
 	"github.com/milvus-io/milvus/pkg/v3/common"
+	"github.com/milvus-io/milvus/pkg/v3/log"
+	"github.com/milvus-io/milvus/pkg/v3/proto/datapb"
 	"github.com/milvus-io/milvus/pkg/v3/util/paramtable"
 	"github.com/milvus-io/milvus/pkg/v3/util/tsoutil"
 )
@@ -56,15 +61,20 @@ type MultiSegmentWriterSuite struct {
 }
 
 type testCompactionAllocator struct {
-	next           int64
-	allocOneCalls  int
-	failAllocOneAt int
+	next            int64
+	allocOneCalls   int
+	failAllocOneAt  int
+	failAllocOneErr error
 }
 
 type closeErrSerializeWriter struct {
 	err error
 }
 
+type closeErrBinlogRecordWriter struct {
+	writtenUncompressed uint64
+}
+
 func (w closeErrSerializeWriter) WriteValue(*storage.Value) error {
 	return nil
 }
@@ -77,6 +87,44 @@ func (w closeErrSerializeWriter) Close() error {
 	return w.err
 }
 
+func (w closeErrBinlogRecordWriter) Write(storage.Record) error {
+	return nil
+}
+
+func (w closeErrBinlogRecordWriter) GetWrittenUncompressed() uint64 {
+	return w.writtenUncompressed
+}
+
+func (w closeErrBinlogRecordWriter) Close() error {
+	return nil
+}
+
+func (w closeErrBinlogRecordWriter) GetLogs() (
+	map[storage.FieldID]*datapb.FieldBinlog,
+	*datapb.FieldBinlog,
+	map[storage.FieldID]*datapb.FieldBinlog,
+	string,
+	[]int64,
+) {
+	return nil, nil, nil, "", nil
+}
+
+func (w closeErrBinlogRecordWriter) GetRowNum() int64 {
+	return 0
+}
+
+func (w closeErrBinlogRecordWriter) FlushChunk() error {
+	return nil
+}
+
+func (w closeErrBinlogRecordWriter) GetBufferUncompressed() uint64 {
+	return w.writtenUncompressed
+}
+
+func (w closeErrBinlogRecordWriter) Schema() *schemapb.CollectionSchema {
+	return nil
+}
+
 func (a *testCompactionAllocator) Alloc(count uint32) (int64, int64, error) {
 	start := a.next
 	a.next += int64(count)
@@ -86,7 +134,10 @@ func (a *testCompactionAllocator) Alloc(count uint32) (int64, int64, error) {
 func (a *testCompactionAllocator) AllocOne() (int64, error) {
 	a.allocOneCalls++
 	if a.failAllocOneAt > 0 && a.allocOneCalls == a.failAllocOneAt {
-		return 0, errors.New("ID is exhausted")
+		if a.failAllocOneErr != nil {
+			return 0, a.failAllocOneErr
+		}
+		return 0, allocator.NewIDExhaustedError(0, 0, 1)
 	}
 	a.next++
 	return a.next, nil
@@ -385,9 +436,58 @@ func (s *MultiSegmentWriterSuite) TestSegmentRotation() {
 	s.GreaterOrEqual(len(finalSegments), 2)
 }
 
-func (s *MultiSegmentWriterSuite) TestCloseAfterRotateAllocFailureDoesNotRecloseClosedWriter() {
+func (s *MultiSegmentWriterSuite) TestSegmentIDExhaustionGrowsCurrentSegment() {
 	schema := s.genSimpleSchema()
-	segmentAlloc := &testCompactionAllocator{next: 1000, failAllocOneAt: 2}
+	segmentAlloc := allocator.NewLocalAllocator(1000, 1001)
+	logAlloc := allocator.NewLocalAllocator(2000, 3000)
+	allocator := NewCompactionAllocator(segmentAlloc, logAlloc)
+	core, recordedLogs := observer.New(zapcore.WarnLevel)
+	log.ReplaceGlobals(zap.New(core), &log.ZapProperties{
+		Core:  core,
+		Level: zap.NewAtomicLevelAt(zapcore.WarnLevel),
+	})
+	s.T().Cleanup(func() {
+		logger, props, err := log.InitLogger(&log.Config{
+			Level:               "debug",
+			Stdout:              true,
+			DisableErrorVerbose: true,
+		})
+		s.Require().NoError(err)
+		log.ReplaceGlobals(logger, props)
+	})
+
+	writer, err := NewMultiSegmentWriter(
+		context.Background(),
+		s.mockBinlogIO,
+		allocator,
+		1,
+		schema,
+		s.params,
+		1000,
+		s.partitionID,
+		s.collectionID,
+		s.channel,
+		1,
+		storage.WithStorageConfig(s.params.StorageConfig),
+	)
+	s.Require().NoError(err)
+
+	s.Require().NoError(writer.WriteValue(s.genTestValue(1)))
+	s.Require().NoError(writer.WriteValue(s.genTestValue(2)))
+	s.Require().NoError(writer.WriteValue(s.genTestValue(3)))
+	s.Require().NoError(writer.Close())
+
+	segments := writer.GetCompactionSegments()
+	s.Require().Len(segments, 1)
+	s.EqualValues(1000, segments[0].SegmentID)
+	s.EqualValues(3, segments[0].NumOfRows)
+	s.Len(recordedLogs.FilterMessage("pre-allocated compaction segment IDs exhausted, continue writing current segment").All(), 1)
+}
+
+func (s *MultiSegmentWriterSuite) TestRotateAllocFailureKeepsCurrentWriterOpen() {
+	schema := s.genSimpleSchema()
+	allocErr := errors.New("rootcoord unavailable")
+	segmentAlloc := &testCompactionAllocator{next: 1000, failAllocOneAt: 2, failAllocOneErr: allocErr}
 	logAlloc := &testCompactionAllocator{next: 2000}
 	allocator := NewCompactionAllocator(segmentAlloc, logAlloc)
 
@@ -410,16 +510,40 @@ func (s *MultiSegmentWriterSuite) TestCloseAfterRotateAllocFailureDoesNotReclose
 	err = writer.WriteValue(s.genTestValue(1))
 	s.Require().NoError(err)
 	err = writer.WriteValue(s.genTestValue(2))
-	s.Require().Error(err)
-	s.Contains(err.Error(), "ID is exhausted")
-	s.Nil(writer.writer)
-	s.Len(writer.GetCompactionSegments(), 1)
+	s.ErrorIs(err, allocErr)
+	s.NotNil(writer.writer)
+	s.Empty(writer.GetCompactionSegments())
 
 	err = writer.Close()
 	s.NoError(err)
 	s.Len(writer.GetCompactionSegments(), 1)
 }
 
+func (s *MultiSegmentWriterSuite) TestLogIDExhaustionDuringRotationReturnsError() {
+	segmentAlloc := &testCompactionAllocator{next: 1000}
+	logAlloc := &testCompactionAllocator{next: 2000}
+	writer := &MultiSegmentWriter{
+		allocator: NewCompactionAllocator(segmentAlloc, logAlloc),
+		writer: &storage.BinlogValueWriter{
+			BinlogRecordWriter: closeErrBinlogRecordWriter{writtenUncompressed: 2},
+			SerializeWriter:    closeErrSerializeWriter{err: allocator.NewIDExhaustedError(0, 0, 1)},
+		},
+		segmentSize:      1,
+		currentSegmentID: 1000,
+		collectionID:     s.collectionID,
+		partitionID:      s.partitionID,
+		channel:          s.channel,
+	}
+
+	var err error
+	s.NotPanics(func() {
+		err = writer.rotateWriterOrGrowCurrent()
+	})
+	s.True(allocator.IsIDExhausted(err))
+	s.Nil(writer.writer)
+	s.Empty(writer.GetCompactionSegments())
+}
+
 func (s *MultiSegmentWriterSuite) TestCloseFailureClearsWriter() {
 	closeErr := errors.New("close failed")
 	writer := &MultiSegmentWriter{

internal/datanode/compactor/segment_writer.go

@@ -22,6 +22,7 @@ import (
 	"math"
 
 	"github.com/apache/arrow/go/v17/arrow/array"
+	"github.com/cockroachdb/errors"
 	"github.com/samber/lo"
 	"go.uber.org/atomic"
 	"go.uber.org/zap"
@@ -71,10 +72,13 @@ type MultiSegmentWriter struct {
 }
 
 type compactionAlloactor struct {
-	segmentAlloc allocator.Interface
-	logIDAlloc   allocator.Interface
+	segmentAlloc             allocator.Interface
+	logIDAlloc               allocator.Interface
+	segmentIDBudgetExhausted bool
 }
 
+var errCompactionSegmentIDsExhausted = errors.New("pre-allocated compaction segment IDs exhausted")
+
 func NewCompactionAllocator(segmentAlloc, logIDAlloc allocator.Interface) *compactionAlloactor {
 	return &compactionAlloactor{
 		segmentAlloc: segmentAlloc,
@@ -83,9 +87,20 @@ func NewCompactionAllocator(segmentAlloc, logIDAlloc allocator.Interface) *compa
 }
 
 func (alloc *compactionAlloactor) allocSegmentID() (typeutil.UniqueID, error) {
+	if alloc.isSegmentIDBudgetExhausted() {
+		return 0, errCompactionSegmentIDsExhausted
+	}
 	return alloc.segmentAlloc.AllocOne()
 }
 
+func (alloc *compactionAlloactor) isSegmentIDBudgetExhausted() bool {
+	return alloc.segmentIDBudgetExhausted
+}
+
+func (alloc *compactionAlloactor) markSegmentIDBudgetExhausted() {
+	alloc.segmentIDBudgetExhausted = true
+}
+
 func NewMultiSegmentWriter(ctx context.Context, binlogIO io.BinlogIO, allocator *compactionAlloactor, segmentSize int64,
 	schema *schemapb.CollectionSchema, params compaction.Params,
 	maxRows int64, partitionID, collectionID int64, channel string, batchSize int, rwOption ...storage.RwOption,
@@ -151,15 +166,17 @@ func (w *MultiSegmentWriter) closeWriter() error {
 }
 
 func (w *MultiSegmentWriter) rotateWriter() error {
-	if err := w.closeWriter(); err != nil {
+	newSegmentID, err := w.allocator.allocSegmentID()
+	if err != nil {
+		if allocator.IsIDExhausted(err) {
+			return errors.Mark(err, errCompactionSegmentIDsExhausted)
+		}
 		return err
 	}
 
-	newSegmentID, err := w.allocator.allocSegmentID()
-	if err != nil {
+	if err := w.closeWriter(); err != nil {
 		return err
 	}
-	w.currentSegmentID = newSegmentID
 
 	chunkSize := w.binLogMaxSize
 
@@ -176,10 +193,41 @@ func (w *MultiSegmentWriter) rotateWriter() error {
 		return err
 	}
 
+	w.currentSegmentID = newSegmentID
 	w.writer = storage.NewBinlogValueWriter(rw, w.batchSize)
 	return nil
 }
 
+func (w *MultiSegmentWriter) rotateWriterOrGrowCurrent() error {
+	if w.writer == nil {
+		return w.rotateWriter()
+	}
+	if w.allocator.isSegmentIDBudgetExhausted() {
+		return nil
+	}
+	if w.writer.GetWrittenUncompressed() < uint64(w.segmentSize) {
+		return nil
+	}
+
+	if err := w.rotateWriter(); err != nil {
+		if !errors.Is(err, errCompactionSegmentIDsExhausted) {
+			return err
+		}
+
+		w.allocator.markSegmentIDBudgetExhausted()
+		writtenUncompressed := w.writer.GetWrittenUncompressed()
+		log.Warn("pre-allocated compaction segment IDs exhausted, continue writing current segment",
+			zap.Int64("collectionID", w.collectionID),
+			zap.Int64("partitionID", w.partitionID),
+			zap.String("channel", w.channel),
+			zap.Int64("segmentID", w.currentSegmentID),
+			zap.Uint64("currentSize", writtenUncompressed),
+			zap.Int64("expectedSegmentSize", w.segmentSize),
+			zap.Error(err))
+	}
+	return nil
+}
+
 func (w *MultiSegmentWriter) GetWrittenUncompressed() uint64 {
 	if w.writer == nil {
 		return 0
@@ -199,19 +247,15 @@ func (w *MultiSegmentWriter) GetCompactionSegments() []*datapb.CompactionSegment
 }
 
 func (w *MultiSegmentWriter) Write(r storage.Record) error {
-	if w.writer == nil || w.writer.GetWrittenUncompressed() >= uint64(w.segmentSize) {
-		if err := w.rotateWriter(); err != nil {
-			return err
-		}
+	if err := w.rotateWriterOrGrowCurrent(); err != nil {
+		return err
 	}
 	return w.writer.Write(r)
 }
 
 func (w *MultiSegmentWriter) WriteValue(v *storage.Value) error {
-	if w.writer == nil || w.writer.GetWrittenUncompressed() >= uint64(w.segmentSize) {
-		if err := w.rotateWriter(); err != nil {
-			return err
-		}
+	if err := w.rotateWriterOrGrowCurrent(); err != nil {
+		return err
 	}
 
 	return w.writer.WriteValue(v)