ddl: extract add-index KV generation helper

pkg/ddl/backfilling_operators.go

@@ -915,7 +915,7 @@ func (w *indexIngestWorker) WriteChunk(rs *IndexRecordChunk) (count int, bytes i
 		indexConditionCheckers = nil
 	}
 	cnt, kvBytes, err := writeChunk(w.ctx, w.writers, w.indexes, indexConditionCheckers, w.copCtx,
-		sc.TimeZone(), sc.ErrCtx(), vars.GetWriteStmtBufs(), rs.Chunk, w.tbl.Meta(), w.tbl.UseNewCollate())
+		sc.TimeZone(), sc.ErrCtx(), vars.GetWriteStmtBufs(), rs.Chunk, w.tbl.GetPhysicalID(), w.tbl.Meta(), w.tbl.UseNewCollate())
 	if err != nil || cnt == 0 {
 		return 0, 0, err
 	}

pkg/ddl/index.go

@@ -2697,6 +2697,58 @@ func getLocalWriterConfig(indexCnt, writerCnt int) *backend.LocalWriterConfig {
 	return writerCfg
 }
 
+func generateIndexKVsForRow(
+	indexes []table.Index,
+	loc *time.Location,
+	errCtx errctx.Context,
+	handle kv.Handle,
+	physicalID int64,
+	idxDataBuf []types.Datum,
+	checkPartialCondition func(int, table.Index) (bool, error),
+	fetchIndexValues func(int, table.Index, []types.Datum) ([]types.Datum, error),
+	restoreData func(int, table.Index) []types.Datum,
+	sink func(int, table.Index, []types.Datum, table.IndexKVGenerator) (int64, error),
+) (int64, error) {
+	var totalBytes int64
+	for i, index := range indexes {
+		if index.Meta().HasCondition() && checkPartialCondition != nil {
+			ok, err := checkPartialCondition(i, index)
+			if err != nil {
+				return totalBytes, errors.Trace(err)
+			}
+			if !ok {
+				continue
+			}
+		}
+		indexedValues, err := fetchIndexValues(i, index, idxDataBuf)
+		if err != nil {
+			return totalBytes, errors.Trace(err)
+		}
+		var rsData []types.Datum
+		if restoreData != nil {
+			rsData = restoreData(i, index)
+		}
+		indexHandle := indexKVHandleForPhysicalTable(index, physicalID, handle)
+		iter := index.GenIndexKVIter(errCtx, loc, indexedValues, indexHandle, rsData)
+		kvBytes, err := sink(i, index, indexedValues, iter)
+		if err != nil {
+			return totalBytes, errors.Trace(err)
+		}
+		totalBytes += kvBytes
+	}
+	return totalBytes, nil
+}
+
+func indexKVHandleForPhysicalTable(index table.Index, physicalID int64, handle kv.Handle) kv.Handle {
+	if index.Meta().Global && index.Meta().GlobalIndexVersion >= model.GlobalIndexVersionV1 {
+		if ph, ok := handle.(kv.PartitionHandle); ok {
+			handle = ph.Handle
+		}
+		return kv.NewPartitionHandle(physicalID, handle)
+	}
+	return handle
+}
+
 func writeChunk(
 	ctx context.Context,
 	writers []ingest.Writer,
@@ -2707,6 +2759,7 @@ func writeChunk(
 	errCtx errctx.Context,
 	writeStmtBufs *variable.WriteStmtBufs,
 	copChunk *chunk.Chunk,
+	physicalID int64,
 	tblInfo *model.TableInfo,
 	useNewCollate bool,
 ) (rowCnt int, bytes int, err error) {
@@ -2757,34 +2810,42 @@ func writeChunk(
 		if err != nil {
 			return 0, totalBytes, errors.Trace(err)
 		}
-		for i, index := range indexes {
-			// If the `IndexRecordChunk.conditionPushed` is true and we have only 1 index, the `indexConditionCheckers`
-			// will not be initialized.
-			if index.Meta().HasCondition() && indexConditionCheckers != nil {
+		var checkPartialCondition func(int, table.Index) (bool, error)
+		if indexConditionCheckers != nil {
+			checkPartialCondition = func(i int, _ table.Index) (bool, error) {
 				ok, err := indexConditionCheckers[i](row)
 				if err != nil {
-					return 0, 0, errors.Trace(err)
-				}
-				if !ok {
-					continue
+					return false, errors.Trace(err)
 				}
+				return ok, nil
 			}
-
-			idxID := index.Meta().ID
-			idxDataBuf = ExtractDatumByOffsets(ectx,
-				row, copCtx.IndexColumnOutputOffsets(idxID), c.ExprColumnInfos, idxDataBuf)
-			idxData := idxDataBuf[:len(index.Meta().Columns)]
-			var rsData []types.Datum
-			if needRestoreForIndexes[i] {
-				rsData = getRestoreData(useNewCollate, c.TableInfo, copCtx.IndexInfo(idxID), c.PrimaryKeyInfo, restoreDataBuf)
-			}
-			kvBytes, err := writeOneKV(ctx, writers[i], index, loc, errCtx, writeStmtBufs, idxData, rsData, h)
-			if err != nil {
-				err = ingest.TryConvertToKeyExistsErr(err, index.Meta(), tblInfo)
-				return 0, totalBytes, errors.Trace(err)
-			}
-			totalBytes += int(kvBytes)
 		}
+		kvBytes, err := generateIndexKVsForRow(
+			indexes, loc, errCtx, h, physicalID, idxDataBuf, checkPartialCondition,
+			func(_ int, index table.Index, buf []types.Datum) ([]types.Datum, error) {
+				idxID := index.Meta().ID
+				idxData := ExtractDatumByOffsets(ectx, row, copCtx.IndexColumnOutputOffsets(idxID), c.ExprColumnInfos, buf)
+				return idxData[:len(index.Meta().Columns)], nil
+			},
+			func(i int, index table.Index) []types.Datum {
+				if !needRestoreForIndexes[i] {
+					return nil
+				}
+				return getRestoreData(useNewCollate, c.TableInfo, copCtx.IndexInfo(index.Meta().ID), c.PrimaryKeyInfo, restoreDataBuf)
+			},
+			func(i int, _ table.Index, _ []types.Datum, iter table.IndexKVGenerator) (int64, error) {
+				kvBytes, err := writeOneKV(ctx, writers[i], writeStmtBufs, iter, h)
+				if err != nil {
+					err = ingest.TryConvertToKeyExistsErr(err, indexes[i].Meta(), tblInfo)
+					return 0, errors.Trace(err)
+				}
+				return kvBytes, nil
+			},
+		)
+		if err != nil {
+			return 0, totalBytes + int(kvBytes), errors.Trace(err)
+		}
+		totalBytes += int(kvBytes)
 		count++
 	}
 	return count, totalBytes, nil
@@ -2804,15 +2865,11 @@ func maxIndexColumnCount(indexes []table.Index) int {
 func writeOneKV(
 	ctx context.Context,
 	writer ingest.Writer,
-	index table.Index,
-	loc *time.Location,
-	errCtx errctx.Context,
 	writeBufs *variable.WriteStmtBufs,
-	idxDt, rsData []types.Datum,
+	iter table.IndexKVGenerator,
 	handle kv.Handle,
 ) (int64, error) {
 	var kvBytes int64
-	iter := index.GenIndexKVIter(errCtx, loc, idxDt, handle, rsData)
 	for iter.Valid() {
 		key, idxVal, _, err := iter.Next(writeBufs.IndexKeyBuf, writeBufs.RowValBuf)
 		if err != nil {

pkg/ddl/reorg_util_test.go

@@ -19,6 +19,7 @@ import (
 	"testing"
 
 	"github.com/docker/go-units"
+	"github.com/pingcap/tidb/pkg/kv"
 	"github.com/pingcap/tidb/pkg/meta/model"
 	"github.com/pingcap/tidb/pkg/store/helper"
 	"github.com/pingcap/tidb/pkg/table/tables"
@@ -84,6 +85,23 @@ func expectedRegionRange(tableID int64) ([]byte, []byte) {
 	return mockCodec{}.EncodeRegionRange(tableStart, tableEnd)
 }
 
+func TestIndexKVHandleForPhysicalTableRewrapsGlobalIndexHandle(t *testing.T) {
+	idx, err := tables.NewIndex(10, &model.TableInfo{ID: 1}, &model.IndexInfo{
+		ID:                 2,
+		Global:             true,
+		GlobalIndexVersion: model.GlobalIndexVersionV1,
+	})
+	require.NoError(t, err)
+
+	handle := indexKVHandleForPhysicalTable(idx, 20, kv.NewPartitionHandle(10, kv.IntHandle(42)))
+	partitionHandle, ok := handle.(kv.PartitionHandle)
+	require.True(t, ok)
+	require.Equal(t, int64(20), partitionHandle.PartitionID)
+	require.True(t, partitionHandle.Handle.Equal(kv.IntHandle(42)))
+	_, nested := partitionHandle.Handle.(kv.PartitionHandle)
+	require.False(t, nested)
+}
+
 func TestEstimateTableSizeByIDUsesMaxApproximateSizes(t *testing.T) {
 	pdCli := &mockPDHTTPClient{
 		regionInfos: []*pdhttp.RegionsInfo{

pkg/table/tables/BUILD.bazel

@@ -80,7 +80,7 @@ go_test(
     ],
     embed = [":tables"],
     flaky = True,
-    shard_count = 44,
+    shard_count = 45,
     deps = [
         "//pkg/ddl",
         "//pkg/domain",

pkg/table/tables/index.go

@@ -206,13 +206,14 @@ func (c *index) GenIndexValue(ec errctx.Context, loc *time.Location, distinct, u
 	return idx, err
 }
 
-// getIndexedValue will produce the result like:
+// BuildMultiValueIndexValueGroups expands indexed values for a multi-valued index.
+// It will produce the result like:
 // 1. If not multi-valued index, return directly.
 // 2. (i1, [m1,m2], i2, ...) ==> [(i1, m1, i2, ...), (i1, m2, i2, ...)]
 // 3. (i1, null, i2, ...) ==> [(i1, null, i2, ...)]
 // 4. (i1, [], i2, ...) ==> nothing.
-func (c *index) getIndexedValue(indexedValues []types.Datum) [][]types.Datum {
-	if !c.idxInfo.MVIndex {
+func BuildMultiValueIndexValueGroups(tblInfo *model.TableInfo, idxInfo *model.IndexInfo, indexedValues []types.Datum) [][]types.Datum {
+	if idxInfo == nil || !idxInfo.MVIndex {
 		return [][]types.Datum{indexedValues}
 	}
 
@@ -224,7 +225,7 @@ func (c *index) getIndexedValue(indexedValues []types.Datum) [][]types.Datum {
 	for !jsonIsNull {
 		val := make([]types.Datum, 0, len(indexedValues))
 		for i, v := range indexedValues {
-			if !c.tblInfo.Columns[c.idxInfo.Columns[i].Offset].FieldType.IsArray() {
+			if !tblInfo.Columns[idxInfo.Columns[i].Offset].FieldType.IsArray() {
 				val = append(val, v)
 			} else {
 				// if the datum type is not JSON, it must come from cleanup index.
@@ -258,6 +259,35 @@ out:
 	return vals
 }
 
+// EncodeRawIndexKeyValues encodes one raw index-key suffix per generated index entry.
+// For multi-valued indexes, it expands the indexed values in the same order as index
+// entry generation, then truncates and encodes each group.
+func EncodeRawIndexKeyValues(
+	useNewCollate bool,
+	loc *time.Location,
+	tblInfo *model.TableInfo,
+	idxInfo *model.IndexInfo,
+	indexedValues []types.Datum,
+) ([][]byte, error) {
+	rawValueGroups := BuildMultiValueIndexValueGroups(tblInfo, idxInfo, indexedValues)
+	rawKeys := make([][]byte, 0, len(rawValueGroups))
+	encoder := codec.NewEncoder(useNewCollate)
+	for _, rawValues := range rawValueGroups {
+		clonedValues := append([]types.Datum(nil), rawValues...)
+		tablecodec.TruncateIndexValues(tblInfo, idxInfo, clonedValues)
+		rawKey, err := encoder.EncodeKey(loc, nil, clonedValues...)
+		if err != nil {
+			return nil, err
+		}
+		rawKeys = append(rawKeys, rawKey)
+	}
+	return rawKeys, nil
+}
+
+func (c *index) getIndexedValue(indexedValues []types.Datum) [][]types.Datum {
+	return BuildMultiValueIndexValueGroups(c.tblInfo, c.idxInfo, indexedValues)
+}
+
 // MeetPartialCondition checks whether the row meets the partial index condition of the index.
 func (c *index) MeetPartialCondition(row []types.Datum) (meet bool, err error) {
 	if c.conditionExpr == nil {

pkg/table/tables/index_test.go

@@ -183,6 +183,56 @@ func buildTableInfo(t *testing.T, sql string) *model.TableInfo {
 	return tblInfo
 }
 
+func TestEncodeRawIndexKeyValues(t *testing.T) {
+	loc := time.UTC
+
+	t.Run("non multi-valued index truncates prefix columns", func(t *testing.T) {
+		tblInfo := buildTableInfo(t, "create table t (a varchar(10), key idx(a(2)))")
+		idxInfo := tblInfo.Indices[0]
+
+		rawKeys, err := tables.EncodeRawIndexKeyValues(true, loc, tblInfo, idxInfo, types.MakeDatums("abcd"))
+		require.NoError(t, err)
+		require.Len(t, rawKeys, 1)
+
+		remain, datum, err := codec.DecodeOne(rawKeys[0])
+		require.NoError(t, err)
+		require.Empty(t, remain)
+		require.Equal(t, "ab", datum.GetString())
+	})
+
+	t.Run("multi-valued index expands unique array elements in order", func(t *testing.T) {
+		jsonType := types.NewFieldType(mysql.TypeJSON)
+		jsonType.SetArray(true)
+		tblInfo := &model.TableInfo{
+			Columns: []*model.ColumnInfo{
+				{ID: 1, Offset: 0, FieldType: *jsonType},
+			},
+		}
+		idxInfo := &model.IndexInfo{
+			MVIndex: true,
+			Columns: []*model.IndexColumn{
+				{Offset: 0},
+			},
+		}
+
+		arrayDatum, err := types.ParseBinaryJSONFromString(`[2, 1, 2]`)
+		require.NoError(t, err)
+		rawKeys, err := tables.EncodeRawIndexKeyValues(true, loc, tblInfo, idxInfo, types.MakeDatums(arrayDatum))
+		require.NoError(t, err)
+		require.Len(t, rawKeys, 2)
+
+		remain, datum, err := codec.DecodeOne(rawKeys[0])
+		require.NoError(t, err)
+		require.Empty(t, remain)
+		require.EqualValues(t, 2, datum.GetInt64())
+
+		remain, datum, err = codec.DecodeOne(rawKeys[1])
+		require.NoError(t, err)
+		require.Empty(t, remain)
+		require.EqualValues(t, 1, datum.GetInt64())
+	})
+}
+
 func TestGenIndexValueFromIndex(t *testing.T) {
 	tblInfo := buildTableInfo(t, "create table a (a int primary key, b int not null, c text, unique key key_b(b));")
 	tblInfo.State = model.StatePublic