feat: add a namespace compactor (#44253)

#44011

---------

Signed-off-by: sunby <sunbingyi1992@gmail.com>

internal/datanode/compactor/merge_sort.go

@@ -31,6 +31,7 @@ func mergeSortMultipleSegments(ctx context.Context,
 	currentTime time.Time,
 	collectionTtl int64,
 	compactionParams compaction.Params,
+	sortByFields []int64,
 ) ([]*datapb.CompactionSegment, error) {
 	_ = tr.RecordSpan()
 
@@ -106,7 +107,7 @@ func mergeSortMultipleSegments(ctx context.Context,
 		log.Warn("compaction only support int64 and varchar pk field")
 	}
 
-	if _, err = storage.MergeSort(compactionParams.BinLogMaxSize, plan.GetSchema(), segmentReaders, writer, predicate, []int64{pkField.FieldID}); err != nil {
+	if _, err = storage.MergeSort(compactionParams.BinLogMaxSize, plan.GetSchema(), segmentReaders, writer, predicate, sortByFields); err != nil {
 		writer.Close()
 		return nil, err
 	}

internal/datanode/compactor/mix_compactor.go

@@ -65,6 +65,7 @@ type mixCompactionTask struct {
 	tr   *timerecord.TimeRecorder
 
 	compactionParams compaction.Params
+	sortByFieldIDs   []int64
 }
 
 var _ Compactor = (*mixCompactionTask)(nil)
@@ -74,6 +75,7 @@ func NewMixCompactionTask(
 	binlogIO io.BinlogIO,
 	plan *datapb.CompactionPlan,
 	compactionParams compaction.Params,
+	sortByFieldIDs []int64,
 ) *mixCompactionTask {
 	ctx1, cancel := context.WithCancel(ctx)
 	return &mixCompactionTask{
@@ -85,6 +87,7 @@ func NewMixCompactionTask(
 		currentTime:      time.Now(),
 		done:             make(chan struct{}, 1),
 		compactionParams: compactionParams,
+		sortByFieldIDs:   sortByFieldIDs,
 	}
 }
 
@@ -366,7 +369,7 @@ func (t *mixCompactionTask) Compact() (*datapb.CompactionPlanResult, error) {
 		// TODO: the implementation of mergeSortMultipleSegments is not correct, also see issue: https://github.com/milvus-io/milvus/issues/43034
 		log.Info("compact by merge sort")
 		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)
+			t.plan.GetSegmentBinlogs(), t.tr, t.currentTime, t.plan.GetCollectionTtl(), t.compactionParams, t.sortByFieldIDs)
 		if err != nil {
 			log.Warn("compact wrong, fail to merge sort segments", zap.Error(err))
 			return nil, err

internal/datanode/compactor/mix_compactor_test.go

@@ -94,7 +94,9 @@ func (s *MixCompactionTaskStorageV1Suite) setupTest() {
 		JsonParams:             params,
 	}
 
-	s.task = NewMixCompactionTask(context.Background(), s.mockBinlogIO, plan, compaction.GenParams())
+	pk, err := typeutil.GetPrimaryFieldSchema(s.meta.GetSchema())
+	s.Require().NoError(err)
+	s.task = NewMixCompactionTask(context.Background(), s.mockBinlogIO, plan, compaction.GenParams(), []int64{pk.FieldID})
 }
 
 func (s *MixCompactionTaskStorageV1Suite) SetupTest() {
@@ -128,7 +130,9 @@ func (s *MixCompactionTaskStorageV1Suite) SetupBM25() {
 		JsonParams:             params,
 	}
 
-	s.task = NewMixCompactionTask(context.Background(), s.mockBinlogIO, plan, compaction.GenParams())
+	pk, err := typeutil.GetPrimaryFieldSchema(s.meta.GetSchema())
+	s.Require().NoError(err)
+	s.task = NewMixCompactionTask(context.Background(), s.mockBinlogIO, plan, compaction.GenParams(), []int64{pk.FieldID})
 }
 
 func (s *MixCompactionTaskStorageV1Suite) SetupSubTest() {

internal/datanode/compactor/namespace_compactor.go

@@ -0,0 +1,45 @@
+package compactor
+
+import (
+	"context"
+	"math"
+
+	"github.com/milvus-io/milvus/internal/compaction"
+	"github.com/milvus-io/milvus/internal/flushcommon/io"
+	"github.com/milvus-io/milvus/pkg/v2/proto/datapb"
+	"github.com/milvus-io/milvus/pkg/v2/util/merr"
+)
+
+// NamespaceCompactor compacts data with the same namespace together
+// Input segments must be sorted by namespace
+type NamespaceCompactor struct {
+	*mixCompactionTask
+}
+
+func checkInputSorted(plan *datapb.CompactionPlan) bool {
+	for _, segment := range plan.GetSegmentBinlogs() {
+		if !segment.IsSorted {
+			return false
+		}
+	}
+	return true
+}
+
+func (c *NamespaceCompactor) Compact() (*datapb.CompactionPlanResult, error) {
+	if !checkInputSorted(c.plan) {
+		return nil, merr.WrapErrIllegalCompactionPlan("input segments must be sorted by namespace")
+	}
+	res, err := c.mixCompactionTask.Compact()
+	if err != nil {
+		return nil, err
+	}
+	// TODO: after compact
+	return res, nil
+}
+
+func NewNamespaceCompactor(ctx context.Context, plan *datapb.CompactionPlan, binlogIO io.BinlogIO, compactionParams compaction.Params, sortByFieldIDs []int64) *NamespaceCompactor {
+	compactionParams.BinLogMaxSize = math.MaxInt64
+	return &NamespaceCompactor{
+		mixCompactionTask: NewMixCompactionTask(ctx, binlogIO, plan, compactionParams, sortByFieldIDs),
+	}
+}

internal/datanode/compactor/namespace_compactor_test.go

@@ -0,0 +1,224 @@
+package compactor
+
+import (
+	"context"
+	"fmt"
+	"math"
+	"testing"
+
+	"github.com/apache/arrow/go/v17/arrow/array"
+	"github.com/milvus-io/milvus-proto/go-api/v2/schemapb"
+	"github.com/milvus-io/milvus/internal/allocator"
+	"github.com/milvus-io/milvus/internal/compaction"
+	"github.com/milvus-io/milvus/internal/flushcommon/metacache"
+	"github.com/milvus-io/milvus/internal/flushcommon/metacache/pkoracle"
+	"github.com/milvus-io/milvus/internal/flushcommon/syncmgr"
+	"github.com/milvus-io/milvus/internal/mocks/flushcommon/mock_util"
+	"github.com/milvus-io/milvus/internal/storage"
+	"github.com/milvus-io/milvus/internal/storagev2/packed"
+	"github.com/milvus-io/milvus/internal/util/initcore"
+	"github.com/milvus-io/milvus/pkg/v2/common"
+	"github.com/milvus-io/milvus/pkg/v2/objectstorage"
+	"github.com/milvus-io/milvus/pkg/v2/proto/datapb"
+	"github.com/milvus-io/milvus/pkg/v2/proto/indexpb"
+	"github.com/milvus-io/milvus/pkg/v2/util/paramtable"
+	"github.com/milvus-io/milvus/pkg/v2/util/tsoutil"
+	"github.com/milvus-io/milvus/pkg/v2/util/typeutil"
+	"github.com/stretchr/testify/mock"
+	"github.com/stretchr/testify/suite"
+)
+
+type NamespaceCompactorTestSuite struct {
+	suite.Suite
+	binlogIO       *mock_util.MockBinlogIO
+	schema         *schemapb.CollectionSchema
+	sortedSegments []*datapb.CompactionSegmentBinlogs
+}
+
+func (s *NamespaceCompactorTestSuite) SetupSuite() {
+	paramtable.Get().Init(paramtable.NewBaseTable())
+	paramtable.Get().Save("common.storageType", "local")
+	paramtable.Get().Save("common.storage.enableV2", "true")
+	initcore.InitStorageV2FileSystem(paramtable.Get())
+
+	s.binlogIO = mock_util.NewMockBinlogIO(s.T())
+	s.binlogIO.EXPECT().Upload(mock.Anything, mock.Anything).Return(nil)
+	s.schema = &schemapb.CollectionSchema{
+		Fields: []*schemapb.FieldSchema{
+			{
+				FieldID:  common.RowIDField,
+				Name:     "row_id",
+				DataType: schemapb.DataType_Int64,
+			},
+			{
+				FieldID:  common.TimeStampField,
+				Name:     "timestamp",
+				DataType: schemapb.DataType_Int64,
+			},
+			{
+				FieldID:      100,
+				Name:         "pk",
+				DataType:     schemapb.DataType_Int64,
+				IsPrimaryKey: true,
+			},
+			{
+				FieldID:  101,
+				Name:     "namespace",
+				DataType: schemapb.DataType_Int64,
+			},
+		},
+	}
+	s.setupSortedSegments()
+}
+
+func (s *NamespaceCompactorTestSuite) setupSortedSegments() {
+	num := 3
+	rows := 10000
+	collectionID := int64(100)
+	partitionID := int64(100)
+	alloc := allocator.NewLocalAllocator(0, math.MaxInt64)
+
+	for i := 0; i < num; i++ {
+		data, err := storage.NewInsertData(s.schema)
+		s.Require().NoError(err)
+		for j := 0; j < rows; j++ {
+			v := map[int64]interface{}{
+				common.RowIDField:     int64(j),
+				common.TimeStampField: int64(tsoutil.ComposeTSByTime(getMilvusBirthday(), 0)),
+				100:                   int64(j),
+				101:                   int64(j),
+			}
+			data.Append(v)
+		}
+		pack := new(syncmgr.SyncPack).WithCollectionID(collectionID).WithPartitionID(partitionID).WithSegmentID(int64(i)).WithChannelName(fmt.Sprintf("by-dev-rootcoord-dml_0_%dv0", collectionID)).WithInsertData([]*storage.InsertData{data})
+		rootPath := paramtable.Get().LocalStorageCfg.Path.GetValue()
+		cm := storage.NewLocalChunkManager(objectstorage.RootPath(rootPath))
+		bfs := pkoracle.NewBloomFilterSet()
+		seg := metacache.NewSegmentInfo(&datapb.SegmentInfo{}, bfs, nil)
+		metacache.UpdateNumOfRows(int64(rows))(seg)
+		mc := metacache.NewMockMetaCache(s.T())
+		mc.EXPECT().Collection().Return(collectionID).Maybe()
+		mc.EXPECT().GetSchema(mock.Anything).Return(s.schema).Maybe()
+		mc.EXPECT().GetSegmentByID(int64(i)).Return(seg, true).Maybe()
+		mc.EXPECT().GetSegmentsBy(mock.Anything, mock.Anything).Return([]*metacache.SegmentInfo{seg}).Maybe()
+		mc.EXPECT().UpdateSegments(mock.Anything, mock.Anything).Run(func(action metacache.SegmentAction, filters ...metacache.SegmentFilter) {
+			action(seg)
+		}).Return().Maybe()
+		bw := syncmgr.NewBulkPackWriterV2(mc, s.schema, cm, alloc, packed.DefaultWriteBufferSize, 0, &indexpb.StorageConfig{
+			StorageType: "local",
+			RootPath:    rootPath,
+		})
+		inserts, _, _, _, _, err := bw.Write(context.Background(), pack)
+		s.Require().NoError(err)
+		s.sortedSegments = append(s.sortedSegments, &datapb.CompactionSegmentBinlogs{
+			SegmentID:      int64(i),
+			FieldBinlogs:   storage.SortFieldBinlogs(inserts),
+			Deltalogs:      []*datapb.FieldBinlog{},
+			StorageVersion: storage.StorageV2,
+			IsSorted:       true,
+		})
+	}
+}
+
+func (s *NamespaceCompactorTestSuite) TestCompactSorted() {
+	plan := &datapb.CompactionPlan{
+		SegmentBinlogs: s.sortedSegments,
+		Schema:         s.schema,
+		PreAllocatedSegmentIDs: &datapb.IDRange{
+			Begin: 0,
+			End:   math.MaxInt64,
+		},
+		PreAllocatedLogIDs: &datapb.IDRange{
+			Begin: 0,
+			End:   math.MaxInt64,
+		},
+		MaxSize: 1024 * 1024 * 1024,
+	}
+	params := compaction.GenParams()
+	sortedByFieldIDs := []int64{101, 100}
+
+	c := NewNamespaceCompactor(context.Background(), plan, s.binlogIO, params, sortedByFieldIDs)
+	result, err := c.Compact()
+
+	s.Require().NoError(err)
+	s.Require().Equal(datapb.CompactionTaskState_completed, result.State)
+	for _, segment := range result.GetSegments() {
+		s.assertSorted(segment, sortedByFieldIDs)
+	}
+}
+
+func (s *NamespaceCompactorTestSuite) assertSorted(segment *datapb.CompactionSegment, sortedByFieldIDs []int64) {
+	reader, err := storage.NewBinlogRecordReader(context.Background(), segment.GetInsertLogs(), s.schema, storage.WithVersion(segment.GetStorageVersion()), storage.WithStorageConfig(&indexpb.StorageConfig{
+		StorageType: "local",
+		RootPath:    paramtable.Get().LocalStorageCfg.Path.GetValue(),
+	}))
+	s.Require().NoError(err)
+	records := make([]storage.Record, 0)
+	for {
+		record, err := reader.Next()
+		if err != nil {
+			break
+		}
+		record.Retain()
+		records = append(records, record)
+	}
+
+	cmps := make([]func(ri, i, rj, j int) int, 0)
+	for _, fieldID := range sortedByFieldIDs {
+		field := typeutil.GetField(s.schema, fieldID)
+		switch field.DataType {
+		case schemapb.DataType_Int64:
+			cmps = append(cmps, func(ri, i, rj, j int) int {
+				vi := records[ri].Column(fieldID).(*array.Int64).Value(i)
+				vj := records[rj].Column(fieldID).(*array.Int64).Value(j)
+				if vi < vj {
+					return -1
+				}
+				if vi > vj {
+					return 1
+				}
+				return 0
+			})
+		case schemapb.DataType_String:
+			cmps = append(cmps, func(ri, i, rj, j int) int {
+				vi := records[ri].Column(fieldID).(*array.String).Value(i)
+				vj := records[rj].Column(fieldID).(*array.String).Value(j)
+				if vi < vj {
+					return -1
+				}
+				if vi > vj {
+					return 1
+				}
+				return 0
+			})
+		default:
+			panic("unsupported data type")
+		}
+	}
+
+	prevri := -1
+	previ := -1
+
+	for ri := 0; ri < len(records); ri++ {
+		for i := 0; i < records[ri].Len(); i++ {
+			if prevri == -1 {
+				prevri = ri
+				previ = i
+				continue
+			}
+			for _, cmp := range cmps {
+				c := cmp(prevri, previ, ri, i)
+				s.Require().True(c <= 0, "not sorted")
+				if c < 0 {
+					break
+				}
+			}
+			prevri = ri
+			previ = i
+		}
+	}
+}
+
+func TestNamespaceCompactorTestSuite(t *testing.T) {
+	suite.Run(t, new(NamespaceCompactorTestSuite))
+}

internal/datanode/services.go

@@ -214,11 +214,16 @@ func (node *DataNode) CompactionV2(ctx context.Context, req *datapb.CompactionPl
 		if req.GetPreAllocatedSegmentIDs() == nil || req.GetPreAllocatedSegmentIDs().GetBegin() == 0 {
 			return merr.Status(merr.WrapErrParameterInvalidMsg("invalid pre-allocated segmentID range")), nil
 		}
+		pk, err := typeutil.GetPrimaryFieldSchema(req.GetSchema())
+		if err != nil {
+			return merr.Status(err), err
+		}
 		task = compactor.NewMixCompactionTask(
 			taskCtx,
 			binlogIO,
 			req,
 			compactionParams,
+			[]int64{pk.GetFieldID()},
 		)
 	case datapb.CompactionType_ClusteringCompaction:
 		if req.GetPreAllocatedSegmentIDs() == nil || req.GetPreAllocatedSegmentIDs().GetBegin() == 0 {