From ded1ed25174bd39d9cc68e8751322058ef83dd02 Mon Sep 17 00:00:00 2001 From: "xiao.li" Date: Tue, 7 Jul 2026 19:26:30 +0800 Subject: [PATCH 1/2] ddl: extract add-index KV generation helper --- pkg/ddl/backfilling_operators.go | 2 +- pkg/ddl/index.go | 114 ++++++++++++----- pkg/ddl/index_cop.go | 27 ---- pkg/ddl/reorg_util_test.go | 18 +++ pkg/executor/admin.go | 8 +- pkg/table/tables/BUILD.bazel | 2 +- pkg/table/tables/index.go | 38 +++++- pkg/table/tables/index_test.go | 50 ++++++++ pkg/table/tables/mutation_checker_test.go | 142 ++++++++++++++++++++++ pkg/table/tables/tables.go | 57 +++++++-- 10 files changed, 381 insertions(+), 77 deletions(-) diff --git a/pkg/ddl/backfilling_operators.go b/pkg/ddl/backfilling_operators.go index 80d8ac9f285e8..e6ea2774ef2bb 100644 --- a/pkg/ddl/backfilling_operators.go +++ b/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 } diff --git a/pkg/ddl/index.go b/pkg/ddl/index.go index f675c705921dd..2137acf037e94 100644 --- a/pkg/ddl/index.go +++ b/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,43 @@ 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 tables.TryGetHandleRestoredData( + useNewCollate, c.TableInfo, c.PrimaryKeyInfo, restoreDataBuf, copCtx.IndexInfo(index.Meta().ID)) + }, + 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 +2866,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 { diff --git a/pkg/ddl/index_cop.go b/pkg/ddl/index_cop.go index 65adf1ee8e7c6..edccaffe048f2 100644 --- a/pkg/ddl/index_cop.go +++ b/pkg/ddl/index_cop.go @@ -145,33 +145,6 @@ func completeErr(err error, idxInfo *model.IndexInfo) error { return errors.Trace(err) } -func getRestoreData(useNewCollate bool, tblInfo *model.TableInfo, targetIdx, pkIdx *model.IndexInfo, handleDts []types.Datum) []types.Datum { - if !useNewCollate || !tblInfo.IsCommonHandle || tblInfo.CommonHandleVersion == 0 { - return nil - } - if pkIdx == nil { - return nil - } - for i, pkIdxCol := range pkIdx.Columns { - pkCol := tblInfo.Columns[pkIdxCol.Offset] - if !types.NeedRestoredDataWithCollate(&pkCol.FieldType, useNewCollate) { - // Since the handle data cannot be null, we can use SetNull to - // indicate that this column does not need to be restored. - handleDts[i].SetNull() - continue - } - tables.TryTruncateRestoredData(&handleDts[i], pkCol, pkIdxCol, targetIdx) - tables.ConvertDatumToTailSpaceCount(&handleDts[i], pkCol) - } - dtToRestored := handleDts[:0] - for _, handleDt := range handleDts { - if !handleDt.IsNull() { - dtToRestored = append(dtToRestored, handleDt) - } - } - return dtToRestored -} - func buildDAGPB(ctx context.Context, exprCtx exprctx.BuildContext, distSQLCtx *distsqlctx.DistSQLContext, pushDownFlags uint64, tblInfo *model.TableInfo, colInfos []*model.ColumnInfo, selectExpr expression.Expression) (*tipb.DAGRequest, bool, error) { conditionPushed := false diff --git a/pkg/ddl/reorg_util_test.go b/pkg/ddl/reorg_util_test.go index d637ed6d6aa9e..a2a3b6d9a5566 100644 --- a/pkg/ddl/reorg_util_test.go +++ b/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{ diff --git a/pkg/executor/admin.go b/pkg/executor/admin.go index 950c5cf40b052..da934ec102894 100644 --- a/pkg/executor/admin.go +++ b/pkg/executor/admin.go @@ -404,10 +404,12 @@ func (e *RecoverIndexExec) fetchRecoverRows(ctx context.Context, srcResult dists return nil, err } e.idxValsBufs[result.scanRowCount] = idxVals - rsData := tables.TryGetHandleRestoredDataWrapper( - e.table, + tblInfo := e.table.Meta() + rsData := tables.TryGetHandleRestoredData( + e.table.UseNewCollate(), + tblInfo, + tables.FindPrimaryIndex(tblInfo), plannercore.GetCommonHandleDatum(e.handleCols, row), - nil, e.index.Meta(), ) e.recoverRows = append(e.recoverRows, recoverRows{handle: handle, idxVals: idxVals, rsData: rsData, skip: true}) diff --git a/pkg/table/tables/BUILD.bazel b/pkg/table/tables/BUILD.bazel index 9775792b76383..3383b156c6c8d 100644 --- a/pkg/table/tables/BUILD.bazel +++ b/pkg/table/tables/BUILD.bazel @@ -80,7 +80,7 @@ go_test( ], embed = [":tables"], flaky = True, - shard_count = 44, + shard_count = 46, deps = [ "//pkg/ddl", "//pkg/domain", diff --git a/pkg/table/tables/index.go b/pkg/table/tables/index.go index 5d433e5d656c8..6887d16308b55 100644 --- a/pkg/table/tables/index.go +++ b/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 { diff --git a/pkg/table/tables/index_test.go b/pkg/table/tables/index_test.go index 3d5e7d8a0e7ed..a29650252982c 100644 --- a/pkg/table/tables/index_test.go +++ b/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 diff --git a/pkg/table/tables/mutation_checker_test.go b/pkg/table/tables/mutation_checker_test.go index 7e83ba2fa251e..bd0754e9a4eeb 100644 --- a/pkg/table/tables/mutation_checker_test.go +++ b/pkg/table/tables/mutation_checker_test.go @@ -399,3 +399,145 @@ func buildIndexKeyValue(index table.Index, rowToInsert []types.Datum, tc types.C } return key, value, nil } + +func TestTryGetHandleRestoredData(t *testing.T) { + t.Run("returns restored common handle data", func(t *testing.T) { + varcharType := types.NewFieldTypeWithCollation(mysql.TypeVarchar, "utf8_unicode_ci", types.UnspecifiedLength) + intType := types.NewFieldType(mysql.TypeLonglong) + charType := types.NewFieldTypeWithCollation(mysql.TypeString, "utf8_unicode_ci", types.UnspecifiedLength) + tableInfo := &model.TableInfo{ + ID: 1, + Name: ast.NewCIStr("t"), + Columns: []*model.ColumnInfo{ + {ID: 1, Name: ast.NewCIStr("a"), Offset: 0, FieldType: *varcharType}, + {ID: 2, Name: ast.NewCIStr("b"), Offset: 1, FieldType: *intType}, + {ID: 3, Name: ast.NewCIStr("c"), Offset: 2, FieldType: *charType}, + }, + PKIsHandle: false, + IsCommonHandle: true, + CommonHandleVersion: 1, + Indices: []*model.IndexInfo{ + { + ID: 1, + Name: ast.NewCIStr("primary"), + Primary: true, + Columns: []*model.IndexColumn{ + {Name: ast.NewCIStr("a"), Offset: 0, Length: types.UnspecifiedLength}, + {Name: ast.NewCIStr("c"), Offset: 2, Length: types.UnspecifiedLength}, + }, + }, + { + ID: 2, + Name: ast.NewCIStr("idx"), + Columns: []*model.IndexColumn{{Name: ast.NewCIStr("b"), Offset: 1, Length: types.UnspecifiedLength}}, + }, + }, + } + + rsData := TryGetHandleRestoredData(true, tableInfo, FindPrimaryIndex(tableInfo), types.MakeDatums("a", "c"), tableInfo.Indices[1]) + require.Len(t, rsData, 2) + require.Equal(t, "a", rsData[0].GetString()) + require.Equal(t, "c", rsData[1].GetString()) + }) + + t.Run("does not mutate input while filtering and converting restored data", func(t *testing.T) { + varcharType := types.NewFieldTypeWithCollation(mysql.TypeVarchar, "utf8_unicode_ci", types.UnspecifiedLength) + charBinType := types.NewFieldTypeWithCollation(mysql.TypeString, "utf8_bin", types.UnspecifiedLength) + varcharBinType := types.NewFieldTypeWithCollation(mysql.TypeVarchar, "utf8_bin", types.UnspecifiedLength) + intType := types.NewFieldType(mysql.TypeLonglong) + tableInfo := &model.TableInfo{ + ID: 1, + Name: ast.NewCIStr("t"), + Columns: []*model.ColumnInfo{ + {ID: 1, Name: ast.NewCIStr("a"), Offset: 0, FieldType: *varcharType}, + {ID: 2, Name: ast.NewCIStr("b"), Offset: 1, FieldType: *charBinType}, + {ID: 3, Name: ast.NewCIStr("c"), Offset: 2, FieldType: *varcharBinType}, + {ID: 4, Name: ast.NewCIStr("d"), Offset: 3, FieldType: *intType}, + }, + PKIsHandle: false, + IsCommonHandle: true, + CommonHandleVersion: 1, + Indices: []*model.IndexInfo{ + { + ID: 1, + Name: ast.NewCIStr("primary"), + Primary: true, + Columns: []*model.IndexColumn{ + {Name: ast.NewCIStr("a"), Offset: 0, Length: types.UnspecifiedLength}, + {Name: ast.NewCIStr("b"), Offset: 1, Length: types.UnspecifiedLength}, + {Name: ast.NewCIStr("c"), Offset: 2, Length: types.UnspecifiedLength}, + }, + }, + { + ID: 2, + Name: ast.NewCIStr("idx"), + Columns: []*model.IndexColumn{{Name: ast.NewCIStr("d"), Offset: 3, Length: types.UnspecifiedLength}}, + }, + }, + } + handleDts := types.MakeDatums("alpha", "ignored ", "keep ") + + rsData := TryGetHandleRestoredData(true, tableInfo, FindPrimaryIndex(tableInfo), handleDts, tableInfo.Indices[1]) + require.Len(t, rsData, 2) + require.Equal(t, "alpha", rsData[0].GetString()) + require.Equal(t, int64(2), rsData[1].GetInt64()) + require.Equal(t, "alpha", handleDts[0].GetString()) + require.Equal(t, "ignored ", handleDts[1].GetString()) + require.Equal(t, "keep ", handleDts[2].GetString()) + }) + + t.Run("uses longer shared secondary index prefix without mutating input", func(t *testing.T) { + varcharType := types.NewFieldTypeWithCollation(mysql.TypeVarchar, "utf8_unicode_ci", types.UnspecifiedLength) + tableInfo := &model.TableInfo{ + ID: 1, + Name: ast.NewCIStr("t"), + Columns: []*model.ColumnInfo{ + {ID: 1, Name: ast.NewCIStr("a"), Offset: 0, FieldType: *varcharType}, + }, + PKIsHandle: false, + IsCommonHandle: true, + CommonHandleVersion: 1, + Indices: []*model.IndexInfo{ + { + ID: 1, + Name: ast.NewCIStr("primary"), + Primary: true, + Columns: []*model.IndexColumn{{Name: ast.NewCIStr("a"), Offset: 0, Length: 2}}, + }, + { + ID: 2, + Name: ast.NewCIStr("idx"), + Columns: []*model.IndexColumn{{Name: ast.NewCIStr("a"), Offset: 0, Length: 4}}, + }, + }, + } + handleDts := types.MakeDatums("abcdef") + + rsData := TryGetHandleRestoredData(true, tableInfo, FindPrimaryIndex(tableInfo), handleDts, tableInfo.Indices[1]) + require.Len(t, rsData, 1) + require.Equal(t, "abcd", rsData[0].GetString()) + require.Equal(t, "abcdef", handleDts[0].GetString()) + }) + + t.Run("returns nil when restored data is not needed", func(t *testing.T) { + varcharType := types.NewFieldTypeWithCollation(mysql.TypeVarchar, "utf8_unicode_ci", types.UnspecifiedLength) + tableInfo := &model.TableInfo{ + ID: 1, + Name: ast.NewCIStr("t"), + Columns: []*model.ColumnInfo{{ID: 1, Name: ast.NewCIStr("a"), Offset: 0, FieldType: *varcharType}}, + PKIsHandle: false, + IsCommonHandle: true, + CommonHandleVersion: 1, + Indices: []*model.IndexInfo{ + { + ID: 1, + Name: ast.NewCIStr("primary"), + Primary: true, + Columns: []*model.IndexColumn{{Name: ast.NewCIStr("a"), Offset: 0, Length: types.UnspecifiedLength}}, + }, + }, + } + + require.Nil(t, TryGetHandleRestoredData(false, tableInfo, FindPrimaryIndex(tableInfo), types.MakeDatums("a"), tableInfo.Indices[0])) + }) +} diff --git a/pkg/table/tables/tables.go b/pkg/table/tables/tables.go index 69ea903197bb3..917e309bc27a9 100644 --- a/pkg/table/tables/tables.go +++ b/pkg/table/tables/tables.go @@ -1784,33 +1784,64 @@ func (t *TableCommon) GetSequenceCommon() *sequenceCommon { return t.sequence } +// TryGetHandleRestoredData returns restored common-handle data if needed. +func TryGetHandleRestoredData( + useNewCollate bool, + tblInfo *model.TableInfo, + pkIdx *model.IndexInfo, + handleDts []types.Datum, + idx *model.IndexInfo, +) []types.Datum { + if !needHandleRestoredData(useNewCollate, tblInfo, pkIdx) { + return nil + } + restoredHandleDts := make([]types.Datum, len(handleDts)) + for i := range handleDts { + handleDts[i].Copy(&restoredHandleDts[i]) + } + for i, pkIdxCol := range pkIdx.Columns { + pkCol := tblInfo.Columns[pkIdxCol.Offset] + if !types.NeedRestoredDataWithCollate(&pkCol.FieldType, useNewCollate) { + // Since the handle data cannot be null, SetNull marks this column as not restored. + restoredHandleDts[i].SetNull() + continue + } + TryTruncateRestoredData(&restoredHandleDts[i], pkCol, pkIdxCol, idx) + ConvertDatumToTailSpaceCount(&restoredHandleDts[i], pkCol) + } + rsData := restoredHandleDts[:0] + for _, handleDt := range restoredHandleDts { + if !handleDt.IsNull() { + rsData = append(rsData, handleDt) + } + } + return rsData +} + +func needHandleRestoredData(useNewCollate bool, tblInfo *model.TableInfo, pkIdx *model.IndexInfo) bool { + return useNewCollate && tblInfo.IsCommonHandle && tblInfo.CommonHandleVersion != 0 && pkIdx != nil +} + // TryGetHandleRestoredDataWrapper tries to get the restored data for handle if // needed. The argument can be a slice or a map. func TryGetHandleRestoredDataWrapper(tbl table.Table, row []types.Datum, rowMap map[int64]types.Datum, idx *model.IndexInfo) []types.Datum { useNewCollate := tbl.UseNewCollate() tblInfo := tbl.Meta() - if !useNewCollate || !tblInfo.IsCommonHandle || tblInfo.CommonHandleVersion == 0 { + pkIdx := FindPrimaryIndex(tblInfo) + if !needHandleRestoredData(useNewCollate, tblInfo, pkIdx) { return nil } - rsData := make([]types.Datum, 0, 4) - pkIdx := FindPrimaryIndex(tblInfo) + handleDts := make([]types.Datum, 0, len(pkIdx.Columns)) for _, pkIdxCol := range pkIdx.Columns { pkCol := tblInfo.Columns[pkIdxCol.Offset] - if !types.NeedRestoredDataWithCollate(&pkCol.FieldType, useNewCollate) { - continue - } - var datum types.Datum if len(rowMap) > 0 { - datum = rowMap[pkCol.ID] + handleDts = append(handleDts, rowMap[pkCol.ID]) } else { - datum = row[pkCol.Offset] + handleDts = append(handleDts, row[pkCol.Offset]) } - TryTruncateRestoredData(&datum, pkCol, pkIdxCol, idx) - ConvertDatumToTailSpaceCount(&datum, pkCol) - rsData = append(rsData, datum) } - return rsData + return TryGetHandleRestoredData(useNewCollate, tblInfo, pkIdx, handleDts, idx) } // TryTruncateRestoredData tries to truncate index values. From 0861e8fac83de87214ccd613d2ca48c1dde69d37 Mon Sep 17 00:00:00 2001 From: "xiao.li" Date: Wed, 8 Jul 2026 16:01:57 +0800 Subject: [PATCH 2/2] ddl: keep restored handle data helper local --- pkg/ddl/index.go | 3 +- pkg/ddl/index_cop.go | 27 ++++ pkg/executor/admin.go | 8 +- pkg/table/tables/BUILD.bazel | 2 +- pkg/table/tables/mutation_checker_test.go | 142 ---------------------- pkg/table/tables/tables.go | 57 ++------- 6 files changed, 45 insertions(+), 194 deletions(-) diff --git a/pkg/ddl/index.go b/pkg/ddl/index.go index 2137acf037e94..2e8d4014039d2 100644 --- a/pkg/ddl/index.go +++ b/pkg/ddl/index.go @@ -2831,8 +2831,7 @@ func writeChunk( if !needRestoreForIndexes[i] { return nil } - return tables.TryGetHandleRestoredData( - useNewCollate, c.TableInfo, c.PrimaryKeyInfo, restoreDataBuf, copCtx.IndexInfo(index.Meta().ID)) + 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) diff --git a/pkg/ddl/index_cop.go b/pkg/ddl/index_cop.go index edccaffe048f2..65adf1ee8e7c6 100644 --- a/pkg/ddl/index_cop.go +++ b/pkg/ddl/index_cop.go @@ -145,6 +145,33 @@ func completeErr(err error, idxInfo *model.IndexInfo) error { return errors.Trace(err) } +func getRestoreData(useNewCollate bool, tblInfo *model.TableInfo, targetIdx, pkIdx *model.IndexInfo, handleDts []types.Datum) []types.Datum { + if !useNewCollate || !tblInfo.IsCommonHandle || tblInfo.CommonHandleVersion == 0 { + return nil + } + if pkIdx == nil { + return nil + } + for i, pkIdxCol := range pkIdx.Columns { + pkCol := tblInfo.Columns[pkIdxCol.Offset] + if !types.NeedRestoredDataWithCollate(&pkCol.FieldType, useNewCollate) { + // Since the handle data cannot be null, we can use SetNull to + // indicate that this column does not need to be restored. + handleDts[i].SetNull() + continue + } + tables.TryTruncateRestoredData(&handleDts[i], pkCol, pkIdxCol, targetIdx) + tables.ConvertDatumToTailSpaceCount(&handleDts[i], pkCol) + } + dtToRestored := handleDts[:0] + for _, handleDt := range handleDts { + if !handleDt.IsNull() { + dtToRestored = append(dtToRestored, handleDt) + } + } + return dtToRestored +} + func buildDAGPB(ctx context.Context, exprCtx exprctx.BuildContext, distSQLCtx *distsqlctx.DistSQLContext, pushDownFlags uint64, tblInfo *model.TableInfo, colInfos []*model.ColumnInfo, selectExpr expression.Expression) (*tipb.DAGRequest, bool, error) { conditionPushed := false diff --git a/pkg/executor/admin.go b/pkg/executor/admin.go index da934ec102894..950c5cf40b052 100644 --- a/pkg/executor/admin.go +++ b/pkg/executor/admin.go @@ -404,12 +404,10 @@ func (e *RecoverIndexExec) fetchRecoverRows(ctx context.Context, srcResult dists return nil, err } e.idxValsBufs[result.scanRowCount] = idxVals - tblInfo := e.table.Meta() - rsData := tables.TryGetHandleRestoredData( - e.table.UseNewCollate(), - tblInfo, - tables.FindPrimaryIndex(tblInfo), + rsData := tables.TryGetHandleRestoredDataWrapper( + e.table, plannercore.GetCommonHandleDatum(e.handleCols, row), + nil, e.index.Meta(), ) e.recoverRows = append(e.recoverRows, recoverRows{handle: handle, idxVals: idxVals, rsData: rsData, skip: true}) diff --git a/pkg/table/tables/BUILD.bazel b/pkg/table/tables/BUILD.bazel index 3383b156c6c8d..e9b4d421d24cf 100644 --- a/pkg/table/tables/BUILD.bazel +++ b/pkg/table/tables/BUILD.bazel @@ -80,7 +80,7 @@ go_test( ], embed = [":tables"], flaky = True, - shard_count = 46, + shard_count = 45, deps = [ "//pkg/ddl", "//pkg/domain", diff --git a/pkg/table/tables/mutation_checker_test.go b/pkg/table/tables/mutation_checker_test.go index bd0754e9a4eeb..7e83ba2fa251e 100644 --- a/pkg/table/tables/mutation_checker_test.go +++ b/pkg/table/tables/mutation_checker_test.go @@ -399,145 +399,3 @@ func buildIndexKeyValue(index table.Index, rowToInsert []types.Datum, tc types.C } return key, value, nil } - -func TestTryGetHandleRestoredData(t *testing.T) { - t.Run("returns restored common handle data", func(t *testing.T) { - varcharType := types.NewFieldTypeWithCollation(mysql.TypeVarchar, "utf8_unicode_ci", types.UnspecifiedLength) - intType := types.NewFieldType(mysql.TypeLonglong) - charType := types.NewFieldTypeWithCollation(mysql.TypeString, "utf8_unicode_ci", types.UnspecifiedLength) - tableInfo := &model.TableInfo{ - ID: 1, - Name: ast.NewCIStr("t"), - Columns: []*model.ColumnInfo{ - {ID: 1, Name: ast.NewCIStr("a"), Offset: 0, FieldType: *varcharType}, - {ID: 2, Name: ast.NewCIStr("b"), Offset: 1, FieldType: *intType}, - {ID: 3, Name: ast.NewCIStr("c"), Offset: 2, FieldType: *charType}, - }, - PKIsHandle: false, - IsCommonHandle: true, - CommonHandleVersion: 1, - Indices: []*model.IndexInfo{ - { - ID: 1, - Name: ast.NewCIStr("primary"), - Primary: true, - Columns: []*model.IndexColumn{ - {Name: ast.NewCIStr("a"), Offset: 0, Length: types.UnspecifiedLength}, - {Name: ast.NewCIStr("c"), Offset: 2, Length: types.UnspecifiedLength}, - }, - }, - { - ID: 2, - Name: ast.NewCIStr("idx"), - Columns: []*model.IndexColumn{{Name: ast.NewCIStr("b"), Offset: 1, Length: types.UnspecifiedLength}}, - }, - }, - } - - rsData := TryGetHandleRestoredData(true, tableInfo, FindPrimaryIndex(tableInfo), types.MakeDatums("a", "c"), tableInfo.Indices[1]) - require.Len(t, rsData, 2) - require.Equal(t, "a", rsData[0].GetString()) - require.Equal(t, "c", rsData[1].GetString()) - }) - - t.Run("does not mutate input while filtering and converting restored data", func(t *testing.T) { - varcharType := types.NewFieldTypeWithCollation(mysql.TypeVarchar, "utf8_unicode_ci", types.UnspecifiedLength) - charBinType := types.NewFieldTypeWithCollation(mysql.TypeString, "utf8_bin", types.UnspecifiedLength) - varcharBinType := types.NewFieldTypeWithCollation(mysql.TypeVarchar, "utf8_bin", types.UnspecifiedLength) - intType := types.NewFieldType(mysql.TypeLonglong) - tableInfo := &model.TableInfo{ - ID: 1, - Name: ast.NewCIStr("t"), - Columns: []*model.ColumnInfo{ - {ID: 1, Name: ast.NewCIStr("a"), Offset: 0, FieldType: *varcharType}, - {ID: 2, Name: ast.NewCIStr("b"), Offset: 1, FieldType: *charBinType}, - {ID: 3, Name: ast.NewCIStr("c"), Offset: 2, FieldType: *varcharBinType}, - {ID: 4, Name: ast.NewCIStr("d"), Offset: 3, FieldType: *intType}, - }, - PKIsHandle: false, - IsCommonHandle: true, - CommonHandleVersion: 1, - Indices: []*model.IndexInfo{ - { - ID: 1, - Name: ast.NewCIStr("primary"), - Primary: true, - Columns: []*model.IndexColumn{ - {Name: ast.NewCIStr("a"), Offset: 0, Length: types.UnspecifiedLength}, - {Name: ast.NewCIStr("b"), Offset: 1, Length: types.UnspecifiedLength}, - {Name: ast.NewCIStr("c"), Offset: 2, Length: types.UnspecifiedLength}, - }, - }, - { - ID: 2, - Name: ast.NewCIStr("idx"), - Columns: []*model.IndexColumn{{Name: ast.NewCIStr("d"), Offset: 3, Length: types.UnspecifiedLength}}, - }, - }, - } - handleDts := types.MakeDatums("alpha", "ignored ", "keep ") - - rsData := TryGetHandleRestoredData(true, tableInfo, FindPrimaryIndex(tableInfo), handleDts, tableInfo.Indices[1]) - require.Len(t, rsData, 2) - require.Equal(t, "alpha", rsData[0].GetString()) - require.Equal(t, int64(2), rsData[1].GetInt64()) - require.Equal(t, "alpha", handleDts[0].GetString()) - require.Equal(t, "ignored ", handleDts[1].GetString()) - require.Equal(t, "keep ", handleDts[2].GetString()) - }) - - t.Run("uses longer shared secondary index prefix without mutating input", func(t *testing.T) { - varcharType := types.NewFieldTypeWithCollation(mysql.TypeVarchar, "utf8_unicode_ci", types.UnspecifiedLength) - tableInfo := &model.TableInfo{ - ID: 1, - Name: ast.NewCIStr("t"), - Columns: []*model.ColumnInfo{ - {ID: 1, Name: ast.NewCIStr("a"), Offset: 0, FieldType: *varcharType}, - }, - PKIsHandle: false, - IsCommonHandle: true, - CommonHandleVersion: 1, - Indices: []*model.IndexInfo{ - { - ID: 1, - Name: ast.NewCIStr("primary"), - Primary: true, - Columns: []*model.IndexColumn{{Name: ast.NewCIStr("a"), Offset: 0, Length: 2}}, - }, - { - ID: 2, - Name: ast.NewCIStr("idx"), - Columns: []*model.IndexColumn{{Name: ast.NewCIStr("a"), Offset: 0, Length: 4}}, - }, - }, - } - handleDts := types.MakeDatums("abcdef") - - rsData := TryGetHandleRestoredData(true, tableInfo, FindPrimaryIndex(tableInfo), handleDts, tableInfo.Indices[1]) - require.Len(t, rsData, 1) - require.Equal(t, "abcd", rsData[0].GetString()) - require.Equal(t, "abcdef", handleDts[0].GetString()) - }) - - t.Run("returns nil when restored data is not needed", func(t *testing.T) { - varcharType := types.NewFieldTypeWithCollation(mysql.TypeVarchar, "utf8_unicode_ci", types.UnspecifiedLength) - tableInfo := &model.TableInfo{ - ID: 1, - Name: ast.NewCIStr("t"), - Columns: []*model.ColumnInfo{{ID: 1, Name: ast.NewCIStr("a"), Offset: 0, FieldType: *varcharType}}, - PKIsHandle: false, - IsCommonHandle: true, - CommonHandleVersion: 1, - Indices: []*model.IndexInfo{ - { - ID: 1, - Name: ast.NewCIStr("primary"), - Primary: true, - Columns: []*model.IndexColumn{{Name: ast.NewCIStr("a"), Offset: 0, Length: types.UnspecifiedLength}}, - }, - }, - } - - require.Nil(t, TryGetHandleRestoredData(false, tableInfo, FindPrimaryIndex(tableInfo), types.MakeDatums("a"), tableInfo.Indices[0])) - }) -} diff --git a/pkg/table/tables/tables.go b/pkg/table/tables/tables.go index 917e309bc27a9..69ea903197bb3 100644 --- a/pkg/table/tables/tables.go +++ b/pkg/table/tables/tables.go @@ -1784,64 +1784,33 @@ func (t *TableCommon) GetSequenceCommon() *sequenceCommon { return t.sequence } -// TryGetHandleRestoredData returns restored common-handle data if needed. -func TryGetHandleRestoredData( - useNewCollate bool, - tblInfo *model.TableInfo, - pkIdx *model.IndexInfo, - handleDts []types.Datum, - idx *model.IndexInfo, -) []types.Datum { - if !needHandleRestoredData(useNewCollate, tblInfo, pkIdx) { - return nil - } - restoredHandleDts := make([]types.Datum, len(handleDts)) - for i := range handleDts { - handleDts[i].Copy(&restoredHandleDts[i]) - } - for i, pkIdxCol := range pkIdx.Columns { - pkCol := tblInfo.Columns[pkIdxCol.Offset] - if !types.NeedRestoredDataWithCollate(&pkCol.FieldType, useNewCollate) { - // Since the handle data cannot be null, SetNull marks this column as not restored. - restoredHandleDts[i].SetNull() - continue - } - TryTruncateRestoredData(&restoredHandleDts[i], pkCol, pkIdxCol, idx) - ConvertDatumToTailSpaceCount(&restoredHandleDts[i], pkCol) - } - rsData := restoredHandleDts[:0] - for _, handleDt := range restoredHandleDts { - if !handleDt.IsNull() { - rsData = append(rsData, handleDt) - } - } - return rsData -} - -func needHandleRestoredData(useNewCollate bool, tblInfo *model.TableInfo, pkIdx *model.IndexInfo) bool { - return useNewCollate && tblInfo.IsCommonHandle && tblInfo.CommonHandleVersion != 0 && pkIdx != nil -} - // TryGetHandleRestoredDataWrapper tries to get the restored data for handle if // needed. The argument can be a slice or a map. func TryGetHandleRestoredDataWrapper(tbl table.Table, row []types.Datum, rowMap map[int64]types.Datum, idx *model.IndexInfo) []types.Datum { useNewCollate := tbl.UseNewCollate() tblInfo := tbl.Meta() - pkIdx := FindPrimaryIndex(tblInfo) - if !needHandleRestoredData(useNewCollate, tblInfo, pkIdx) { + if !useNewCollate || !tblInfo.IsCommonHandle || tblInfo.CommonHandleVersion == 0 { return nil } - handleDts := make([]types.Datum, 0, len(pkIdx.Columns)) + rsData := make([]types.Datum, 0, 4) + pkIdx := FindPrimaryIndex(tblInfo) for _, pkIdxCol := range pkIdx.Columns { pkCol := tblInfo.Columns[pkIdxCol.Offset] + if !types.NeedRestoredDataWithCollate(&pkCol.FieldType, useNewCollate) { + continue + } + var datum types.Datum if len(rowMap) > 0 { - handleDts = append(handleDts, rowMap[pkCol.ID]) + datum = rowMap[pkCol.ID] } else { - handleDts = append(handleDts, row[pkCol.Offset]) + datum = row[pkCol.Offset] } + TryTruncateRestoredData(&datum, pkCol, pkIdxCol, idx) + ConvertDatumToTailSpaceCount(&datum, pkCol) + rsData = append(rsData, datum) } - return TryGetHandleRestoredData(useNewCollate, tblInfo, pkIdx, handleDts, idx) + return rsData } // TryTruncateRestoredData tries to truncate index values.