enhance: [cp3.0]rewrite storage.MergeSort as a proper k-way merge

internal/datanode/compactor/merge_sort.go

@@ -133,6 +133,17 @@ func mergeSortMultipleSegments(ctx context.Context,
 	}
 
 	if _, err = storage.MergeSort(compactionParams.BinLogMaxSize, writerSchema, segmentReaders, writer, predicate, sortByFields); err != nil {
+		// segmentReaders[i] is built from binlogs[i], so the reader index the
+		// error names, when it names one, indexes into this list.
+		segmentIDs := make([]int64, len(binlogs))
+		for i, s := range binlogs {
+			segmentIDs[i] = s.GetSegmentID()
+		}
+		log.Warn(ctx, "compact wrong, failed to merge sort segments",
+			mlog.Int64("collectionID", collectionID),
+			mlog.Int64s("segmentIDsByReaderIndex", segmentIDs),
+			mlog.Int64s("sortByFields", sortByFields),
+			mlog.Err(err))
 		if closeErr := writer.Close(); closeErr != nil {
 			log.Warn(ctx, "failed to close writer after merge sort error", mlog.Err(closeErr))
 		}

internal/datanode/compactor/mix_compactor.go

@@ -477,31 +477,25 @@ func (t *mixCompactionTask) Compact() (*datapb.CompactionPlanResult, error) {
 		return nil, err
 	}
 
-	sortMergeAppicable := t.compactionParams.UseMergeSort
-	if sortMergeAppicable {
-		for _, segment := range t.plan.GetSegmentBinlogs() {
-			if !segment.GetIsSorted() && !segment.GetIsSortedByNamespace() {
-				sortMergeAppicable = false
-				break
-			}
-		}
-
-		if len(t.plan.GetSegmentBinlogs()) > t.compactionParams.MaxSegmentMergeSort {
-			// sort merge is not applicable if there is only one segment or too many segments
-			sortMergeAppicable = false
-		}
-	}
+	useMergeSort := canMergeSort(t.plan, t.compactionParams)
 
 	var res []*datapb.CompactionSegment
 	var err error
-	if sortMergeAppicable {
+	if useMergeSort {
 		mlog.Info(context.TODO(), "compact by merge sort")
 		writerOpts := t.buildWriterOptions(ctx)
 		res, err = mergeSortMultipleSegments(ctxTimeout, t.plan, t.collectionID, t.partitionID, t.maxRows, t.binlogIO,
 			t.plan.GetSegmentBinlogs(), t.tr, t.currentTime, t.plan.GetCollectionTtl(), t.compactionParams,
 			writerOpts, t.lobContext, t.sortByFieldIDs)
 		if err != nil {
-			mlog.Warn(context.TODO(), "compact wrong, fail to merge sort segments", mlog.Err(err))
+			// Compactor-boundary catch-all: mergeSortMultipleSegments can fail
+			// before it ever reaches the merge sort step (reader/writer
+			// construction, deltalog composition, ...), and those paths don't
+			// log on their own, so this line is their only record.
+			mlog.Warn(ctx, "compact wrong, merge sort compaction failed (compactor boundary)",
+				mlog.Int64("planID", t.GetPlanID()),
+				mlog.Int64("collectionID", t.collectionID),
+				mlog.Err(err))
 			return nil, err
 		}
 	} else {
@@ -583,6 +577,37 @@ func (t *mixCompactionTask) Compact() (*datapb.CompactionPlanResult, error) {
 	return planResult, nil
 }
 
+// canMergeSort reports whether this plan is eligible for storage.MergeSort,
+// which merges without sorting and so requires every input to already be
+// ordered by the plan's merge key. A plan that is not eligible falls back to
+// mergeSplit, which does not assume ordering.
+func canMergeSort(plan *datapb.CompactionPlan, params compaction.Params) bool {
+	if !params.UseMergeSort {
+		return false
+	}
+	// The two sorted flags record which order the compactors that write sorted
+	// output used. datanode/services.go derives this plan's merge key from the
+	// same EnableNamespace setting -- [pk], or [partitionKey, pk] -- so the flag
+	// that must be set is the matching one, not either one. It also rejects the
+	// plan outright when a namespace-enabled collection has no partition key
+	// (namespace.mode=partition), which is why reading the setting is enough
+	// here rather than inspecting the merge key itself.
+	namespaceEnabled := plan.GetSchema().GetEnableNamespace()
+	for _, segment := range plan.GetSegmentBinlogs() {
+		sortedByMergeKey := segment.GetIsSorted()
+		if namespaceEnabled {
+			sortedByMergeKey = segment.GetIsSortedByNamespace()
+		}
+		if !sortedByMergeKey {
+			return false
+		}
+	}
+	// Each reader holds a live record, so memory grows with the reader count. A
+	// single segment is allowed: merge sort keeps the output flagged sorted,
+	// whereas mergeSplit emits it unsorted and needs a follow-up sort compaction.
+	return len(plan.GetSegmentBinlogs()) <= params.MaxSegmentMergeSort
+}
+
 func (t *mixCompactionTask) Complete() {
 	t.done <- struct{}{}
 }

internal/datanode/compactor/mix_compactor_test.go

@@ -1525,3 +1525,109 @@ func BenchmarkMixCompactor(b *testing.B) {
 
 	s.TearDownTest()
 }
+
+func TestCanMergeSort(t *testing.T) {
+	params := compaction.Params{UseMergeSort: true, MaxSegmentMergeSort: 30}
+	// EnableNamespace is false here, so the merge key is [pk] and IsSorted is
+	// the flag canMergeSort requires. TestCanMergeSortMatchesMergeKey covers
+	// the namespace-enabled half.
+	namespaceDisabledSchema := &schemapb.CollectionSchema{}
+
+	t.Run("disabled by param", func(t *testing.T) {
+		plan := &datapb.CompactionPlan{
+			Schema:         namespaceDisabledSchema,
+			SegmentBinlogs: []*datapb.CompactionSegmentBinlogs{{SegmentID: 1, IsSorted: true}},
+		}
+		assert.False(t, canMergeSort(plan, compaction.Params{UseMergeSort: false, MaxSegmentMergeSort: 30}))
+	})
+
+	t.Run("unsorted segment rejected", func(t *testing.T) {
+		plan := &datapb.CompactionPlan{
+			Schema:         namespaceDisabledSchema,
+			SegmentBinlogs: []*datapb.CompactionSegmentBinlogs{{SegmentID: 1}},
+		}
+		assert.False(t, canMergeSort(plan, params))
+	})
+
+	t.Run("single sorted segment allowed", func(t *testing.T) {
+		plan := &datapb.CompactionPlan{
+			Schema:         namespaceDisabledSchema,
+			SegmentBinlogs: []*datapb.CompactionSegmentBinlogs{{SegmentID: 1, IsSorted: true}},
+		}
+		assert.True(t, canMergeSort(plan, params))
+	})
+
+	t.Run("too many segments rejected", func(t *testing.T) {
+		segs := make([]*datapb.CompactionSegmentBinlogs, params.MaxSegmentMergeSort+1)
+		for i := range segs {
+			segs[i] = &datapb.CompactionSegmentBinlogs{SegmentID: int64(i), IsSorted: true}
+		}
+		plan := &datapb.CompactionPlan{Schema: namespaceDisabledSchema, SegmentBinlogs: segs}
+		assert.False(t, canMergeSort(plan, params))
+	})
+
+	t.Run("exactly max segments allowed", func(t *testing.T) {
+		segs := make([]*datapb.CompactionSegmentBinlogs, params.MaxSegmentMergeSort)
+		for i := range segs {
+			segs[i] = &datapb.CompactionSegmentBinlogs{SegmentID: int64(i), IsSorted: true}
+		}
+		plan := &datapb.CompactionPlan{Schema: namespaceDisabledSchema, SegmentBinlogs: segs}
+		assert.True(t, canMergeSort(plan, params))
+	})
+
+	t.Run("later unsorted segment rejected", func(t *testing.T) {
+		plan := &datapb.CompactionPlan{
+			Schema: namespaceDisabledSchema,
+			SegmentBinlogs: []*datapb.CompactionSegmentBinlogs{
+				{SegmentID: 1, IsSorted: true},
+				{SegmentID: 2},
+			},
+		}
+		assert.False(t, canMergeSort(plan, params))
+	})
+}
+
+func TestCanMergeSortMatchesMergeKey(t *testing.T) {
+	params := compaction.Params{UseMergeSort: true, MaxSegmentMergeSort: 30}
+
+	t.Run("namespace enabled rejects pk-only sorted segment", func(t *testing.T) {
+		plan := &datapb.CompactionPlan{
+			Schema: &schemapb.CollectionSchema{EnableNamespace: true},
+			SegmentBinlogs: []*datapb.CompactionSegmentBinlogs{
+				{SegmentID: 1, IsSorted: true, IsSortedByNamespace: false},
+			},
+		}
+		assert.False(t, canMergeSort(plan, params))
+	})
+
+	t.Run("namespace enabled accepts namespace-sorted segment", func(t *testing.T) {
+		plan := &datapb.CompactionPlan{
+			Schema: &schemapb.CollectionSchema{EnableNamespace: true},
+			SegmentBinlogs: []*datapb.CompactionSegmentBinlogs{
+				{SegmentID: 1, IsSorted: false, IsSortedByNamespace: true},
+			},
+		}
+		assert.True(t, canMergeSort(plan, params))
+	})
+
+	t.Run("namespace disabled rejects namespace-sorted segment", func(t *testing.T) {
+		plan := &datapb.CompactionPlan{
+			Schema: &schemapb.CollectionSchema{EnableNamespace: false},
+			SegmentBinlogs: []*datapb.CompactionSegmentBinlogs{
+				{SegmentID: 1, IsSorted: false, IsSortedByNamespace: true},
+			},
+		}
+		assert.False(t, canMergeSort(plan, params))
+	})
+
+	t.Run("mixed flags rejected when one does not match", func(t *testing.T) {
+		plan := &datapb.CompactionPlan{
+			Schema: &schemapb.CollectionSchema{EnableNamespace: true},
+			SegmentBinlogs: []*datapb.CompactionSegmentBinlogs{
+				{SegmentID: 1, IsSortedByNamespace: true},
+				{SegmentID: 2, IsSorted: true},
+			},
+		}
+		assert.False(t, canMergeSort(plan, params))
+	})
+}

internal/datanode/compactor/namespace_compactor_test.go

@@ -47,6 +47,10 @@ func (s *NamespaceCompactorTestSuite) SetupSuite() {
 	s.binlogIO = mock_util.NewMockBinlogIO(s.T())
 	s.binlogIO.EXPECT().Upload(mock.Anything, mock.Anything).Return(nil).Maybe()
 	s.schema = &schemapb.CollectionSchema{
+		// NamespaceCompactor is only constructed inside the namespaceEnabled
+		// branch of datanode/services.go, and the merge key below is
+		// [partitionKey, pk] accordingly.
+		EnableNamespace: true,
 		Fields: []*schemapb.FieldSchema{
 			{
 				FieldID:  common.RowIDField,
@@ -65,9 +69,10 @@ func (s *NamespaceCompactorTestSuite) SetupSuite() {
 				IsPrimaryKey: true,
 			},
 			{
-				FieldID:  101,
-				Name:     "namespace",
-				DataType: schemapb.DataType_Int64,
+				FieldID:        101,
+				Name:           "namespace",
+				DataType:       schemapb.DataType_Int64,
+				IsPartitionKey: true,
 			},
 		},
 	}