ddl: support pre-split index regions before creating index

pkg/ddl/BUILD.bazel

@@ -43,6 +43,7 @@ go_library(
         "index.go",
         "index_cop.go",
         "index_merge_tmp.go",
+        "index_presplit.go",
         "job_scheduler.go",
         "job_submitter.go",
         "job_worker.go",
@@ -152,6 +153,7 @@ go_library(
         "//pkg/util/context",
         "//pkg/util/dbterror",
         "//pkg/util/dbterror/exeerrors",
+        "//pkg/util/dbterror/plannererrors",
         "//pkg/util/domainutil",
         "//pkg/util/engine",
         "//pkg/util/execdetails",

pkg/ddl/backfilling.go

@@ -244,7 +244,6 @@ type reorgBackfillTask struct {
 	startKey kv.Key
 	endKey   kv.Key
 	jobID    int64
-	sqlQuery string
 	priority int
 }
 

pkg/ddl/executor.go

@@ -4616,6 +4616,10 @@ func (e *executor) CreatePrimaryKey(ctx sessionctx.Context, ti ast.Ident, indexN
 		}
 	}
 
+	splitOpt, err := buildIndexPresplitOpt(indexOption)
+	if err != nil {
+		return errors.Trace(err)
+	}
 	sqlMode := ctx.GetSessionVars().SQLMode
 	// global is set to  'false' is just there to be backwards compatible,
 	// to avoid unmarshal issues, it is now part of indexOption.
@@ -4642,6 +4646,7 @@ func (e *executor) CreatePrimaryKey(ctx sessionctx.Context, ti ast.Ident, indexN
 			SQLMode:                 sqlMode,
 			Global:                  false,
 			IsPK:                    true,
+			SplitOpt:                splitOpt,
 		}},
 		OpType: model.OpAddIndex,
 	}
@@ -4905,29 +4910,9 @@ func (e *executor) createIndex(ctx sessionctx.Context, ti ast.Ident, keyType ast
 		return e.addHypoIndexIntoCtx(ctx, ti.Schema, ti.Name, indexInfo)
 	}
 
-	if indexOption != nil && indexOption.SplitOpt != nil {
-		opt := indexOption.SplitOpt
-		lowers := make([]string, 0, len(opt.Lower))
-		for _, expL := range opt.Lower {
-			var sb strings.Builder
-			rCtx := format.NewRestoreCtx(format.DefaultRestoreFlags, &sb)
-			err := expL.Restore(rCtx)
-			if err != nil {
-				return errors.Trace(err)
-			}
-			lowers = append(lowers, sb.String())
-		}
-		uppers := make([]string, 0, len(opt.Upper))
-		for _, expU := range opt.Upper {
-			var sb strings.Builder
-			rCtx := format.NewRestoreCtx(format.DefaultRestoreFlags, &sb)
-			err := expU.Restore(rCtx)
-			if err != nil {
-				return errors.Trace(err)
-			}
-			uppers = append(uppers, sb.String())
-		}
-
+	splitOpt, err := buildIndexPresplitOpt(indexOption)
+	if err != nil {
+		return errors.Trace(err)
 	}
 
 	// global is set to  'false' is just there to be backwards compatible,
@@ -4950,6 +4935,7 @@ func (e *executor) createIndex(ctx sessionctx.Context, ti ast.Ident, keyType ast
 			IndexOption:             indexOption,
 			HiddenCols:              hiddenCols,
 			Global:                  global,
+			SplitOpt:                splitOpt,
 		}},
 		OpType: model.OpAddIndex,
 	}
@@ -5003,6 +4989,56 @@ func newReorgMetaFromVariables(job *model.Job, sctx sessionctx.Context) (*model.
 	return reorgMeta, nil
 }
 
+func buildIndexPresplitOpt(indexOpt *ast.IndexOption) (*model.IndexArgSplitOpt, error) {
+	if indexOpt == nil {
+		return nil, nil
+	}
+	opt := indexOpt.SplitOpt
+	if opt == nil {
+		return nil, nil
+	}
+	lowers := make([]string, 0, len(opt.Lower))
+	for _, expL := range opt.Lower {
+		var sb strings.Builder
+		rCtx := format.NewRestoreCtx(format.DefaultRestoreFlags, &sb)
+		err := expL.Restore(rCtx)
+		if err != nil {
+			return nil, errors.Trace(err)
+		}
+		lowers = append(lowers, sb.String())
+	}
+	uppers := make([]string, 0, len(opt.Upper))
+	for _, expU := range opt.Upper {
+		var sb strings.Builder
+		rCtx := format.NewRestoreCtx(format.DefaultRestoreFlags, &sb)
+		err := expU.Restore(rCtx)
+		if err != nil {
+			return nil, errors.Trace(err)
+		}
+		uppers = append(uppers, sb.String())
+	}
+	valLists := make([][]string, 0, len(opt.ValueLists))
+	for _, lst := range opt.ValueLists {
+		values := make([]string, 0, len(lst))
+		for _, exp := range lst {
+			var sb strings.Builder
+			rCtx := format.NewRestoreCtx(format.DefaultRestoreFlags, &sb)
+			err := exp.Restore(rCtx)
+			if err != nil {
+				return nil, errors.Trace(err)
+			}
+			values = append(values, sb.String())
+		}
+		valLists = append(valLists, values)
+	}
+	return &model.IndexArgSplitOpt{
+		Lower:      lowers,
+		Upper:      uppers,
+		Num:        opt.Num,
+		ValueLists: valLists,
+	}, nil
+}
+
 // LastReorgMetaFastReorgDisabled is used for test.
 var LastReorgMetaFastReorgDisabled bool
 

pkg/ddl/index_presplit.go

@@ -274,17 +274,17 @@ func evalSplitDatumFromArgs(
 	idxInfo *model.IndexInfo,
 	idxArg *model.IndexArg,
 ) (*splitArgs, error) {
-	opt := idxArg.IndexOption.SplitOpt
+	opt := idxArg.SplitOpt
 	if opt == nil {
 		return nil, nil
 	}
 	if len(opt.ValueLists) > 0 {
 		indexValues := make([][]types.Datum, 0, len(opt.ValueLists))
-		for i, valuesItem := range opt.ValueLists {
-			if len(valuesItem) > len(idxInfo.Columns) {
+		for i, valueList := range opt.ValueLists {
+			if len(valueList) > len(idxInfo.Columns) {
 				return nil, plannererrors.ErrWrongValueCountOnRow.GenWithStackByArgs(i + 1)
 			}
-			values, err := evalConstExprNodes(evalCtx, valuesItem, tblInfo, idxInfo)
+			values, err := evalConstExprNodes(evalCtx, valueList, tblInfo, idxInfo)
 			if err != nil {
 				return nil, err
 			}
@@ -294,7 +294,7 @@ func evalSplitDatumFromArgs(
 	}
 
 	// Split index regions by lower, upper value.
-	checkLowerUpperValue := func(valuesItem []ast.ExprNode, name string) ([]types.Datum, error) {
+	checkLowerUpperValue := func(valuesItem []string, name string) ([]types.Datum, error) {
 		if len(valuesItem) == 0 {
 			return nil, errors.Errorf("Split index `%v` region %s value count should more than 0", idxInfo.Name, name)
 		}
@@ -328,15 +328,16 @@ func evalSplitDatumFromArgs(
 
 func evalConstExprNodes(
 	evalCtx exprctx.EvalContext,
-	valuesItem []ast.ExprNode,
+	valueList []string,
 	tblInfo *model.TableInfo,
 	idxInfo *model.IndexInfo,
 ) ([]types.Datum, error) {
-	values := make([]types.Datum, 0, len(valuesItem))
-	for j, valueItem := range valuesItem {
+	values := make([]types.Datum, 0, len(valueList))
+	for j, value := range valueList {
 		colOffset := idxInfo.Columns[j].Offset
 		col := tblInfo.Columns[colOffset]
-		switch x := valueItem.(type) {
+		valExpr := ast.NewValueExpr(value, "", "")
+		switch x := valExpr.(type) {
 		case *driver.ValueExpr:
 			constant := &expression.Constant{
 				Value:   x.Datum,

pkg/meta/model/job_args.go

@@ -1323,19 +1323,15 @@ type IndexArg struct {
 	IfExist  bool  `json:"if_exist,omitempty"`
 	IsGlobal bool  `json:"is_global,omitempty"`
 
-	// For index presplit
-	BetweenLower []string
-	BetweenUpper []string
-	BetweenCnt   int
-	ByData       [][]string
+	SplitOpt *IndexArgSplitOpt `json:"split_opt,omitempty"`
 }
 
 // IndexArgSplitOpt is a field of IndexArg used by index presplit.
 type IndexArgSplitOpt struct {
-	BetweenLower []string
-	BetweenUpper []string
-	BetweenCnt   int
-	ByData       [][]string
+	Lower      []string   `json:"lower"`
+	Upper      []string   `json:"upper"`
+	Num        int64      `json:"num"`
+	ValueLists [][]string `json:"value_lists"`
 }
 
 // ModifyIndexArgs is the argument for add/drop/rename index jobs,

pkg/util/BUILD.bazel

@@ -17,6 +17,7 @@ go_library(
         "rlimit_windows.go",
         "security.go",
         "session_pool.go",
+        "split.go",
         "tokenlimiter.go",
         "urls.go",
         "util.go",

tests/realtikvtest/addindextest3/functional_test.go

@@ -149,17 +149,62 @@ func TestAddIndexPresplitIndexRegions(t *testing.T) {
 	tk := testkit.NewTestKit(t, store)
 	tk.MustExec("use test")
 
-	tk.MustExec("create table t (a int primary key, b int);")
+	tk.MustExec("create table t (a int primary key, b int, c int);")
 	for i := 0; i < 10; i++ {
-		insertSQL := fmt.Sprintf("insert into t values (%d, %d);", 10000*i, 10000*i)
+		insertSQL := fmt.Sprintf("insert into t values (%[1]d, %[1]d, %[1]d);", 10000*i)
 		tk.MustExec(insertSQL)
 	}
-	ret := tk.MustQuery("show table t regions;").Rows()
-	require.Len(t, ret, 1)
+	retRows := tk.MustQuery("show table t regions;").Rows()
+	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 = (by (10000), (20000), (30000));")
-	ret = tk.MustQuery("show table t regions;").Rows()
-	require.Len(t, ret, 1)
-	// TODO(tangenta): test between by and partitions.
+	retRows = tk.MustQuery("show table t regions;").Rows()
+	idxRegionCnt := 0
+	for _, r := range retRows {
+		startKey := r[2].(string)
+		if strings.Contains(startKey, "_i_1") {
+			idxRegionCnt++
+		}
+	}
+	require.Equal(t, 3, idxRegionCnt)
+
+	tk.MustExec("alter table t add index idx2(c) pre_split_regions = (between (0) and (100000) regions 3);")
+	retRows = tk.MustQuery("show table t regions;").Rows()
+	idxRegionCnt = 0
+	for _, r := range retRows {
+		startKey := r[2].(string)
+		if strings.Contains(startKey, "_i_2") {
+			idxRegionCnt++
+		}
+	}
+	require.Equal(t, 3, idxRegionCnt)
+
+	// Test partition tables.
+	tk.MustExec("drop table t;")
+	tk.MustExec("create table t (a int primary key, b int, c int) partition by hash(a) partitions 4;")
+	for i := 0; i < 10; i++ {
+		insertSQL := fmt.Sprintf("insert into t values (%[1]d, %[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));")
+	retRows = tk.MustQuery("show table t regions;").Rows()
+	idxRegionCnt = 0
+	for _, r := range retRows {
+		startKey := r[2].(string)
+		if strings.Contains(startKey, "_i_1") {
+			idxRegionCnt++
+		}
+	}
+	require.Equal(t, 9, idxRegionCnt)
+	tk.MustExec("alter table t add index idx2(c) pre_split_regions = (between (0) and (100000) regions 3);")
+	retRows = tk.MustQuery("show table t regions;").Rows()
+	idxRegionCnt = 0
+	for _, r := range retRows {
+		startKey := r[2].(string)
+		if strings.Contains(startKey, "_i_2") {
+			idxRegionCnt++
+		}
+	}
+	require.Equal(t, 12, idxRegionCnt)
 }