Repository: pingcap/tidb
File: pkg/executor/importer/import.go (Go)

Code before:
	if e.InImportInto && isAutoDetectingFormat && e.Format != DataFormatCSV {
		if err2 = e.checkNonCSVFormatOptions(); err2 != nil {
			return err2
		}
	}

	e.dataFiles = dataFiles
	e.totalRealSize = totalRealSize
	e.TotalFileSize = totalSize

	return nil
}

// CalResourceParams calculates resource related parameters according to the total
// file size and target node cpu count.
func (e *LoadDataController) CalResourceParams(ctx context.Context, ksCodec []byte) error {
	start := time.Now()
	targetNodeCPUCnt, err := handle.GetCPUCountOfNode(ctx)
	if err != nil {
		return err
	}
	factors, err := handle.GetScheduleTuneFactors(ctx, e.Keyspace)
	if err != nil {
		return err
	}
	totalSize := e.totalRealSize
	failpoint.InjectCall("mockImportDataSize", &totalSize)
	numOfIndexGenKV := GetNumOfIndexGenKV(e.TableInfo)
	var indexSizeRatio float64
	if numOfIndexGenKV > 0 {
		indexSizeRatio, err = e.sampleIndexSizeRatio(ctx, ksCodec)
		if err != nil {
			e.logger.Warn("meet error when sampling index size ratio", zap.Error(err))
		}
	}
	cal := scheduler.NewRCCalc(totalSize, targetNodeCPUCnt, indexSizeRatio, factors)
	e.ThreadCnt = cal.CalcRequiredSlots()
	e.MaxNodeCnt = cal.CalcMaxNodeCountForImportInto()
	e.DistSQLScanConcurrency = scheduler.CalcDistSQLConcurrency(e.ThreadCnt, e.MaxNodeCnt, targetNodeCPUCnt)
	e.logger.Info("auto calculate resource related params",
		zap.Int("thread", e.ThreadCnt),
		zap.Int("maxNode", e.MaxNodeCnt),
		zap.Int("distsqlScanConcurrency", e.DistSQLScanConcurrency),
		zap.Int("targetNodeCPU", targetNodeCPUCnt),
		zap.String("totalFileSize", units.BytesSize(float64(e.TotalFileSize))),
		zap.String("totalRealSize", units.BytesSize(float64(totalSize))),
		zap.Int("fileCount", len(e.dataFiles)),
		zap.Int("numOfIndexGenKV", numOfIndexGenKV),
		zap.Float64("indexSizeRatio", indexSizeRatio),
		zap.Float64("amplifyFactor", factors.AmplifyFactor),
		zap.Duration("costTime", time.Since(start)),

Review comment (bug):
`CalResourceParams` now uses `e.totalRealSize` exclusively. If `InitDataFiles` hasn't been called (e.g. unit tests construct a controller with only `Plan.TotalFileSize` set), `e.totalRealSize` will be 0 and resource calculation will be wrong. Consider either (1) enforcing/validating that `InitDataFiles` has run (return an error if `e.dataFiles` is nil/empty), or (2) falling back to `e.TotalFileSize` when `e.totalRealSize == 0`.
```suggestion
	totalSize := e.totalRealSize
	if totalSize == 0 && e.TotalFileSize > 0 {
		// Fallback to the planned total file size when real size has not been initialized,
		// e.g. when InitDataFiles has not been called (such as in some unit tests).
		totalSize = e.TotalFileSize
	}
```

Code after:
		}
	}
	var totalSize, totalRealSize int64
	for _, dfile := range dataFiles {
		totalSize += dfile.FileSize
		realSize := dfile.RealSize
		failpoint.Inject("amplifyRealSize", func(val failpoint.Value) {
			factor := int64(val.(int))
			realSize *= factor
		})
		totalRealSize += realSize
	}

	e.dataFiles = dataFiles
	e.TotalFileSize = totalSize
	e.TotalRealSize = totalRealSize

	return nil
}

// CalResourceParams calculates resource related parameters according to the total
// file size and target node cpu count.
func (e *LoadDataController) CalResourceParams(ctx context.Context, ksCodec []byte) error {
	start := time.Now()
	targetNodeCPUCnt, err := handle.GetCPUCountOfNode(ctx)
	if err != nil {
		return err
	}
	factors, err := handle.GetScheduleTuneFactors(ctx, e.Keyspace)
	if err != nil {
		return err
	}
	totalSize := e.TotalRealSize
	numOfIndexGenKV := GetNumOfIndexGenKV(e.TableInfo)
	var indexSizeRatio float64
	if numOfIndexGenKV > 0 {
		indexSizeRatio, err = e.sampleIndexSizeRatio(ctx, ksCodec)
		if err != nil {
			e.logger.Warn("meet error when sampling index size ratio", zap.Error(err))
		}
	}
	cal := scheduler.NewRCCalc(totalSize, targetNodeCPUCnt, indexSizeRatio, factors)
	e.ThreadCnt = cal.CalcRequiredSlots()
	e.MaxNodeCnt = cal.CalcMaxNodeCountForImportInto()
	e.DistSQLScanConcurrency = scheduler.CalcDistSQLConcurrency(e.ThreadCnt, e.MaxNodeCnt, targetNodeCPUCnt)
	e.logger.Info("auto calculate resource related params",
		zap.Int("thread", e.ThreadCnt),
		zap.Int("maxNode", e.MaxNodeCnt),
		zap.Int("distsqlScanConcurrency", e.DistSQLScanConcurrency),
		zap.Int("targetNodeCPU", targetNodeCPUCnt),
		zap.String("totalFileSize", units.BytesSize(float64(e.TotalFileSize))),