import-into: support running with 1 thread when using global sort (pingcap#58572)

joccau · zeminzhou · commit 579734dbac7f · 2025-05-06T14:00:04.000+08:00
close pingcap#58680
diff --git a/pkg/disttask/importinto/job.go b/pkg/disttask/importinto/job.go
@@ -28,7 +28,6 @@ import (
 	"github.com/pingcap/tidb/pkg/domain/infosync"
 	"github.com/pingcap/tidb/pkg/executor/importer"
 	"github.com/pingcap/tidb/pkg/kv"
-	"github.com/pingcap/tidb/pkg/lightning/checkpoints"
 	"github.com/pingcap/tidb/pkg/metrics"
 	"github.com/pingcap/tidb/pkg/sessionctx"
 	"github.com/pingcap/tidb/pkg/util/logutil"
@@ -39,20 +38,20 @@ import (
 // SubmitStandaloneTask submits a task to the distribute framework that only runs on the current node.
 // when import from server-disk, pass engine checkpoints too, as scheduler might run on another
 // node where we can't access the data files.
-func SubmitStandaloneTask(ctx context.Context, plan *importer.Plan, stmt string, ecp map[int32]*checkpoints.EngineCheckpoint) (int64, *proto.TaskBase, error) {
+func SubmitStandaloneTask(ctx context.Context, plan *importer.Plan, stmt string, chunkMap map[int32][]importer.Chunk) (int64, *proto.TaskBase, error) {
 	serverInfo, err := infosync.GetServerInfo()
 	if err != nil {
 		return 0, nil, err
 	}
-	return doSubmitTask(ctx, plan, stmt, serverInfo, toChunkMap(ecp))
+	return doSubmitTask(ctx, plan, stmt, serverInfo, chunkMap)
 }
 
 // SubmitTask submits a task to the distribute framework that runs on all managed nodes.
 func SubmitTask(ctx context.Context, plan *importer.Plan, stmt string) (int64, *proto.TaskBase, error) {
 	return doSubmitTask(ctx, plan, stmt, nil, nil)
 }
 
-func doSubmitTask(ctx context.Context, plan *importer.Plan, stmt string, instance *infosync.ServerInfo, chunkMap map[int32][]Chunk) (int64, *proto.TaskBase, error) {
+func doSubmitTask(ctx context.Context, plan *importer.Plan, stmt string, instance *infosync.ServerInfo, chunkMap map[int32][]importer.Chunk) (int64, *proto.TaskBase, error) {
 	var instances []*infosync.ServerInfo
 	if instance != nil {
 		instances = append(instances, instance)
diff --git a/pkg/disttask/importinto/planner.go b/pkg/disttask/importinto/planner.go
@@ -54,7 +54,7 @@ type LogicalPlan struct {
 	Plan              importer.Plan
 	Stmt              string
 	EligibleInstances []*infosync.ServerInfo
-	ChunkMap          map[int32][]Chunk
+	ChunkMap          map[int32][]importer.Chunk
 }
 
 // ToTaskMeta converts the logical plan to task meta.
@@ -156,7 +156,7 @@ func (p *LogicalPlan) ToPhysicalPlan(planCtx planner.PlanCtx) (*planner.Physical
 type ImportSpec struct {
 	ID     int32
 	Plan   importer.Plan
-	Chunks []Chunk
+	Chunks []importer.Chunk
 }
 
 // ToSubtaskMeta converts the import spec to subtask meta.
@@ -257,7 +257,7 @@ func buildController(plan *importer.Plan, stmt string) (*importer.LoadDataContro
 }
 
 func generateImportSpecs(pCtx planner.PlanCtx, p *LogicalPlan) ([]planner.PipelineSpec, error) {
-	var chunkMap map[int32][]Chunk
+	var chunkMap map[int32][]importer.Chunk
 	if len(p.ChunkMap) > 0 {
 		chunkMap = p.ChunkMap
 	} else {
@@ -270,21 +270,21 @@ func generateImportSpecs(pCtx planner.PlanCtx, p *LogicalPlan) ([]planner.Pipeli
 		}
 
 		controller.SetExecuteNodeCnt(pCtx.ExecuteNodesCnt)
-		engineCheckpoints, err2 := controller.PopulateChunks(pCtx.Ctx)
+		chunkMap, err2 = controller.PopulateChunks(pCtx.Ctx)
 		if err2 != nil {
 			return nil, err2
 		}
-		chunkMap = toChunkMap(engineCheckpoints)
 	}
+
 	importSpecs := make([]planner.PipelineSpec, 0, len(chunkMap))
-	for id := range chunkMap {
+	for id, chunks := range chunkMap {
 		if id == common.IndexEngineID {
 			continue
 		}
 		importSpec := &ImportSpec{
 			ID:     id,
 			Plan:   p.Plan,
-			Chunks: chunkMap[id],
+			Chunks: chunks,
 		}
 		importSpecs = append(importSpecs, importSpec)
 	}
diff --git a/pkg/disttask/importinto/planner_test.go b/pkg/disttask/importinto/planner_test.go
@@ -44,7 +44,7 @@ func TestLogicalPlan(t *testing.T) {
 		Plan:              importer.Plan{},
 		Stmt:              `IMPORT INTO db.tb FROM 'gs://test-load/*.csv?endpoint=xxx'`,
 		EligibleInstances: []*infosync.ServerInfo{{ID: "1"}},
-		ChunkMap:          map[int32][]Chunk{1: {{Path: "gs://test-load/1.csv"}}},
+		ChunkMap:          map[int32][]importer.Chunk{1: {{Path: "gs://test-load/1.csv"}}},
 	}
 	bs, err := logicalPlan.ToTaskMeta()
 	require.NoError(t, err)
@@ -65,7 +65,7 @@ func TestToPhysicalPlan(t *testing.T) {
 		},
 		Stmt:              `IMPORT INTO db.tb FROM 'gs://test-load/*.csv?endpoint=xxx'`,
 		EligibleInstances: []*infosync.ServerInfo{{ID: "1"}},
-		ChunkMap:          map[int32][]Chunk{chunkID: {{Path: "gs://test-load/1.csv"}}},
+		ChunkMap:          map[int32][]importer.Chunk{chunkID: {{Path: "gs://test-load/1.csv"}}},
 	}
 	planCtx := planner.PlanCtx{
 		NextTaskStep: proto.ImportStepImport,
diff --git a/pkg/disttask/importinto/proto.go b/pkg/disttask/importinto/proto.go
@@ -23,7 +23,6 @@ import (
 	"github.com/pingcap/tidb/pkg/executor/importer"
 	"github.com/pingcap/tidb/pkg/lightning/backend"
 	"github.com/pingcap/tidb/pkg/lightning/backend/external"
-	"github.com/pingcap/tidb/pkg/lightning/mydump"
 	"github.com/pingcap/tidb/pkg/lightning/verification"
 	"github.com/pingcap/tidb/pkg/meta/autoid"
 )
@@ -44,7 +43,7 @@ type TaskMeta struct {
 	// files to the framework scheduler which might run on another instance.
 	// we use a map from engine ID to chunks since we need support split_file for CSV,
 	// so need to split them into engines before passing to scheduler.
-	ChunkMap map[int32][]Chunk
+	ChunkMap map[int32][]importer.Chunk
 }
 
 // ImportStepMeta is the meta of import step.
@@ -53,7 +52,7 @@ type TaskMeta struct {
 type ImportStepMeta struct {
 	// this is the engine ID, not the id in tidb_background_subtask table.
 	ID       int32
-	Chunks   []Chunk
+	Chunks   []importer.Chunk
 	Checksum map[int64]Checksum // see KVGroupChecksum for definition of map key.
 	Result   Result
 	// MaxIDs stores the max id that have been used during encoding for each allocator type.
@@ -142,7 +141,7 @@ func (sv *SharedVars) mergeIndexSummary(indexID int64, summary *external.WriterS
 // TaskExecutor will split the subtask into minimal tasks(Chunks -> Chunk)
 type importStepMinimalTask struct {
 	Plan       importer.Plan
-	Chunk      Chunk
+	Chunk      importer.Chunk
 	SharedVars *SharedVars
 	panicked   *atomic.Bool
 }
@@ -158,19 +157,6 @@ func (t *importStepMinimalTask) String() string {
 	return fmt.Sprintf("chunk:%s:%d", t.Chunk.Path, t.Chunk.Offset)
 }
 
-// Chunk records the chunk information.
-type Chunk struct {
-	Path         string
-	FileSize     int64
-	Offset       int64
-	EndOffset    int64
-	PrevRowIDMax int64
-	RowIDMax     int64
-	Type         mydump.SourceType
-	Compression  mydump.Compression
-	Timestamp    int64
-}
-
 // Checksum records the checksum information.
 type Checksum struct {
 	Sum  uint64
diff --git a/pkg/disttask/importinto/scheduler.go b/pkg/disttask/importinto/scheduler.go
@@ -35,7 +35,6 @@ import (
 	"github.com/pingcap/tidb/pkg/errno"
 	"github.com/pingcap/tidb/pkg/executor/importer"
 	"github.com/pingcap/tidb/pkg/kv"
-	"github.com/pingcap/tidb/pkg/lightning/checkpoints"
 	"github.com/pingcap/tidb/pkg/lightning/common"
 	"github.com/pingcap/tidb/pkg/lightning/config"
 	"github.com/pingcap/tidb/pkg/lightning/metric"
@@ -548,18 +547,6 @@ func updateMeta(task *proto.Task, taskMeta *TaskMeta) error {
 	return nil
 }
 
-// todo: converting back and forth, we should unify struct and remove this function later.
-func toChunkMap(engineCheckpoints map[int32]*checkpoints.EngineCheckpoint) map[int32][]Chunk {
-	chunkMap := make(map[int32][]Chunk, len(engineCheckpoints))
-	for id, ecp := range engineCheckpoints {
-		chunkMap[id] = make([]Chunk, 0, len(ecp.Chunks))
-		for _, chunkCheckpoint := range ecp.Chunks {
-			chunkMap[id] = append(chunkMap[id], toChunk(*chunkCheckpoint))
-		}
-	}
-	return chunkMap
-}
-
 func getStepOfEncode(globalSort bool) proto.Step {
 	if globalSort {
 		return proto.ImportStepEncodeAndSort
diff --git a/pkg/disttask/importinto/scheduler_testkit_test.go b/pkg/disttask/importinto/scheduler_testkit_test.go
@@ -69,7 +69,7 @@ func TestSchedulerExtLocalSort(t *testing.T) {
 		},
 		Stmt:              `IMPORT INTO db.tb FROM 'gs://test-load/*.csv?endpoint=xxx'`,
 		EligibleInstances: []*infosync.ServerInfo{{ID: "1"}},
-		ChunkMap:          map[int32][]importinto.Chunk{1: {{Path: "gs://test-load/1.csv"}}},
+		ChunkMap:          map[int32][]importer.Chunk{1: {{Path: "gs://test-load/1.csv"}}},
 	}
 	bs, err := logicalPlan.ToTaskMeta()
 	require.NoError(t, err)
@@ -209,7 +209,7 @@ func TestSchedulerExtGlobalSort(t *testing.T) {
 		},
 		Stmt:              `IMPORT INTO db.tb FROM 'gs://test-load/*.csv?endpoint=xxx'`,
 		EligibleInstances: []*infosync.ServerInfo{{ID: "1"}},
-		ChunkMap: map[int32][]importinto.Chunk{
+		ChunkMap: map[int32][]importer.Chunk{
 			1: {{Path: "gs://test-load/1.csv"}},
 			2: {{Path: "gs://test-load/2.csv"}},
 		},
diff --git a/pkg/disttask/importinto/wrapper.go b/pkg/disttask/importinto/wrapper.go
@@ -15,11 +15,12 @@
 package importinto
 
 import (
+	"github.com/pingcap/tidb/pkg/executor/importer"
 	"github.com/pingcap/tidb/pkg/lightning/checkpoints"
 	"github.com/pingcap/tidb/pkg/lightning/mydump"
 )
 
-func toChunkCheckpoint(chunk Chunk) checkpoints.ChunkCheckpoint {
+func toChunkCheckpoint(chunk importer.Chunk) checkpoints.ChunkCheckpoint {
 	return checkpoints.ChunkCheckpoint{
 		Key: checkpoints.ChunkCheckpointKey{
 			Path:   chunk.Path,
@@ -41,8 +42,8 @@ func toChunkCheckpoint(chunk Chunk) checkpoints.ChunkCheckpoint {
 	}
 }
 
-func toChunk(chunkCheckpoint checkpoints.ChunkCheckpoint) Chunk {
-	return Chunk{
+func toChunk(chunkCheckpoint checkpoints.ChunkCheckpoint) importer.Chunk {
+	return importer.Chunk{
 		Path:         chunkCheckpoint.FileMeta.Path,
 		FileSize:     chunkCheckpoint.FileMeta.FileSize,
 		Offset:       chunkCheckpoint.Chunk.Offset,
diff --git a/pkg/executor/import_into.go b/pkg/executor/import_into.go
@@ -214,11 +214,11 @@ func (e *ImportIntoExec) submitTask(ctx context.Context) (int64, *proto.TaskBase
 	logutil.Logger(ctx).Info("get job importer", zap.Stringer("param", e.controller.Parameters),
 		zap.Bool("dist-task-enabled", vardef.EnableDistTask.Load()))
 	if importFromServer {
-		ecp, err2 := e.controller.PopulateChunks(ctx)
+		chunkMap, err2 := e.controller.PopulateChunks(ctx)
 		if err2 != nil {
 			return 0, nil, err2
 		}
-		return importinto.SubmitStandaloneTask(ctx, e.controller.Plan, e.stmt, ecp)
+		return importinto.SubmitStandaloneTask(ctx, e.controller.Plan, e.stmt, chunkMap)
 	}
 	// if tidb_enable_dist_task=true, we import distributively, otherwise we import on current node.
 	if vardef.EnableDistTask.Load() {
diff --git a/pkg/executor/importer/import.go b/pkg/executor/importer/import.go
@@ -1014,6 +1014,7 @@ func (e *LoadDataController) InitDataStore(ctx context.Context) error {
 	}
 	return nil
 }
+
 func (*LoadDataController) initExternalStore(ctx context.Context, u *url.URL, target string) (storage.ExternalStorage, error) {
 	b, err2 := storage.ParseBackendFromURL(u, nil)
 	if err2 != nil {
diff --git a/pkg/executor/importer/importer_testkit_test.go b/pkg/executor/importer/importer_testkit_test.go
@@ -370,6 +370,7 @@ func TestPopulateChunks(t *testing.T) {
 	engines, err := ti.PopulateChunks(ctx)
 	require.NoError(t, err)
 	require.Len(t, engines, 3)
-	require.Len(t, engines[0].Chunks, 2)
-	require.Len(t, engines[1].Chunks, 1)
+	require.Len(t, engines[0], 2)
+	require.Len(t, engines[1], 1)
+	require.Len(t, engines[common.IndexEngineID], 0)
 }
diff --git a/pkg/executor/importer/precheck.go b/pkg/executor/importer/precheck.go
@@ -46,7 +46,7 @@ var GetEtcdClient = getEtcdClient
 //   - when import from file
 //     1. there is no active job on the target table
 //     2. the total file size > 0
-//     3. if global sort, thread count >= 16 and have required privileges
+//     3. if global sort, check required privileges
 //   - target table should be empty
 //   - no CDC or PiTR tasks running
 //
@@ -63,11 +63,6 @@ func (e *LoadDataController) CheckRequirements(ctx context.Context, conn sqlexec
 		if err := e.checkTotalFileSize(); err != nil {
 			return err
 		}
-		// run global sort with < 8 thread might OOM on ingest step
-		// TODO: remove this limit after control memory usage.
-		if e.IsGlobalSort() && e.ThreadCnt < 8 {
-			return exeerrors.ErrLoadDataPreCheckFailed.FastGenByArgs("global sort requires at least 8 threads")
-		}
 	}
 	if err := e.checkTableEmpty(ctx, conn); err != nil {
 		return err
diff --git a/pkg/executor/importer/precheck_test.go b/pkg/executor/importer/precheck_test.go
@@ -104,16 +104,9 @@ func TestCheckRequirements(t *testing.T) {
 
 	// make checkTotalFileSize pass
 	c.TotalFileSize = 1
-	// global sort with thread count < 8
-	c.ThreadCnt = 7
-	c.CloudStorageURI = "s3://test"
-	err = c.CheckRequirements(ctx, conn)
-	require.ErrorIs(t, err, exeerrors.ErrLoadDataPreCheckFailed)
-	require.ErrorContains(t, err, "global sort requires at least 8 threads")
-
-	// reset fields, make global sort thread check pass
 	c.ThreadCnt = 1
 	c.CloudStorageURI = ""
+
 	// non-empty table
 	_, err = conn.Execute(ctx, "insert into test.t values(1)")
 	require.NoError(t, err)
@@ -178,8 +171,8 @@ func TestCheckRequirements(t *testing.T) {
 	require.NoError(t, err)
 	require.NoError(t, c.CheckRequirements(ctx, conn))
 
-	// with global sort
-	c.Plan.ThreadCnt = 8
+	// with global sort with threadCnt < 8
+	c.Plan.ThreadCnt = 2
 	c.Plan.CloudStorageURI = ":"
 	require.ErrorIs(t, c.CheckRequirements(ctx, conn), exeerrors.ErrLoadDataInvalidURI)
 	c.Plan.CloudStorageURI = "sdsdsdsd://sdsdsdsd"
diff --git a/pkg/executor/importer/table_import.go b/pkg/executor/importer/table_import.go
diff --git a/pkg/meta/model/table.go b/pkg/meta/model/table.go

Original file line number	Diff line number	Diff line change
`@@ -54,7 +54,7 @@ type LogicalPlan struct {`
`54`	`54`	`Plan importer.Plan`
`55`	`55`	`Stmt string`
`56`	`56`	`EligibleInstances []*infosync.ServerInfo`
`57`		`- ChunkMap map[int32][]Chunk`
	`57`	`+ ChunkMap map[int32][]importer.Chunk`
`58`	`58`	`}`
`59`	`59`
`60`	`60`	`// ToTaskMeta converts the logical plan to task meta.`
`@@ -156,7 +156,7 @@ func (p LogicalPlan) ToPhysicalPlan(planCtx planner.PlanCtx) (planner.Physical`
`156`	`156`	`type ImportSpec struct {`
`157`	`157`	`ID int32`
`158`	`158`	`Plan importer.Plan`
`159`		`- Chunks []Chunk`
	`159`	`+ Chunks []importer.Chunk`
`160`	`160`	`}`
`161`	`161`
`162`	`162`	`// ToSubtaskMeta converts the import spec to subtask meta.`
`@@ -257,7 +257,7 @@ func buildController(plan importer.Plan, stmt string) (importer.LoadDataContro`
`257`	`257`	`}`
`258`	`258`
`259`	`259`	`func generateImportSpecs(pCtx planner.PlanCtx, p *LogicalPlan) ([]planner.PipelineSpec, error) {`
`260`		`- var chunkMap map[int32][]Chunk`
	`260`	`+ var chunkMap map[int32][]importer.Chunk`
`261`	`261`	`if len(p.ChunkMap) > 0 {`
`262`	`262`	`chunkMap = p.ChunkMap`
`263`	`263`	`} else {`
`@@ -270,21 +270,21 @@ func generateImportSpecs(pCtx planner.PlanCtx, p *LogicalPlan) ([]planner.Pipeli`
`270`	`270`	`}`
`271`	`271`
`272`	`272`	`controller.SetExecuteNodeCnt(pCtx.ExecuteNodesCnt)`
`273`		`- engineCheckpoints, err2 := controller.PopulateChunks(pCtx.Ctx)`
	`273`	`+ chunkMap, err2 = controller.PopulateChunks(pCtx.Ctx)`
`274`	`274`	`if err2 != nil {`
`275`	`275`	`return nil, err2`
`276`	`276`	`}`
`277`		`- chunkMap = toChunkMap(engineCheckpoints)`
`278`	`277`	`}`
	`278`	`+`
`279`	`279`	`importSpecs := make([]planner.PipelineSpec, 0, len(chunkMap))`
`280`		`- for id := range chunkMap {`
	`280`	`+ for id, chunks := range chunkMap {`
`281`	`281`	`if id == common.IndexEngineID {`
`282`	`282`	`continue`
`283`	`283`	`}`
`284`	`284`	`importSpec := &ImportSpec{`
`285`	`285`	`ID: id,`
`286`	`286`	`Plan: p.Plan,`
`287`		`- Chunks: chunkMap[id],`
	`287`	`+ Chunks: chunks,`
`288`	`288`	`}`
`289`	`289`	`importSpecs = append(importSpecs, importSpec)`
`290`	`290`	`}`
Original file line number	Diff line number	Diff line change
`@@ -214,11 +214,11 @@ func (e ImportIntoExec) submitTask(ctx context.Context) (int64, proto.TaskBase`
`214`	`214`	`logutil.Logger(ctx).Info("get job importer", zap.Stringer("param", e.controller.Parameters),`
`215`	`215`	`zap.Bool("dist-task-enabled", vardef.EnableDistTask.Load()))`
`216`	`216`	`if importFromServer {`
`217`		`- ecp, err2 := e.controller.PopulateChunks(ctx)`
	`217`	`+ chunkMap, err2 := e.controller.PopulateChunks(ctx)`
`218`	`218`	`if err2 != nil {`
`219`	`219`	`return 0, nil, err2`
`220`	`220`	`}`
`221`		`- return importinto.SubmitStandaloneTask(ctx, e.controller.Plan, e.stmt, ecp)`
	`221`	`+ return importinto.SubmitStandaloneTask(ctx, e.controller.Plan, e.stmt, chunkMap)`
`222`	`222`	`}`
`223`	`223`	`// if tidb_enable_dist_task=true, we import distributively, otherwise we import on current node.`
`224`	`224`	`if vardef.EnableDistTask.Load() {`
Original file line number	Diff line number	Diff line change
`@@ -1014,6 +1014,7 @@ func (e *LoadDataController) InitDataStore(ctx context.Context) error {`
`1014`	`1014`	`}`
`1015`	`1015`	`return nil`
`1016`	`1016`	`}`
	`1017`	`+`
`1017`	`1018`	`func (LoadDataController) initExternalStore(ctx context.Context, u url.URL, target string) (storage.ExternalStorage, error) {`
`1018`	`1019`	`b, err2 := storage.ParseBackendFromURL(u, nil)`
`1019`	`1020`	`if err2 != nil {`