Fix large files import

Signed-off-by: bigsheeper <yihao.dai@zilliz.com>

configs/milvus.yaml

@@ -403,7 +403,7 @@ dataCoord:
     enableLevelZero: false
   enableCompaction: true # Enable data segment compaction
   compaction:
-    enableAutoCompaction: true
+    enableAutoCompaction: false
     rpcTimeout: 10 # compaction rpc request timeout in seconds
     maxParallelTaskNum: 10 # max parallel compaction task number
     indexBasedCompaction: true
@@ -475,8 +475,8 @@ dataNode:
     # suggest to set it bigger on large collection numbers to avoid blocking
     updateChannelCheckpointMaxParallel: 1000
   import:
-    maxConcurrentTaskNum: 10 # The maximum number of import/pre-import tasks allowed to run concurrently on a datanode.
-    bufferSize: 64 # The maximum buffer size in MB that each import/pre-import task is permitted to utilize.
+    maxConcurrentTaskNum: 20 # The maximum number of import/pre-import tasks allowed to run concurrently on a datanode.
+    bufferSize: 256 # The maximum buffer size in MB that each import/pre-import task is permitted to utilize.
 
 # Configures the system log output.
 log:

deployments/docker/dev/docker-compose.yml

@@ -46,7 +46,7 @@ services:
       MINIO_ACCESS_KEY: minioadmin
       MINIO_SECRET_KEY: minioadmin
     volumes:
-      - ${DOCKER_VOLUME_DIRECTORY:-.}/volumes/minio:/minio_data
+      - /data/milvus-data/volumes/minio:/minio_data
     command: minio server /minio_data --console-address ":9001"
     healthcheck:
       test:

internal/datacoord/import_checker.go

@@ -195,7 +195,7 @@ func (c *importChecker) checkImportState(requestID int64) {
 			return
 		}
 	}
-	if AreAllTasksFinished(tasks, c.meta) {
+	if AreAllTasksFinished(tasks, c.meta, c.imeta) {
 		return
 	}
 	for _, task := range tasks {

internal/datacoord/import_checker_test.go

@@ -49,7 +49,7 @@ func TestImportChecker(t *testing.T) {
 		return rand.Int63(), nil
 	})
 
-	imeta, err := NewImportMeta(catalog)
+	imeta, err := NewImportMeta(nil, catalog)
 	assert.NoError(t, err)
 
 	pit1 := &preImportTask{

internal/datacoord/import_meta.go

@@ -17,6 +17,10 @@
 package datacoord
 
 import (
+	"context"
+	"github.com/milvus-io/milvus-proto/go-api/v2/schemapb"
+	"github.com/milvus-io/milvus/internal/datacoord/broker"
+	"github.com/milvus-io/milvus/pkg/util/retry"
 	"sync"
 	"time"
 
@@ -39,7 +43,19 @@ type importMeta struct {
 	catalog metastore.DataCoordCatalog
 }
 
-func NewImportMeta(catalog metastore.DataCoordCatalog) (ImportMeta, error) {
+func NewImportMeta(broker broker.Broker, catalog metastore.DataCoordCatalog) (ImportMeta, error) {
+	// TODO: dyh, remove dependency of broker
+	getSchema := func(collectionID int64) (*schemapb.CollectionSchema, error) {
+		var schema *schemapb.CollectionSchema
+		ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
+		defer cancel()
+		err := retry.Do(ctx, func() error {
+			resp, err := broker.DescribeCollectionInternal(ctx, collectionID)
+			schema = resp.GetSchema()
+			return err
+		})
+		return schema, err
+	}
 	restoredTasks, err := catalog.ListImportTasks()
 	if err != nil {
 		return nil, err
@@ -48,9 +64,29 @@ func NewImportMeta(catalog metastore.DataCoordCatalog) (ImportMeta, error) {
 	for _, t := range restoredTasks {
 		switch task := t.(type) {
 		case *datapb.PreImportTask:
-			tasks[task.GetTaskID()] = &preImportTask{PreImportTask: task} // TODO: dyh, resolve empty schema
+			catalog.DropImportTask(task.GetTaskID())
+			continue
+			schema, err := getSchema(task.GetCollectionID())
+			if err != nil {
+				return nil, err
+			}
+			tasks[task.GetTaskID()] = &preImportTask{
+				PreImportTask:  task,
+				schema:         schema,
+				lastActiveTime: time.Now(),
+			}
 		case *datapb.ImportTaskV2:
-			tasks[task.GetTaskID()] = &importTask{ImportTaskV2: task}
+			catalog.DropImportTask(task.GetTaskID())
+			continue
+			schema, err := getSchema(task.GetCollectionID())
+			if err != nil {
+				return nil, err
+			}
+			tasks[task.GetTaskID()] = &importTask{
+				ImportTaskV2:   task,
+				schema:         schema,
+				lastActiveTime: time.Now(),
+			}
 		}
 	}
 	return &importMeta{

internal/datacoord/import_meta_test.go

@@ -39,7 +39,7 @@ func TestImportMeta_Restore(t *testing.T) {
 		},
 	}, nil)
 
-	meta, err := NewImportMeta(catalog)
+	meta, err := NewImportMeta(nil, catalog)
 	assert.NoError(t, err)
 
 	tasks := meta.GetBy()
@@ -58,7 +58,7 @@ func TestImportMeta_Normal(t *testing.T) {
 	catalog.EXPECT().SaveImportTask(mock.Anything).Return(nil)
 	catalog.EXPECT().DropImportTask(mock.Anything).Return(nil)
 
-	meta, err := NewImportMeta(catalog)
+	meta, err := NewImportMeta(nil, catalog)
 	assert.NoError(t, err)
 
 	task1 := &importTask{
@@ -117,7 +117,7 @@ func TestImportMeta_Failed(t *testing.T) {
 	catalog.EXPECT().SaveImportTask(mock.Anything).Return(mockErr)
 	catalog.EXPECT().DropImportTask(mock.Anything).Return(mockErr)
 
-	meta, err := NewImportMeta(catalog)
+	meta, err := NewImportMeta(nil, catalog)
 	assert.NoError(t, err)
 
 	task1 := &importTask{

internal/datacoord/import_scheduler.go

@@ -18,6 +18,7 @@ package datacoord
 
 import (
 	"github.com/milvus-io/milvus/internal/proto/internalpb"
+	"sort"
 	"sync"
 	"time"
 
@@ -84,27 +85,37 @@ func (s *importScheduler) Close() {
 }
 
 func (s *importScheduler) process() {
-	tasks := s.imeta.GetBy()
-	for _, task := range tasks {
-		switch task.GetState() {
-		case internalpb.ImportState_Pending:
-			switch task.GetType() {
-			case PreImportTaskType:
-				s.processPendingPreImport(task)
-			case ImportTaskType:
-				s.processPendingImport(task)
-			}
-		case internalpb.ImportState_InProgress:
-			switch task.GetType() {
-			case PreImportTaskType:
-				s.processInProgressPreImport(task)
-			case ImportTaskType:
-				s.processInProgressImport(task)
+	all := s.imeta.GetBy()
+	tasksByReq := lo.GroupBy(all, func(t ImportTask) int64 {
+		return t.GetRequestID()
+	})
+	requests := lo.Keys(tasksByReq)
+	sort.Slice(requests, func(i, j int) bool {
+		return requests[i] < requests[j]
+	})
+	for _, request := range requests {
+		tasks := tasksByReq[request]
+		for _, task := range tasks {
+			switch task.GetState() {
+			case internalpb.ImportState_Pending:
+				switch task.GetType() {
+				case PreImportTaskType:
+					s.processPendingPreImport(task)
+				case ImportTaskType:
+					s.processPendingImport(task)
+				}
+			case internalpb.ImportState_InProgress:
+				switch task.GetType() {
+				case PreImportTaskType:
+					s.processInProgressPreImport(task)
+				case ImportTaskType:
+					s.processInProgressImport(task)
+				}
+			case internalpb.ImportState_Completed:
+				s.processCompleted(task)
+			case internalpb.ImportState_Failed:
+				s.processFailed(task)
 			}
-		case internalpb.ImportState_Completed:
-			s.processCompleted(task)
-		case internalpb.ImportState_Failed:
-			s.processFailed(task)
 		}
 	}
 }
@@ -137,13 +148,13 @@ func (s *importScheduler) getIdleNode() int64 {
 			return nodeID
 		}
 	}
+	//log.Warn("no datanode can be scheduled")
 	return NullNodeID
 }
 
 func (s *importScheduler) processPendingPreImport(task ImportTask) {
 	nodeID := s.getIdleNode()
 	if nodeID == NullNodeID {
-		log.Warn("no datanode can be scheduled", WrapLogFields(task)...)
 		return
 	}
 	log.Info("processing pending preimport task...", WrapLogFields(task)...)
@@ -165,7 +176,6 @@ func (s *importScheduler) processPendingPreImport(task ImportTask) {
 func (s *importScheduler) processPendingImport(task ImportTask) {
 	nodeID := s.getIdleNode()
 	if nodeID == NullNodeID {
-		log.Warn("no datanode can be scheduled", WrapLogFields(task)...)
 		return
 	}
 	log.Info("processing pending import task...", WrapLogFields(task)...)
@@ -232,8 +242,11 @@ func (s *importScheduler) processInProgressImport(task ImportTask) {
 		return
 	}
 	for _, info := range resp.GetImportSegmentsInfo() {
-		// TODO: dyh, check if segment info changed to save meta op
-		op1 := UpdateBinlogsOperator(info.GetSegmentID(), info.GetBinlogs(), info.GetStatslogs(), nil)
+		segment := s.meta.GetSegment(info.GetSegmentID())
+		if info.GetImportedRows() <= segment.GetNumOfRows() {
+			continue // rows not changed, no need to update
+		}
+		op1 := ReplaceBinlogsOperator(info.GetSegmentID(), info.GetBinlogs(), info.GetStatslogs(), nil)
 		op2 := UpdateNumOfRows(info.GetSegmentID(), info.GetImportedRows())
 		err = s.meta.UpdateSegmentsInfo(op1, op2)
 		if err != nil {

internal/datacoord/import_scheduler_test.go

@@ -62,7 +62,7 @@ func (s *ImportSchedulerSuite) SetupTest() {
 		ID:     s.collectionID,
 		Schema: newTestSchema(),
 	})
-	s.imeta, err = NewImportMeta(s.catalog)
+	s.imeta, err = NewImportMeta(nil, s.catalog)
 	s.NoError(err)
 	s.scheduler = NewImportScheduler(s.meta, s.cluster, s.alloc, s.imeta).(*importScheduler)
 }

internal/datacoord/import_util.go

@@ -18,6 +18,8 @@ package datacoord
 
 import (
 	"context"
+	"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
+	"github.com/milvus-io/milvus/pkg/log"
 	"sort"
 	"time"
 
@@ -245,20 +247,32 @@ func AddImportSegment(cluster Cluster, meta *meta, segmentID int64) error {
 	return err
 }
 
-func AreAllTasksFinished(tasks []ImportTask, meta *meta) bool {
+func AreAllTasksFinished(tasks []ImportTask, meta *meta, imeta ImportMeta) bool {
+	var finished = true
 	for _, task := range tasks {
 		if task.GetState() != internalpb.ImportState_Completed {
 			return false
 		}
 		segmentIDs := task.(*importTask).GetSegmentIDs()
+		validSegments := make([]int64, 0)
 		for _, segmentID := range segmentIDs {
 			segment := meta.GetSegment(segmentID)
+			if segment.GetState() == commonpb.SegmentState_Dropped {
+				continue // this segment has been compacted
+			}
+			validSegments = append(validSegments, segmentID)
 			if segment.GetIsImporting() {
-				return false
+				finished = false
+			}
+		}
+		if len(validSegments) < len(segmentIDs) {
+			err := imeta.Update(task.GetTaskID(), UpdateSegmentIDs(validSegments))
+			if err != nil {
+				log.Warn("update segmentIDs failed", WrapLogFields(task, zap.Error(err))...)
 			}
 		}
 	}
-	return true
+	return finished
 }
 
 func GetImportProgress(requestID int64, imeta ImportMeta, meta *meta) (int64, internalpb.ImportState, string) {