enhance: Support L0 import (milvus-io#33514)

issue: milvus-io#33157 --------- Signed-off-by: bigsheeper <[email protected]>
bigsheeper · Jun 7, 2024 · 3540eee · 3540eee
1 parent 60baaed
commit 3540eee
Show file tree

Hide file tree

Showing 30 changed files with 1,425 additions and 389 deletions.
diff --git a/internal/datacoord/import_checker.go b/internal/datacoord/import_checker.go
@@ -28,6 +28,7 @@ import (
 	"github.com/milvus-io/milvus/internal/datacoord/broker"
 	"github.com/milvus-io/milvus/internal/proto/datapb"
 	"github.com/milvus-io/milvus/internal/proto/internalpb"
+	"github.com/milvus-io/milvus/internal/util/importutilv2"
 	"github.com/milvus-io/milvus/pkg/log"
 	"github.com/milvus-io/milvus/pkg/metrics"
 	"github.com/milvus-io/milvus/pkg/util/tsoutil"
@@ -253,7 +254,7 @@ func (c *importChecker) checkImportingJob(job ImportJob) {
 
 	// Verify completion of index building for imported segments.
 	unindexed := c.meta.indexMeta.GetUnindexedSegments(job.GetCollectionID(), segmentIDs)
-	if Params.DataCoordCfg.WaitForIndex.GetAsBool() && len(unindexed) > 0 {
+	if Params.DataCoordCfg.WaitForIndex.GetAsBool() && len(unindexed) > 0 && !importutilv2.IsL0Import(job.GetOptions()) {
 		log.Debug("waiting for import segments building index...", zap.Int64s("unindexed", unindexed))
 		return
 	}

diff --git a/internal/datacoord/import_scheduler.go b/internal/datacoord/import_scheduler.go
@@ -304,13 +304,13 @@ func (s *importScheduler) processInProgressImport(task ImportTask) {
 	if resp.GetState() == datapb.ImportTaskStateV2_Completed {
 		for _, info := range resp.GetImportSegmentsInfo() {
 			// try to parse path and fill logID
-			err = binlog.CompressBinLogs(info.GetBinlogs(), info.GetStatslogs())
+			err = binlog.CompressBinLogs(info.GetBinlogs(), info.GetDeltalogs(), info.GetStatslogs())
 			if err != nil {
-				log.Warn("fail to CompressFieldBinlogs for import binlogs",
+				log.Warn("fail to CompressBinLogs for import binlogs",
 					WrapTaskLog(task, zap.Int64("segmentID", info.GetSegmentID()), zap.Error(err))...)
 				return
 			}
-			op1 := UpdateBinlogsOperator(info.GetSegmentID(), info.GetBinlogs(), info.GetStatslogs(), nil)
+			op1 := UpdateBinlogsOperator(info.GetSegmentID(), info.GetBinlogs(), info.GetStatslogs(), info.GetDeltalogs())
 			op2 := UpdateStatusOperator(info.GetSegmentID(), commonpb.SegmentState_Flushed)
 			err = s.meta.UpdateSegmentsInfo(op1, op2)
 			if err != nil {

diff --git a/internal/datacoord/import_util.go b/internal/datacoord/import_util.go
@@ -30,6 +30,7 @@ import (
 	"github.com/milvus-io/milvus/internal/proto/datapb"
 	"github.com/milvus-io/milvus/internal/proto/internalpb"
 	"github.com/milvus-io/milvus/internal/storage"
+	"github.com/milvus-io/milvus/internal/util/importutilv2"
 	"github.com/milvus-io/milvus/pkg/log"
 	"github.com/milvus-io/milvus/pkg/util/merr"
 	"github.com/milvus-io/milvus/pkg/util/paramtable"
@@ -96,7 +97,7 @@ func NewImportTasks(fileGroups [][]*datapb.ImportFileStats,
 				FileStats:    group,
 			},
 		}
-		segments, err := AssignSegments(task, manager)
+		segments, err := AssignSegments(job, task, manager)
 		if err != nil {
 			return nil, err
 		}
@@ -106,7 +107,7 @@ func NewImportTasks(fileGroups [][]*datapb.ImportFileStats,
 	return tasks, nil
 }
 
-func AssignSegments(task ImportTask, manager Manager) ([]int64, error) {
+func AssignSegments(job ImportJob, task ImportTask, manager Manager) ([]int64, error) {
 	// merge hashed sizes
 	hashedDataSize := make(map[string]map[int64]int64) // vchannel->(partitionID->size)
 	for _, fileStats := range task.GetFileStats() {
@@ -120,15 +121,24 @@ func AssignSegments(task ImportTask, manager Manager) ([]int64, error) {
 		}
 	}
 
+	isL0Import := importutilv2.IsL0Import(job.GetOptions())
+
 	segmentMaxSize := paramtable.Get().DataCoordCfg.SegmentMaxSize.GetAsInt64() * 1024 * 1024
+	if isL0Import {
+		segmentMaxSize = paramtable.Get().DataNodeCfg.FlushDeleteBufferBytes.GetAsInt64()
+	}
+	segmentLevel := datapb.SegmentLevel_L1
+	if isL0Import {
+		segmentLevel = datapb.SegmentLevel_L0
+	}
 
 	// alloc new segments
 	segments := make([]int64, 0)
 	addSegment := func(vchannel string, partitionID int64, size int64) error {
 		ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
 		defer cancel()
 		for size > 0 {
-			segmentInfo, err := manager.AllocImportSegment(ctx, task.GetTaskID(), task.GetCollectionID(), partitionID, vchannel)
+			segmentInfo, err := manager.AllocImportSegment(ctx, task.GetTaskID(), task.GetCollectionID(), partitionID, vchannel, segmentLevel)
 			if err != nil {
 				return err
 			}
@@ -215,7 +225,12 @@ func RegroupImportFiles(job ImportJob, files []*datapb.ImportFileStats) [][]*dat
 		return nil
 	}
 
+	isL0Import := importutilv2.IsL0Import(job.GetOptions())
 	segmentMaxSize := paramtable.Get().DataCoordCfg.SegmentMaxSize.GetAsInt() * 1024 * 1024
+	if isL0Import {
+		segmentMaxSize = paramtable.Get().DataNodeCfg.FlushDeleteBufferBytes.GetAsInt()
+	}
+
 	threshold := paramtable.Get().DataCoordCfg.MaxSizeInMBPerImportTask.GetAsInt() * 1024 * 1024
 	maxSizePerFileGroup := segmentMaxSize * len(job.GetPartitionIDs()) * len(job.GetVchannels())
 	if maxSizePerFileGroup > threshold {

diff --git a/internal/datacoord/import_util_test.go b/internal/datacoord/import_util_test.go
@@ -96,15 +96,16 @@ func TestImportUtil_NewImportTasks(t *testing.T) {
 		return id, id + n, nil
 	})
 	manager := NewMockManager(t)
-	manager.EXPECT().AllocImportSegment(mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything).
-		RunAndReturn(func(ctx context.Context, taskID int64, collectionID int64, partitionID int64, vchannel string) (*SegmentInfo, error) {
+	manager.EXPECT().AllocImportSegment(mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything).
+		RunAndReturn(func(ctx context.Context, taskID int64, collectionID int64, partitionID int64, vchannel string, level datapb.SegmentLevel) (*SegmentInfo, error) {
 			return &SegmentInfo{
 				SegmentInfo: &datapb.SegmentInfo{
 					ID:            rand.Int63(),
 					CollectionID:  collectionID,
 					PartitionID:   partitionID,
 					InsertChannel: vchannel,
 					IsImporting:   true,
+					Level:         level,
 				},
 			}, nil
 		})

diff --git a/internal/datacoord/mock_segment_manager.go b/internal/datacoord/mock_segment_manager.go
diff --git a/internal/datacoord/segment_manager.go b/internal/datacoord/segment_manager.go
@@ -76,7 +76,7 @@ type Manager interface {
 
 	// AllocSegment allocates rows and record the allocation.
 	AllocSegment(ctx context.Context, collectionID, partitionID UniqueID, channelName string, requestRows int64) ([]*Allocation, error)
-	AllocImportSegment(ctx context.Context, taskID int64, collectionID UniqueID, partitionID UniqueID, channelName string) (*SegmentInfo, error)
+	AllocImportSegment(ctx context.Context, taskID int64, collectionID UniqueID, partitionID UniqueID, channelName string, level datapb.SegmentLevel) (*SegmentInfo, error)
 	// DropSegment drops the segment from manager.
 	DropSegment(ctx context.Context, segmentID UniqueID)
 	// FlushImportSegments set importing segment state to Flushed.
@@ -350,7 +350,7 @@ func (s *SegmentManager) genExpireTs(ctx context.Context) (Timestamp, error) {
 }
 
 func (s *SegmentManager) AllocImportSegment(ctx context.Context, taskID int64, collectionID UniqueID,
-	partitionID UniqueID, channelName string,
+	partitionID UniqueID, channelName string, level datapb.SegmentLevel,
 ) (*SegmentInfo, error) {
 	log := log.Ctx(ctx)
 	ctx, sp := otel.Tracer(typeutil.DataCoordRole).Start(ctx, "open-Segment")
@@ -378,7 +378,7 @@ func (s *SegmentManager) AllocImportSegment(ctx context.Context, taskID int64, c
 		NumOfRows:      0,
 		State:          commonpb.SegmentState_Importing,
 		MaxRowNum:      0,
-		Level:          datapb.SegmentLevel_L1,
+		Level:          level,
 		LastExpireTime: math.MaxUint64,
 		StartPosition:  position,
 		DmlPosition:    position,
@@ -394,9 +394,10 @@ func (s *SegmentManager) AllocImportSegment(ctx context.Context, taskID int64, c
 	s.segments = append(s.segments, id)
 	log.Info("add import segment done",
 		zap.Int64("taskID", taskID),
-		zap.Int64("CollectionID", segmentInfo.CollectionID),
-		zap.Int64("SegmentID", segmentInfo.ID),
-		zap.String("Channel", segmentInfo.InsertChannel))
+		zap.Int64("collectionID", segmentInfo.CollectionID),
+		zap.Int64("segmentID", segmentInfo.ID),
+		zap.String("channel", segmentInfo.InsertChannel),
+		zap.String("level", level.String()))
 
 	return segment, nil
 }