mirror of
https://gitee.com/milvus-io/milvus.git
synced 2025-12-06 17:18:35 +08:00
feat: add a namespace compactor (#44253)
#44011 --------- Signed-off-by: sunby <sunbingyi1992@gmail.com>
This commit is contained in:
Bingyi Sun
GitHub
parent
9add663a08
commit
08a2678f07
@ -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
|
||||
}
|
||||
|
||||
@ -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
|
||||
|
||||
@ -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() {
|
||||
|
||||
45
internal/datanode/compactor/namespace_compactor.go
Normal file
45
internal/datanode/compactor/namespace_compactor.go
Normal file
@ -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),
|
||||
}
|
||||
}
|
||||
224
internal/datanode/compactor/namespace_compactor_test.go
Normal file
224
internal/datanode/compactor/namespace_compactor_test.go
Normal file
@ -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))
|
||||
}
|
||||
@ -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 {
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user