parser: support pre-split global index add special comment support fo…

…r pre_split index option (#58408)

ref #57551, ref #57552

pkg/ddl/index_presplit.go

@@ -40,8 +40,6 @@ import (
 	"go.uber.org/zap"
 )
 
-// TODO(tangenta): support global index.
-// Wrap the job.Query to with special comments.
 func preSplitIndexRegions(
 	ctx context.Context,
 	sctx sessionctx.Context,
@@ -120,6 +118,11 @@ func getSplitIdxKeysFromValueList(
 		return getSplitIdxPhysicalKeysFromValueList(sctx, tblInfo, idxInfo, tblInfo.ID, byRows, destKeys)
 	}
 
+	if idxInfo.Global {
+		destKeys = make([][]byte, 0, len(byRows)+1)
+		return getSplitIdxPhysicalKeysFromValueList(sctx, tblInfo, idxInfo, tblInfo.ID, byRows, destKeys)
+	}
+
 	destKeys = make([][]byte, 0, (len(byRows)+1)*len(pi.Definitions))
 	for _, p := range pi.Definitions {
 		destKeys, err = getSplitIdxPhysicalKeysFromValueList(sctx, tblInfo, idxInfo, p.ID, byRows, destKeys)
@@ -305,6 +308,20 @@ func evalSplitDatumFromArgs(
 		return &splitArgs{byRows: indexValues}, nil
 	}
 
+	if len(opt.Lower) == 0 && len(opt.Upper) == 0 && opt.Num > 0 {
+		lowerVals := make([]types.Datum, 0, len(idxInfo.Columns))
+		upperVals := make([]types.Datum, 0, len(idxInfo.Columns))
+		for i := 0; i < len(idxInfo.Columns); i++ {
+			lowerVals = append(lowerVals, types.MinNotNullDatum())
+			upperVals = append(upperVals, types.MaxValueDatum())
+		}
+		return &splitArgs{
+			betweenLower: lowerVals,
+			betweenUpper: upperVals,
+			regionsCnt:   int(opt.Num),
+		}, nil
+	}
+
 	// Split index regions by lower, upper value.
 	checkLowerUpperValue := func(valuesItem []string, name string) ([]types.Datum, error) {
 		if len(valuesItem) == 0 {

pkg/parser/ast/ddl.go

@@ -828,16 +828,22 @@ func (n *IndexOption) Restore(ctx *format.RestoreCtx) error {
 		if hasPrevOption {
 			ctx.WritePlain(" ")
 		}
-		ctx.WriteKeyWord("PRE_SPLIT_REGIONS")
-		ctx.WritePlain(" = ")
-		if n.SplitOpt.Num != 0 && len(n.SplitOpt.Lower) == 0 {
-			ctx.WritePlainf("%d", n.SplitOpt.Num)
-		} else {
-			ctx.WritePlain("(")
-			if err := n.SplitOpt.Restore(ctx); err != nil {
-				return errors.Annotate(err, "An error occurred while splicing IndexOption SplitOpt")
+		err := ctx.WriteWithSpecialComments(tidb.FeatureIDPresplit, func() error {
+			ctx.WriteKeyWord("PRE_SPLIT_REGIONS")
+			ctx.WritePlain(" = ")
+			if n.SplitOpt.Num != 0 && len(n.SplitOpt.Lower) == 0 {
+				ctx.WritePlainf("%d", n.SplitOpt.Num)
+			} else {
+				ctx.WritePlain("(")
+				if err := n.SplitOpt.Restore(ctx); err != nil {
+					return errors.Annotate(err, "An error occurred while splicing IndexOption SplitOpt")
+				}
+				ctx.WritePlain(")")
 			}
-			ctx.WritePlain(")")
+			return nil
+		})
+		if err != nil {
+			return err
 		}
 	}
 	return nil

pkg/parser/ast/ddl_test.go

@@ -934,6 +934,34 @@ func TestTableOptionTTLRestoreWithTTLEnableOffFlag(t *testing.T) {
 	}
 }
 
+func TestPresplitIndexSpecialComments(t *testing.T) {
+	specialCmtFlag := format.DefaultRestoreFlags | format.RestoreTiDBSpecialComment
+	cases := []struct {
+		sourceSQL string
+		flags     format.RestoreFlags
+		expectSQL string
+	}{
+		{"ALTER TABLE t ADD INDEX (a) PRE_SPLIT_REGIONS = 4", specialCmtFlag, "ALTER TABLE `t` ADD INDEX(`a`) /*T![pre_split] PRE_SPLIT_REGIONS = 4 */"},
+		{"ALTER TABLE t ADD INDEX (a) PRE_SPLIT_REGIONS 4", specialCmtFlag, "ALTER TABLE `t` ADD INDEX(`a`) /*T![pre_split] PRE_SPLIT_REGIONS = 4 */"},
+		{"ALTER TABLE t ADD PRIMARY KEY (a) CLUSTERED PRE_SPLIT_REGIONS = 4", specialCmtFlag, "ALTER TABLE `t` ADD PRIMARY KEY(`a`) /*T![clustered_index] CLUSTERED */ /*T![pre_split] PRE_SPLIT_REGIONS = 4 */"},
+		{"ALTER TABLE t ADD PRIMARY KEY (a) PRE_SPLIT_REGIONS = 4 NONCLUSTERED", specialCmtFlag, "ALTER TABLE `t` ADD PRIMARY KEY(`a`) /*T![clustered_index] NONCLUSTERED */ /*T![pre_split] PRE_SPLIT_REGIONS = 4 */"},
+		{"ALTER TABLE t ADD INDEX (a) PRE_SPLIT_REGIONS = (between (1, 'a') and (2, 'b') regions 4);", specialCmtFlag, "ALTER TABLE `t` ADD INDEX(`a`) /*T![pre_split] PRE_SPLIT_REGIONS = (BETWEEN (1,_UTF8MB4'a') AND (2,_UTF8MB4'b') REGIONS 4) */"},
+		{"ALTER TABLE t ADD INDEX idx(a) pre_split_regions = 100, ADD INDEX idx2(b) pre_split_regions = (by(1),(2),(3))", specialCmtFlag, "ALTER TABLE `t` ADD INDEX `idx`(`a`) /*T![pre_split] PRE_SPLIT_REGIONS = 100 */, ADD INDEX `idx2`(`b`) /*T![pre_split] PRE_SPLIT_REGIONS = (BY (1),(2),(3)) */"},
+		{"ALTER TABLE t ADD INDEX (a) comment 'a' PRE_SPLIT_REGIONS = (between (1, 'a') and (2, 'b') regions 4);", specialCmtFlag, "ALTER TABLE `t` ADD INDEX(`a`) COMMENT 'a' /*T![pre_split] PRE_SPLIT_REGIONS = (BETWEEN (1,_UTF8MB4'a') AND (2,_UTF8MB4'b') REGIONS 4) */"},
+	}
+
+	extractNodeFunc := func(node Node) Node {
+		return node
+	}
+
+	for _, ca := range cases {
+		testCases := []NodeRestoreTestCase{
+			{ca.sourceSQL, ca.expectSQL},
+		}
+		runNodeRestoreTestWithFlags(t, testCases, "%s", extractNodeFunc, ca.flags)
+	}
+}
+
 func TestResourceGroupDDLStmtRestore(t *testing.T) {
 	createTestCases := []NodeRestoreTestCase{
 		{

pkg/parser/tidb/features.go

@@ -34,6 +34,8 @@ const (
 	FeatureIDResourceGroup = "resource_group"
 	// FeatureIDGlobalIndex is the `Global Index` feature.
 	FeatureIDGlobalIndex = "global_index"
+	// FeatureIDPresplit is the pre-split feature.
+	FeatureIDPresplit = "pre_split"
 )
 
 var featureIDs = map[string]struct{}{
@@ -45,6 +47,7 @@ var featureIDs = map[string]struct{}{
 	FeatureIDPlacement:      {},
 	FeatureIDTTL:            {},
 	FeatureIDGlobalIndex:    {},
+	FeatureIDPresplit:       {},
 }
 
 // CanParseFeature is used to check if a feature can be parsed.

tests/realtikvtest/addindextest3/functional_test.go

@@ -170,6 +170,14 @@ func TestAddIndexPresplitIndexRegions(t *testing.T) {
 			splitKeyHex = nil
 		}
 	}
+	var idxID int64
+	nextIdxID := func() int64 {
+		idxID++
+		return idxID
+	}
+	resetIdxID := func() {
+		idxID = 0
+	}
 
 	tk.MustExec("create table t (a int primary key, b int);")
 	for i := 0; i < 10; i++ {
@@ -180,39 +188,57 @@ func TestAddIndexPresplitIndexRegions(t *testing.T) {
 	require.Len(t, retRows, 1)
 	tk.MustExec("set @@global.tidb_ddl_enable_fast_reorg = off;")
 	tk.MustExec("set @@global.tidb_enable_dist_task = off;")
+	tk.MustExec("alter table t add index idx(b) pre_split_regions = 4;")
+	checkSplitKeys(nextIdxID(), 3, true)
+	tk.MustExec("drop index idx on t;")
 	tk.MustExec("alter table t add index idx(b) pre_split_regions = (by (10000), (20000), (30000));")
-	checkSplitKeys(1, 3, true)
+	checkSplitKeys(nextIdxID(), 3, true)
+	tk.MustExec("drop index idx on t;")
+	tk.MustExec("alter table t add index idx(b) /*T![pre_split] pre_split_regions = (by (10000), (20000), (30000)) */;")
+	checkSplitKeys(nextIdxID(), 3, true)
 	tk.MustExec("drop index idx on t;")
 	tk.MustExec("alter table t add index idx(b) pre_split_regions = (between (0) and (10 * 10000) regions 3);")
-	checkSplitKeys(2, 2, true)
+	checkSplitKeys(nextIdxID(), 2, true)
 	tk.MustExec("drop index idx on t;")
 	tk.MustExec("set @@global.tidb_ddl_enable_fast_reorg = on;")
 
 	tk.MustExec("alter table t add index idx(b) pre_split_regions = (by (10000), (20000), (30000));")
-	checkSplitKeys(3, 0, false)
-	checkSplitKeys(tablecodec.TempIndexPrefix|3, 3, true)
+	nextID := nextIdxID()
+	checkSplitKeys(nextID, 0, false)
+	checkSplitKeys(tablecodec.TempIndexPrefix|nextID, 3, true)
 
 	tk.MustExec("set @@global.tidb_ddl_enable_fast_reorg = off;")
 
 	// Test partition tables.
+	resetIdxID()
 	tk.MustExec("drop table t;")
 	tk.MustExec("create table t (a int primary key, b int) partition by hash(a) partitions 4;")
 	for i := 0; i < 10; i++ {
 		insertSQL := fmt.Sprintf("insert into t values (%[1]d, %[1]d);", 10000*i)
 		tk.MustExec(insertSQL)
 	}
 	tk.MustExec("alter table t add index idx(b) pre_split_regions = (by (10000), (20000), (30000));")
-	checkSplitKeys(1, 3*4, true)
+	checkSplitKeys(nextIdxID(), 3*4, true)
 	tk.MustExec("drop index idx on t;")
 	tk.MustExec("alter table t add index idx(b) pre_split_regions = (between (0) and (10 * 10000) regions 3);")
-	checkSplitKeys(2, 2*4, true)
+	checkSplitKeys(nextIdxID(), 2*4, true)
 	tk.MustExec("drop index idx on t;")
 	tk.MustExec("set @@global.tidb_ddl_enable_fast_reorg = on;")
 	tk.MustExec("alter table t add index idx(b) pre_split_regions = (by (10000), (20000), (30000));")
-	checkSplitKeys(3, 0, false)
+	checkSplitKeys(nextIdxID(), 0, false)
 	checkSplitKeys(tablecodec.TempIndexPrefix|3, 12, true)
 	tk.MustExec("drop index idx on t;")
 	tk.MustExec("set @@global.tidb_ddl_enable_fast_reorg = off;")
+
+	resetIdxID()
+	tk.MustExec("drop table t;")
+	tk.MustExec("set @@global.tidb_ddl_enable_fast_reorg = on;")
+	tk.MustExec("set @@global.tidb_enable_dist_task = off;")
+	tk.MustExec("create table t (a int, b int) partition by range (b)" +
+		" (partition p0 values less than (10), " +
+		"  partition p1 values less than (maxvalue));")
+	tk.MustExec("alter table t add unique index p_a (a) global pre_split_regions = (by (5), (15));")
+	checkSplitKeys(tablecodec.TempIndexPrefix|nextIdxID(), 2, true)
 }
 
 func TestAddIndexPresplitFunctional(t *testing.T) {