[#61999] improve dumpling string key handling

Author: takaidohigasi

dumpling/export/dump.go

@@ -14,6 +14,7 @@ import (
 	"slices"
 	"strconv"
 	"strings"
+	"sync"
 	"sync/atomic"
 	"time"
 
@@ -53,6 +54,9 @@ var errEmptyHandleVals = errors.New("empty handleVals for TiDB table")
 // see https://docs.pingcap.com/zh/tidb/dev/system-variables#tidb_enable_paging-%E4%BB%8E-v540-%E7%89%88%E6%9C%AC%E5%BC%80%E5%A7%8B%E5%BC%95%E5%85%A5
 var enablePagingVersion = semver.New("6.2.0")
 
+// Maximum number of chunks to prevent infinite loops during boundary sampling
+const maxChunkLimit = 1000000
+
 // Dumper is the dump progress structure
 type Dumper struct {
 	tctx      *tcontext.Context
@@ -68,7 +72,10 @@ type Dumper struct {
 	totalTables                   int64
 	charsetAndDefaultCollationMap map[string]string
 
-	speedRecorder *SpeedRecorder
+	speedRecorder      *SpeedRecorder
+	adaptiveChunkSizer *AdaptiveChunkSizer
+
+	chunkedTables sync.Map
 }
 
 // NewDumper returns a new Dumper
@@ -93,6 +100,7 @@ func NewDumper(ctx context.Context, conf *Config) (*Dumper, error) {
 		cancelCtx:                 cancelFn,
 		selectTiDBTableRegionFunc: selectTiDBTableRegion,
 		speedRecorder:             NewSpeedRecorder(),
+		adaptiveChunkSizer:        NewAdaptiveChunkSizer(int64(conf.Rows)),
 	}
 
 	var err error
@@ -347,27 +355,25 @@ func (d *Dumper) startWriters(tctx *tcontext.Context, wg *errgroup.Group, taskCh
 		if err != nil {
 			return nil, func() {}, err
 		}
-		writer := NewWriter(tctx, int64(i), conf, conn, d.extStore, d.metrics)
+		writer := NewWriter(tctx, int64(i), conf, conn, d.extStore, d.metrics, d.adaptiveChunkSizer)
 		writer.rebuildConnFn = rebuildConnFn
 		writer.setFinishTableCallBack(func(task Task) {
-			if _, ok := task.(*TaskTableData); ok {
-				IncCounter(d.metrics.finishedTablesCounter)
-				// FIXME: actually finishing the last chunk doesn't means this table is 'finished'.
-				//  We can call this table is 'finished' if all its chunks are finished.
-				//  Comment this log now to avoid ambiguity.
-				// tctx.L().Debug("finished dumping table data",
-				//	zap.String("database", td.Meta.DatabaseName()),
-				//	zap.String("table", td.Meta.TableName()))
-				failpoint.Inject("EnableLogProgress", func() {
-					time.Sleep(1 * time.Second)
-					tctx.L().Debug("EnableLogProgress, sleep 1s")
-				})
-			}
+			// this is called when a file is finished.
 		})
 		writer.setFinishTaskCallBack(func(task Task) {
 			IncGauge(d.metrics.taskChannelCapacity)
 			if td, ok := task.(*TaskTableData); ok {
 				d.metrics.completedChunks.Add(1)
+				if val, ok := d.chunkedTables.Load(td.Meta.ChunkKey()); ok {
+					chunkStats := val.(*tableChunkStat)
+					finishedChunks := chunkStats.finished.Add(1)
+					if chunkStats.finalized.Load() && finishedChunks == chunkStats.sent.Load() {
+						IncCounter(d.metrics.finishedTablesCounter)
+						d.chunkedTables.Delete(td.Meta.ChunkKey())
+					}
+				} else {
+					IncCounter(d.metrics.finishedTablesCounter)
+				}
 				tctx.L().Debug("finish dumping table data task",
 					zap.String("database", td.Meta.DatabaseName()),
 					zap.String("table", td.Meta.TableName()),
@@ -619,7 +625,7 @@ func (d *Dumper) dumpTableData(tctx *tcontext.Context, conn *BaseConn, meta Tabl
 	}
 
 	// Update total rows
-	fieldName, _ := pickupPossibleField(tctx, meta, conn)
+	fieldName, _, _ := pickupPossibleField(tctx, meta, conn)
 	c := estimateCount(tctx, meta.DatabaseName(), meta.TableName(), conn, fieldName, conf)
 	AddCounter(d.metrics.estimateTotalRowsCounter, float64(c))
 
@@ -729,9 +735,10 @@ func (d *Dumper) sequentialDumpTable(tctx *tcontext.Context, conn *BaseConn, met
 }
 
 // concurrentDumpTable tries to split table into several chunks to dump
-func (d *Dumper) concurrentDumpTable(tctx *tcontext.Context, conn *BaseConn, meta TableMeta, taskChan chan<- Task) error {
+func (d *Dumper) concurrentDumpTable(tctx *tcontext.Context, conn *BaseConn, meta TableMeta, taskChan chan<- Task) (err error) {
 	conf := d.conf
 	db, tbl := meta.DatabaseName(), meta.TableName()
+
 	if conf.ServerInfo.ServerType == version.ServerTypeTiDB &&
 		conf.ServerInfo.ServerVersion != nil &&
 		(conf.ServerInfo.ServerVersion.Compare(*tableSampleVersion) >= 0 ||
@@ -751,7 +758,7 @@ func (d *Dumper) concurrentDumpTable(tctx *tcontext.Context, conn *BaseConn, met
 		return err
 	}
 
-	field, err := pickupPossibleField(tctx, meta, conn)
+	field, isStringField, err := pickupPossibleField(tctx, meta, conn)
 	if err != nil || field == "" {
 		// skip split chunk logic if not found proper field
 		tctx.L().Info("fallback to sequential dump due to no proper field. This won't influence the whole dump process",
@@ -774,6 +781,15 @@ func (d *Dumper) concurrentDumpTable(tctx *tcontext.Context, conn *BaseConn, met
 		return d.dumpWholeTableDirectly(tctx, meta, taskChan, "", orderByClause, 0, 1)
 	}
 
+	// Handle string-based chunking
+	if isStringField {
+		tctx.L().Info("using string-based chunking",
+			zap.String("database", db),
+			zap.String("table", tbl),
+			zap.String("field", field))
+		return d.concurrentDumpStringField(tctx, conn, meta, taskChan, field, orderByClause, count)
+	}
+
 	minv, maxv, err := d.selectMinAndMaxIntValue(tctx, conn, db, tbl, field)
 	if err != nil {
 		tctx.L().Info("fallback to sequential dump due to cannot get bounding values. This won't influence the whole dump process",
@@ -798,6 +814,15 @@ func (d *Dumper) concurrentDumpTable(tctx *tcontext.Context, conn *BaseConn, met
 
 	chunkIndex := 0
 	nullValueCondition := ""
+	chunkStats := newTableChunkStat()
+	d.chunkedTables.Store(meta.ChunkKey(), chunkStats)
+	defer func() {
+		chunkStats.finalized.Store(true)
+		if chunkStats.finished.Load() == chunkStats.sent.Load() {
+			IncCounter(d.metrics.finishedTablesCounter)
+			d.chunkedTables.Delete(meta.ChunkKey())
+		}
+	}()
 	if conf.Where == "" {
 		nullValueCondition = fmt.Sprintf("`%s` IS NULL OR ", escapeString(field))
 	}
@@ -820,6 +845,11 @@ func (d *Dumper) concurrentDumpTable(tctx *tcontext.Context, conn *BaseConn, met
 }
 
 func (d *Dumper) sendTaskToChan(tctx *tcontext.Context, task Task, taskChan chan<- Task) (ctxDone bool) {
+	if td, ok := task.(*TaskTableData); ok {
+		if val, ok := d.chunkedTables.Load(td.Meta.ChunkKey()); ok {
+			val.(*tableChunkStat).sent.Add(1)
+		}
+	}
 	select {
 	case <-tctx.Done():
 		return true
@@ -831,6 +861,20 @@ func (d *Dumper) sendTaskToChan(tctx *tcontext.Context, task Task, taskChan chan
 	}
 }
 
+type tableChunkStat struct {
+	sent      *gatomic.Int32
+	finished  *gatomic.Int32
+	finalized *gatomic.Bool
+}
+
+func newTableChunkStat() *tableChunkStat {
+	return &tableChunkStat{
+		sent:      gatomic.NewInt32(0),
+		finished:  gatomic.NewInt32(0),
+		finalized: gatomic.NewBool(false),
+	}
+}
+
 func (d *Dumper) selectMinAndMaxIntValue(tctx *tcontext.Context, conn *BaseConn, db, tbl, field string) (minv, maxv *big.Int, err error) {
 	conf, zero := d.conf, &big.Int{}
 	query := fmt.Sprintf("SELECT MIN(`%s`),MAX(`%s`) FROM `%s`.`%s`",
@@ -923,6 +967,16 @@ func (d *Dumper) concurrentDumpTiDBPartitionTables(tctx *tcontext.Context, conn
 		totalChunk += len(handleVals) + 1
 		cachedHandleVals[i] = handleVals
 	}
+
+	chunkStats := newTableChunkStat()
+	d.chunkedTables.Store(meta.ChunkKey(), chunkStats)
+	defer func() {
+		chunkStats.finalized.Store(true)
+		if chunkStats.finished.Load() == chunkStats.sent.Load() {
+			IncCounter(d.metrics.finishedTablesCounter)
+			d.chunkedTables.Delete(meta.ChunkKey())
+		}
+	}()
 	for i, partition := range partitions {
 		err := d.sendConcurrentDumpTiDBTasks(tctx, meta, taskChan, handleColNames, cachedHandleVals[i], partition, startChunkIdx, totalChunk)
 		if err != nil {
@@ -1709,3 +1763,70 @@ func (d *Dumper) newTaskTableData(meta TableMeta, data TableDataIR, currentChunk
 	d.metrics.totalChunks.Add(1)
 	return NewTaskTableData(meta, data, currentChunk, totalChunks)
 }
+
+// extractOrderByColumns extracts column names from ORDER BY clause
+// Input: "ORDER BY `item_id`,`photo_index`"
+// Output: ["`item_id`", "`photo_index`"]
+func extractOrderByColumns(orderByClause string) []string {
+	// Remove "ORDER BY " prefix
+	columnsStr := strings.TrimPrefix(orderByClause, "ORDER BY ")
+
+	// Split by comma and trim spaces
+	columns := strings.Split(columnsStr, ",")
+	for i, col := range columns {
+		columns[i] = strings.TrimSpace(col)
+	}
+
+	return columns
+}
+
+// checkRowNumberSupport tests if the database supports ROW_NUMBER() window function
+// ROW_NUMBER() is supported by:
+// - MySQL 8.0+ (window functions introduced)
+// Not supported by:
+// - MySQL 5.7 and earlier (no window functions)
+// - Older database versions
+func checkRowNumberSupport(tctx *tcontext.Context, conn *BaseConn, database, table string) bool {
+	// First, set the database context to avoid "No database selected" error
+	useDbQuery := fmt.Sprintf("USE `%s`", escapeString(database))
+	err := conn.QuerySQL(tctx, func(rows *sql.Rows) error {
+		// No rows expected from USE statement
+		return nil
+	}, func() {}, useDbQuery)
+
+	if err != nil {
+		// Can't set database context, fallback to OFFSET-based approach
+		tctx.L().Debug("Cannot set database context for ROW_NUMBER() check", zap.Error(err))
+		return false
+	}
+
+	// Test ROW_NUMBER() against the actual table to detect support
+	// This will return error 1305 (function not found) on older MySQL versions
+	testQuery := fmt.Sprintf("SELECT ROW_NUMBER() FROM `%s` LIMIT 1", escapeString(table))
+
+	var supported bool
+	err = conn.QuerySQL(tctx, func(rows *sql.Rows) error {
+		// If we get here, ROW_NUMBER() syntax is recognized (even if query fails for other reasons)
+		supported = true
+		return nil
+	}, func() {}, testQuery)
+
+	if err != nil {
+		// Check if it's the expected "function not found" error (1305)
+		if mysqlErr, ok := errors.Cause(err).(*mysql.MySQLError); ok && mysqlErr.Number == 1305 {
+			// Error 1305: FUNCTION does not exist - ROW_NUMBER() not supported
+			tctx.L().Info("ROW_NUMBER() not supported by this database version",
+				zap.Uint16("errorCode", mysqlErr.Number))
+			return false
+		}
+		// Other errors might indicate syntax issues or connection problems
+		// Still treat as not supported for safety
+		tctx.L().Debug("ROW_NUMBER() check failed with unexpected error", zap.Error(err))
+		return false
+	}
+
+	tctx.L().Info("ROW_NUMBER() support check completed",
+		zap.Bool("supported", supported))
+
+	return supported
+}

dumpling/export/ir.go

@@ -35,6 +35,7 @@ type TableMeta interface {
 	ShowCreateView() string
 	AvgRowLength() uint64
 	HasImplicitRowID() bool
+	ChunkKey() string
 }
 
 // SQLRowIter is the iterator on a collection of sql.Row.

dumpling/export/metadata.go

@@ -26,6 +26,10 @@ type globalMetadata struct {
 	storage storage.ExternalStorage
 }
 
+func (t *tableMeta) ChunkKey() string {
+	return t.DatabaseName() + "." + t.TableName()
+}
+
 const (
 	metadataPath       = "metadata"
 	metadataTimeLayout = time.DateTime