ddl: support region split policy

Signed-off-by: xhe <[email protected]>

Author: xhebox

errors.toml

@@ -806,6 +806,11 @@ error = '''
 Incorrect table name '%-.100s'
 '''
 
+["ddl:1105"]
+error = '''
+It is recommended to add a region split strategy to the new index to avoid write hotspots
+'''
+
 ["ddl:1111"]
 error = '''
 Invalid use of group function

pkg/ddl/BUILD.bazel

@@ -184,6 +184,7 @@ go_library(
         "//pkg/util/memory",
         "//pkg/util/ppcpuusage",
         "//pkg/util/ranger",
+        "//pkg/util/regionsplit",
         "//pkg/util/rowDecoder",
         "//pkg/util/rowcodec",
         "//pkg/util/set",

pkg/ddl/executor.go

@@ -1040,6 +1040,24 @@ func (e *executor) CreateTable(ctx sessionctx.Context, s *ast.CreateTableStmt) (
 	if err = checkTableInfoValidWithStmt(metaBuildCtx, tbInfo, s); err != nil {
 		return err
 	}
+
+	// Process region split policies from CREATE TABLE
+	if len(s.SplitIndex) > 0 {
+		for _, splitOpt := range s.SplitIndex {
+			policy, indexName, isPK, err := normalizeSplitPolicy(splitOpt, tbInfo)
+			if err != nil {
+				return errors.Trace(err)
+			}
+			if indexName == "" || (isPK && tbInfo.PKIsHandle) {
+				tbInfo.TableSplitPolicy = policy
+			} else {
+				indexInfo := tbInfo.FindIndexByName(indexName)
+				indexInfo.RegionSplitPolicy = policy
+			}
+		}
+
+	}
+
 	if err = checkTableForeignKeysValid(ctx, is, schema.Name.L, tbInfo); err != nil {
 		return err
 	}
@@ -1933,6 +1951,8 @@ func (e *executor) AlterTable(ctx context.Context, sctx sessionctx.Context, stmt
 		case ast.AlterTableDisableKeys, ast.AlterTableEnableKeys:
 			// Nothing to do now, see https://github.com/pingcap/tidb/issues/1051
 			// MyISAM specific
+		case ast.AlterTableSplitIndex:
+			err = e.AlterTableSetRegionSplitPolicy(sctx, ident, spec.SplitIndex)
 		case ast.AlterTableRemoveTTL:
 			// the parser makes sure we have only one `ast.AlterTableRemoveTTL` in an alter statement
 			err = e.AlterTableRemoveTTL(sctx, ident)
@@ -5082,6 +5102,10 @@ func (e *executor) createIndex(ctx sessionctx.Context, ti ast.Ident, keyType ast
 		OpType: model.OpAddIndex,
 	}
 
+	// Check if we should warn about missing region split policy
+	// If other indexes already have region split policy, warn user to set policy for new index
+	checkAndWarnMissingRegionSplitPolicy(ctx, tblInfo, indexName)
+
 	err = e.doDDLJob2(ctx, job, args)
 	// key exists, but if_not_exists flags is true, so we ignore this error.
 	if dbterror.ErrDupKeyName.Equal(err) && ifNotExists {
@@ -7175,3 +7199,56 @@ func checkColumnReferencedByPartialCondition(t *model.TableInfo, colName ast.CIS
 
 	return nil
 }
+
+// AlterTableSetRegionSplitPolicy sets persistent region split policy for table
+func (e *executor) AlterTableSetRegionSplitPolicy(ctx sessionctx.Context, ident ast.Ident, splitOpt *ast.SplitIndexOption) error {
+	schema, tb, err := e.getSchemaAndTableByIdent(ident)
+	if err != nil {
+		return errors.Trace(err)
+	}
+
+	meta := tb.Meta()
+
+	policy, indexName, isPK, err := normalizeSplitPolicy(splitOpt, tb.Meta())
+	if err != nil {
+		return errors.Trace(err)
+	}
+
+	job := &model.Job{
+		Version:        model.GetJobVerInUse(),
+		SchemaID:       schema.ID,
+		TableID:        meta.ID,
+		SchemaName:     schema.Name.L,
+		TableName:      meta.Name.L,
+		Type:           model.ActionAlterTableSetRegionSplitPolicy,
+		BinlogInfo:     &model.HistoryInfo{},
+		CDCWriteSource: ctx.GetSessionVars().CDCWriteSource,
+		SQLMode:        ctx.GetSessionVars().SQLMode,
+	}
+
+	args := &model.AlterTableSetRegionSplitPolicyArgs{
+		IndexName:  indexName,
+		PrimaryKey: isPK,
+		Policy:     policy,
+	}
+	return e.doDDLJob2(ctx, job, args)
+}
+
+// checkAndWarnMissingRegionSplitPolicy checks if table has other indexes with region split policy,
+// and warns user to set policy for newly created index.
+func checkAndWarnMissingRegionSplitPolicy(ctx sessionctx.Context, tblInfo *model.TableInfo, newIndexName ast.CIStr) {
+	// Check if any existing index has RegionSplitPolicy
+	hasExistingPolicy := false
+	for _, idx := range tblInfo.Indices {
+		if idx.RegionSplitPolicy != nil {
+			hasExistingPolicy = true
+			break
+		}
+	}
+
+	// If existing indexes have region split policy, warn about new index
+	if hasExistingPolicy {
+		ctx.GetSessionVars().StmtCtx.AppendWarning(errors.NewNoStackError(
+			"It is recommended to add a region split strategy to the new index '" + newIndexName.O + "' to avoid write hotspots"))
+	}
+}

pkg/ddl/job_worker.go

@@ -1082,6 +1082,8 @@ func (w *worker) runOneJobStep(
 		ver, err = onRefreshMeta(jobCtx, job)
 	case model.ActionAlterTableAffinity:
 		ver, err = onAlterTableAffinity(jobCtx, job)
+	case model.ActionAlterTableSetRegionSplitPolicy:
+		ver, err = w.onAlterTableSetRegionSplitPolicy(jobCtx, job)
 	default:
 		// Invalid job, cancel it.
 		job.State = model.JobStateCancelled

pkg/ddl/split_region.go

@@ -16,17 +16,24 @@ package ddl
 
 import (
 	"context"
+	"strconv"
+	"strings"
 
 	"github.com/pingcap/errors"
 	"github.com/pingcap/tidb/pkg/ddl/logutil"
 	"github.com/pingcap/tidb/pkg/kv"
 	"github.com/pingcap/tidb/pkg/meta/autoid"
 	"github.com/pingcap/tidb/pkg/meta/model"
+	"github.com/pingcap/tidb/pkg/parser/ast"
+	"github.com/pingcap/tidb/pkg/parser/format"
 	"github.com/pingcap/tidb/pkg/parser/mysql"
 	"github.com/pingcap/tidb/pkg/sessionctx"
 	"github.com/pingcap/tidb/pkg/sessionctx/vardef"
+	"github.com/pingcap/tidb/pkg/table/tables"
 	"github.com/pingcap/tidb/pkg/tablecodec"
 	"github.com/pingcap/tidb/pkg/types"
+	"github.com/pingcap/tidb/pkg/util/dbterror"
+	"github.com/pingcap/tidb/pkg/util/regionsplit"
 	tikverr "github.com/tikv/client-go/v2/error"
 	"go.uber.org/zap"
 )
@@ -36,11 +43,19 @@ const GlobalScatterGroupID int64 = -1
 
 func splitPartitionTableRegion(ctx sessionctx.Context, store kv.SplittableStore, tbInfo *model.TableInfo, parts []model.PartitionDefinition, scatterScope string) {
 	// Max partition count is 8192, should we sample and just choose some partitions to split?
-	var regionIDs []uint64
 	ctxWithTimeout, cancel := context.WithTimeout(context.Background(), ctx.GetSessionVars().GetSplitRegionTimeout())
 	defer cancel()
 	ctxWithTimeout = kv.WithInternalSourceType(ctxWithTimeout, kv.InternalTxnDDL)
-	if shardingBits(tbInfo) > 0 && tbInfo.PreSplitRegions > 0 {
+
+	var regionIDs []uint64
+	if hasSplitPolicies(tbInfo) {
+		regionIDs = append(regionIDs,
+			applySplitPoliciesForTable(ctxWithTimeout, ctx, store, tbInfo, tbInfo.ID)...)
+		for _, def := range parts {
+			regionIDs = append(regionIDs,
+				applySplitPoliciesForTable(ctxWithTimeout, ctx, store, tbInfo, def.ID)...)
+		}
+	} else if shardingBits(tbInfo) > 0 && tbInfo.PreSplitRegions > 0 {
 		regionIDs = make([]uint64, 0, len(parts)*(len(tbInfo.Indices)+1))
 		scatter, tableID := getScatterConfig(scatterScope, tbInfo.ID)
 		// Try to split global index region here.
@@ -63,8 +78,11 @@ func splitTableRegion(ctx sessionctx.Context, store kv.SplittableStore, tbInfo *
 	ctxWithTimeout, cancel := context.WithTimeout(context.Background(), ctx.GetSessionVars().GetSplitRegionTimeout())
 	defer cancel()
 	ctxWithTimeout = kv.WithInternalSourceType(ctxWithTimeout, kv.InternalTxnDDL)
+
 	var regionIDs []uint64
-	if shardingBits(tbInfo) > 0 && tbInfo.PreSplitRegions > 0 {
+	if hasSplitPolicies(tbInfo) {
+		regionIDs = applySplitPoliciesForTable(ctxWithTimeout, ctx, store, tbInfo, tbInfo.ID)
+	} else if shardingBits(tbInfo) > 0 && tbInfo.PreSplitRegions > 0 {
 		regionIDs = preSplitPhysicalTableByShardRowID(ctxWithTimeout, store, tbInfo, tbInfo.ID, scatterScope)
 	} else {
 		regionIDs = append(regionIDs, SplitRecordRegion(ctxWithTimeout, store, tbInfo.ID, tbInfo.ID, scatterScope))
@@ -200,3 +218,205 @@ func WaitScatterRegionFinish(ctx context.Context, store kv.SplittableStore, regi
 		}
 	}
 }
+
+// waitScatterIfNeeded waits for scatter to finish based on scatterScope
+// Extracted from multiple places to avoid code duplication
+func waitScatterIfNeeded(ctx context.Context, store kv.SplittableStore,
+	scatterScope string, regionIDs []uint64) {
+	if scatterScope == vardef.ScatterOff {
+		return // No need to wait
+	}
+
+	if len(regionIDs) == 0 {
+		return
+	}
+
+	WaitScatterRegionFinish(ctx, store, regionIDs...)
+}
+
+func hasSplitPolicies(tbInfo *model.TableInfo) bool {
+	if tbInfo.TableSplitPolicy != nil {
+		return true
+	}
+	for _, idx := range tbInfo.Indices {
+		if idx.RegionSplitPolicy != nil {
+			return true
+		}
+	}
+	return false
+}
+
+func applySplitPoliciesForTable(ctx context.Context, sctx sessionctx.Context, store kv.SplittableStore,
+	tbInfo *model.TableInfo, physicalTableID int64) []uint64 {
+	var keys [][]byte
+
+	sc := sctx.GetSessionVars().StmtCtx
+	scatterScope := sctx.GetSessionVars().ScatterRegion
+	scatter, tableID := getScatterConfig(scatterScope, tbInfo.ID)
+
+	// apply table policy
+	for p := tbInfo.TableSplitPolicy; p != nil; p = nil {
+		policy := tbInfo.TableSplitPolicy
+		lower, err := parseValuesToDatums(policy.Lower)
+		if err != nil {
+			logutil.DDLLogger().Warn("failed to parse lower bound for table policy",
+				zap.String("table", tbInfo.Name.O), zap.Error(err))
+			break
+		}
+		upper, err := parseValuesToDatums(policy.Upper)
+		if err != nil {
+			logutil.DDLLogger().Warn("failed to parse upper bound for table policy",
+				zap.String("table", tbInfo.Name.O), zap.Error(err))
+			break
+		}
+
+		handleCols := regionsplit.BuildHandleColsForSplit(tbInfo)
+		keys, err = regionsplit.GetSplitTableKeys(sc, tbInfo, handleCols, physicalTableID, lower, upper, int(policy.Regions), keys, dbterror.ErrInvalidSplitRegionRanges)
+		if err != nil {
+			logutil.DDLLogger().Warn("failed to generate split keys for table policy",
+				zap.String("table", tbInfo.Name.O), zap.Error(err))
+		}
+	}
+
+	// 2. Apply index policies (including PRIMARY)
+	for _, idx := range tbInfo.Indices {
+		if tbInfo.GetPartitionInfo() != nil &&
+			((idx.Global && tbInfo.ID != physicalTableID) || (!idx.Global && tbInfo.ID == physicalTableID)) {
+			continue
+		}
+
+		if idx.RegionSplitPolicy == nil {
+			continue
+		}
+
+		policy := idx.RegionSplitPolicy
+		lower, err := parseValuesToDatums(policy.Lower)
+		if err != nil {
+			logutil.DDLLogger().Warn("failed to parse lower bound for index policy",
+				zap.String("table", tbInfo.Name.O),
+				zap.String("index", idx.Name.O),
+				zap.Error(err))
+			continue
+		}
+		upper, err := parseValuesToDatums(policy.Upper)
+		if err != nil {
+			logutil.DDLLogger().Warn("failed to parse upper bound for index policy",
+				zap.String("table", tbInfo.Name.O),
+				zap.String("index", idx.Name.O),
+				zap.Error(err))
+			continue
+		}
+
+		keys, err = regionsplit.GetSplitIndexKeys(sc, tbInfo, idx, physicalTableID, lower, upper, int(policy.Regions), keys, dbterror.ErrInvalidSplitRegionRanges)
+		if err != nil {
+			logutil.DDLLogger().Warn("failed to generate split keys for index policy",
+				zap.String("table", tbInfo.Name.O),
+				zap.String("index", idx.Name.O),
+				zap.Error(err))
+		}
+	}
+
+	regionIDs := []uint64{}
+	if len(keys) > 0 {
+		ids, err := store.SplitRegions(ctx, keys, scatter, &tableID)
+		if err != nil {
+			logutil.DDLLogger().Warn("split regions failed", zap.Error(err))
+		} else {
+			regionIDs = ids
+		}
+	}
+	return regionIDs
+}
+
+func parseValuesToDatums(values []string) ([]types.Datum, error) {
+	datums := make([]types.Datum, len(values))
+	for i, val := range values {
+		// Try to parse as different types
+		// First try integer
+		if intVal, err := strconv.ParseInt(val, 10, 64); err == nil {
+			datums[i] = types.NewIntDatum(intVal)
+			continue
+		}
+		// Then try float
+		if floatVal, err := strconv.ParseFloat(val, 64); err == nil {
+			datums[i] = types.NewFloat64Datum(floatVal)
+			continue
+		}
+		// Default to string (remove quotes if present)
+		strVal := val
+		if len(strVal) >= 2 && strVal[0] == '\'' && strVal[len(strVal)-1] == '\'' {
+			strVal = strVal[1 : len(strVal)-1]
+		}
+		datums[i] = types.NewStringDatum(strVal)
+	}
+	return datums, nil
+}
+
+func normalizeSplitPolicy(splitOpt *ast.SplitIndexOption, tbInfo *model.TableInfo) (*model.RegionSplitPolicy, string, bool, error) {
+	if tbInfo.HasClusteredIndex() && splitOpt.TableLevel {
+		// cannot specify both SPLIT BETWEEN for CLUSTERED table
+		// it is for _tidb_rowid
+		return nil, "", false, dbterror.ErrInvalidRegionSplitPolicy
+	}
+
+	indexName := ""
+	if !splitOpt.TableLevel {
+		pkName := strings.ToLower(mysql.PrimaryKeyName)
+		indexName = splitOpt.IndexName.L
+
+		// fill primary key name
+		isPK := splitOpt.PrimaryKey
+		if isPK && indexName == "" {
+			indexName = pkName
+		}
+
+		if isPK != (pkName == indexName) {
+			// specified pk, but incorrect name, or reverse
+			return nil, "", false, dbterror.ErrWrongNameForIndex.GenWithStackByArgs(indexName)
+		}
+	}
+
+	if indexName != "" {
+		var colen int
+		if tbInfo.PKIsHandle && splitOpt.PrimaryKey {
+			colen = 1
+		} else if tbInfo.IsCommonHandle && splitOpt.PrimaryKey {
+			pk := tables.FindPrimaryIndex(tbInfo)
+			colen = len(pk.Columns)
+			indexName = pk.Name.L
+		} else {
+			var idx *model.IndexInfo
+			idx = tbInfo.FindIndexByName(indexName)
+			if idx == nil {
+				return nil, "", false, dbterror.ErrWrongNameForIndex.GenWithStackByArgs(indexName)
+			}
+			colen = len(idx.Columns)
+		}
+		if colen != len(splitOpt.SplitOpt.Upper) || colen != len(splitOpt.SplitOpt.Lower) {
+			return nil, "", false, dbterror.ErrInvalidSplitRegionRanges.GenWithStackByArgs("length of index columns and split values differ")
+		}
+	}
+
+	var buf strings.Builder
+	restoreCtx := format.NewRestoreCtx(format.DefaultRestoreFlags, &buf)
+
+	policy := &model.RegionSplitPolicy{
+		Regions: splitOpt.SplitOpt.Num,
+	}
+
+	policy.Lower = make([]string, len(splitOpt.SplitOpt.Lower))
+	for i, expr := range splitOpt.SplitOpt.Lower {
+		buf.Reset()
+		_ = expr.Restore(restoreCtx)
+		policy.Lower[i] = buf.String()
+	}
+
+	policy.Upper = make([]string, len(splitOpt.SplitOpt.Upper))
+	for i, expr := range splitOpt.SplitOpt.Upper {
+		buf.Reset()
+		_ = expr.Restore(restoreCtx)
+		policy.Upper[i] = buf.String()
+	}
+
+	return policy, indexName, splitOpt.PrimaryKey, nil
+}