From 8c906d45193f3a09774d838e443c76769f2f0a72 Mon Sep 17 00:00:00 2001 From: Ruihao Chen Date: Thu, 2 Jul 2026 15:44:10 +0800 Subject: [PATCH 01/35] ddl, importer: capture collation mode for background encoding --- pkg/ddl/backfilling.go | 45 ++++++++-------- pkg/ddl/backfilling_dist_executor.go | 4 +- pkg/ddl/backfilling_dist_scheduler.go | 8 ++- pkg/ddl/backfilling_operators.go | 11 ++-- pkg/ddl/backfilling_txn_executor.go | 4 +- pkg/ddl/executor.go | 2 + pkg/ddl/export_test.go | 3 +- pkg/ddl/index.go | 11 ++-- pkg/ddl/index_cop.go | 11 ++-- pkg/ddl/index_merge_tmp.go | 2 +- pkg/ddl/job_scheduler.go | 25 --------- pkg/dxf/importinto/planner.go | 7 ++- pkg/dxf/importinto/task_executor.go | 23 +++++++-- pkg/executor/importer/import.go | 30 ++++++++++- pkg/executor/importer/sampler.go | 26 ++++++++-- pkg/executor/importer/table_import.go | 68 +++++++++++++++++++++---- pkg/expression/expression.go | 18 ++++++- pkg/lightning/backend/encode/encode.go | 3 ++ pkg/lightning/backend/kv/base.go | 7 ++- pkg/lightning/backend/kv/sql2kv.go | 8 +++ pkg/meta/model/reorg.go | 18 +++++++ pkg/planner/core/expression_rewriter.go | 1 + pkg/table/tables/partition.go | 13 +++-- pkg/table/tables/tables.go | 19 ++++++- 24 files changed, 273 insertions(+), 94 deletions(-) diff --git a/pkg/ddl/backfilling.go b/pkg/ddl/backfilling.go index a2d6aa02b0fb4..60d8d5d67fe20 100644 --- a/pkg/ddl/backfilling.go +++ b/pkg/ddl/backfilling.go @@ -47,6 +47,7 @@ import ( "github.com/pingcap/tidb/pkg/table" "github.com/pingcap/tidb/pkg/tablecodec" "github.com/pingcap/tidb/pkg/util" + "github.com/pingcap/tidb/pkg/util/collate" contextutil "github.com/pingcap/tidb/pkg/util/context" "github.com/pingcap/tidb/pkg/util/dbterror" decoder "github.com/pingcap/tidb/pkg/util/rowDecoder" @@ -159,14 +160,15 @@ type backfillTaskContext struct { type backfillCtx struct { id int *ddlCtx - warnings contextutil.WarnHandlerExt - loc *time.Location - exprCtx exprctx.BuildContext - tblCtx table.MutateContext - schemaName string - table table.Table - batchCnt int - jobContext *ReorgContext + warnings contextutil.WarnHandlerExt + loc *time.Location + exprCtx exprctx.BuildContext + tblCtx table.MutateContext + schemaName string + table table.Table + batchCnt int + useNewCollate bool + jobContext *ReorgContext metricCounter prometheus.Counter conflictCounter prometheus.Counter @@ -213,18 +215,20 @@ func newBackfillCtx(id int, rInfo *reorgInfo, schemaName string, tbl table.Table } batchCnt := rInfo.ReorgMeta.GetBatchSize() + useNewCollate := rInfo.ReorgMeta.GetUseNewCollateOrDefault(collate.NewCollationEnabled()) metricTableID := backfillMetricsTableID(rInfo, label) return &backfillCtx{ - id: id, - ddlCtx: rInfo.jobCtx.oldDDLCtx, - warnings: warnHandler, - exprCtx: exprCtx, - tblCtx: tblCtx, - loc: exprCtx.GetEvalCtx().Location(), - schemaName: schemaName, - table: tbl, - batchCnt: batchCnt, - jobContext: jobCtx, + id: id, + ddlCtx: rInfo.jobCtx.oldDDLCtx, + warnings: warnHandler, + exprCtx: exprCtx, + tblCtx: tblCtx, + loc: exprCtx.GetEvalCtx().Location(), + schemaName: schemaName, + table: tbl, + batchCnt: batchCnt, + useNewCollate: useNewCollate, + jobContext: jobCtx, metricCounter: getBackfillTotalByTableID( metricTableID, label, schemaName, tbl.Meta().Name.String(), colOrIdxName), conflictCounter: getBackfillTotalByTableID( @@ -719,12 +723,13 @@ func sendTasks( return nil } -func makeupDecodeColMap(dbName ast.CIStr, t table.Table) (map[int64]decoder.Column, error) { +func makeupDecodeColMap(dbName ast.CIStr, t table.Table, useNewCollate bool) (map[int64]decoder.Column, error) { writableColInfos := make([]*model.ColumnInfo, 0, len(t.WritableCols())) for _, col := range t.WritableCols() { writableColInfos = append(writableColInfos, col.ColumnInfo) } - exprCols, _, err := expression.ColumnInfos2ColumnsAndNames(newReorgExprCtx(), dbName, t.Meta().Name, writableColInfos, t.Meta()) + exprCols, _, err := expression.ColumnInfos2ColumnsAndNamesWithCollate( + newReorgExprCtx(), dbName, t.Meta().Name, writableColInfos, t.Meta(), useNewCollate) if err != nil { return nil, err } diff --git a/pkg/ddl/backfilling_dist_executor.go b/pkg/ddl/backfilling_dist_executor.go index 0403d4ad0e854..3f0436ff56c23 100644 --- a/pkg/ddl/backfilling_dist_executor.go +++ b/pkg/ddl/backfilling_dist_executor.go @@ -137,9 +137,9 @@ func (s *backfillDistExecutor) newBackfillStepExecutor( store := s.TaskRuntime.Store() sessPool := sess.NewSessionPool(s.TaskRuntime.SysSessionPool()) - // TODO getTableByTxn is using DDL ctx which is never cancelled except when shutdown. + // TODO This is using DDL ctx which is never cancelled except when shutdown. // we should move this operation out of GetStepExecutor, and put into Init. - _, tblIface, err := getTableByTxn(ddlObj.ctx, store, jobMeta.SchemaID, jobMeta.TableID) + tblIface, err := getUserTableFromTaskStore(ddlObj.ctx, store, jobMeta) if err != nil { return nil, err } diff --git a/pkg/ddl/backfilling_dist_scheduler.go b/pkg/ddl/backfilling_dist_scheduler.go index 46a1a98d34505..9a9ac1486fe28 100644 --- a/pkg/ddl/backfilling_dist_scheduler.go +++ b/pkg/ddl/backfilling_dist_scheduler.go @@ -47,7 +47,9 @@ import ( "github.com/pingcap/tidb/pkg/sessionctx/vardef" "github.com/pingcap/tidb/pkg/store/helper" "github.com/pingcap/tidb/pkg/table" + "github.com/pingcap/tidb/pkg/table/tables" "github.com/pingcap/tidb/pkg/util/backoff" + "github.com/pingcap/tidb/pkg/util/collate" "github.com/tikv/client-go/v2/oracle" "github.com/tikv/client-go/v2/tikv" "go.uber.org/zap" @@ -200,7 +202,11 @@ func getUserTableFromTaskStore( return nil, err } // we don't touch table data during add-index, a fake Allocators is enough. - tbl, err := table.TableFromMeta(autoid.NewAllocators(tblInfo.SepAutoInc()), tblInfo) + tbl, err := tables.TableFromMetaWithCollate( + job.ReorgMeta.GetUseNewCollateOrDefault(collate.NewCollationEnabled()), + autoid.NewAllocators(tblInfo.SepAutoInc()), + tblInfo, + ) if err != nil { return nil, err } diff --git a/pkg/ddl/backfilling_operators.go b/pkg/ddl/backfilling_operators.go index 4af509e6baf51..a1d9ac715a4b2 100644 --- a/pkg/ddl/backfilling_operators.go +++ b/pkg/ddl/backfilling_operators.go @@ -49,6 +49,7 @@ import ( "github.com/pingcap/tidb/pkg/table/tables" "github.com/pingcap/tidb/pkg/tablecodec" "github.com/pingcap/tidb/pkg/util/chunk" + "github.com/pingcap/tidb/pkg/util/collate" contextutil "github.com/pingcap/tidb/pkg/util/context" "github.com/pingcap/tidb/pkg/util/dbterror" "github.com/pingcap/tidb/pkg/util/execdetails" @@ -109,9 +110,10 @@ func NewAddIndexIngestPipeline( concurrency int, collector execute.Collector, ) (*operator.AsyncPipeline, error) { + useNewCollate := reorgMeta.GetUseNewCollateOrDefault(collate.NewCollationEnabled()) indexes := make([]table.Index, 0, len(idxInfos)) for _, idxInfo := range idxInfos { - index, err := tables.NewIndex(tbl.GetPhysicalID(), tbl.Meta(), idxInfo) + index, err := tables.NewIndexWithCollate(useNewCollate, tbl.GetPhysicalID(), tbl.Meta(), idxInfo) if err != nil { return nil, err } @@ -167,9 +169,10 @@ func NewWriteIndexToExternalStoragePipeline( collector execute.Collector, tikvCodec tikv.Codec, ) (*operator.AsyncPipeline, error) { + useNewCollate := reorgMeta.GetUseNewCollateOrDefault(collate.NewCollationEnabled()) indexes := make([]table.Index, 0, len(idxInfos)) for _, idxInfo := range idxInfos { - index, err := tables.NewIndex(tbl.GetPhysicalID(), tbl.Meta(), idxInfo) + index, err := tables.NewIndexWithCollate(useNewCollate, tbl.GetPhysicalID(), tbl.Meta(), idxInfo) if err != nil { return nil, err } @@ -914,7 +917,9 @@ func (w *indexIngestWorker) WriteChunk(rs *IndexRecordChunk) (count int, bytes i // skip running the checker in TiDB side. 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()) + useNewCollate := w.reorgMeta.GetUseNewCollateOrDefault(collate.NewCollationEnabled()) + cnt, kvBytes, err := writeChunk(w.ctx, w.writers, w.indexes, indexConditionCheckers, w.copCtx, + sc.TimeZone(), sc.ErrCtx(), vars.GetWriteStmtBufs(), rs.Chunk, w.tbl.Meta(), useNewCollate) if err != nil || cnt == 0 { return 0, 0, err } diff --git a/pkg/ddl/backfilling_txn_executor.go b/pkg/ddl/backfilling_txn_executor.go index 46c2121ce6acf..65322b100c30c 100644 --- a/pkg/ddl/backfilling_txn_executor.go +++ b/pkg/ddl/backfilling_txn_executor.go @@ -33,6 +33,7 @@ import ( "github.com/pingcap/tidb/pkg/sessionctx/stmtctx" "github.com/pingcap/tidb/pkg/sessionctx/vardef" "github.com/pingcap/tidb/pkg/table" + "github.com/pingcap/tidb/pkg/util/collate" contextutil "github.com/pingcap/tidb/pkg/util/context" "github.com/pingcap/tidb/pkg/util/execdetails" "github.com/pingcap/tidb/pkg/util/intest" @@ -82,7 +83,8 @@ type txnBackfillExecutor struct { func newTxnBackfillExecutor(ctx context.Context, info *reorgInfo, sessPool *sess.Pool, tp backfillerType, tbl table.PhysicalTable, jobCtx *ReorgContext) (backfillExecutor, error) { - decColMap, err := makeupDecodeColMap(info.dbInfo.Name, tbl) + useNewCollate := info.ReorgMeta.GetUseNewCollateOrDefault(collate.NewCollationEnabled()) + decColMap, err := makeupDecodeColMap(info.dbInfo.Name, tbl, useNewCollate) if err != nil { return nil, err } diff --git a/pkg/ddl/executor.go b/pkg/ddl/executor.go index ea8f170051857..804be37b7bcca 100644 --- a/pkg/ddl/executor.go +++ b/pkg/ddl/executor.go @@ -7535,6 +7535,7 @@ func getJobCheckInterval(action model.ActionType, i int) (time.Duration, bool) { // NewDDLReorgMeta create a DDL ReorgMeta. func NewDDLReorgMeta(ctx sessionctx.Context) *model.DDLReorgMeta { tzName, tzOffset := ddlutil.GetTimeZone(ctx) + useNewCollate := collate.NewCollationEnabled() return &model.DDLReorgMeta{ SQLMode: ctx.GetSessionVars().SQLMode, Warnings: make(map[errors.ErrorID]*terror.Error), @@ -7542,6 +7543,7 @@ func NewDDLReorgMeta(ctx sessionctx.Context) *model.DDLReorgMeta { Location: &model.TimeZoneLocation{Name: tzName, Offset: tzOffset}, ResourceGroupName: ctx.GetSessionVars().StmtCtx.ResourceGroupName, Version: model.CurrentReorgMetaVersion, + UseNewCollate: &useNewCollate, } } diff --git a/pkg/ddl/export_test.go b/pkg/ddl/export_test.go index 40d13e07a14e7..2b55667c3e4f7 100644 --- a/pkg/ddl/export_test.go +++ b/pkg/ddl/export_test.go @@ -33,6 +33,7 @@ import ( "github.com/pingcap/tidb/pkg/table" "github.com/pingcap/tidb/pkg/types" "github.com/pingcap/tidb/pkg/util/chunk" + "github.com/pingcap/tidb/pkg/util/collate" "github.com/pingcap/tidb/pkg/util/mock" ) @@ -80,6 +81,6 @@ func ConvertRowToHandleAndIndexDatum( c := copCtx.GetBase() idxData := ddl.ExtractDatumByOffsets(ctx, row, copCtx.IndexColumnOutputOffsets(idxID), c.ExprColumnInfos, idxDataBuf) handleData := ddl.ExtractDatumByOffsets(ctx, row, c.HandleOutputOffsets, c.ExprColumnInfos, handleDataBuf) - handle, err := ddl.BuildHandle(handleData, c.TableInfo, c.PrimaryKeyInfo, time.Local, errctx.StrictNoWarningContext) + handle, err := ddl.BuildHandle(collate.NewCollationEnabled(), handleData, c.TableInfo, c.PrimaryKeyInfo, time.Local, errctx.StrictNoWarningContext) return handle, idxData, err } diff --git a/pkg/ddl/index.go b/pkg/ddl/index.go index 863fabffbdde8..9beef955ec0cb 100644 --- a/pkg/ddl/index.go +++ b/pkg/ddl/index.go @@ -82,7 +82,6 @@ import ( "github.com/pingcap/tidb/pkg/util" "github.com/pingcap/tidb/pkg/util/backoff" "github.com/pingcap/tidb/pkg/util/chunk" - "github.com/pingcap/tidb/pkg/util/collate" "github.com/pingcap/tidb/pkg/util/dbterror" "github.com/pingcap/tidb/pkg/util/engine" "github.com/pingcap/tidb/pkg/util/generatedexpr" @@ -2406,7 +2405,7 @@ func newAddIndexTxnWorker( continue } indexInfo := model.FindIndexInfoByID(t.Meta().Indices, elem.ID) - index, err := tables.NewIndex(t.GetPhysicalID(), t.Meta(), indexInfo) + index, err := tables.NewIndexWithCollate(bfCtx.useNewCollate, t.GetPhysicalID(), t.Meta(), indexInfo) if err != nil { return nil, err } @@ -2479,7 +2478,7 @@ func (w *baseIndexWorker) getIndexRecord(idxInfo *model.IndexInfo, handle kv.Han idxVal[j] = idxColumnVal } - rsData := tables.TryGetHandleRestoredDataWrapper(collate.NewCollationEnabled(), w.table.Meta(), nil, w.rowMap, idxInfo) + rsData := tables.TryGetHandleRestoredDataWrapper(w.useNewCollate, w.table.Meta(), nil, w.rowMap, idxInfo) idxRecord := &indexRecord{handle: handle, key: recordKey, vals: idxVal, rsData: rsData} return idxRecord, nil } @@ -2709,6 +2708,7 @@ func writeChunk( writeStmtBufs *variable.WriteStmtBufs, copChunk *chunk.Chunk, tblInfo *model.TableInfo, + useNewCollate bool, ) (rowCnt int, bytes int, err error) { iter := chunk.NewIterator4Chunk(copChunk) c := copCtx.GetBase() @@ -2733,7 +2733,6 @@ func writeChunk( }() needRestoreForIndexes := make([]bool, len(indexes)) restore, pkNeedRestore := false, false - useNewCollate := collate.NewCollationEnabled() if c.PrimaryKeyInfo != nil && c.TableInfo.IsCommonHandle && c.TableInfo.CommonHandleVersion != 0 { pkNeedRestore = tables.NeedRestoredData(useNewCollate, c.PrimaryKeyInfo.Columns, c.TableInfo.Columns) } @@ -2754,7 +2753,7 @@ func writeChunk( restoreDataBuf[i] = *datum.Clone() } } - h, err := BuildHandle(handleDataBuf, c.TableInfo, c.PrimaryKeyInfo, loc, errCtx) + h, err := BuildHandle(useNewCollate, handleDataBuf, c.TableInfo, c.PrimaryKeyInfo, loc, errCtx) if err != nil { return 0, totalBytes, errors.Trace(err) } @@ -2777,7 +2776,7 @@ func writeChunk( idxData := idxDataBuf[:len(index.Meta().Columns)] var rsData []types.Datum if needRestoreForIndexes[i] { - rsData = getRestoreData(c.TableInfo, copCtx.IndexInfo(idxID), c.PrimaryKeyInfo, restoreDataBuf) + 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 { diff --git a/pkg/ddl/index_cop.go b/pkg/ddl/index_cop.go index baafbf4f5149d..65adf1ee8e7c6 100644 --- a/pkg/ddl/index_cop.go +++ b/pkg/ddl/index_cop.go @@ -36,7 +36,6 @@ import ( "github.com/pingcap/tidb/pkg/types" "github.com/pingcap/tidb/pkg/util/chunk" "github.com/pingcap/tidb/pkg/util/codec" - "github.com/pingcap/tidb/pkg/util/collate" "github.com/pingcap/tidb/pkg/util/logutil" "github.com/pingcap/tidb/pkg/util/timeutil" "github.com/pingcap/tipb/go-tipb" @@ -146,8 +145,8 @@ func completeErr(err error, idxInfo *model.IndexInfo) error { return errors.Trace(err) } -func getRestoreData(tblInfo *model.TableInfo, targetIdx, pkIdx *model.IndexInfo, handleDts []types.Datum) []types.Datum { - if !collate.NewCollationEnabled() || !tblInfo.IsCommonHandle || tblInfo.CommonHandleVersion == 0 { +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 { @@ -155,7 +154,7 @@ func getRestoreData(tblInfo *model.TableInfo, targetIdx, pkIdx *model.IndexInfo, } for i, pkIdxCol := range pkIdx.Columns { pkCol := tblInfo.Columns[pkIdxCol.Offset] - if !types.NeedRestoredData(&pkCol.FieldType) { + 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() @@ -257,11 +256,11 @@ func ExtractDatumByOffsets(ctx expression.EvalContext, row chunk.Row, offsets [] } // BuildHandle is exported for test. -func BuildHandle(pkDts []types.Datum, tblInfo *model.TableInfo, +func BuildHandle(useNewCollate bool, pkDts []types.Datum, tblInfo *model.TableInfo, pkInfo *model.IndexInfo, loc *time.Location, errCtx errctx.Context) (kv.Handle, error) { if tblInfo.IsCommonHandle { tablecodec.TruncateIndexValues(tblInfo, pkInfo, pkDts) - handleBytes, err := codec.EncodeKey(loc, nil, pkDts...) + handleBytes, err := codec.NewEncoder(useNewCollate).EncodeKey(loc, nil, pkDts...) err = errCtx.HandleError(err) if err != nil { return nil, err diff --git a/pkg/ddl/index_merge_tmp.go b/pkg/ddl/index_merge_tmp.go index e35efb15a23a9..63bec58a51b9c 100644 --- a/pkg/ddl/index_merge_tmp.go +++ b/pkg/ddl/index_merge_tmp.go @@ -148,7 +148,7 @@ func newMergeTempIndexWorker(bfCtx *backfillCtx, t table.PhysicalTable, elements allIndexes := make([]table.Index, 0, len(elements)) for _, elem := range elements { indexInfo := model.FindIndexInfoByID(t.Meta().Indices, elem.ID) - index, err := tables.NewIndex(t.GetPhysicalID(), t.Meta(), indexInfo) + index, err := tables.NewIndexWithCollate(bfCtx.useNewCollate, t.GetPhysicalID(), t.Meta(), indexInfo) if err != nil { return nil, err } diff --git a/pkg/ddl/job_scheduler.go b/pkg/ddl/job_scheduler.go index 7cbbd6f8a7913..752e48c92ceed 100644 --- a/pkg/ddl/job_scheduler.go +++ b/pkg/ddl/job_scheduler.go @@ -40,13 +40,11 @@ import ( "github.com/pingcap/tidb/pkg/ddl/util" "github.com/pingcap/tidb/pkg/kv" "github.com/pingcap/tidb/pkg/meta" - "github.com/pingcap/tidb/pkg/meta/autoid" "github.com/pingcap/tidb/pkg/meta/metadef" "github.com/pingcap/tidb/pkg/meta/model" "github.com/pingcap/tidb/pkg/metrics" "github.com/pingcap/tidb/pkg/owner" "github.com/pingcap/tidb/pkg/sessionctx/vardef" - "github.com/pingcap/tidb/pkg/table" tidbutil "github.com/pingcap/tidb/pkg/util" "github.com/pingcap/tidb/pkg/util/dbterror" "github.com/pingcap/tidb/pkg/util/etcd" @@ -691,29 +689,6 @@ func (s *jobScheduler) workerPoolExhausted() bool { s.bgJobWorkerPool.available() == 0 } -func getTableByTxn(ctx context.Context, store kv.Storage, schemaID, tableID int64) (*model.DBInfo, table.Table, error) { - var tbl table.Table - var dbInfo *model.DBInfo - err := kv.RunInNewTxn(ctx, store, false, func(_ context.Context, txn kv.Transaction) error { - t := meta.NewMutator(txn) - var err1 error - dbInfo, err1 = t.GetDatabase(schemaID) - if err1 != nil { - return errors.Trace(err1) - } - tblInfo, err1 := getTableInfo(t, tableID, schemaID) - if err1 != nil { - return errors.Trace(err1) - } - // This tableInfo should never interact with the autoid allocator, - // so we can use the autoid.Allocators{} here. - // TODO(tangenta): Use model.TableInfo instead of tables.Table. - tbl, err1 = table.TableFromMeta(autoid.Allocators{}, tblInfo) - return errors.Trace(err1) - }) - return dbInfo, tbl, err -} - func updateDDLJob2Table( ctx context.Context, se *sess.Session, diff --git a/pkg/dxf/importinto/planner.go b/pkg/dxf/importinto/planner.go index adc4eff13ceec..6619dcf9543ca 100644 --- a/pkg/dxf/importinto/planner.go +++ b/pkg/dxf/importinto/planner.go @@ -39,6 +39,7 @@ import ( "github.com/pingcap/tidb/pkg/objstore/storeapi" "github.com/pingcap/tidb/pkg/parser/mysql" "github.com/pingcap/tidb/pkg/table/tables" + "github.com/pingcap/tidb/pkg/util/collate" "github.com/pingcap/tidb/pkg/util/logutil" "go.uber.org/zap" ) @@ -351,7 +352,11 @@ func (*PostProcessSpec) ToSubtaskMeta(planCtx planner.PlanCtx) ([]byte, error) { func buildControllerForPlan(p *LogicalPlan) (*importer.LoadDataController, error) { plan, stmt := &p.Plan, p.Stmt idAlloc := kv.NewPanickingAllocators(plan.TableInfo.SepAutoInc()) - tbl, err := tables.TableFromMeta(idAlloc, plan.TableInfo) + tbl, err := tables.TableFromMetaWithCollate( + plan.GetUseNewCollateOrDefault(collate.NewCollationEnabled()), + idAlloc, + plan.TableInfo, + ) if err != nil { return nil, err } diff --git a/pkg/dxf/importinto/task_executor.go b/pkg/dxf/importinto/task_executor.go index 534bc2e830707..76ec9d8d47df6 100644 --- a/pkg/dxf/importinto/task_executor.go +++ b/pkg/dxf/importinto/task_executor.go @@ -51,6 +51,7 @@ import ( "github.com/pingcap/tidb/pkg/objstore/storeapi" "github.com/pingcap/tidb/pkg/resourcemanager/pool/workerpool" "github.com/pingcap/tidb/pkg/table/tables" + "github.com/pingcap/tidb/pkg/util/collate" "github.com/pingcap/tidb/pkg/util/dbterror/exeerrors" "github.com/pingcap/tidb/pkg/util/logutil" "go.uber.org/zap" @@ -106,7 +107,11 @@ func getTableImporter( logger *zap.Logger, ) (*importer.TableImporter, error) { idAlloc := kv.NewPanickingAllocators(taskMeta.Plan.TableInfo.SepAutoInc()) - tbl, err := tables.TableFromMeta(idAlloc, taskMeta.Plan.TableInfo) + tbl, err := tables.TableFromMetaWithCollate( + taskMeta.Plan.GetUseNewCollateOrDefault(collate.NewCollationEnabled()), + idAlloc, + taskMeta.Plan.TableInfo, + ) if err != nil { return nil, err } @@ -123,9 +128,21 @@ func getTableImporter( } failpoint.Inject("createTableImporterForTest", func() { - failpoint.Return(importer.NewTableImporterForTest(ctx, controller, strconv.FormatInt(taskID, 10), store)) + failpoint.Return(importer.NewTableImporterForTest( + ctx, + controller, + strconv.FormatInt(taskID, 10), + store, + importer.WithEncodingTable(tbl), + )) }) - return importer.NewTableImporter(ctx, controller, strconv.FormatInt(taskID, 10), store) + return importer.NewTableImporter( + ctx, + controller, + strconv.FormatInt(taskID, 10), + store, + importer.WithEncodingTable(tbl), + ) } func (s *importStepExecutor) Init(ctx context.Context) (err error) { diff --git a/pkg/executor/importer/import.go b/pkg/executor/importer/import.go index 86fd170494952..054a8b5b24ba7 100644 --- a/pkg/executor/importer/import.go +++ b/pkg/executor/importer/import.go @@ -65,6 +65,7 @@ import ( "github.com/pingcap/tidb/pkg/table" tidbutil "github.com/pingcap/tidb/pkg/util" "github.com/pingcap/tidb/pkg/util/chunk" + "github.com/pingcap/tidb/pkg/util/collate" contextutil "github.com/pingcap/tidb/pkg/util/context" "github.com/pingcap/tidb/pkg/util/cpu" "github.com/pingcap/tidb/pkg/util/dbterror" @@ -335,6 +336,10 @@ type Plan struct { ManualRecovery bool // the keyspace name when submitting this job, only for import-into Keyspace string + // UseNewCollate captures the collation mode used by background import workers + // when encoding keys. Nil means old metadata and should fall back to the + // caller-provided default. + UseNewCollate *bool `json:"use_new_collate,omitempty"` } // GetOnDupKeyMode returns the conflict handling mode. @@ -349,6 +354,20 @@ func (p *Plan) GetOnDupKeyMode() OnDupKeyMode { return p.OnDupKey } +// GetUseNewCollateOrDefault returns the captured collation mode, or defaultVal +// for import metadata generated before the field existed. +func (p *Plan) GetUseNewCollateOrDefault(defaultVal bool) bool { + if p == nil || p.UseNewCollate == nil { + return defaultVal + } + return *p.UseNewCollate +} + +// SetUseNewCollate stores the collation mode used by background import workers. +func (p *Plan) SetUseNewCollate(useNewCollate bool) { + p.UseNewCollate = &useNewCollate +} + // ASTArgs is the arguments for ast.LoadDataStmt. // TODO: remove this struct and use the struct which can be serialized. type ASTArgs struct { @@ -553,6 +572,7 @@ func NewImportPlan(ctx context.Context, userSctx sessionctx.Context, plan *plann User: userSctx.GetSessionVars().User.String(), Keyspace: userSctx.GetStore().GetKeyspace(), } + p.SetUseNewCollate(collate.NewCollationEnabled()) if err := p.initOptions(ctx, userSctx, plan.Options); err != nil { return nil, err } @@ -1924,6 +1944,7 @@ func createColAssignSimpleExprs( assignments []*ast.Assignment, ctx expression.BuildContext, mu *sync.Mutex, + useNewCollate bool, ) (_ []expression.Expression, _ []contextutil.SQLWarn, retErr error) { if mu != nil { mu.Lock() @@ -1932,7 +1953,7 @@ func createColAssignSimpleExprs( res := make([]expression.Expression, 0, len(assignments)) var allWarnings []contextutil.SQLWarn for _, assign := range assignments { - newExpr, err := expression.BuildSimpleExpr(ctx, assign.Expr) + newExpr, err := expression.BuildSimpleExpr(ctx, assign.Expr, expression.WithUseNewCollate(useNewCollate)) // col assign expr warnings is static, we should generate it for each row processed. // so we save it and clear it here. if ctx.GetEvalCtx().WarningCount() > 0 { @@ -1948,7 +1969,12 @@ func createColAssignSimpleExprs( // CreateColAssignSimpleExprs creates the column assignment expressions using `expression.BuildContext`. func (e *LoadDataController) CreateColAssignSimpleExprs(ctx expression.BuildContext) (_ []expression.Expression, _ []contextutil.SQLWarn, retErr error) { - return createColAssignSimpleExprs(e.ColumnAssignments, ctx, &e.colAssignMu) + return createColAssignSimpleExprs( + e.ColumnAssignments, + ctx, + &e.colAssignMu, + e.GetUseNewCollateOrDefault(collate.NewCollationEnabled()), + ) } func (e *LoadDataController) getLocalBackendCfg(keyspace, pdAddr, dataDir string) ingestctrl.BackendConfig { diff --git a/pkg/executor/importer/sampler.go b/pkg/executor/importer/sampler.go index e5f6bda1cc6ad..cda668d7f12dd 100644 --- a/pkg/executor/importer/sampler.go +++ b/pkg/executor/importer/sampler.go @@ -39,6 +39,7 @@ import ( plannercore "github.com/pingcap/tidb/pkg/planner/core" "github.com/pingcap/tidb/pkg/table" "github.com/pingcap/tidb/pkg/table/tables" + "github.com/pingcap/tidb/pkg/util/collate" contextutil "github.com/pingcap/tidb/pkg/util/context" "github.com/pingcap/tidb/pkg/util/dbterror/exeerrors" "go.uber.org/zap" @@ -82,6 +83,7 @@ type KVSizeSampleConfig struct { FieldNullDef []string LineFieldsInfo plannercore.LineFieldsInfo IgnoreLines uint64 + UseNewCollate *bool ColumnsAndUserVars []*ast.ColumnNameOrUserVar ColumnAssignments []*ast.Assignment @@ -103,6 +105,13 @@ func (r *SampledKVSizeResult) TotalKVSize() int64 { return int64(r.DataKVSize + r.IndexKVSize) } +func (cfg *KVSizeSampleConfig) getUseNewCollateOrDefault(defaultVal bool) bool { + if cfg == nil || cfg.UseNewCollate == nil { + return defaultVal + } + return *cfg.UseNewCollate +} + // SampleFileImportKVSize samples source rows with nextgen's KV encoder and returns // the sampled source bytes and encoded KV sizes. func SampleFileImportKVSize( @@ -155,6 +164,7 @@ func (e *LoadDataController) buildKVSizeSampleConfig() *KVSizeSampleConfig { FieldNullDef: append([]string(nil), e.FieldNullDef...), LineFieldsInfo: e.LineFieldsInfo, IgnoreLines: e.IgnoreLines, + UseNewCollate: e.UseNewCollate, ColumnsAndUserVars: e.ColumnsAndUserVars, ColumnAssignments: e.ColumnAssignments, } @@ -209,7 +219,12 @@ func validateKVSizeSampleConfig(cfg *KVSizeSampleConfig) error { func (s *kvSizeSampler) CreateColAssignSimpleExprs( ctx expression.BuildContext, ) (_ []expression.Expression, _ []contextutil.SQLWarn, retErr error) { - return createColAssignSimpleExprs(s.cfg.ColumnAssignments, ctx, &s.colAssignMu) + return createColAssignSimpleExprs( + s.cfg.ColumnAssignments, + ctx, + &s.colAssignMu, + s.cfg.getUseNewCollateOrDefault(collate.NewCollationEnabled()), + ) } func (s *kvSizeSampler) generateCSVConfig() *config.CSVConfig { @@ -281,9 +296,10 @@ func (s *kvSizeSampler) getKVEncoder( SysVars: s.cfg.ImportantSysVars, AutoRandomSeed: chunk.PrevRowIDMax, }, - Path: chunk.Path, - Table: encTable, - Logger: log.Logger{Logger: logger.With(zap.String("path", chunk.Path))}, + Path: chunk.Path, + Table: encTable, + Logger: log.Logger{Logger: logger.With(zap.String("path", chunk.Path))}, + UseNewCollate: s.cfg.UseNewCollate, } return newTableKVEncoderInner(cfg, s, s.fieldMappings, s.insertColumns) } @@ -331,7 +347,7 @@ func (s *kvSizeSampler) sampleOneFile( ParquetMeta: file.ParquetMeta, } idAlloc := kv.NewPanickingAllocators(s.table.Meta().SepAutoInc()) - tbl, err := tables.TableFromMeta(idAlloc, s.table.Meta()) + tbl, err := tables.TableFromMetaWithCollate(s.cfg.getUseNewCollateOrDefault(collate.NewCollationEnabled()), idAlloc, s.table.Meta()) if err != nil { return 0, 0, 0, errors.Annotatef(err, "failed to tables.TableFromMeta %s", s.table.Meta().Name) } diff --git a/pkg/executor/importer/table_import.go b/pkg/executor/importer/table_import.go index a96b4c5337619..c52b435f935c3 100644 --- a/pkg/executor/importer/table_import.go +++ b/pkg/executor/importer/table_import.go @@ -56,6 +56,7 @@ import ( "github.com/pingcap/tidb/pkg/table" "github.com/pingcap/tidb/pkg/table/tables" tidbutil "github.com/pingcap/tidb/pkg/util" + "github.com/pingcap/tidb/pkg/util/collate" "github.com/pingcap/tidb/pkg/util/etcd" "github.com/pingcap/tidb/pkg/util/logutil" "github.com/pingcap/tidb/pkg/util/promutil" @@ -91,6 +92,44 @@ var ( defaultMaxEngineSize = int64(5 * config.DefaultBatchSize) ) +type tableImporterOptions struct { + encTable table.Table +} + +// TableImporterOption configures NewTableImporter. +type TableImporterOption func(*tableImporterOptions) + +// WithEncodingTable sets the table used to encode rows and track allocated IDs. +func WithEncodingTable(tbl table.Table) TableImporterOption { + return func(opts *tableImporterOptions) { + opts.encTable = tbl + } +} + +func getTableImporterOptions(options []TableImporterOption) tableImporterOptions { + opts := tableImporterOptions{} + for _, opt := range options { + opt(&opts) + } + return opts +} + +func newEncodingTable(e *LoadDataController) (table.Table, error) { + idAlloc := kv.NewPanickingAllocators(e.Table.Meta().SepAutoInc()) + tbl, err := tables.TableFromMetaWithCollate(e.GetUseNewCollateOrDefault(collate.NewCollationEnabled()), idAlloc, e.Table.Meta()) + if err != nil { + return nil, errors.Annotatef(err, "failed to tables.TableFromMeta %s", e.Table.Meta().Name) + } + return tbl, nil +} + +func getEncodingTable(e *LoadDataController, opts tableImporterOptions) (table.Table, error) { + if opts.encTable != nil { + return opts.encTable, nil + } + return newEncodingTable(e) +} + // Chunk records the chunk information. type Chunk struct { Path string @@ -188,11 +227,12 @@ func NewTableImporter( e *LoadDataController, id string, kvStore tidbkv.Storage, + options ...TableImporterOption, ) (ti *TableImporter, err error) { - idAlloc := kv.NewPanickingAllocators(e.Table.Meta().SepAutoInc()) - tbl, err := tables.TableFromMeta(idAlloc, e.Table.Meta()) + opts := getTableImporterOptions(options) + tbl, err := getEncodingTable(e, opts) if err != nil { - return nil, errors.Annotatef(err, "failed to tables.TableFromMeta %s", e.Table.Meta().Name) + return nil, err } tidbCfg := tidb.GetGlobalConfig() @@ -284,12 +324,18 @@ func (s *storeHelper) GetTiKVCodec() tikv.Codec { var _ ingestctrl.StoreHelper = (*storeHelper)(nil) // NewTableImporterForTest creates a new table importer for test. -func NewTableImporterForTest(ctx context.Context, e *LoadDataController, id string, kvStore tidbkv.Storage) (*TableImporter, error) { +func NewTableImporterForTest( + ctx context.Context, + e *LoadDataController, + id string, + kvStore tidbkv.Storage, + options ...TableImporterOption, +) (*TableImporter, error) { helper := &storeHelper{kvStore: kvStore} - idAlloc := kv.NewPanickingAllocators(e.Table.Meta().SepAutoInc()) - tbl, err := tables.TableFromMeta(idAlloc, e.Table.Meta()) + opts := getTableImporterOptions(options) + tbl, err := getEncodingTable(e, opts) if err != nil { - return nil, errors.Annotatef(err, "failed to tables.TableFromMeta %s", e.Table.Meta().Name) + return nil, err } tidbCfg := tidb.GetGlobalConfig() @@ -393,9 +439,10 @@ func (e *LoadDataController) getKVEncoder(logger *zap.Logger, chunk *Chunk, encT SysVars: e.ImportantSysVars, AutoRandomSeed: chunk.PrevRowIDMax, }, - Path: chunk.Path, - Table: encTable, - Logger: log.Logger{Logger: logger.With(zap.String("path", chunk.Path))}, + Path: chunk.Path, + Table: encTable, + Logger: log.Logger{Logger: logger.With(zap.String("path", chunk.Path))}, + UseNewCollate: e.UseNewCollate, } return NewTableKVEncoder(cfg, e) } @@ -410,6 +457,7 @@ func (ti *TableImporter) GetKVEncoderForDupResolve() (*TableKVEncoder, error) { Table: ti.encTable, Logger: log.Logger{Logger: ti.logger}, UseIdentityAutoRowID: true, + UseNewCollate: ti.UseNewCollate, } return NewTableKVEncoderForDupResolve(cfg, ti.LoadDataController) } diff --git a/pkg/expression/expression.go b/pkg/expression/expression.go index a64d9573faec9..aa396e8d5a33a 100644 --- a/pkg/expression/expression.go +++ b/pkg/expression/expression.go @@ -30,6 +30,7 @@ import ( "github.com/pingcap/tidb/pkg/planner/cascades/base" "github.com/pingcap/tidb/pkg/types" "github.com/pingcap/tidb/pkg/util/chunk" + "github.com/pingcap/tidb/pkg/util/collate" "github.com/pingcap/tidb/pkg/util/generatedexpr" "github.com/pingcap/tidb/pkg/util/size" "github.com/pingcap/tidb/pkg/util/zeropool" @@ -1106,6 +1107,18 @@ func TableInfo2SchemaAndNames(ctx BuildContext, dbName ast.CIStr, tbl *model.Tab // ColumnInfos2ColumnsAndNames converts the ColumnInfo to the *Column and NameSlice. func ColumnInfos2ColumnsAndNames(ctx BuildContext, dbName, tblName ast.CIStr, colInfos []*model.ColumnInfo, tblInfo *model.TableInfo) ([]*Column, types.NameSlice, error) { + return ColumnInfos2ColumnsAndNamesWithCollate(ctx, dbName, tblName, colInfos, tblInfo, collate.NewCollationEnabled()) +} + +// ColumnInfos2ColumnsAndNamesWithCollate converts the ColumnInfo to the *Column +// and NameSlice with a fixed collation mode. +func ColumnInfos2ColumnsAndNamesWithCollate( + ctx BuildContext, + dbName, tblName ast.CIStr, + colInfos []*model.ColumnInfo, + tblInfo *model.TableInfo, + useNewCollate bool, +) ([]*Column, types.NameSlice, error) { columns := make([]*Column, 0, len(colInfos)) names := make([]*types.FieldName, 0, len(colInfos)) for i, col := range colInfos { @@ -1146,7 +1159,10 @@ func ColumnInfos2ColumnsAndNames(ctx BuildContext, dbName, tblName ast.CIStr, co if err != nil { return nil, nil, errors.Trace(err) } - e, err := BuildSimpleExpr(ctx, expr, WithInputSchemaAndNames(mockSchema, names, tblInfo), WithAllowCastArray(true)) + e, err := BuildSimpleExpr(ctx, expr, + WithInputSchemaAndNames(mockSchema, names, tblInfo), + WithAllowCastArray(true), + WithUseNewCollate(useNewCollate)) if err != nil { return nil, nil, errors.Trace(err) } diff --git a/pkg/lightning/backend/encode/encode.go b/pkg/lightning/backend/encode/encode.go index 931067181122e..3aed962304343 100644 --- a/pkg/lightning/backend/encode/encode.go +++ b/pkg/lightning/backend/encode/encode.go @@ -34,6 +34,9 @@ type EncodingConfig struct { // when encoding. // default false, in this case we will do sharding automatically if needed. UseIdentityAutoRowID bool + // UseNewCollate captures the collation mode used while encoding. Nil means + // old metadata and should fall back to the current process default. + UseNewCollate *bool } // EncodingBuilder consists of operations to handle encoding backend row data formats from source. diff --git a/pkg/lightning/backend/kv/base.go b/pkg/lightning/backend/kv/base.go index 05ac73a7eaf96..7de2a8354a2bf 100644 --- a/pkg/lightning/backend/kv/base.go +++ b/pkg/lightning/backend/kv/base.go @@ -32,6 +32,7 @@ import ( "github.com/pingcap/tidb/pkg/types" "github.com/pingcap/tidb/pkg/util/chunk" "github.com/pingcap/tidb/pkg/util/codec" + "github.com/pingcap/tidb/pkg/util/collate" "github.com/pingcap/tidb/pkg/util/redact" "go.uber.org/zap" "go.uber.org/zap/zapcore" @@ -167,7 +168,11 @@ func NewBaseKVEncoder(config *encode.EncodingConfig) (*BaseKVEncoder, error) { } // collect expressions for evaluating stored generated columns - genCols, err := CollectGeneratedColumns(se, meta, cols) + useNewCollate := collate.NewCollationEnabled() + if config.UseNewCollate != nil { + useNewCollate = *config.UseNewCollate + } + genCols, err := CollectGeneratedColumnsWithCollate(se, meta, cols, useNewCollate) if err != nil { return nil, errors.Annotate(err, "failed to parse generated column expressions") } diff --git a/pkg/lightning/backend/kv/sql2kv.go b/pkg/lightning/backend/kv/sql2kv.go index 31666a97fa947..6d65ce94792bf 100644 --- a/pkg/lightning/backend/kv/sql2kv.go +++ b/pkg/lightning/backend/kv/sql2kv.go @@ -35,6 +35,7 @@ import ( "github.com/pingcap/tidb/pkg/table" "github.com/pingcap/tidb/pkg/tablecodec" "github.com/pingcap/tidb/pkg/types" + "github.com/pingcap/tidb/pkg/util/collate" ) type tableKVEncoder struct { @@ -69,6 +70,12 @@ func NewTableKVEncoder( // CollectGeneratedColumns collects all expressions required to evaluate the // results of all generated columns. The returning slice is in evaluation order. func CollectGeneratedColumns(se *Session, meta *model.TableInfo, cols []*table.Column) ([]GeneratedCol, error) { + return CollectGeneratedColumnsWithCollate(se, meta, cols, collate.NewCollationEnabled()) +} + +// CollectGeneratedColumnsWithCollate collects all expressions required to +// evaluate generated columns with a fixed collation mode. +func CollectGeneratedColumnsWithCollate(se *Session, meta *model.TableInfo, cols []*table.Column, useNewCollate bool) ([]GeneratedCol, error) { hasGenCol := false for _, col := range cols { if col.GeneratedExpr != nil { @@ -112,6 +119,7 @@ func CollectGeneratedColumns(se *Session, meta *model.TableInfo, cols []*table.C col.GeneratedExpr.Internal(), expression.WithInputSchemaAndNames(schema, names, meta), expression.WithAllowCastArray(true), + expression.WithUseNewCollate(useNewCollate), ) if err != nil { return nil, err diff --git a/pkg/meta/model/reorg.go b/pkg/meta/model/reorg.go index e904c903dfc33..03a84398bd6a8 100644 --- a/pkg/meta/model/reorg.go +++ b/pkg/meta/model/reorg.go @@ -95,6 +95,10 @@ type DDLReorgMeta struct { MaxNodeCount int `json:"max_node_count"` AnalyzeState int8 `json:"analyze_state"` Stage ReorgStage `json:"stage"` + // UseNewCollate captures the collation mode used by background reorg workers + // when encoding keys. Nil means the metadata was generated before this field + // existed and should fall back to the caller-provided default. + UseNewCollate *bool `json:"use_new_collate,omitempty"` // These two variables are used to control the concurrency and batch size of the reorganization process. // They can be adjusted dynamically through `admin alter ddl jobs` command. // Note: Don't get or set these two variables directly, use the functions instead. @@ -151,6 +155,20 @@ func (dm *DDLReorgMeta) SetMaxWriteSpeed(maxWriteSpeed int) { dm.MaxWriteSpeed.Store(int64(maxWriteSpeed)) } +// GetUseNewCollateOrDefault returns the captured collation mode, or defaultVal +// for reorg metadata generated before the field existed. +func (dm *DDLReorgMeta) GetUseNewCollateOrDefault(defaultVal bool) bool { + if dm == nil || dm.UseNewCollate == nil { + return defaultVal + } + return *dm.UseNewCollate +} + +// SetUseNewCollate stores the collation mode used by background reorg workers. +func (dm *DDLReorgMeta) SetUseNewCollate(useNewCollate bool) { + dm.UseNewCollate = &useNewCollate +} + const ( // ReorgMetaVersion0 is the minimum version of DDLReorgMeta. ReorgMetaVersion0 = int64(0) diff --git a/pkg/planner/core/expression_rewriter.go b/pkg/planner/core/expression_rewriter.go index f380aa2bf0147..af503d6ce9037 100644 --- a/pkg/planner/core/expression_rewriter.go +++ b/pkg/planner/core/expression_rewriter.go @@ -282,6 +282,7 @@ func (b *PlanBuilder) getExpressionRewriter(ctx context.Context, p base.LogicalP rewriter.planCtx.aggrMap = nil rewriter.planCtx.insertPlan = nil rewriter.planCtx.rollExpand = b.currentBlockExpand + rewriter.useNewCollate = collate.NewCollationEnabled() return } diff --git a/pkg/table/tables/partition.go b/pkg/table/tables/partition.go index 8ee3895538d3a..0a6f31856d63c 100644 --- a/pkg/table/tables/partition.go +++ b/pkg/table/tables/partition.go @@ -45,6 +45,7 @@ import ( "github.com/pingcap/tidb/pkg/util" "github.com/pingcap/tidb/pkg/util/chunk" "github.com/pingcap/tidb/pkg/util/codec" + "github.com/pingcap/tidb/pkg/util/collate" "github.com/pingcap/tidb/pkg/util/dbterror" "github.com/pingcap/tidb/pkg/util/hack" "github.com/pingcap/tidb/pkg/util/logutil" @@ -181,7 +182,8 @@ func newPartitionedTable(tbl *TableCommon, tblInfo *model.TableInfo) (table.Part } else { tblInfo.Indices = origIndices } - err := initTableCommonWithIndices(&t.TableCommon, tblInfo, p.ID, tbl.Columns, tbl.allocs, tbl.Constraints) + err := initTableCommonWithIndicesAndCollate( + &t.TableCommon, tbl.encoder.UseNewCollate(), tblInfo, p.ID, tbl.Columns, tbl.allocs, tbl.Constraints) if err != nil { return nil, errors.Trace(err) } @@ -273,7 +275,8 @@ func newPartitionedTable(tbl *TableCommon, tblInfo *model.TableInfo) (table.Part func initPartition(t *partitionedTable, def model.PartitionDefinition) (*partition, error) { var newPart partition - err := initTableCommonWithIndices(&newPart.TableCommon, t.meta, def.ID, t.Columns, t.allocs, t.Constraints) + err := initTableCommonWithIndicesAndCollate( + &newPart.TableCommon, t.encoder.UseNewCollate(), t.meta, def.ID, t.Columns, t.allocs, t.Constraints) if err != nil { return nil, err } @@ -1676,7 +1679,11 @@ func GetReorganizedPartitionedTable(t table.Table) (table.PartitionedTable, erro return nil, err } var tc TableCommon - tc.initTableCommon(tblInfo, tblInfo.ID, t.Cols(), t.Allocators(nil), constraints) + useNewCollate := collate.NewCollationEnabled() + if pt, ok := t.(*partitionedTable); ok { + useNewCollate = pt.encoder.UseNewCollate() + } + tc.initTableCommonWithCollate(useNewCollate, tblInfo, tblInfo.ID, t.Cols(), t.Allocators(nil), constraints) // and rebuild the partitioning structure return newPartitionedTable(&tc, tblInfo) diff --git a/pkg/table/tables/tables.go b/pkg/table/tables/tables.go index da525127a579a..bf815b209d29d 100644 --- a/pkg/table/tables/tables.go +++ b/pkg/table/tables/tables.go @@ -166,6 +166,12 @@ func MockTableFromMeta(tblInfo *model.TableInfo) table.Table { // TableFromMeta creates a Table instance from model.TableInfo. func TableFromMeta(allocs autoid.Allocators, tblInfo *model.TableInfo) (table.Table, error) { + return TableFromMetaWithCollate(collate.NewCollationEnabled(), allocs, tblInfo) +} + +// TableFromMetaWithCollate creates a Table instance from model.TableInfo with a +// fixed collation mode. +func TableFromMetaWithCollate(useNewCollate bool, allocs autoid.Allocators, tblInfo *model.TableInfo) (table.Table, error) { if tblInfo.State == model.StateNone { return nil, table.ErrTableStateCantNone.GenWithStackByArgs(tblInfo.Name) } @@ -216,7 +222,7 @@ func TableFromMeta(allocs autoid.Allocators, tblInfo *model.TableInfo) (table.Ta return nil, err } var t TableCommon - t.initTableCommon(tblInfo, tblInfo.ID, columns, allocs, constraints) + t.initTableCommonWithCollate(useNewCollate, tblInfo, tblInfo.ID, columns, allocs, constraints) if tblInfo.GetPartitionInfo() == nil { if err := initTableIndices(&t); err != nil { return nil, err @@ -243,6 +249,10 @@ func buildGeneratedExpr(tblInfo *model.TableInfo, genExpr string) (ast.ExprNode, // initTableCommon initializes a TableCommon struct. func (t *TableCommon) initTableCommon(tblInfo *model.TableInfo, physicalTableID int64, cols []*table.Column, allocs autoid.Allocators, constraints []*table.Constraint) { + t.initTableCommonWithCollate(collate.NewCollationEnabled(), tblInfo, physicalTableID, cols, allocs, constraints) +} + +func (t *TableCommon) initTableCommonWithCollate(useNewCollate bool, tblInfo *model.TableInfo, physicalTableID int64, cols []*table.Column, allocs autoid.Allocators, constraints []*table.Constraint) { t.tableID = tblInfo.ID t.physicalTableID = physicalTableID t.allocs = allocs @@ -251,7 +261,7 @@ func (t *TableCommon) initTableCommon(tblInfo *model.TableInfo, physicalTableID t.Constraints = constraints t.recordPrefix = tablecodec.GenTableRecordPrefix(physicalTableID) t.indexPrefix = tablecodec.GenTableIndexPrefix(physicalTableID) - t.encoder = codec.NewEncoder(collate.NewCollationEnabled()) + t.encoder = codec.NewEncoder(useNewCollate) if tblInfo.IsSequence() { t.sequence = &sequenceCommon{meta: tblInfo.Sequence} } @@ -315,6 +325,11 @@ func initTableCommonWithIndices(t *TableCommon, tblInfo *model.TableInfo, physic return initTableIndices(t) } +func initTableCommonWithIndicesAndCollate(t *TableCommon, useNewCollate bool, tblInfo *model.TableInfo, physicalTableID int64, cols []*table.Column, allocs autoid.Allocators, constraints []*table.Constraint) error { + t.initTableCommonWithCollate(useNewCollate, tblInfo, physicalTableID, cols, allocs, constraints) + return initTableIndices(t) +} + // Indices implements table.Table Indices interface. func (t *TableCommon) Indices() []table.Index { return t.indices From 556f001936ce866e97987663da16b1aaf94c4a3f Mon Sep 17 00:00:00 2001 From: Ruihao Chen Date: Thu, 2 Jul 2026 18:07:28 +0800 Subject: [PATCH 02/35] table: simplify collation-aware table init helpers --- pkg/table/tables/partition.go | 6 +++--- pkg/table/tables/tables.go | 20 +++++--------------- 2 files changed, 8 insertions(+), 18 deletions(-) diff --git a/pkg/table/tables/partition.go b/pkg/table/tables/partition.go index 0a6f31856d63c..f8caac477c493 100644 --- a/pkg/table/tables/partition.go +++ b/pkg/table/tables/partition.go @@ -182,7 +182,7 @@ func newPartitionedTable(tbl *TableCommon, tblInfo *model.TableInfo) (table.Part } else { tblInfo.Indices = origIndices } - err := initTableCommonWithIndicesAndCollate( + err := initTableCommonWithIndices( &t.TableCommon, tbl.encoder.UseNewCollate(), tblInfo, p.ID, tbl.Columns, tbl.allocs, tbl.Constraints) if err != nil { return nil, errors.Trace(err) @@ -275,7 +275,7 @@ func newPartitionedTable(tbl *TableCommon, tblInfo *model.TableInfo) (table.Part func initPartition(t *partitionedTable, def model.PartitionDefinition) (*partition, error) { var newPart partition - err := initTableCommonWithIndicesAndCollate( + err := initTableCommonWithIndices( &newPart.TableCommon, t.encoder.UseNewCollate(), t.meta, def.ID, t.Columns, t.allocs, t.Constraints) if err != nil { return nil, err @@ -1683,7 +1683,7 @@ func GetReorganizedPartitionedTable(t table.Table) (table.PartitionedTable, erro if pt, ok := t.(*partitionedTable); ok { useNewCollate = pt.encoder.UseNewCollate() } - tc.initTableCommonWithCollate(useNewCollate, tblInfo, tblInfo.ID, t.Cols(), t.Allocators(nil), constraints) + tc.initTableCommon(useNewCollate, tblInfo, tblInfo.ID, t.Cols(), t.Allocators(nil), constraints) // and rebuild the partitioning structure return newPartitionedTable(&tc, tblInfo) diff --git a/pkg/table/tables/tables.go b/pkg/table/tables/tables.go index bf815b209d29d..9a7c3e2c8adfb 100644 --- a/pkg/table/tables/tables.go +++ b/pkg/table/tables/tables.go @@ -142,7 +142,7 @@ func MockTableFromMeta(tblInfo *model.TableInfo) table.Table { return nil } var t TableCommon - t.initTableCommon(tblInfo, tblInfo.ID, columns, autoid.NewAllocators(false), constraints) + t.initTableCommon(collate.NewCollationEnabled(), tblInfo, tblInfo.ID, columns, autoid.NewAllocators(false), constraints) if tblInfo.TableCacheStatusType != model.TableCacheStatusDisable { ret, err := newCachedTable(&t) if err != nil { @@ -222,7 +222,7 @@ func TableFromMetaWithCollate(useNewCollate bool, allocs autoid.Allocators, tblI return nil, err } var t TableCommon - t.initTableCommonWithCollate(useNewCollate, tblInfo, tblInfo.ID, columns, allocs, constraints) + t.initTableCommon(useNewCollate, tblInfo, tblInfo.ID, columns, allocs, constraints) if tblInfo.GetPartitionInfo() == nil { if err := initTableIndices(&t); err != nil { return nil, err @@ -247,12 +247,7 @@ func buildGeneratedExpr(tblInfo *model.TableInfo, genExpr string) (ast.ExprNode, return expr, nil } -// initTableCommon initializes a TableCommon struct. -func (t *TableCommon) initTableCommon(tblInfo *model.TableInfo, physicalTableID int64, cols []*table.Column, allocs autoid.Allocators, constraints []*table.Constraint) { - t.initTableCommonWithCollate(collate.NewCollationEnabled(), tblInfo, physicalTableID, cols, allocs, constraints) -} - -func (t *TableCommon) initTableCommonWithCollate(useNewCollate bool, tblInfo *model.TableInfo, physicalTableID int64, cols []*table.Column, allocs autoid.Allocators, constraints []*table.Constraint) { +func (t *TableCommon) initTableCommon(useNewCollate bool, tblInfo *model.TableInfo, physicalTableID int64, cols []*table.Column, allocs autoid.Allocators, constraints []*table.Constraint) { t.tableID = tblInfo.ID t.physicalTableID = physicalTableID t.allocs = allocs @@ -320,13 +315,8 @@ func asIndex(idx table.Index) *index { return idx.(*index) } -func initTableCommonWithIndices(t *TableCommon, tblInfo *model.TableInfo, physicalTableID int64, cols []*table.Column, allocs autoid.Allocators, constraints []*table.Constraint) error { - t.initTableCommon(tblInfo, physicalTableID, cols, allocs, constraints) - return initTableIndices(t) -} - -func initTableCommonWithIndicesAndCollate(t *TableCommon, useNewCollate bool, tblInfo *model.TableInfo, physicalTableID int64, cols []*table.Column, allocs autoid.Allocators, constraints []*table.Constraint) error { - t.initTableCommonWithCollate(useNewCollate, tblInfo, physicalTableID, cols, allocs, constraints) +func initTableCommonWithIndices(t *TableCommon, useNewCollate bool, tblInfo *model.TableInfo, physicalTableID int64, cols []*table.Column, allocs autoid.Allocators, constraints []*table.Constraint) error { + t.initTableCommon(useNewCollate, tblInfo, physicalTableID, cols, allocs, constraints) return initTableIndices(t) } From ed3adf7fff8a0e0c84c8b53625f22bc976b3aa76 Mon Sep 17 00:00:00 2001 From: Ruihao Chen Date: Thu, 2 Jul 2026 19:38:48 +0800 Subject: [PATCH 03/35] ddl, importer: use table collation mode in background paths --- lightning/pkg/errormanager/errormanager.go | 6 ++-- pkg/ddl/backfilling.go | 4 +-- pkg/ddl/backfilling_operators.go | 10 ++----- pkg/ddl/backfilling_txn_executor.go | 4 +-- pkg/ddl/index.go | 4 +-- pkg/ddl/index_merge_tmp.go | 2 +- pkg/ddl/reorg.go | 6 ++-- pkg/dxf/importinto/conflictedkv/handler.go | 2 +- pkg/executor/admin.go | 9 ++++-- pkg/executor/batch_checker.go | 2 +- pkg/executor/importer/import.go | 5 ++-- pkg/executor/importer/sampler.go | 11 +++---- pkg/executor/importer/table_import.go | 9 +++--- pkg/infoschema/tables.go | 11 +++++++ pkg/lightning/backend/kv/base.go | 3 +- pkg/lightning/backend/kv/kv2sql.go | 2 +- pkg/table/table.go | 3 ++ pkg/table/tables/mutation_checker_test.go | 3 +- pkg/table/tables/tables.go | 24 ++++++++++----- pkg/tablecodec/tablecodec.go | 8 ++++- pkg/util/rowDecoder/decoder.go | 34 ++++++++++++---------- 21 files changed, 94 insertions(+), 68 deletions(-) diff --git a/lightning/pkg/errormanager/errormanager.go b/lightning/pkg/errormanager/errormanager.go index 3ef21f9fb0435..ef3ddf4e4cd6f 100644 --- a/lightning/pkg/errormanager/errormanager.go +++ b/lightning/pkg/errormanager/errormanager.go @@ -616,7 +616,7 @@ func (em *ErrorManager) ReplaceConflictKeys( return errors.Trace(err) } decodedData, _, err := tables.DecodeRawRowData(encoder.SessionCtx.GetExprCtx(), - tbl.Meta(), overwrittenHandle, tbl.Cols(), overwritten) + tbl, overwrittenHandle, tbl.Cols(), overwritten) if err != nil { return errors.Trace(err) } @@ -796,7 +796,7 @@ func (em *ErrorManager) ReplaceConflictKeys( return errors.Trace(err) } decodedData, _, err := tables.DecodeRawRowData(encoder.SessionCtx.GetExprCtx(), - tbl.Meta(), handle, tbl.Cols(), latestValue) + tbl, handle, tbl.Cols(), latestValue) if err != nil { return errors.Trace(err) } @@ -825,7 +825,7 @@ func (em *ErrorManager) ReplaceConflictKeys( return errors.Trace(err) } decodedData, _, err := tables.DecodeRawRowData(encoder.SessionCtx.GetExprCtx(), - tbl.Meta(), handle, tbl.Cols(), rawValue) + tbl, handle, tbl.Cols(), rawValue) if err != nil { return errors.Trace(err) } diff --git a/pkg/ddl/backfilling.go b/pkg/ddl/backfilling.go index 60d8d5d67fe20..8fd3f62f3b207 100644 --- a/pkg/ddl/backfilling.go +++ b/pkg/ddl/backfilling.go @@ -47,7 +47,6 @@ import ( "github.com/pingcap/tidb/pkg/table" "github.com/pingcap/tidb/pkg/tablecodec" "github.com/pingcap/tidb/pkg/util" - "github.com/pingcap/tidb/pkg/util/collate" contextutil "github.com/pingcap/tidb/pkg/util/context" "github.com/pingcap/tidb/pkg/util/dbterror" decoder "github.com/pingcap/tidb/pkg/util/rowDecoder" @@ -215,7 +214,6 @@ func newBackfillCtx(id int, rInfo *reorgInfo, schemaName string, tbl table.Table } batchCnt := rInfo.ReorgMeta.GetBatchSize() - useNewCollate := rInfo.ReorgMeta.GetUseNewCollateOrDefault(collate.NewCollationEnabled()) metricTableID := backfillMetricsTableID(rInfo, label) return &backfillCtx{ id: id, @@ -227,7 +225,7 @@ func newBackfillCtx(id int, rInfo *reorgInfo, schemaName string, tbl table.Table schemaName: schemaName, table: tbl, batchCnt: batchCnt, - useNewCollate: useNewCollate, + useNewCollate: tbl.UseNewCollate(), jobContext: jobCtx, metricCounter: getBackfillTotalByTableID( metricTableID, label, schemaName, tbl.Meta().Name.String(), colOrIdxName), diff --git a/pkg/ddl/backfilling_operators.go b/pkg/ddl/backfilling_operators.go index a1d9ac715a4b2..80d8ac9f285e8 100644 --- a/pkg/ddl/backfilling_operators.go +++ b/pkg/ddl/backfilling_operators.go @@ -49,7 +49,6 @@ import ( "github.com/pingcap/tidb/pkg/table/tables" "github.com/pingcap/tidb/pkg/tablecodec" "github.com/pingcap/tidb/pkg/util/chunk" - "github.com/pingcap/tidb/pkg/util/collate" contextutil "github.com/pingcap/tidb/pkg/util/context" "github.com/pingcap/tidb/pkg/util/dbterror" "github.com/pingcap/tidb/pkg/util/execdetails" @@ -110,10 +109,9 @@ func NewAddIndexIngestPipeline( concurrency int, collector execute.Collector, ) (*operator.AsyncPipeline, error) { - useNewCollate := reorgMeta.GetUseNewCollateOrDefault(collate.NewCollationEnabled()) indexes := make([]table.Index, 0, len(idxInfos)) for _, idxInfo := range idxInfos { - index, err := tables.NewIndexWithCollate(useNewCollate, tbl.GetPhysicalID(), tbl.Meta(), idxInfo) + index, err := tables.NewIndexWithCollate(tbl.UseNewCollate(), tbl.GetPhysicalID(), tbl.Meta(), idxInfo) if err != nil { return nil, err } @@ -169,10 +167,9 @@ func NewWriteIndexToExternalStoragePipeline( collector execute.Collector, tikvCodec tikv.Codec, ) (*operator.AsyncPipeline, error) { - useNewCollate := reorgMeta.GetUseNewCollateOrDefault(collate.NewCollationEnabled()) indexes := make([]table.Index, 0, len(idxInfos)) for _, idxInfo := range idxInfos { - index, err := tables.NewIndexWithCollate(useNewCollate, tbl.GetPhysicalID(), tbl.Meta(), idxInfo) + index, err := tables.NewIndexWithCollate(tbl.UseNewCollate(), tbl.GetPhysicalID(), tbl.Meta(), idxInfo) if err != nil { return nil, err } @@ -917,9 +914,8 @@ func (w *indexIngestWorker) WriteChunk(rs *IndexRecordChunk) (count int, bytes i // skip running the checker in TiDB side. indexConditionCheckers = nil } - useNewCollate := w.reorgMeta.GetUseNewCollateOrDefault(collate.NewCollationEnabled()) cnt, kvBytes, err := writeChunk(w.ctx, w.writers, w.indexes, indexConditionCheckers, w.copCtx, - sc.TimeZone(), sc.ErrCtx(), vars.GetWriteStmtBufs(), rs.Chunk, w.tbl.Meta(), useNewCollate) + sc.TimeZone(), sc.ErrCtx(), vars.GetWriteStmtBufs(), rs.Chunk, w.tbl.Meta(), w.tbl.UseNewCollate()) if err != nil || cnt == 0 { return 0, 0, err } diff --git a/pkg/ddl/backfilling_txn_executor.go b/pkg/ddl/backfilling_txn_executor.go index 65322b100c30c..5cb14c875ec4c 100644 --- a/pkg/ddl/backfilling_txn_executor.go +++ b/pkg/ddl/backfilling_txn_executor.go @@ -33,7 +33,6 @@ import ( "github.com/pingcap/tidb/pkg/sessionctx/stmtctx" "github.com/pingcap/tidb/pkg/sessionctx/vardef" "github.com/pingcap/tidb/pkg/table" - "github.com/pingcap/tidb/pkg/util/collate" contextutil "github.com/pingcap/tidb/pkg/util/context" "github.com/pingcap/tidb/pkg/util/execdetails" "github.com/pingcap/tidb/pkg/util/intest" @@ -83,8 +82,7 @@ type txnBackfillExecutor struct { func newTxnBackfillExecutor(ctx context.Context, info *reorgInfo, sessPool *sess.Pool, tp backfillerType, tbl table.PhysicalTable, jobCtx *ReorgContext) (backfillExecutor, error) { - useNewCollate := info.ReorgMeta.GetUseNewCollateOrDefault(collate.NewCollationEnabled()) - decColMap, err := makeupDecodeColMap(info.dbInfo.Name, tbl, useNewCollate) + decColMap, err := makeupDecodeColMap(info.dbInfo.Name, tbl, tbl.UseNewCollate()) if err != nil { return nil, err } diff --git a/pkg/ddl/index.go b/pkg/ddl/index.go index 9beef955ec0cb..9e31dad256d99 100644 --- a/pkg/ddl/index.go +++ b/pkg/ddl/index.go @@ -2405,7 +2405,7 @@ func newAddIndexTxnWorker( continue } indexInfo := model.FindIndexInfoByID(t.Meta().Indices, elem.ID) - index, err := tables.NewIndexWithCollate(bfCtx.useNewCollate, t.GetPhysicalID(), t.Meta(), indexInfo) + index, err := tables.NewIndexWithCollate(t.UseNewCollate(), t.GetPhysicalID(), t.Meta(), indexInfo) if err != nil { return nil, err } @@ -2478,7 +2478,7 @@ func (w *baseIndexWorker) getIndexRecord(idxInfo *model.IndexInfo, handle kv.Han idxVal[j] = idxColumnVal } - rsData := tables.TryGetHandleRestoredDataWrapper(w.useNewCollate, w.table.Meta(), nil, w.rowMap, idxInfo) + rsData := tables.TryGetHandleRestoredDataWrapper(w.table, nil, w.rowMap, idxInfo) idxRecord := &indexRecord{handle: handle, key: recordKey, vals: idxVal, rsData: rsData} return idxRecord, nil } diff --git a/pkg/ddl/index_merge_tmp.go b/pkg/ddl/index_merge_tmp.go index 63bec58a51b9c..27864a5d14645 100644 --- a/pkg/ddl/index_merge_tmp.go +++ b/pkg/ddl/index_merge_tmp.go @@ -148,7 +148,7 @@ func newMergeTempIndexWorker(bfCtx *backfillCtx, t table.PhysicalTable, elements allIndexes := make([]table.Index, 0, len(elements)) for _, elem := range elements { indexInfo := model.FindIndexInfoByID(t.Meta().Indices, elem.ID) - index, err := tables.NewIndexWithCollate(bfCtx.useNewCollate, t.GetPhysicalID(), t.Meta(), indexInfo) + index, err := tables.NewIndexWithCollate(t.UseNewCollate(), t.GetPhysicalID(), t.Meta(), indexInfo) if err != nil { return nil, err } diff --git a/pkg/ddl/reorg.go b/pkg/ddl/reorg.go index a190a92ec4183..2f7dad05b6526 100644 --- a/pkg/ddl/reorg.go +++ b/pkg/ddl/reorg.go @@ -802,7 +802,7 @@ func GetTableMaxHandle(ctx *ReorgContext, store kv.Storage, startTS uint64, tbl row := chk.GetRow(0) if tblInfo.IsCommonHandle { pkIdx := tables.FindPrimaryIndex(tblInfo) - maxHandle, err = buildCommonHandleFromChunkRow(time.UTC, tblInfo, pkIdx, handleCols, row) + maxHandle, err = buildCommonHandleFromChunkRow(tbl.UseNewCollate(), time.UTC, tblInfo, pkIdx, handleCols, row) return maxHandle, false, err } return kv.IntHandle(row.GetInt64(0)), false, nil @@ -847,7 +847,7 @@ func buildHandleCols(tbl table.PhysicalTable) []*model.ColumnInfo { return handleCols } -func buildCommonHandleFromChunkRow(loc *time.Location, tblInfo *model.TableInfo, idxInfo *model.IndexInfo, +func buildCommonHandleFromChunkRow(useNewCollate bool, loc *time.Location, tblInfo *model.TableInfo, idxInfo *model.IndexInfo, cols []*model.ColumnInfo, row chunk.Row) (kv.Handle, error) { fieldTypes := make([]*types.FieldType, 0, len(cols)) for _, col := range cols { @@ -857,7 +857,7 @@ func buildCommonHandleFromChunkRow(loc *time.Location, tblInfo *model.TableInfo, tablecodec.TruncateIndexValues(tblInfo, idxInfo, datumRow) var handleBytes []byte - handleBytes, err := codec.EncodeKey(loc, nil, datumRow...) + handleBytes, err := codec.NewEncoder(useNewCollate).EncodeKey(loc, nil, datumRow...) if err != nil { return nil, err } diff --git a/pkg/dxf/importinto/conflictedkv/handler.go b/pkg/dxf/importinto/conflictedkv/handler.go index 15bef0c297013..9c06d9ebb95af 100644 --- a/pkg/dxf/importinto/conflictedkv/handler.go +++ b/pkg/dxf/importinto/conflictedkv/handler.go @@ -147,7 +147,7 @@ func (h *BaseHandler) encodeAndHandleRow(ctx context.Context, handle tidbkv.Handle, val []byte) (err error) { tblMeta := h.targetTable.Meta() decodedData, _, err := tables.DecodeRawRowData(h.encoder.SessionCtx.GetExprCtx(), - tblMeta, handle, h.targetTable.Cols(), val) + h.targetTable, handle, h.targetTable.Cols(), val) if err != nil { return errors.Trace(err) } diff --git a/pkg/executor/admin.go b/pkg/executor/admin.go index d2f73236304df..950c5cf40b052 100644 --- a/pkg/executor/admin.go +++ b/pkg/executor/admin.go @@ -35,7 +35,6 @@ import ( "github.com/pingcap/tidb/pkg/types" "github.com/pingcap/tidb/pkg/util" "github.com/pingcap/tidb/pkg/util/chunk" - "github.com/pingcap/tidb/pkg/util/collate" "github.com/pingcap/tidb/pkg/util/logutil" "github.com/pingcap/tidb/pkg/util/ranger" "github.com/pingcap/tidb/pkg/util/timeutil" @@ -364,7 +363,6 @@ func (e *RecoverIndexExec) fetchRecoverRows(ctx context.Context, srcResult dists idxValLen := len(e.index.Meta().Columns) result.scanRowCount = 0 - useNewCollate := collate.NewCollationEnabled() for { err := srcResult.Next(ctx, e.srcChunk) if err != nil { @@ -406,7 +404,12 @@ func (e *RecoverIndexExec) fetchRecoverRows(ctx context.Context, srcResult dists return nil, err } e.idxValsBufs[result.scanRowCount] = idxVals - rsData := tables.TryGetHandleRestoredDataWrapper(useNewCollate, e.table.Meta(), plannercore.GetCommonHandleDatum(e.handleCols, row), nil, e.index.Meta()) + 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}) result.scanRowCount++ result.currentHandle = handle diff --git a/pkg/executor/batch_checker.go b/pkg/executor/batch_checker.go index a2d8a8a40b4fa..81209bd6dc1a8 100644 --- a/pkg/executor/batch_checker.go +++ b/pkg/executor/batch_checker.go @@ -300,7 +300,7 @@ func getOldRow(ctx context.Context, sctx sessionctx.Context, txn kv.Transaction, } cols := t.WritableCols() - oldRow, oldRowMap, err := tables.DecodeRawRowData(sctx.GetExprCtx(), t.Meta(), handle, cols, oldValue) + oldRow, oldRowMap, err := tables.DecodeRawRowData(sctx.GetExprCtx(), t, handle, cols, oldValue) if err != nil { return nil, err } diff --git a/pkg/executor/importer/import.go b/pkg/executor/importer/import.go index 054a8b5b24ba7..8710d143bed91 100644 --- a/pkg/executor/importer/import.go +++ b/pkg/executor/importer/import.go @@ -65,7 +65,6 @@ import ( "github.com/pingcap/tidb/pkg/table" tidbutil "github.com/pingcap/tidb/pkg/util" "github.com/pingcap/tidb/pkg/util/chunk" - "github.com/pingcap/tidb/pkg/util/collate" contextutil "github.com/pingcap/tidb/pkg/util/context" "github.com/pingcap/tidb/pkg/util/cpu" "github.com/pingcap/tidb/pkg/util/dbterror" @@ -572,7 +571,7 @@ func NewImportPlan(ctx context.Context, userSctx sessionctx.Context, plan *plann User: userSctx.GetSessionVars().User.String(), Keyspace: userSctx.GetStore().GetKeyspace(), } - p.SetUseNewCollate(collate.NewCollationEnabled()) + p.SetUseNewCollate(tbl.UseNewCollate()) if err := p.initOptions(ctx, userSctx, plan.Options); err != nil { return nil, err } @@ -1973,7 +1972,7 @@ func (e *LoadDataController) CreateColAssignSimpleExprs(ctx expression.BuildCont e.ColumnAssignments, ctx, &e.colAssignMu, - e.GetUseNewCollateOrDefault(collate.NewCollationEnabled()), + e.Table.UseNewCollate(), ) } diff --git a/pkg/executor/importer/sampler.go b/pkg/executor/importer/sampler.go index cda668d7f12dd..41c3f6cdc755e 100644 --- a/pkg/executor/importer/sampler.go +++ b/pkg/executor/importer/sampler.go @@ -39,7 +39,6 @@ import ( plannercore "github.com/pingcap/tidb/pkg/planner/core" "github.com/pingcap/tidb/pkg/table" "github.com/pingcap/tidb/pkg/table/tables" - "github.com/pingcap/tidb/pkg/util/collate" contextutil "github.com/pingcap/tidb/pkg/util/context" "github.com/pingcap/tidb/pkg/util/dbterror/exeerrors" "go.uber.org/zap" @@ -156,6 +155,7 @@ func (e *LoadDataController) sampleKVSize( } func (e *LoadDataController) buildKVSizeSampleConfig() *KVSizeSampleConfig { + useNewCollate := e.Table.UseNewCollate() return &KVSizeSampleConfig{ Format: e.Format, SQLMode: e.SQLMode, @@ -164,7 +164,7 @@ func (e *LoadDataController) buildKVSizeSampleConfig() *KVSizeSampleConfig { FieldNullDef: append([]string(nil), e.FieldNullDef...), LineFieldsInfo: e.LineFieldsInfo, IgnoreLines: e.IgnoreLines, - UseNewCollate: e.UseNewCollate, + UseNewCollate: &useNewCollate, ColumnsAndUserVars: e.ColumnsAndUserVars, ColumnAssignments: e.ColumnAssignments, } @@ -223,7 +223,7 @@ func (s *kvSizeSampler) CreateColAssignSimpleExprs( s.cfg.ColumnAssignments, ctx, &s.colAssignMu, - s.cfg.getUseNewCollateOrDefault(collate.NewCollationEnabled()), + s.table.UseNewCollate(), ) } @@ -289,6 +289,7 @@ func (s *kvSizeSampler) getKVEncoder( chunk *Chunk, encTable table.Table, ) (*TableKVEncoder, error) { + useNewCollate := encTable.UseNewCollate() cfg := &encode.EncodingConfig{ SessionOptions: encode.SessionOptions{ SQLMode: s.cfg.SQLMode, @@ -299,7 +300,7 @@ func (s *kvSizeSampler) getKVEncoder( Path: chunk.Path, Table: encTable, Logger: log.Logger{Logger: logger.With(zap.String("path", chunk.Path))}, - UseNewCollate: s.cfg.UseNewCollate, + UseNewCollate: &useNewCollate, } return newTableKVEncoderInner(cfg, s, s.fieldMappings, s.insertColumns) } @@ -347,7 +348,7 @@ func (s *kvSizeSampler) sampleOneFile( ParquetMeta: file.ParquetMeta, } idAlloc := kv.NewPanickingAllocators(s.table.Meta().SepAutoInc()) - tbl, err := tables.TableFromMetaWithCollate(s.cfg.getUseNewCollateOrDefault(collate.NewCollationEnabled()), idAlloc, s.table.Meta()) + tbl, err := tables.TableFromMetaWithCollate(s.table.UseNewCollate(), idAlloc, s.table.Meta()) if err != nil { return 0, 0, 0, errors.Annotatef(err, "failed to tables.TableFromMeta %s", s.table.Meta().Name) } diff --git a/pkg/executor/importer/table_import.go b/pkg/executor/importer/table_import.go index c52b435f935c3..007811b259bb7 100644 --- a/pkg/executor/importer/table_import.go +++ b/pkg/executor/importer/table_import.go @@ -56,7 +56,6 @@ import ( "github.com/pingcap/tidb/pkg/table" "github.com/pingcap/tidb/pkg/table/tables" tidbutil "github.com/pingcap/tidb/pkg/util" - "github.com/pingcap/tidb/pkg/util/collate" "github.com/pingcap/tidb/pkg/util/etcd" "github.com/pingcap/tidb/pkg/util/logutil" "github.com/pingcap/tidb/pkg/util/promutil" @@ -116,7 +115,7 @@ func getTableImporterOptions(options []TableImporterOption) tableImporterOptions func newEncodingTable(e *LoadDataController) (table.Table, error) { idAlloc := kv.NewPanickingAllocators(e.Table.Meta().SepAutoInc()) - tbl, err := tables.TableFromMetaWithCollate(e.GetUseNewCollateOrDefault(collate.NewCollationEnabled()), idAlloc, e.Table.Meta()) + tbl, err := tables.TableFromMetaWithCollate(e.Table.UseNewCollate(), idAlloc, e.Table.Meta()) if err != nil { return nil, errors.Annotatef(err, "failed to tables.TableFromMeta %s", e.Table.Meta().Name) } @@ -432,6 +431,7 @@ func (ti *TableImporter) getKVEncoder(chunk *Chunk) (*TableKVEncoder, error) { } func (e *LoadDataController) getKVEncoder(logger *zap.Logger, chunk *Chunk, encTable table.Table) (*TableKVEncoder, error) { + useNewCollate := encTable.UseNewCollate() cfg := &encode.EncodingConfig{ SessionOptions: encode.SessionOptions{ SQLMode: e.SQLMode, @@ -442,13 +442,14 @@ func (e *LoadDataController) getKVEncoder(logger *zap.Logger, chunk *Chunk, encT Path: chunk.Path, Table: encTable, Logger: log.Logger{Logger: logger.With(zap.String("path", chunk.Path))}, - UseNewCollate: e.UseNewCollate, + UseNewCollate: &useNewCollate, } return NewTableKVEncoder(cfg, e) } // GetKVEncoderForDupResolve get the KV encoder for duplicate resolution. func (ti *TableImporter) GetKVEncoderForDupResolve() (*TableKVEncoder, error) { + useNewCollate := ti.encTable.UseNewCollate() cfg := &encode.EncodingConfig{ SessionOptions: encode.SessionOptions{ SQLMode: ti.SQLMode, @@ -457,7 +458,7 @@ func (ti *TableImporter) GetKVEncoderForDupResolve() (*TableKVEncoder, error) { Table: ti.encTable, Logger: log.Logger{Logger: ti.logger}, UseIdentityAutoRowID: true, - UseNewCollate: ti.UseNewCollate, + UseNewCollate: &useNewCollate, } return NewTableKVEncoderForDupResolve(cfg, ti.LoadDataController) } diff --git a/pkg/infoschema/tables.go b/pkg/infoschema/tables.go index 3072d29e07c87..017f0f6178d0a 100644 --- a/pkg/infoschema/tables.go +++ b/pkg/infoschema/tables.go @@ -55,6 +55,7 @@ import ( "github.com/pingcap/tidb/pkg/table" "github.com/pingcap/tidb/pkg/types" "github.com/pingcap/tidb/pkg/util" + "github.com/pingcap/tidb/pkg/util/collate" "github.com/pingcap/tidb/pkg/util/deadlockhistory" "github.com/pingcap/tidb/pkg/util/execdetails" "github.com/pingcap/tidb/pkg/util/logutil" @@ -2647,6 +2648,11 @@ func (it *infoschemaTable) Meta() *model.TableInfo { return it.meta } +// UseNewCollate implements table.Table UseNewCollate interface. +func (it *infoschemaTable) UseNewCollate() bool { + return collate.NewCollationEnabled() +} + // GetPhysicalID implements table.Table GetPhysicalID interface. func (it *infoschemaTable) GetPhysicalID() int64 { return it.meta.ID @@ -2745,6 +2751,11 @@ func (vt *VirtualTable) Meta() *model.TableInfo { return nil } +// UseNewCollate implements table.Table UseNewCollate interface. +func (vt *VirtualTable) UseNewCollate() bool { + return collate.NewCollationEnabled() +} + // GetPhysicalID implements table.Table GetPhysicalID interface. func (vt *VirtualTable) GetPhysicalID() int64 { return 0 diff --git a/pkg/lightning/backend/kv/base.go b/pkg/lightning/backend/kv/base.go index 7de2a8354a2bf..5d45483413edc 100644 --- a/pkg/lightning/backend/kv/base.go +++ b/pkg/lightning/backend/kv/base.go @@ -32,7 +32,6 @@ import ( "github.com/pingcap/tidb/pkg/types" "github.com/pingcap/tidb/pkg/util/chunk" "github.com/pingcap/tidb/pkg/util/codec" - "github.com/pingcap/tidb/pkg/util/collate" "github.com/pingcap/tidb/pkg/util/redact" "go.uber.org/zap" "go.uber.org/zap/zapcore" @@ -168,7 +167,7 @@ func NewBaseKVEncoder(config *encode.EncodingConfig) (*BaseKVEncoder, error) { } // collect expressions for evaluating stored generated columns - useNewCollate := collate.NewCollationEnabled() + useNewCollate := config.Table.UseNewCollate() if config.UseNewCollate != nil { useNewCollate = *config.UseNewCollate } diff --git a/pkg/lightning/backend/kv/kv2sql.go b/pkg/lightning/backend/kv/kv2sql.go index 0a535166b1902..3b172e9fdae06 100644 --- a/pkg/lightning/backend/kv/kv2sql.go +++ b/pkg/lightning/backend/kv/kv2sql.go @@ -54,7 +54,7 @@ func (t *TableKVDecoder) DecodeHandleFromIndex(indexInfo *model.IndexInfo, key, // DecodeRawRowData decodes raw row data into a datum slice and a (columnID:columnValue) map. func (t *TableKVDecoder) DecodeRawRowData(h kv.Handle, value []byte) ([]types.Datum, map[int64]types.Datum, error) { - return tables.DecodeRawRowData(t.se.GetExprCtx(), t.tbl.Meta(), h, t.tbl.Cols(), value) + return tables.DecodeRawRowData(t.se.GetExprCtx(), t.tbl, h, t.tbl.Cols(), value) } // DecodeRawRowDataAsStr decodes raw row data into a string. diff --git a/pkg/table/table.go b/pkg/table/table.go index 575f2e18c7a24..130bb18c1d119 100644 --- a/pkg/table/table.go +++ b/pkg/table/table.go @@ -459,6 +459,9 @@ type Table interface { // Meta returns TableInfo. Meta() *model.TableInfo + // UseNewCollate returns the collation mode used for table key encoding. + UseNewCollate() bool + // Type returns the type of table Type() Type diff --git a/pkg/table/tables/mutation_checker_test.go b/pkg/table/tables/mutation_checker_test.go index e3b0e3076301f..7e83ba2fa251e 100644 --- a/pkg/table/tables/mutation_checker_test.go +++ b/pkg/table/tables/mutation_checker_test.go @@ -392,8 +392,7 @@ func buildIndexKeyValue(index table.Index, rowToInsert []types.Datum, tc types.C if err != nil { return nil, nil, err } - useNewCollate := table.encoder.UseNewCollate() - rsData := TryGetHandleRestoredDataWrapper(useNewCollate, table.meta, rowToInsert, nil, indexInfo) + rsData := TryGetHandleRestoredDataWrapper(table, rowToInsert, nil, indexInfo) value, err := index.GenIndexValue(errctx.StrictNoWarningContext, tc.Location(), distinct, false, indexedValues, handle, rsData, nil) if err != nil { return nil, nil, err diff --git a/pkg/table/tables/tables.go b/pkg/table/tables/tables.go index 9a7c3e2c8adfb..ed955c2af78fc 100644 --- a/pkg/table/tables/tables.go +++ b/pkg/table/tables/tables.go @@ -1039,7 +1039,7 @@ func (t *TableCommon) addIndices(sctx table.MutateContext, recordID kv.Handle, r } dupErr = kv.GenKeyExistsErr(colStrVals, fmt.Sprintf("%s.%s", v.TableMeta().Name.String(), v.Meta().Name.String())) } - rsData := TryGetHandleRestoredDataWrapper(t.encoder.UseNewCollate(), t.meta, r, nil, v.Meta()) + rsData := TryGetHandleRestoredDataWrapper(t, r, nil, v.Meta()) if dupHandle, err := asIndex(v).create(sctx, txn, indexVals, recordID, rsData, false, opt); err != nil { if kv.ErrKeyExists.Equal(err) { return dupHandle, dupErr @@ -1064,7 +1064,7 @@ func RowWithCols(t table.Table, ctx sessionctx.Context, h kv.Handle, cols []*tab if err != nil { return nil, err } - v, _, err := DecodeRawRowData(ctx.GetExprCtx(), t.Meta(), h, cols, value) + v, _, err := DecodeRawRowData(ctx.GetExprCtx(), t, h, cols, value) if err != nil { return nil, err } @@ -1084,8 +1084,10 @@ func containFullColInHandle(meta *model.TableInfo, col *table.Column) (containFu } // DecodeRawRowData decodes raw row data into a datum slice and a (columnID:columnValue) map. -func DecodeRawRowData(ctx expression.BuildContext, meta *model.TableInfo, h kv.Handle, cols []*table.Column, +func DecodeRawRowData(ctx expression.BuildContext, tbl table.Table, h kv.Handle, cols []*table.Column, value []byte) ([]types.Datum, map[int64]types.Datum, error) { + meta := tbl.Meta() + useNewCollate := tbl.UseNewCollate() v := make([]types.Datum, len(cols)) colTps := make(map[int64]*types.FieldType, len(cols)) prefixCols := make(map[int64]struct{}) @@ -1101,7 +1103,7 @@ func DecodeRawRowData(ctx expression.BuildContext, meta *model.TableInfo, h kv.H } continue } - if col.IsCommonHandleColumn(meta) && !types.NeedRestoredData(&col.FieldType) { + if col.IsCommonHandleColumn(meta) && !types.NeedRestoredDataWithCollate(&col.FieldType, useNewCollate) { if containFullCol, idxInHandle := containFullColInHandle(meta, col); containFullCol { dtBytes := h.EncodedCol(idxInHandle) _, dt, err := codec.DecodeOne(dtBytes) @@ -1128,7 +1130,8 @@ func DecodeRawRowData(ctx expression.BuildContext, meta *model.TableInfo, h kv.H if col == nil { continue } - if col.IsPKHandleColumn(meta) || (col.IsCommonHandleColumn(meta) && !types.NeedRestoredData(&col.FieldType)) { + if col.IsPKHandleColumn(meta) || + (col.IsCommonHandleColumn(meta) && !types.NeedRestoredDataWithCollate(&col.FieldType, useNewCollate)) { if _, isPrefix := prefixCols[col.ID]; !isPrefix { continue } @@ -1319,7 +1322,7 @@ func (t *TableCommon) removeRowIndices(ctx table.MutateContext, txn kv.Transacti } func (t *TableCommon) buildIndexForRow(ctx table.MutateContext, h kv.Handle, vals []types.Datum, newData []types.Datum, idx *index, txn kv.Transaction, untouched bool, opt *table.CreateIdxOpt) error { - rsData := TryGetHandleRestoredDataWrapper(t.encoder.UseNewCollate(), t.meta, newData, nil, idx.Meta()) + rsData := TryGetHandleRestoredDataWrapper(t, newData, nil, idx.Meta()) if _, err := idx.create(ctx, txn, vals, h, rsData, untouched, opt); err != nil { if kv.ErrKeyExists.Equal(err) { // Make error message consistent with MySQL. @@ -1525,6 +1528,11 @@ func (t *TableCommon) Type() table.Type { return table.NormalTable } +// UseNewCollate implements table.Table UseNewCollate interface. +func (t *TableCommon) UseNewCollate() bool { + return t.encoder.UseNewCollate() +} + func (t *TableCommon) canSkip(col *table.Column, value *types.Datum) bool { return CanSkip(t.encoder.UseNewCollate(), t.Meta(), col, value) } @@ -1782,8 +1790,10 @@ func (t *TableCommon) GetSequenceCommon() *sequenceCommon { // TryGetHandleRestoredDataWrapper tries to get the restored data for handle if // needed. The argument can be a slice or a map. -func TryGetHandleRestoredDataWrapper(useNewCollate bool, tblInfo *model.TableInfo, +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 { return nil } diff --git a/pkg/tablecodec/tablecodec.go b/pkg/tablecodec/tablecodec.go index 410c1dc1a6fdb..eed1f7c51635a 100644 --- a/pkg/tablecodec/tablecodec.go +++ b/pkg/tablecodec/tablecodec.go @@ -557,6 +557,12 @@ func DecodeRowToDatumMap(b []byte, cols map[int64]*types.FieldType, loc *time.Lo // DecodeHandleToDatumMap decodes a handle into datum map. func DecodeHandleToDatumMap(handle kv.Handle, handleColIDs []int64, + cols map[int64]*types.FieldType, loc *time.Location, row map[int64]types.Datum) (map[int64]types.Datum, error) { + return DecodeHandleToDatumMapWithCollate(collate.NewCollationEnabled(), handle, handleColIDs, cols, loc, row) +} + +// DecodeHandleToDatumMapWithCollate decodes a handle into datum map with a fixed collation mode. +func DecodeHandleToDatumMapWithCollate(useNewCollate bool, handle kv.Handle, handleColIDs []int64, cols map[int64]*types.FieldType, loc *time.Location, row map[int64]types.Datum) (map[int64]types.Datum, error) { if handle == nil || len(handleColIDs) == 0 { return row, nil @@ -569,7 +575,7 @@ func DecodeHandleToDatumMap(handle kv.Handle, handleColIDs []int64, if !ok { continue } - if types.NeedRestoredData(ft) { + if types.NeedRestoredDataWithCollate(ft, useNewCollate) { continue } d, err := decodeHandleToDatum(handle, ft, idx) diff --git a/pkg/util/rowDecoder/decoder.go b/pkg/util/rowDecoder/decoder.go index c81c9fea62045..6f024197fab6b 100644 --- a/pkg/util/rowDecoder/decoder.go +++ b/pkg/util/rowDecoder/decoder.go @@ -38,13 +38,14 @@ type Column struct { // RowDecoder decodes a byte slice into datums and eval the generated column value. type RowDecoder struct { - tbl table.Table - mutRow chunk.MutRow - colMap map[int64]Column - colTypes map[int64]*types.FieldType - defaultVals []types.Datum - cols []*table.Column - pkCols []int64 + tbl table.Table + mutRow chunk.MutRow + colMap map[int64]Column + colTypes map[int64]*types.FieldType + defaultVals []types.Datum + cols []*table.Column + pkCols []int64 + useNewCollate bool } // NewRowDecoder returns a new RowDecoder. @@ -70,13 +71,14 @@ func NewRowDecoder(tbl table.Table, cols []*table.Column, decodeColMap map[int64 pkCols = []int64{model.ExtraHandleID} } return &RowDecoder{ - tbl: tbl, - mutRow: chunk.MutRowFromTypes(tps), - colMap: decodeColMap, - colTypes: colFieldMap, - defaultVals: make([]types.Datum, len(cols)), - cols: cols, - pkCols: pkCols, + tbl: tbl, + mutRow: chunk.MutRowFromTypes(tps), + colMap: decodeColMap, + colTypes: colFieldMap, + defaultVals: make([]types.Datum, len(cols)), + cols: cols, + pkCols: pkCols, + useNewCollate: tbl.UseNewCollate(), } } @@ -91,7 +93,7 @@ func (rd *RowDecoder) DecodeAndEvalRowWithMap(ctx exprctx.BuildContext, handle k if err != nil { return nil, err } - row, err = tablecodec.DecodeHandleToDatumMap(handle, rd.pkCols, rd.colTypes, decodeLoc, row) + row, err = tablecodec.DecodeHandleToDatumMapWithCollate(rd.useNewCollate, handle, rd.pkCols, rd.colTypes, decodeLoc, row) if err != nil { return nil, err } @@ -148,7 +150,7 @@ func (rd *RowDecoder) DecodeTheExistedColumnMap(ctx exprctx.BuildContext, handle if err != nil { return nil, err } - row, err = tablecodec.DecodeHandleToDatumMap(handle, rd.pkCols, rd.colTypes, decodeLoc, row) + row, err = tablecodec.DecodeHandleToDatumMapWithCollate(rd.useNewCollate, handle, rd.pkCols, rd.colTypes, decodeLoc, row) if err != nil { return nil, err } From dec0b7bcbc401b0a7bf9c4b1f3577bf8126f30ed Mon Sep 17 00:00:00 2001 From: Ruihao Chen Date: Thu, 2 Jul 2026 20:34:47 +0800 Subject: [PATCH 04/35] ddl: avoid table collation mode in txn backfill --- pkg/ddl/backfilling.go | 10 ++++++---- pkg/ddl/backfilling_txn_executor.go | 2 +- pkg/ddl/index.go | 2 +- pkg/ddl/index_merge_tmp.go | 2 +- 4 files changed, 9 insertions(+), 7 deletions(-) diff --git a/pkg/ddl/backfilling.go b/pkg/ddl/backfilling.go index 8fd3f62f3b207..fe7bbf7464232 100644 --- a/pkg/ddl/backfilling.go +++ b/pkg/ddl/backfilling.go @@ -47,6 +47,7 @@ import ( "github.com/pingcap/tidb/pkg/table" "github.com/pingcap/tidb/pkg/tablecodec" "github.com/pingcap/tidb/pkg/util" + "github.com/pingcap/tidb/pkg/util/collate" contextutil "github.com/pingcap/tidb/pkg/util/context" "github.com/pingcap/tidb/pkg/util/dbterror" decoder "github.com/pingcap/tidb/pkg/util/rowDecoder" @@ -214,6 +215,7 @@ func newBackfillCtx(id int, rInfo *reorgInfo, schemaName string, tbl table.Table } batchCnt := rInfo.ReorgMeta.GetBatchSize() + useNewCollate := rInfo.ReorgMeta.GetUseNewCollateOrDefault(collate.NewCollationEnabled()) metricTableID := backfillMetricsTableID(rInfo, label) return &backfillCtx{ id: id, @@ -225,7 +227,7 @@ func newBackfillCtx(id int, rInfo *reorgInfo, schemaName string, tbl table.Table schemaName: schemaName, table: tbl, batchCnt: batchCnt, - useNewCollate: tbl.UseNewCollate(), + useNewCollate: useNewCollate, jobContext: jobCtx, metricCounter: getBackfillTotalByTableID( metricTableID, label, schemaName, tbl.Meta().Name.String(), colOrIdxName), @@ -721,13 +723,13 @@ func sendTasks( return nil } -func makeupDecodeColMap(dbName ast.CIStr, t table.Table, useNewCollate bool) (map[int64]decoder.Column, error) { +func makeupDecodeColMap(dbName ast.CIStr, t table.Table) (map[int64]decoder.Column, error) { writableColInfos := make([]*model.ColumnInfo, 0, len(t.WritableCols())) for _, col := range t.WritableCols() { writableColInfos = append(writableColInfos, col.ColumnInfo) } - exprCols, _, err := expression.ColumnInfos2ColumnsAndNamesWithCollate( - newReorgExprCtx(), dbName, t.Meta().Name, writableColInfos, t.Meta(), useNewCollate) + exprCols, _, err := expression.ColumnInfos2ColumnsAndNames( + newReorgExprCtx(), dbName, t.Meta().Name, writableColInfos, t.Meta()) if err != nil { return nil, err } diff --git a/pkg/ddl/backfilling_txn_executor.go b/pkg/ddl/backfilling_txn_executor.go index 5cb14c875ec4c..46c2121ce6acf 100644 --- a/pkg/ddl/backfilling_txn_executor.go +++ b/pkg/ddl/backfilling_txn_executor.go @@ -82,7 +82,7 @@ type txnBackfillExecutor struct { func newTxnBackfillExecutor(ctx context.Context, info *reorgInfo, sessPool *sess.Pool, tp backfillerType, tbl table.PhysicalTable, jobCtx *ReorgContext) (backfillExecutor, error) { - decColMap, err := makeupDecodeColMap(info.dbInfo.Name, tbl, tbl.UseNewCollate()) + decColMap, err := makeupDecodeColMap(info.dbInfo.Name, tbl) if err != nil { return nil, err } diff --git a/pkg/ddl/index.go b/pkg/ddl/index.go index 9e31dad256d99..5b520d79ce503 100644 --- a/pkg/ddl/index.go +++ b/pkg/ddl/index.go @@ -2405,7 +2405,7 @@ func newAddIndexTxnWorker( continue } indexInfo := model.FindIndexInfoByID(t.Meta().Indices, elem.ID) - index, err := tables.NewIndexWithCollate(t.UseNewCollate(), t.GetPhysicalID(), t.Meta(), indexInfo) + index, err := tables.NewIndexWithCollate(bfCtx.useNewCollate, t.GetPhysicalID(), t.Meta(), indexInfo) if err != nil { return nil, err } diff --git a/pkg/ddl/index_merge_tmp.go b/pkg/ddl/index_merge_tmp.go index 27864a5d14645..63bec58a51b9c 100644 --- a/pkg/ddl/index_merge_tmp.go +++ b/pkg/ddl/index_merge_tmp.go @@ -148,7 +148,7 @@ func newMergeTempIndexWorker(bfCtx *backfillCtx, t table.PhysicalTable, elements allIndexes := make([]table.Index, 0, len(elements)) for _, elem := range elements { indexInfo := model.FindIndexInfoByID(t.Meta().Indices, elem.ID) - index, err := tables.NewIndexWithCollate(t.UseNewCollate(), t.GetPhysicalID(), t.Meta(), indexInfo) + index, err := tables.NewIndexWithCollate(bfCtx.useNewCollate, t.GetPhysicalID(), t.Meta(), indexInfo) if err != nil { return nil, err } From f8a0520edd93f41aecacea27bb86af9e30f3d924 Mon Sep 17 00:00:00 2001 From: Ruihao Chen Date: Thu, 2 Jul 2026 23:43:52 +0800 Subject: [PATCH 05/35] ddl: use local collation mode in txn add index --- pkg/ddl/backfilling.go | 40 +++++++++++++++++--------------------- pkg/ddl/index.go | 2 +- pkg/ddl/index_merge_tmp.go | 2 +- 3 files changed, 20 insertions(+), 24 deletions(-) diff --git a/pkg/ddl/backfilling.go b/pkg/ddl/backfilling.go index fe7bbf7464232..0ec15ce683c56 100644 --- a/pkg/ddl/backfilling.go +++ b/pkg/ddl/backfilling.go @@ -47,7 +47,6 @@ import ( "github.com/pingcap/tidb/pkg/table" "github.com/pingcap/tidb/pkg/tablecodec" "github.com/pingcap/tidb/pkg/util" - "github.com/pingcap/tidb/pkg/util/collate" contextutil "github.com/pingcap/tidb/pkg/util/context" "github.com/pingcap/tidb/pkg/util/dbterror" decoder "github.com/pingcap/tidb/pkg/util/rowDecoder" @@ -160,15 +159,14 @@ type backfillTaskContext struct { type backfillCtx struct { id int *ddlCtx - warnings contextutil.WarnHandlerExt - loc *time.Location - exprCtx exprctx.BuildContext - tblCtx table.MutateContext - schemaName string - table table.Table - batchCnt int - useNewCollate bool - jobContext *ReorgContext + warnings contextutil.WarnHandlerExt + loc *time.Location + exprCtx exprctx.BuildContext + tblCtx table.MutateContext + schemaName string + table table.Table + batchCnt int + jobContext *ReorgContext metricCounter prometheus.Counter conflictCounter prometheus.Counter @@ -215,20 +213,18 @@ func newBackfillCtx(id int, rInfo *reorgInfo, schemaName string, tbl table.Table } batchCnt := rInfo.ReorgMeta.GetBatchSize() - useNewCollate := rInfo.ReorgMeta.GetUseNewCollateOrDefault(collate.NewCollationEnabled()) metricTableID := backfillMetricsTableID(rInfo, label) return &backfillCtx{ - id: id, - ddlCtx: rInfo.jobCtx.oldDDLCtx, - warnings: warnHandler, - exprCtx: exprCtx, - tblCtx: tblCtx, - loc: exprCtx.GetEvalCtx().Location(), - schemaName: schemaName, - table: tbl, - batchCnt: batchCnt, - useNewCollate: useNewCollate, - jobContext: jobCtx, + id: id, + ddlCtx: rInfo.jobCtx.oldDDLCtx, + warnings: warnHandler, + exprCtx: exprCtx, + tblCtx: tblCtx, + loc: exprCtx.GetEvalCtx().Location(), + schemaName: schemaName, + table: tbl, + batchCnt: batchCnt, + jobContext: jobCtx, metricCounter: getBackfillTotalByTableID( metricTableID, label, schemaName, tbl.Meta().Name.String(), colOrIdxName), conflictCounter: getBackfillTotalByTableID( diff --git a/pkg/ddl/index.go b/pkg/ddl/index.go index 5b520d79ce503..f675c705921dd 100644 --- a/pkg/ddl/index.go +++ b/pkg/ddl/index.go @@ -2405,7 +2405,7 @@ func newAddIndexTxnWorker( continue } indexInfo := model.FindIndexInfoByID(t.Meta().Indices, elem.ID) - index, err := tables.NewIndexWithCollate(bfCtx.useNewCollate, t.GetPhysicalID(), t.Meta(), indexInfo) + index, err := tables.NewIndex(t.GetPhysicalID(), t.Meta(), indexInfo) if err != nil { return nil, err } diff --git a/pkg/ddl/index_merge_tmp.go b/pkg/ddl/index_merge_tmp.go index 63bec58a51b9c..e35efb15a23a9 100644 --- a/pkg/ddl/index_merge_tmp.go +++ b/pkg/ddl/index_merge_tmp.go @@ -148,7 +148,7 @@ func newMergeTempIndexWorker(bfCtx *backfillCtx, t table.PhysicalTable, elements allIndexes := make([]table.Index, 0, len(elements)) for _, elem := range elements { indexInfo := model.FindIndexInfoByID(t.Meta().Indices, elem.ID) - index, err := tables.NewIndexWithCollate(bfCtx.useNewCollate, t.GetPhysicalID(), t.Meta(), indexInfo) + index, err := tables.NewIndex(t.GetPhysicalID(), t.Meta(), indexInfo) if err != nil { return nil, err } From b9ed895be1cc2071fa755a4374ba5e8316005e3e Mon Sep 17 00:00:00 2001 From: Ruihao Chen Date: Fri, 3 Jul 2026 00:00:03 +0800 Subject: [PATCH 06/35] ddl: keep collation snapshot for merge temp index --- pkg/ddl/backfilling.go | 24 ++++++++++++++---------- pkg/ddl/index_merge_tmp.go | 2 +- 2 files changed, 15 insertions(+), 11 deletions(-) diff --git a/pkg/ddl/backfilling.go b/pkg/ddl/backfilling.go index 0ec15ce683c56..ce1478f310ae3 100644 --- a/pkg/ddl/backfilling.go +++ b/pkg/ddl/backfilling.go @@ -47,6 +47,7 @@ import ( "github.com/pingcap/tidb/pkg/table" "github.com/pingcap/tidb/pkg/tablecodec" "github.com/pingcap/tidb/pkg/util" + "github.com/pingcap/tidb/pkg/util/collate" contextutil "github.com/pingcap/tidb/pkg/util/context" "github.com/pingcap/tidb/pkg/util/dbterror" decoder "github.com/pingcap/tidb/pkg/util/rowDecoder" @@ -168,6 +169,7 @@ type backfillCtx struct { batchCnt int jobContext *ReorgContext + useNewCollate bool metricCounter prometheus.Counter conflictCounter prometheus.Counter } @@ -213,18 +215,20 @@ func newBackfillCtx(id int, rInfo *reorgInfo, schemaName string, tbl table.Table } batchCnt := rInfo.ReorgMeta.GetBatchSize() + useNewCollate := rInfo.ReorgMeta.GetUseNewCollateOrDefault(collate.NewCollationEnabled()) metricTableID := backfillMetricsTableID(rInfo, label) return &backfillCtx{ - id: id, - ddlCtx: rInfo.jobCtx.oldDDLCtx, - warnings: warnHandler, - exprCtx: exprCtx, - tblCtx: tblCtx, - loc: exprCtx.GetEvalCtx().Location(), - schemaName: schemaName, - table: tbl, - batchCnt: batchCnt, - jobContext: jobCtx, + id: id, + ddlCtx: rInfo.jobCtx.oldDDLCtx, + warnings: warnHandler, + exprCtx: exprCtx, + tblCtx: tblCtx, + loc: exprCtx.GetEvalCtx().Location(), + schemaName: schemaName, + table: tbl, + batchCnt: batchCnt, + jobContext: jobCtx, + useNewCollate: useNewCollate, metricCounter: getBackfillTotalByTableID( metricTableID, label, schemaName, tbl.Meta().Name.String(), colOrIdxName), conflictCounter: getBackfillTotalByTableID( diff --git a/pkg/ddl/index_merge_tmp.go b/pkg/ddl/index_merge_tmp.go index e35efb15a23a9..63bec58a51b9c 100644 --- a/pkg/ddl/index_merge_tmp.go +++ b/pkg/ddl/index_merge_tmp.go @@ -148,7 +148,7 @@ func newMergeTempIndexWorker(bfCtx *backfillCtx, t table.PhysicalTable, elements allIndexes := make([]table.Index, 0, len(elements)) for _, elem := range elements { indexInfo := model.FindIndexInfoByID(t.Meta().Indices, elem.ID) - index, err := tables.NewIndex(t.GetPhysicalID(), t.Meta(), indexInfo) + index, err := tables.NewIndexWithCollate(bfCtx.useNewCollate, t.GetPhysicalID(), t.Meta(), indexInfo) if err != nil { return nil, err } From 0fb3ffd7c9ca131330166e064b16fa838e80c49b Mon Sep 17 00:00:00 2001 From: Ruihao Chen Date: Fri, 3 Jul 2026 00:06:47 +0800 Subject: [PATCH 07/35] ddl: keep txn merge temp index on local collation --- pkg/ddl/backfilling.go | 24 ++++++++++-------------- pkg/ddl/index_merge_tmp.go | 2 +- 2 files changed, 11 insertions(+), 15 deletions(-) diff --git a/pkg/ddl/backfilling.go b/pkg/ddl/backfilling.go index ce1478f310ae3..0ec15ce683c56 100644 --- a/pkg/ddl/backfilling.go +++ b/pkg/ddl/backfilling.go @@ -47,7 +47,6 @@ import ( "github.com/pingcap/tidb/pkg/table" "github.com/pingcap/tidb/pkg/tablecodec" "github.com/pingcap/tidb/pkg/util" - "github.com/pingcap/tidb/pkg/util/collate" contextutil "github.com/pingcap/tidb/pkg/util/context" "github.com/pingcap/tidb/pkg/util/dbterror" decoder "github.com/pingcap/tidb/pkg/util/rowDecoder" @@ -169,7 +168,6 @@ type backfillCtx struct { batchCnt int jobContext *ReorgContext - useNewCollate bool metricCounter prometheus.Counter conflictCounter prometheus.Counter } @@ -215,20 +213,18 @@ func newBackfillCtx(id int, rInfo *reorgInfo, schemaName string, tbl table.Table } batchCnt := rInfo.ReorgMeta.GetBatchSize() - useNewCollate := rInfo.ReorgMeta.GetUseNewCollateOrDefault(collate.NewCollationEnabled()) metricTableID := backfillMetricsTableID(rInfo, label) return &backfillCtx{ - id: id, - ddlCtx: rInfo.jobCtx.oldDDLCtx, - warnings: warnHandler, - exprCtx: exprCtx, - tblCtx: tblCtx, - loc: exprCtx.GetEvalCtx().Location(), - schemaName: schemaName, - table: tbl, - batchCnt: batchCnt, - jobContext: jobCtx, - useNewCollate: useNewCollate, + id: id, + ddlCtx: rInfo.jobCtx.oldDDLCtx, + warnings: warnHandler, + exprCtx: exprCtx, + tblCtx: tblCtx, + loc: exprCtx.GetEvalCtx().Location(), + schemaName: schemaName, + table: tbl, + batchCnt: batchCnt, + jobContext: jobCtx, metricCounter: getBackfillTotalByTableID( metricTableID, label, schemaName, tbl.Meta().Name.String(), colOrIdxName), conflictCounter: getBackfillTotalByTableID( diff --git a/pkg/ddl/index_merge_tmp.go b/pkg/ddl/index_merge_tmp.go index 63bec58a51b9c..e35efb15a23a9 100644 --- a/pkg/ddl/index_merge_tmp.go +++ b/pkg/ddl/index_merge_tmp.go @@ -148,7 +148,7 @@ func newMergeTempIndexWorker(bfCtx *backfillCtx, t table.PhysicalTable, elements allIndexes := make([]table.Index, 0, len(elements)) for _, elem := range elements { indexInfo := model.FindIndexInfoByID(t.Meta().Indices, elem.ID) - index, err := tables.NewIndexWithCollate(bfCtx.useNewCollate, t.GetPhysicalID(), t.Meta(), indexInfo) + index, err := tables.NewIndex(t.GetPhysicalID(), t.Meta(), indexInfo) if err != nil { return nil, err } From 02ed166208413bfc48f0be8af218958e8ae50442 Mon Sep 17 00:00:00 2001 From: Ruihao Chen Date: Fri, 3 Jul 2026 00:36:38 +0800 Subject: [PATCH 08/35] importer: derive encoding collation from table --- pkg/ddl/backfilling.go | 3 +- pkg/dxf/importinto/task_executor.go | 2 -- pkg/executor/importer/sampler.go | 18 ++-------- pkg/executor/importer/table_import.go | 47 +++----------------------- pkg/lightning/backend/encode/encode.go | 3 -- pkg/lightning/backend/kv/base.go | 6 +--- pkg/lightning/backend/kv/kv2sql.go | 3 +- pkg/lightning/backend/kv/sql2kv.go | 9 ++--- 8 files changed, 14 insertions(+), 77 deletions(-) diff --git a/pkg/ddl/backfilling.go b/pkg/ddl/backfilling.go index 0ec15ce683c56..a2d6aa02b0fb4 100644 --- a/pkg/ddl/backfilling.go +++ b/pkg/ddl/backfilling.go @@ -724,8 +724,7 @@ func makeupDecodeColMap(dbName ast.CIStr, t table.Table) (map[int64]decoder.Colu for _, col := range t.WritableCols() { writableColInfos = append(writableColInfos, col.ColumnInfo) } - exprCols, _, err := expression.ColumnInfos2ColumnsAndNames( - newReorgExprCtx(), dbName, t.Meta().Name, writableColInfos, t.Meta()) + exprCols, _, err := expression.ColumnInfos2ColumnsAndNames(newReorgExprCtx(), dbName, t.Meta().Name, writableColInfos, t.Meta()) if err != nil { return nil, err } diff --git a/pkg/dxf/importinto/task_executor.go b/pkg/dxf/importinto/task_executor.go index 76ec9d8d47df6..b97fe6fcd4e85 100644 --- a/pkg/dxf/importinto/task_executor.go +++ b/pkg/dxf/importinto/task_executor.go @@ -133,7 +133,6 @@ func getTableImporter( controller, strconv.FormatInt(taskID, 10), store, - importer.WithEncodingTable(tbl), )) }) return importer.NewTableImporter( @@ -141,7 +140,6 @@ func getTableImporter( controller, strconv.FormatInt(taskID, 10), store, - importer.WithEncodingTable(tbl), ) } diff --git a/pkg/executor/importer/sampler.go b/pkg/executor/importer/sampler.go index 41c3f6cdc755e..073ac642d30a7 100644 --- a/pkg/executor/importer/sampler.go +++ b/pkg/executor/importer/sampler.go @@ -82,7 +82,6 @@ type KVSizeSampleConfig struct { FieldNullDef []string LineFieldsInfo plannercore.LineFieldsInfo IgnoreLines uint64 - UseNewCollate *bool ColumnsAndUserVars []*ast.ColumnNameOrUserVar ColumnAssignments []*ast.Assignment @@ -104,13 +103,6 @@ func (r *SampledKVSizeResult) TotalKVSize() int64 { return int64(r.DataKVSize + r.IndexKVSize) } -func (cfg *KVSizeSampleConfig) getUseNewCollateOrDefault(defaultVal bool) bool { - if cfg == nil || cfg.UseNewCollate == nil { - return defaultVal - } - return *cfg.UseNewCollate -} - // SampleFileImportKVSize samples source rows with nextgen's KV encoder and returns // the sampled source bytes and encoded KV sizes. func SampleFileImportKVSize( @@ -155,7 +147,6 @@ func (e *LoadDataController) sampleKVSize( } func (e *LoadDataController) buildKVSizeSampleConfig() *KVSizeSampleConfig { - useNewCollate := e.Table.UseNewCollate() return &KVSizeSampleConfig{ Format: e.Format, SQLMode: e.SQLMode, @@ -164,7 +155,6 @@ func (e *LoadDataController) buildKVSizeSampleConfig() *KVSizeSampleConfig { FieldNullDef: append([]string(nil), e.FieldNullDef...), LineFieldsInfo: e.LineFieldsInfo, IgnoreLines: e.IgnoreLines, - UseNewCollate: &useNewCollate, ColumnsAndUserVars: e.ColumnsAndUserVars, ColumnAssignments: e.ColumnAssignments, } @@ -289,7 +279,6 @@ func (s *kvSizeSampler) getKVEncoder( chunk *Chunk, encTable table.Table, ) (*TableKVEncoder, error) { - useNewCollate := encTable.UseNewCollate() cfg := &encode.EncodingConfig{ SessionOptions: encode.SessionOptions{ SQLMode: s.cfg.SQLMode, @@ -297,10 +286,9 @@ func (s *kvSizeSampler) getKVEncoder( SysVars: s.cfg.ImportantSysVars, AutoRandomSeed: chunk.PrevRowIDMax, }, - Path: chunk.Path, - Table: encTable, - Logger: log.Logger{Logger: logger.With(zap.String("path", chunk.Path))}, - UseNewCollate: &useNewCollate, + Path: chunk.Path, + Table: encTable, + Logger: log.Logger{Logger: logger.With(zap.String("path", chunk.Path))}, } return newTableKVEncoderInner(cfg, s, s.fieldMappings, s.insertColumns) } diff --git a/pkg/executor/importer/table_import.go b/pkg/executor/importer/table_import.go index 007811b259bb7..655075ac91efd 100644 --- a/pkg/executor/importer/table_import.go +++ b/pkg/executor/importer/table_import.go @@ -91,28 +91,6 @@ var ( defaultMaxEngineSize = int64(5 * config.DefaultBatchSize) ) -type tableImporterOptions struct { - encTable table.Table -} - -// TableImporterOption configures NewTableImporter. -type TableImporterOption func(*tableImporterOptions) - -// WithEncodingTable sets the table used to encode rows and track allocated IDs. -func WithEncodingTable(tbl table.Table) TableImporterOption { - return func(opts *tableImporterOptions) { - opts.encTable = tbl - } -} - -func getTableImporterOptions(options []TableImporterOption) tableImporterOptions { - opts := tableImporterOptions{} - for _, opt := range options { - opt(&opts) - } - return opts -} - func newEncodingTable(e *LoadDataController) (table.Table, error) { idAlloc := kv.NewPanickingAllocators(e.Table.Meta().SepAutoInc()) tbl, err := tables.TableFromMetaWithCollate(e.Table.UseNewCollate(), idAlloc, e.Table.Meta()) @@ -122,13 +100,6 @@ func newEncodingTable(e *LoadDataController) (table.Table, error) { return tbl, nil } -func getEncodingTable(e *LoadDataController, opts tableImporterOptions) (table.Table, error) { - if opts.encTable != nil { - return opts.encTable, nil - } - return newEncodingTable(e) -} - // Chunk records the chunk information. type Chunk struct { Path string @@ -226,10 +197,8 @@ func NewTableImporter( e *LoadDataController, id string, kvStore tidbkv.Storage, - options ...TableImporterOption, ) (ti *TableImporter, err error) { - opts := getTableImporterOptions(options) - tbl, err := getEncodingTable(e, opts) + tbl, err := newEncodingTable(e) if err != nil { return nil, err } @@ -328,11 +297,9 @@ func NewTableImporterForTest( e *LoadDataController, id string, kvStore tidbkv.Storage, - options ...TableImporterOption, ) (*TableImporter, error) { helper := &storeHelper{kvStore: kvStore} - opts := getTableImporterOptions(options) - tbl, err := getEncodingTable(e, opts) + tbl, err := newEncodingTable(e) if err != nil { return nil, err } @@ -431,7 +398,6 @@ func (ti *TableImporter) getKVEncoder(chunk *Chunk) (*TableKVEncoder, error) { } func (e *LoadDataController) getKVEncoder(logger *zap.Logger, chunk *Chunk, encTable table.Table) (*TableKVEncoder, error) { - useNewCollate := encTable.UseNewCollate() cfg := &encode.EncodingConfig{ SessionOptions: encode.SessionOptions{ SQLMode: e.SQLMode, @@ -439,17 +405,15 @@ func (e *LoadDataController) getKVEncoder(logger *zap.Logger, chunk *Chunk, encT SysVars: e.ImportantSysVars, AutoRandomSeed: chunk.PrevRowIDMax, }, - Path: chunk.Path, - Table: encTable, - Logger: log.Logger{Logger: logger.With(zap.String("path", chunk.Path))}, - UseNewCollate: &useNewCollate, + Path: chunk.Path, + Table: encTable, + Logger: log.Logger{Logger: logger.With(zap.String("path", chunk.Path))}, } return NewTableKVEncoder(cfg, e) } // GetKVEncoderForDupResolve get the KV encoder for duplicate resolution. func (ti *TableImporter) GetKVEncoderForDupResolve() (*TableKVEncoder, error) { - useNewCollate := ti.encTable.UseNewCollate() cfg := &encode.EncodingConfig{ SessionOptions: encode.SessionOptions{ SQLMode: ti.SQLMode, @@ -458,7 +422,6 @@ func (ti *TableImporter) GetKVEncoderForDupResolve() (*TableKVEncoder, error) { Table: ti.encTable, Logger: log.Logger{Logger: ti.logger}, UseIdentityAutoRowID: true, - UseNewCollate: &useNewCollate, } return NewTableKVEncoderForDupResolve(cfg, ti.LoadDataController) } diff --git a/pkg/lightning/backend/encode/encode.go b/pkg/lightning/backend/encode/encode.go index 3aed962304343..931067181122e 100644 --- a/pkg/lightning/backend/encode/encode.go +++ b/pkg/lightning/backend/encode/encode.go @@ -34,9 +34,6 @@ type EncodingConfig struct { // when encoding. // default false, in this case we will do sharding automatically if needed. UseIdentityAutoRowID bool - // UseNewCollate captures the collation mode used while encoding. Nil means - // old metadata and should fall back to the current process default. - UseNewCollate *bool } // EncodingBuilder consists of operations to handle encoding backend row data formats from source. diff --git a/pkg/lightning/backend/kv/base.go b/pkg/lightning/backend/kv/base.go index 5d45483413edc..b05c4dc784731 100644 --- a/pkg/lightning/backend/kv/base.go +++ b/pkg/lightning/backend/kv/base.go @@ -167,11 +167,7 @@ func NewBaseKVEncoder(config *encode.EncodingConfig) (*BaseKVEncoder, error) { } // collect expressions for evaluating stored generated columns - useNewCollate := config.Table.UseNewCollate() - if config.UseNewCollate != nil { - useNewCollate = *config.UseNewCollate - } - genCols, err := CollectGeneratedColumnsWithCollate(se, meta, cols, useNewCollate) + genCols, err := CollectGeneratedColumns(se, config.Table) if err != nil { return nil, errors.Annotate(err, "failed to parse generated column expressions") } diff --git a/pkg/lightning/backend/kv/kv2sql.go b/pkg/lightning/backend/kv/kv2sql.go index 3b172e9fdae06..fa589c8fb2829 100644 --- a/pkg/lightning/backend/kv/kv2sql.go +++ b/pkg/lightning/backend/kv/kv2sql.go @@ -134,8 +134,7 @@ func NewTableKVDecoder( return nil, err } - cols := tbl.Cols() - genCols, err := CollectGeneratedColumns(se, tbl.Meta(), cols) + genCols, err := CollectGeneratedColumns(se, tbl) if err != nil { return nil, err } diff --git a/pkg/lightning/backend/kv/sql2kv.go b/pkg/lightning/backend/kv/sql2kv.go index 6d65ce94792bf..07968f245aa3b 100644 --- a/pkg/lightning/backend/kv/sql2kv.go +++ b/pkg/lightning/backend/kv/sql2kv.go @@ -35,7 +35,6 @@ import ( "github.com/pingcap/tidb/pkg/table" "github.com/pingcap/tidb/pkg/tablecodec" "github.com/pingcap/tidb/pkg/types" - "github.com/pingcap/tidb/pkg/util/collate" ) type tableKVEncoder struct { @@ -69,13 +68,11 @@ func NewTableKVEncoder( // CollectGeneratedColumns collects all expressions required to evaluate the // results of all generated columns. The returning slice is in evaluation order. -func CollectGeneratedColumns(se *Session, meta *model.TableInfo, cols []*table.Column) ([]GeneratedCol, error) { - return CollectGeneratedColumnsWithCollate(se, meta, cols, collate.NewCollationEnabled()) +func CollectGeneratedColumns(se *Session, tbl table.Table) ([]GeneratedCol, error) { + return collectGeneratedColumns(se, tbl.Meta(), tbl.Cols(), tbl.UseNewCollate()) } -// CollectGeneratedColumnsWithCollate collects all expressions required to -// evaluate generated columns with a fixed collation mode. -func CollectGeneratedColumnsWithCollate(se *Session, meta *model.TableInfo, cols []*table.Column, useNewCollate bool) ([]GeneratedCol, error) { +func collectGeneratedColumns(se *Session, meta *model.TableInfo, cols []*table.Column, useNewCollate bool) ([]GeneratedCol, error) { hasGenCol := false for _, col := range cols { if col.GeneratedExpr != nil { From d891385e0dc89a0eb51c9e1598a4b262161f16a6 Mon Sep 17 00:00:00 2001 From: Ruihao Chen Date: Fri, 3 Jul 2026 03:23:24 +0800 Subject: [PATCH 09/35] tablecodec: keep handle decode on local collation --- pkg/tablecodec/tablecodec.go | 8 +------- pkg/util/rowDecoder/decoder.go | 34 ++++++++++++++++------------------ 2 files changed, 17 insertions(+), 25 deletions(-) diff --git a/pkg/tablecodec/tablecodec.go b/pkg/tablecodec/tablecodec.go index eed1f7c51635a..410c1dc1a6fdb 100644 --- a/pkg/tablecodec/tablecodec.go +++ b/pkg/tablecodec/tablecodec.go @@ -557,12 +557,6 @@ func DecodeRowToDatumMap(b []byte, cols map[int64]*types.FieldType, loc *time.Lo // DecodeHandleToDatumMap decodes a handle into datum map. func DecodeHandleToDatumMap(handle kv.Handle, handleColIDs []int64, - cols map[int64]*types.FieldType, loc *time.Location, row map[int64]types.Datum) (map[int64]types.Datum, error) { - return DecodeHandleToDatumMapWithCollate(collate.NewCollationEnabled(), handle, handleColIDs, cols, loc, row) -} - -// DecodeHandleToDatumMapWithCollate decodes a handle into datum map with a fixed collation mode. -func DecodeHandleToDatumMapWithCollate(useNewCollate bool, handle kv.Handle, handleColIDs []int64, cols map[int64]*types.FieldType, loc *time.Location, row map[int64]types.Datum) (map[int64]types.Datum, error) { if handle == nil || len(handleColIDs) == 0 { return row, nil @@ -575,7 +569,7 @@ func DecodeHandleToDatumMapWithCollate(useNewCollate bool, handle kv.Handle, han if !ok { continue } - if types.NeedRestoredDataWithCollate(ft, useNewCollate) { + if types.NeedRestoredData(ft) { continue } d, err := decodeHandleToDatum(handle, ft, idx) diff --git a/pkg/util/rowDecoder/decoder.go b/pkg/util/rowDecoder/decoder.go index 6f024197fab6b..c81c9fea62045 100644 --- a/pkg/util/rowDecoder/decoder.go +++ b/pkg/util/rowDecoder/decoder.go @@ -38,14 +38,13 @@ type Column struct { // RowDecoder decodes a byte slice into datums and eval the generated column value. type RowDecoder struct { - tbl table.Table - mutRow chunk.MutRow - colMap map[int64]Column - colTypes map[int64]*types.FieldType - defaultVals []types.Datum - cols []*table.Column - pkCols []int64 - useNewCollate bool + tbl table.Table + mutRow chunk.MutRow + colMap map[int64]Column + colTypes map[int64]*types.FieldType + defaultVals []types.Datum + cols []*table.Column + pkCols []int64 } // NewRowDecoder returns a new RowDecoder. @@ -71,14 +70,13 @@ func NewRowDecoder(tbl table.Table, cols []*table.Column, decodeColMap map[int64 pkCols = []int64{model.ExtraHandleID} } return &RowDecoder{ - tbl: tbl, - mutRow: chunk.MutRowFromTypes(tps), - colMap: decodeColMap, - colTypes: colFieldMap, - defaultVals: make([]types.Datum, len(cols)), - cols: cols, - pkCols: pkCols, - useNewCollate: tbl.UseNewCollate(), + tbl: tbl, + mutRow: chunk.MutRowFromTypes(tps), + colMap: decodeColMap, + colTypes: colFieldMap, + defaultVals: make([]types.Datum, len(cols)), + cols: cols, + pkCols: pkCols, } } @@ -93,7 +91,7 @@ func (rd *RowDecoder) DecodeAndEvalRowWithMap(ctx exprctx.BuildContext, handle k if err != nil { return nil, err } - row, err = tablecodec.DecodeHandleToDatumMapWithCollate(rd.useNewCollate, handle, rd.pkCols, rd.colTypes, decodeLoc, row) + row, err = tablecodec.DecodeHandleToDatumMap(handle, rd.pkCols, rd.colTypes, decodeLoc, row) if err != nil { return nil, err } @@ -150,7 +148,7 @@ func (rd *RowDecoder) DecodeTheExistedColumnMap(ctx exprctx.BuildContext, handle if err != nil { return nil, err } - row, err = tablecodec.DecodeHandleToDatumMapWithCollate(rd.useNewCollate, handle, rd.pkCols, rd.colTypes, decodeLoc, row) + row, err = tablecodec.DecodeHandleToDatumMap(handle, rd.pkCols, rd.colTypes, decodeLoc, row) if err != nil { return nil, err } From 390f262643926bc6f115fc23b817a65e2a1de726 Mon Sep 17 00:00:00 2001 From: Ruihao Chen Date: Fri, 3 Jul 2026 09:41:15 +0800 Subject: [PATCH 10/35] importer: capture import collation from keyspace --- pkg/executor/importer/import.go | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/pkg/executor/importer/import.go b/pkg/executor/importer/import.go index 8710d143bed91..686331aef015a 100644 --- a/pkg/executor/importer/import.go +++ b/pkg/executor/importer/import.go @@ -65,6 +65,7 @@ import ( "github.com/pingcap/tidb/pkg/table" tidbutil "github.com/pingcap/tidb/pkg/util" "github.com/pingcap/tidb/pkg/util/chunk" + "github.com/pingcap/tidb/pkg/util/collate" contextutil "github.com/pingcap/tidb/pkg/util/context" "github.com/pingcap/tidb/pkg/util/cpu" "github.com/pingcap/tidb/pkg/util/dbterror" @@ -571,7 +572,7 @@ func NewImportPlan(ctx context.Context, userSctx sessionctx.Context, plan *plann User: userSctx.GetSessionVars().User.String(), Keyspace: userSctx.GetStore().GetKeyspace(), } - p.SetUseNewCollate(tbl.UseNewCollate()) + p.SetUseNewCollate(collate.NewCollationEnabled()) if err := p.initOptions(ctx, userSctx, plan.Options); err != nil { return nil, err } From c50e3925ec032061ac8ff786081327a9c93fe2e0 Mon Sep 17 00:00:00 2001 From: Ruihao Chen Date: Fri, 3 Jul 2026 10:29:43 +0800 Subject: [PATCH 11/35] table: initialize indices from table common --- pkg/table/tables/partition.go | 5 ++--- pkg/table/tables/tables.go | 39 ++++++++++++++++++----------------- 2 files changed, 22 insertions(+), 22 deletions(-) diff --git a/pkg/table/tables/partition.go b/pkg/table/tables/partition.go index f8caac477c493..abeca65bd1d90 100644 --- a/pkg/table/tables/partition.go +++ b/pkg/table/tables/partition.go @@ -127,7 +127,7 @@ func newPartitionedTable(tbl *TableCommon, tblInfo *model.TableInfo) (table.Part return initEvalBuffer(ret) }, } - if err := initTableIndices(&ret.TableCommon); err != nil { + if err := ret.initTableIndices(); err != nil { return nil, errors.Trace(err) } origIndices := ret.meta.Indices @@ -1678,12 +1678,11 @@ func GetReorganizedPartitionedTable(t table.Table) (table.PartitionedTable, erro if err != nil { return nil, err } - var tc TableCommon useNewCollate := collate.NewCollationEnabled() if pt, ok := t.(*partitionedTable); ok { useNewCollate = pt.encoder.UseNewCollate() } - tc.initTableCommon(useNewCollate, tblInfo, tblInfo.ID, t.Cols(), t.Allocators(nil), constraints) + tc := initTableCommon(useNewCollate, tblInfo, tblInfo.ID, t.Cols(), t.Allocators(nil), constraints) // and rebuild the partitioning structure return newPartitionedTable(&tc, tblInfo) diff --git a/pkg/table/tables/tables.go b/pkg/table/tables/tables.go index ed955c2af78fc..f25522aeaae72 100644 --- a/pkg/table/tables/tables.go +++ b/pkg/table/tables/tables.go @@ -141,8 +141,7 @@ func MockTableFromMeta(tblInfo *model.TableInfo) table.Table { if err != nil { return nil } - var t TableCommon - t.initTableCommon(collate.NewCollationEnabled(), tblInfo, tblInfo.ID, columns, autoid.NewAllocators(false), constraints) + t := initTableCommon(collate.NewCollationEnabled(), tblInfo, tblInfo.ID, columns, autoid.NewAllocators(false), constraints) if tblInfo.TableCacheStatusType != model.TableCacheStatusDisable { ret, err := newCachedTable(&t) if err != nil { @@ -151,7 +150,7 @@ func MockTableFromMeta(tblInfo *model.TableInfo) table.Table { return ret } if tblInfo.GetPartitionInfo() == nil { - if err := initTableIndices(&t); err != nil { + if err := t.initTableIndices(); err != nil { return nil } return &t @@ -221,10 +220,9 @@ func TableFromMetaWithCollate(useNewCollate bool, allocs autoid.Allocators, tblI if err != nil { return nil, err } - var t TableCommon - t.initTableCommon(useNewCollate, tblInfo, tblInfo.ID, columns, allocs, constraints) + t := initTableCommon(useNewCollate, tblInfo, tblInfo.ID, columns, allocs, constraints) if tblInfo.GetPartitionInfo() == nil { - if err := initTableIndices(&t); err != nil { + if err := t.initTableIndices(); err != nil { return nil, err } if tblInfo.TableCacheStatusType != model.TableCacheStatusDisable { @@ -247,24 +245,27 @@ func buildGeneratedExpr(tblInfo *model.TableInfo, genExpr string) (ast.ExprNode, return expr, nil } -func (t *TableCommon) initTableCommon(useNewCollate bool, tblInfo *model.TableInfo, physicalTableID int64, cols []*table.Column, allocs autoid.Allocators, constraints []*table.Constraint) { - t.tableID = tblInfo.ID - t.physicalTableID = physicalTableID - t.allocs = allocs - t.meta = tblInfo - t.Columns = cols - t.Constraints = constraints - t.recordPrefix = tablecodec.GenTableRecordPrefix(physicalTableID) - t.indexPrefix = tablecodec.GenTableIndexPrefix(physicalTableID) - t.encoder = codec.NewEncoder(useNewCollate) +func initTableCommon(useNewCollate bool, tblInfo *model.TableInfo, physicalTableID int64, cols []*table.Column, allocs autoid.Allocators, constraints []*table.Constraint) TableCommon { + t := TableCommon{ + tableID: tblInfo.ID, + physicalTableID: physicalTableID, + allocs: allocs, + meta: tblInfo, + Columns: cols, + Constraints: constraints, + recordPrefix: tablecodec.GenTableRecordPrefix(physicalTableID), + indexPrefix: tablecodec.GenTableIndexPrefix(physicalTableID), + encoder: codec.NewEncoder(useNewCollate), + } if tblInfo.IsSequence() { t.sequence = &sequenceCommon{meta: tblInfo.Sequence} } t.ResetColumnsCache() + return t } // initTableIndices initializes the indices of the TableCommon. -func initTableIndices(t *TableCommon) error { +func (t *TableCommon) initTableIndices() error { tblInfo := t.meta for _, idxInfo := range tblInfo.Indices { if idxInfo.State == model.StateNone { @@ -316,8 +317,8 @@ func asIndex(idx table.Index) *index { } func initTableCommonWithIndices(t *TableCommon, useNewCollate bool, tblInfo *model.TableInfo, physicalTableID int64, cols []*table.Column, allocs autoid.Allocators, constraints []*table.Constraint) error { - t.initTableCommon(useNewCollate, tblInfo, physicalTableID, cols, allocs, constraints) - return initTableIndices(t) + *t = initTableCommon(useNewCollate, tblInfo, physicalTableID, cols, allocs, constraints) + return t.initTableIndices() } // Indices implements table.Table Indices interface. From 65d42164a764f1a945e6575bfe8efd917c08c05b Mon Sep 17 00:00:00 2001 From: Ruihao Chen Date: Fri, 3 Jul 2026 10:43:32 +0800 Subject: [PATCH 12/35] table: use table collation accessor in helpers --- pkg/executor/importer/import.go | 9 +++++---- pkg/lightning/backend/kv/sql2kv.go | 8 +++----- pkg/meta/model/reorg.go | 9 +++++---- pkg/table/tables/partition.go | 7 +------ 4 files changed, 14 insertions(+), 19 deletions(-) diff --git a/pkg/executor/importer/import.go b/pkg/executor/importer/import.go index 686331aef015a..67bfc2f8c12ff 100644 --- a/pkg/executor/importer/import.go +++ b/pkg/executor/importer/import.go @@ -336,9 +336,10 @@ type Plan struct { ManualRecovery bool // the keyspace name when submitting this job, only for import-into Keyspace string - // UseNewCollate captures the collation mode used by background import workers - // when encoding keys. Nil means old metadata and should fall back to the - // caller-provided default. + // UseNewCollate captures the submitting keyspace collation mode for key encoding. + // It is needed because an import task may execute in another keyspace with a + // different collation setting. Nil means old metadata and should fall back to + // the caller-provided default. UseNewCollate *bool `json:"use_new_collate,omitempty"` } @@ -363,7 +364,7 @@ func (p *Plan) GetUseNewCollateOrDefault(defaultVal bool) bool { return *p.UseNewCollate } -// SetUseNewCollate stores the collation mode used by background import workers. +// SetUseNewCollate stores the submitting keyspace collation mode for key encoding. func (p *Plan) SetUseNewCollate(useNewCollate bool) { p.UseNewCollate = &useNewCollate } diff --git a/pkg/lightning/backend/kv/sql2kv.go b/pkg/lightning/backend/kv/sql2kv.go index 07968f245aa3b..140c8eadccd30 100644 --- a/pkg/lightning/backend/kv/sql2kv.go +++ b/pkg/lightning/backend/kv/sql2kv.go @@ -69,10 +69,8 @@ func NewTableKVEncoder( // CollectGeneratedColumns collects all expressions required to evaluate the // results of all generated columns. The returning slice is in evaluation order. func CollectGeneratedColumns(se *Session, tbl table.Table) ([]GeneratedCol, error) { - return collectGeneratedColumns(se, tbl.Meta(), tbl.Cols(), tbl.UseNewCollate()) -} - -func collectGeneratedColumns(se *Session, meta *model.TableInfo, cols []*table.Column, useNewCollate bool) ([]GeneratedCol, error) { + meta := tbl.Meta() + cols := tbl.Cols() hasGenCol := false for _, col := range cols { if col.GeneratedExpr != nil { @@ -116,7 +114,7 @@ func collectGeneratedColumns(se *Session, meta *model.TableInfo, cols []*table.C col.GeneratedExpr.Internal(), expression.WithInputSchemaAndNames(schema, names, meta), expression.WithAllowCastArray(true), - expression.WithUseNewCollate(useNewCollate), + expression.WithUseNewCollate(tbl.UseNewCollate()), ) if err != nil { return nil, err diff --git a/pkg/meta/model/reorg.go b/pkg/meta/model/reorg.go index 03a84398bd6a8..34147d350616f 100644 --- a/pkg/meta/model/reorg.go +++ b/pkg/meta/model/reorg.go @@ -95,9 +95,10 @@ type DDLReorgMeta struct { MaxNodeCount int `json:"max_node_count"` AnalyzeState int8 `json:"analyze_state"` Stage ReorgStage `json:"stage"` - // UseNewCollate captures the collation mode used by background reorg workers - // when encoding keys. Nil means the metadata was generated before this field - // existed and should fall back to the caller-provided default. + // UseNewCollate captures the submitting keyspace collation mode for key encoding. + // It is needed because a reorg task may execute in another keyspace with a + // different collation setting. Nil means old metadata and should fall back to + // the caller-provided default. UseNewCollate *bool `json:"use_new_collate,omitempty"` // These two variables are used to control the concurrency and batch size of the reorganization process. // They can be adjusted dynamically through `admin alter ddl jobs` command. @@ -164,7 +165,7 @@ func (dm *DDLReorgMeta) GetUseNewCollateOrDefault(defaultVal bool) bool { return *dm.UseNewCollate } -// SetUseNewCollate stores the collation mode used by background reorg workers. +// SetUseNewCollate stores the submitting keyspace collation mode for key encoding. func (dm *DDLReorgMeta) SetUseNewCollate(useNewCollate bool) { dm.UseNewCollate = &useNewCollate } diff --git a/pkg/table/tables/partition.go b/pkg/table/tables/partition.go index abeca65bd1d90..d2b6dae2384bd 100644 --- a/pkg/table/tables/partition.go +++ b/pkg/table/tables/partition.go @@ -45,7 +45,6 @@ import ( "github.com/pingcap/tidb/pkg/util" "github.com/pingcap/tidb/pkg/util/chunk" "github.com/pingcap/tidb/pkg/util/codec" - "github.com/pingcap/tidb/pkg/util/collate" "github.com/pingcap/tidb/pkg/util/dbterror" "github.com/pingcap/tidb/pkg/util/hack" "github.com/pingcap/tidb/pkg/util/logutil" @@ -1678,11 +1677,7 @@ func GetReorganizedPartitionedTable(t table.Table) (table.PartitionedTable, erro if err != nil { return nil, err } - useNewCollate := collate.NewCollationEnabled() - if pt, ok := t.(*partitionedTable); ok { - useNewCollate = pt.encoder.UseNewCollate() - } - tc := initTableCommon(useNewCollate, tblInfo, tblInfo.ID, t.Cols(), t.Allocators(nil), constraints) + tc := initTableCommon(t.UseNewCollate(), tblInfo, tblInfo.ID, t.Cols(), t.Allocators(nil), constraints) // and rebuild the partitioning structure return newPartitionedTable(&tc, tblInfo) From d986cd8171174c02c3fe21cd4cc9f5092cd1e9b2 Mon Sep 17 00:00:00 2001 From: Ruihao Chen Date: Fri, 3 Jul 2026 11:01:47 +0800 Subject: [PATCH 13/35] table: clarify table common construction --- pkg/table/tables/partition.go | 12 +++++------- pkg/table/tables/tables.go | 11 +++-------- 2 files changed, 8 insertions(+), 15 deletions(-) diff --git a/pkg/table/tables/partition.go b/pkg/table/tables/partition.go index d2b6dae2384bd..0f2cd754e9a56 100644 --- a/pkg/table/tables/partition.go +++ b/pkg/table/tables/partition.go @@ -181,9 +181,8 @@ func newPartitionedTable(tbl *TableCommon, tblInfo *model.TableInfo) (table.Part } else { tblInfo.Indices = origIndices } - err := initTableCommonWithIndices( - &t.TableCommon, tbl.encoder.UseNewCollate(), tblInfo, p.ID, tbl.Columns, tbl.allocs, tbl.Constraints) - if err != nil { + t.TableCommon = newTableCommon(tblInfo, p.ID, tbl.Columns, tbl.allocs, tbl.Constraints, tbl.encoder.UseNewCollate()) + if err := t.initTableIndices(); err != nil { return nil, errors.Trace(err) } t.table = ret @@ -274,9 +273,8 @@ func newPartitionedTable(tbl *TableCommon, tblInfo *model.TableInfo) (table.Part func initPartition(t *partitionedTable, def model.PartitionDefinition) (*partition, error) { var newPart partition - err := initTableCommonWithIndices( - &newPart.TableCommon, t.encoder.UseNewCollate(), t.meta, def.ID, t.Columns, t.allocs, t.Constraints) - if err != nil { + newPart.TableCommon = newTableCommon(t.meta, def.ID, t.Columns, t.allocs, t.Constraints, t.encoder.UseNewCollate()) + if err := newPart.initTableIndices(); err != nil { return nil, err } newPart.table = t @@ -1677,7 +1675,7 @@ func GetReorganizedPartitionedTable(t table.Table) (table.PartitionedTable, erro if err != nil { return nil, err } - tc := initTableCommon(t.UseNewCollate(), tblInfo, tblInfo.ID, t.Cols(), t.Allocators(nil), constraints) + tc := newTableCommon(tblInfo, tblInfo.ID, t.Cols(), t.Allocators(nil), constraints, t.UseNewCollate()) // and rebuild the partitioning structure return newPartitionedTable(&tc, tblInfo) diff --git a/pkg/table/tables/tables.go b/pkg/table/tables/tables.go index f25522aeaae72..ba1cb844c443b 100644 --- a/pkg/table/tables/tables.go +++ b/pkg/table/tables/tables.go @@ -141,7 +141,7 @@ func MockTableFromMeta(tblInfo *model.TableInfo) table.Table { if err != nil { return nil } - t := initTableCommon(collate.NewCollationEnabled(), tblInfo, tblInfo.ID, columns, autoid.NewAllocators(false), constraints) + t := newTableCommon(tblInfo, tblInfo.ID, columns, autoid.NewAllocators(false), constraints, collate.NewCollationEnabled()) if tblInfo.TableCacheStatusType != model.TableCacheStatusDisable { ret, err := newCachedTable(&t) if err != nil { @@ -220,7 +220,7 @@ func TableFromMetaWithCollate(useNewCollate bool, allocs autoid.Allocators, tblI if err != nil { return nil, err } - t := initTableCommon(useNewCollate, tblInfo, tblInfo.ID, columns, allocs, constraints) + t := newTableCommon(tblInfo, tblInfo.ID, columns, allocs, constraints, useNewCollate) if tblInfo.GetPartitionInfo() == nil { if err := t.initTableIndices(); err != nil { return nil, err @@ -245,7 +245,7 @@ func buildGeneratedExpr(tblInfo *model.TableInfo, genExpr string) (ast.ExprNode, return expr, nil } -func initTableCommon(useNewCollate bool, tblInfo *model.TableInfo, physicalTableID int64, cols []*table.Column, allocs autoid.Allocators, constraints []*table.Constraint) TableCommon { +func newTableCommon(tblInfo *model.TableInfo, physicalTableID int64, cols []*table.Column, allocs autoid.Allocators, constraints []*table.Constraint, useNewCollate bool) TableCommon { t := TableCommon{ tableID: tblInfo.ID, physicalTableID: physicalTableID, @@ -316,11 +316,6 @@ func asIndex(idx table.Index) *index { return idx.(*index) } -func initTableCommonWithIndices(t *TableCommon, useNewCollate bool, tblInfo *model.TableInfo, physicalTableID int64, cols []*table.Column, allocs autoid.Allocators, constraints []*table.Constraint) error { - *t = initTableCommon(useNewCollate, tblInfo, physicalTableID, cols, allocs, constraints) - return t.initTableIndices() -} - // Indices implements table.Table Indices interface. func (t *TableCommon) Indices() []table.Index { return t.indices From f787d1a13d6eff51506a2dfec2decb9b70b2a05c Mon Sep 17 00:00:00 2001 From: Ruihao Chen Date: Fri, 3 Jul 2026 11:30:32 +0800 Subject: [PATCH 14/35] ddl: use reorg collation in cop context --- pkg/ddl/backfilling_txn_executor.go | 4 +- pkg/ddl/copr/copr_ctx.go | 101 ++++++++++++++++++++++-- pkg/planner/core/expression_rewriter.go | 9 ++- 3 files changed, 106 insertions(+), 8 deletions(-) diff --git a/pkg/ddl/backfilling_txn_executor.go b/pkg/ddl/backfilling_txn_executor.go index 46c2121ce6acf..3b136a174c3f5 100644 --- a/pkg/ddl/backfilling_txn_executor.go +++ b/pkg/ddl/backfilling_txn_executor.go @@ -33,6 +33,7 @@ import ( "github.com/pingcap/tidb/pkg/sessionctx/stmtctx" "github.com/pingcap/tidb/pkg/sessionctx/vardef" "github.com/pingcap/tidb/pkg/table" + "github.com/pingcap/tidb/pkg/util/collate" contextutil "github.com/pingcap/tidb/pkg/util/context" "github.com/pingcap/tidb/pkg/util/execdetails" "github.com/pingcap/tidb/pkg/util/intest" @@ -136,12 +137,13 @@ func NewReorgCopContext( tc, ec := evalCtx.TypeCtx(), evalCtx.ErrCtx() pushDownFlags := stmtctx.PushDownFlagsWithTypeFlagsAndErrLevels(tc.Flags(), ec.LevelMap()) - return copr.NewCopContext( + return copr.NewCopContextWithCollate( exprCtx, pushDownFlags, tblInfo, allIdxInfo, requestSource, + reorgMeta.GetUseNewCollateOrDefault(collate.NewCollationEnabled()), ) } diff --git a/pkg/ddl/copr/copr_ctx.go b/pkg/ddl/copr/copr_ctx.go index d39ac039a1b94..cad42057fabfa 100644 --- a/pkg/ddl/copr/copr_ctx.go +++ b/pkg/ddl/copr/copr_ctx.go @@ -24,6 +24,7 @@ import ( "github.com/pingcap/tidb/pkg/parser/ast" "github.com/pingcap/tidb/pkg/table/tables" "github.com/pingcap/tidb/pkg/types" + "github.com/pingcap/tidb/pkg/util/collate" ) // CopContext contains the information that is needed when building a coprocessor request. @@ -46,6 +47,7 @@ type CopContextBase struct { ExprCtx exprctx.BuildContext PushDownFlags uint64 RequestSource string + UseNewCollate bool ColumnInfos []*model.ColumnInfo FieldTypes []*types.FieldType @@ -80,6 +82,24 @@ func NewCopContextBase( tblInfo *model.TableInfo, idxCols []*model.IndexColumn, requestSource string, +) (*CopContextBase, error) { + return newCopContextBase( + exprCtx, + pushDownFlags, + tblInfo, + idxCols, + requestSource, + collate.NewCollationEnabled(), + ) +} + +func newCopContextBase( + exprCtx exprctx.BuildContext, + pushDownFlags uint64, + tblInfo *model.TableInfo, + idxCols []*model.IndexColumn, + requestSource string, + useNewCollate bool, ) (*CopContextBase, error) { var err error usedColumnIDs := make(map[int64]struct{}, len(idxCols)) @@ -125,8 +145,14 @@ func NewCopContextBase( handleIDs = []int64{extra.ID} } - expColInfos, _, err := expression.ColumnInfos2ColumnsAndNames(exprCtx, - ast.CIStr{} /* unused */, tblInfo.Name, colInfos, tblInfo) + expColInfos, _, err := expression.ColumnInfos2ColumnsAndNamesWithCollate( + exprCtx, + ast.CIStr{}, // unused + tblInfo.Name, + colInfos, + tblInfo, + useNewCollate, + ) if err != nil { return nil, err } @@ -139,6 +165,7 @@ func NewCopContextBase( ExprCtx: exprCtx, PushDownFlags: pushDownFlags, RequestSource: requestSource, + UseNewCollate: useNewCollate, ColumnInfos: colInfos, FieldTypes: fieldTps, ExprColumnInfos: expColInfos, @@ -155,11 +182,37 @@ func NewCopContext( tblInfo *model.TableInfo, allIdxInfo []*model.IndexInfo, requestSource string, +) (CopContext, error) { + return NewCopContextWithCollate( + exprCtx, + pushDownFlags, + tblInfo, + allIdxInfo, + requestSource, + collate.NewCollationEnabled(), + ) +} + +// NewCopContextWithCollate creates a CopContext with a fixed collation mode. +func NewCopContextWithCollate( + exprCtx exprctx.BuildContext, + pushDownFlags uint64, + tblInfo *model.TableInfo, + allIdxInfo []*model.IndexInfo, + requestSource string, + useNewCollate bool, ) (CopContext, error) { if len(allIdxInfo) == 1 { - return NewCopContextSingleIndex(exprCtx, pushDownFlags, tblInfo, allIdxInfo[0], requestSource) + return newCopContextSingleIndex( + exprCtx, + pushDownFlags, + tblInfo, + allIdxInfo[0], + requestSource, + useNewCollate, + ) } - return NewCopContextMultiIndex(exprCtx, pushDownFlags, tblInfo, allIdxInfo, requestSource) + return newCopContextMultiIndex(exprCtx, pushDownFlags, tblInfo, allIdxInfo, requestSource, useNewCollate) } // NewCopContextSingleIndex creates a CopContextSingleIndex. @@ -169,6 +222,24 @@ func NewCopContextSingleIndex( tblInfo *model.TableInfo, idxInfo *model.IndexInfo, requestSource string, +) (*CopContextSingleIndex, error) { + return newCopContextSingleIndex( + exprCtx, + pushDownFlags, + tblInfo, + idxInfo, + requestSource, + collate.NewCollationEnabled(), + ) +} + +func newCopContextSingleIndex( + exprCtx exprctx.BuildContext, + pushDownFlags uint64, + tblInfo *model.TableInfo, + idxInfo *model.IndexInfo, + requestSource string, + useNewCollate bool, ) (*CopContextSingleIndex, error) { cols := idxInfo.Columns neededCols, err := tables.ExtractColumnsFromCondition(exprCtx, idxInfo, tblInfo, false) @@ -178,7 +249,7 @@ func NewCopContextSingleIndex( cols = append(cols, neededCols...) cols = tables.DedupIndexColumns(cols) - base, err := NewCopContextBase(exprCtx, pushDownFlags, tblInfo, cols, requestSource) + base, err := newCopContextBase(exprCtx, pushDownFlags, tblInfo, cols, requestSource, useNewCollate) if err != nil { return nil, err } @@ -235,6 +306,24 @@ func NewCopContextMultiIndex( tblInfo *model.TableInfo, allIdxInfo []*model.IndexInfo, requestSource string, +) (*CopContextMultiIndex, error) { + return newCopContextMultiIndex( + exprCtx, + pushDownFlags, + tblInfo, + allIdxInfo, + requestSource, + collate.NewCollationEnabled(), + ) +} + +func newCopContextMultiIndex( + exprCtx exprctx.BuildContext, + pushDownFlags uint64, + tblInfo *model.TableInfo, + allIdxInfo []*model.IndexInfo, + requestSource string, + useNewCollate bool, ) (*CopContextMultiIndex, error) { approxColLen := 0 for _, idxInfo := range allIdxInfo { @@ -252,7 +341,7 @@ func NewCopContextMultiIndex( } allIdxCols = tables.DedupIndexColumns(allIdxCols) - base, err := NewCopContextBase(exprCtx, pushDownFlags, tblInfo, allIdxCols, requestSource) + base, err := newCopContextBase(exprCtx, pushDownFlags, tblInfo, allIdxCols, requestSource, useNewCollate) if err != nil { return nil, err } diff --git a/pkg/planner/core/expression_rewriter.go b/pkg/planner/core/expression_rewriter.go index af503d6ce9037..d13c90448a645 100644 --- a/pkg/planner/core/expression_rewriter.go +++ b/pkg/planner/core/expression_rewriter.go @@ -156,7 +156,14 @@ func buildSimpleExpr(ctx expression.BuildContext, node ast.ExprNode, opts ...exp } if tbl := options.SourceTable; tbl != nil && rewriter.schema == nil { - cols, names, err := expression.ColumnInfos2ColumnsAndNames(ctx, options.SourceTableDB, tbl.Name, tbl.Cols(), tbl) + cols, names, err := expression.ColumnInfos2ColumnsAndNamesWithCollate( + ctx, + options.SourceTableDB, + tbl.Name, + tbl.Cols(), + tbl, + options.UseNewCollate, + ) if err != nil { return nil, err } From ec4757bc4d38b54448552323ee23fdfe74a08c73 Mon Sep 17 00:00:00 2001 From: Ruihao Chen Date: Fri, 3 Jul 2026 11:35:12 +0800 Subject: [PATCH 15/35] ddl: simplify cop context base construction --- pkg/ddl/copr/copr_ctx.go | 21 ++------------------- 1 file changed, 2 insertions(+), 19 deletions(-) diff --git a/pkg/ddl/copr/copr_ctx.go b/pkg/ddl/copr/copr_ctx.go index cad42057fabfa..bd99e17ab4d3e 100644 --- a/pkg/ddl/copr/copr_ctx.go +++ b/pkg/ddl/copr/copr_ctx.go @@ -82,23 +82,6 @@ func NewCopContextBase( tblInfo *model.TableInfo, idxCols []*model.IndexColumn, requestSource string, -) (*CopContextBase, error) { - return newCopContextBase( - exprCtx, - pushDownFlags, - tblInfo, - idxCols, - requestSource, - collate.NewCollationEnabled(), - ) -} - -func newCopContextBase( - exprCtx exprctx.BuildContext, - pushDownFlags uint64, - tblInfo *model.TableInfo, - idxCols []*model.IndexColumn, - requestSource string, useNewCollate bool, ) (*CopContextBase, error) { var err error @@ -249,7 +232,7 @@ func newCopContextSingleIndex( cols = append(cols, neededCols...) cols = tables.DedupIndexColumns(cols) - base, err := newCopContextBase(exprCtx, pushDownFlags, tblInfo, cols, requestSource, useNewCollate) + base, err := NewCopContextBase(exprCtx, pushDownFlags, tblInfo, cols, requestSource, useNewCollate) if err != nil { return nil, err } @@ -341,7 +324,7 @@ func newCopContextMultiIndex( } allIdxCols = tables.DedupIndexColumns(allIdxCols) - base, err := newCopContextBase(exprCtx, pushDownFlags, tblInfo, allIdxCols, requestSource, useNewCollate) + base, err := NewCopContextBase(exprCtx, pushDownFlags, tblInfo, allIdxCols, requestSource, useNewCollate) if err != nil { return nil, err } From 3a6ac4569a45aa90d60bbc954fc94903435305d9 Mon Sep 17 00:00:00 2001 From: Ruihao Chen Date: Fri, 3 Jul 2026 11:40:51 +0800 Subject: [PATCH 16/35] ddl: simplify cop context collation constructor --- pkg/ddl/backfilling_txn_executor.go | 2 +- pkg/ddl/copr/copr_ctx.go | 20 +------------------- 2 files changed, 2 insertions(+), 20 deletions(-) diff --git a/pkg/ddl/backfilling_txn_executor.go b/pkg/ddl/backfilling_txn_executor.go index 3b136a174c3f5..f2a72e3fc1382 100644 --- a/pkg/ddl/backfilling_txn_executor.go +++ b/pkg/ddl/backfilling_txn_executor.go @@ -137,7 +137,7 @@ func NewReorgCopContext( tc, ec := evalCtx.TypeCtx(), evalCtx.ErrCtx() pushDownFlags := stmtctx.PushDownFlagsWithTypeFlagsAndErrLevels(tc.Flags(), ec.LevelMap()) - return copr.NewCopContextWithCollate( + return copr.NewCopContext( exprCtx, pushDownFlags, tblInfo, diff --git a/pkg/ddl/copr/copr_ctx.go b/pkg/ddl/copr/copr_ctx.go index bd99e17ab4d3e..c0341366f7aa9 100644 --- a/pkg/ddl/copr/copr_ctx.go +++ b/pkg/ddl/copr/copr_ctx.go @@ -158,31 +158,13 @@ func NewCopContextBase( }, nil } -// NewCopContext creates a CopContext. +// NewCopContext creates a CopContext with a fixed collation mode. func NewCopContext( exprCtx exprctx.BuildContext, pushDownFlags uint64, tblInfo *model.TableInfo, allIdxInfo []*model.IndexInfo, requestSource string, -) (CopContext, error) { - return NewCopContextWithCollate( - exprCtx, - pushDownFlags, - tblInfo, - allIdxInfo, - requestSource, - collate.NewCollationEnabled(), - ) -} - -// NewCopContextWithCollate creates a CopContext with a fixed collation mode. -func NewCopContextWithCollate( - exprCtx exprctx.BuildContext, - pushDownFlags uint64, - tblInfo *model.TableInfo, - allIdxInfo []*model.IndexInfo, - requestSource string, useNewCollate bool, ) (CopContext, error) { if len(allIdxInfo) == 1 { From ef0767333485d59b13471094d969730526e19d99 Mon Sep 17 00:00:00 2001 From: Ruihao Chen Date: Fri, 3 Jul 2026 11:43:48 +0800 Subject: [PATCH 17/35] ddl: inline cop context index constructors --- pkg/ddl/copr/copr_ctx.go | 43 ++++------------------------------- pkg/ddl/copr/copr_ctx_test.go | 5 +++- 2 files changed, 8 insertions(+), 40 deletions(-) diff --git a/pkg/ddl/copr/copr_ctx.go b/pkg/ddl/copr/copr_ctx.go index c0341366f7aa9..26cab165d69e4 100644 --- a/pkg/ddl/copr/copr_ctx.go +++ b/pkg/ddl/copr/copr_ctx.go @@ -24,7 +24,6 @@ import ( "github.com/pingcap/tidb/pkg/parser/ast" "github.com/pingcap/tidb/pkg/table/tables" "github.com/pingcap/tidb/pkg/types" - "github.com/pingcap/tidb/pkg/util/collate" ) // CopContext contains the information that is needed when building a coprocessor request. @@ -168,7 +167,7 @@ func NewCopContext( useNewCollate bool, ) (CopContext, error) { if len(allIdxInfo) == 1 { - return newCopContextSingleIndex( + return NewCopContextSingleIndex( exprCtx, pushDownFlags, tblInfo, @@ -177,33 +176,16 @@ func NewCopContext( useNewCollate, ) } - return newCopContextMultiIndex(exprCtx, pushDownFlags, tblInfo, allIdxInfo, requestSource, useNewCollate) + return NewCopContextMultiIndex(exprCtx, pushDownFlags, tblInfo, allIdxInfo, requestSource, useNewCollate) } -// NewCopContextSingleIndex creates a CopContextSingleIndex. +// NewCopContextSingleIndex creates a CopContextSingleIndex with a fixed collation mode. func NewCopContextSingleIndex( exprCtx exprctx.BuildContext, pushDownFlags uint64, tblInfo *model.TableInfo, idxInfo *model.IndexInfo, requestSource string, -) (*CopContextSingleIndex, error) { - return newCopContextSingleIndex( - exprCtx, - pushDownFlags, - tblInfo, - idxInfo, - requestSource, - collate.NewCollationEnabled(), - ) -} - -func newCopContextSingleIndex( - exprCtx exprctx.BuildContext, - pushDownFlags uint64, - tblInfo *model.TableInfo, - idxInfo *model.IndexInfo, - requestSource string, useNewCollate bool, ) (*CopContextSingleIndex, error) { cols := idxInfo.Columns @@ -264,30 +246,13 @@ func (c *CopContextSingleIndex) GetCondition() (expression.Expression, error) { return expr, nil } -// NewCopContextMultiIndex creates a CopContextMultiIndex. +// NewCopContextMultiIndex creates a CopContextMultiIndex with a fixed collation mode. func NewCopContextMultiIndex( exprCtx exprctx.BuildContext, pushDownFlags uint64, tblInfo *model.TableInfo, allIdxInfo []*model.IndexInfo, requestSource string, -) (*CopContextMultiIndex, error) { - return newCopContextMultiIndex( - exprCtx, - pushDownFlags, - tblInfo, - allIdxInfo, - requestSource, - collate.NewCollationEnabled(), - ) -} - -func newCopContextMultiIndex( - exprCtx exprctx.BuildContext, - pushDownFlags uint64, - tblInfo *model.TableInfo, - allIdxInfo []*model.IndexInfo, - requestSource string, useNewCollate bool, ) (*CopContextMultiIndex, error) { approxColLen := 0 diff --git a/pkg/ddl/copr/copr_ctx_test.go b/pkg/ddl/copr/copr_ctx_test.go index 0e9a9caa0e8a2..80b4a676f63c9 100644 --- a/pkg/ddl/copr/copr_ctx_test.go +++ b/pkg/ddl/copr/copr_ctx_test.go @@ -110,7 +110,10 @@ func TestNewCopContextSingleIndex(t *testing.T) { copCtx, err := NewCopContextSingleIndex( sctx.GetExprCtx(), sctx.GetSessionVars().StmtCtx.PushDownFlags(), - mockTableInfo, mockIdxInfo, "", + mockTableInfo, + mockIdxInfo, + "", + false, ) require.NoError(t, err) base := copCtx.GetBase() From a755f2db07645c07f0bf21a071abdb6d46b6f90c Mon Sep 17 00:00:00 2001 From: Ruihao Chen Date: Mon, 6 Jul 2026 12:02:44 +0800 Subject: [PATCH 18/35] planner: keep simple expression collation local --- pkg/planner/core/expression_rewriter.go | 10 +--------- 1 file changed, 1 insertion(+), 9 deletions(-) diff --git a/pkg/planner/core/expression_rewriter.go b/pkg/planner/core/expression_rewriter.go index d13c90448a645..f380aa2bf0147 100644 --- a/pkg/planner/core/expression_rewriter.go +++ b/pkg/planner/core/expression_rewriter.go @@ -156,14 +156,7 @@ func buildSimpleExpr(ctx expression.BuildContext, node ast.ExprNode, opts ...exp } if tbl := options.SourceTable; tbl != nil && rewriter.schema == nil { - cols, names, err := expression.ColumnInfos2ColumnsAndNamesWithCollate( - ctx, - options.SourceTableDB, - tbl.Name, - tbl.Cols(), - tbl, - options.UseNewCollate, - ) + cols, names, err := expression.ColumnInfos2ColumnsAndNames(ctx, options.SourceTableDB, tbl.Name, tbl.Cols(), tbl) if err != nil { return nil, err } @@ -289,7 +282,6 @@ func (b *PlanBuilder) getExpressionRewriter(ctx context.Context, p base.LogicalP rewriter.planCtx.aggrMap = nil rewriter.planCtx.insertPlan = nil rewriter.planCtx.rollExpand = b.currentBlockExpand - rewriter.useNewCollate = collate.NewCollationEnabled() return } From 8ce0b066b3d6b41a1ce85387df907073f5581a2b Mon Sep 17 00:00:00 2001 From: Ruihao Chen Date: Mon, 6 Jul 2026 12:12:30 +0800 Subject: [PATCH 19/35] bazel: add collation deps for background encoding --- pkg/dxf/importinto/BUILD.bazel | 1 + pkg/executor/importer/BUILD.bazel | 1 + pkg/infoschema/BUILD.bazel | 1 + 3 files changed, 3 insertions(+) diff --git a/pkg/dxf/importinto/BUILD.bazel b/pkg/dxf/importinto/BUILD.bazel index 1cf63949bd63c..39f544ea50307 100644 --- a/pkg/dxf/importinto/BUILD.bazel +++ b/pkg/dxf/importinto/BUILD.bazel @@ -74,6 +74,7 @@ go_library( "//pkg/util", "//pkg/util/backoff", "//pkg/util/chunk", + "//pkg/util/collate", "//pkg/util/dbterror/exeerrors", "//pkg/util/disttask", "//pkg/util/logutil", diff --git a/pkg/executor/importer/BUILD.bazel b/pkg/executor/importer/BUILD.bazel index 005fe5fc0adc9..8eaba81cfc195 100644 --- a/pkg/executor/importer/BUILD.bazel +++ b/pkg/executor/importer/BUILD.bazel @@ -68,6 +68,7 @@ go_library( "//pkg/util", "//pkg/util/cdcutil", "//pkg/util/chunk", + "//pkg/util/collate", "//pkg/util/context", "//pkg/util/cpu", "//pkg/util/dbterror", diff --git a/pkg/infoschema/BUILD.bazel b/pkg/infoschema/BUILD.bazel index 7a200c9b43a0a..ff71d94bb0064 100644 --- a/pkg/infoschema/BUILD.bazel +++ b/pkg/infoschema/BUILD.bazel @@ -50,6 +50,7 @@ go_library( "//pkg/types", "//pkg/util", "//pkg/util/chunk", + "//pkg/util/collate", "//pkg/util/dbterror", "//pkg/util/deadlockhistory", "//pkg/util/domainutil", From 0a38136736eb2fcd53a085c152962b59872c29db Mon Sep 17 00:00:00 2001 From: Ruihao Chen Date: Mon, 6 Jul 2026 12:25:04 +0800 Subject: [PATCH 20/35] importer: initialize table in chunk process test --- pkg/executor/importer/chunk_process_testkit_test.go | 1 + 1 file changed, 1 insertion(+) diff --git a/pkg/executor/importer/chunk_process_testkit_test.go b/pkg/executor/importer/chunk_process_testkit_test.go index c47dc56ef17bb..6a89be5c738f1 100644 --- a/pkg/executor/importer/chunk_process_testkit_test.go +++ b/pkg/executor/importer/chunk_process_testkit_test.go @@ -92,6 +92,7 @@ func TestFileChunkProcess(t *testing.T) { }, &importer.LoadDataController{ ASTArgs: &importer.ASTArgs{}, + Table: table, InsertColumns: table.VisibleCols(), FieldMappings: fieldMappings, }, From 0b27b769c7dc2e92ebbe40de97bbc0ad580a4e21 Mon Sep 17 00:00:00 2001 From: Ruihao Chen Date: Mon, 6 Jul 2026 14:11:05 +0800 Subject: [PATCH 21/35] Update pkg/executor/importer/import.go Co-authored-by: D3Hunter --- pkg/executor/importer/import.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pkg/executor/importer/import.go b/pkg/executor/importer/import.go index 67bfc2f8c12ff..33e33ea8c0eae 100644 --- a/pkg/executor/importer/import.go +++ b/pkg/executor/importer/import.go @@ -358,7 +358,7 @@ func (p *Plan) GetOnDupKeyMode() OnDupKeyMode { // GetUseNewCollateOrDefault returns the captured collation mode, or defaultVal // for import metadata generated before the field existed. func (p *Plan) GetUseNewCollateOrDefault(defaultVal bool) bool { - if p == nil || p.UseNewCollate == nil { + if p.UseNewCollate == nil { return defaultVal } return *p.UseNewCollate From b574e4a23fd814520503df20e401abec5050887a Mon Sep 17 00:00:00 2001 From: Ruihao Chen Date: Mon, 6 Jul 2026 15:12:40 +0800 Subject: [PATCH 22/35] expression: honor captured collation in simple exprs --- pkg/expression/builtin.go | 6 +- pkg/expression/builtin_compare.go | 32 +++-- pkg/expression/builtin_compare_vec.go | 4 +- .../builtin_compare_vec_generated.go | 14 +-- pkg/expression/builtin_other.go | 9 +- pkg/expression/builtin_other_vec_generated.go | 6 +- pkg/expression/builtin_string.go | 15 ++- pkg/expression/builtin_string_vec.go | 2 +- pkg/expression/collation_context.go | 49 ++++++++ pkg/expression/generator/compare_vec.go | 4 +- pkg/expression/generator/other_vec.go | 7 +- pkg/expression/util.go | 5 +- pkg/planner/core/expression_rewriter.go | 1 + pkg/planner/core/expression_test.go | 109 ++++++++++++++++++ 14 files changed, 214 insertions(+), 49 deletions(-) create mode 100644 pkg/expression/collation_context.go diff --git a/pkg/expression/builtin.go b/pkg/expression/builtin.go index c5e83893ed005..9edf9f12eb61c 100644 --- a/pkg/expression/builtin.go +++ b/pkg/expression/builtin.go @@ -139,7 +139,7 @@ func newBaseBuiltinFunc(ctx BuildContext, funcName string, args []Expression, tp tp: tp, } bf.SetCharsetAndCollation(ec.Charset, ec.Collation) - bf.setCollator(collate.GetCollator(ec.Collation)) + bf.setCollator(getCollatorFromBuildContext(ctx, ec.Collation)) bf.SetCoercibility(ec.Coer) bf.SetRepertoire(ec.Repe) adjustNullFlagForReturnType(ctx.GetEvalCtx(), funcName, args, bf) @@ -230,7 +230,7 @@ func newBaseBuiltinFuncWithTp(ctx BuildContext, funcName string, args []Expressi tp: fieldType, } bf.SetCharsetAndCollation(ec.Charset, ec.Collation) - bf.setCollator(collate.GetCollator(ec.Collation)) + bf.setCollator(getCollatorFromBuildContext(ctx, ec.Collation)) bf.SetCoercibility(ec.Coer) bf.SetRepertoire(ec.Repe) // note this function must be called after wrap cast function to the args @@ -290,7 +290,7 @@ func newBaseBuiltinFuncWithFieldTypes(ctx BuildContext, funcName string, args [] tp: fieldType, } bf.SetCharsetAndCollation(ec.Charset, ec.Collation) - bf.setCollator(collate.GetCollator(ec.Collation)) + bf.setCollator(getCollatorFromBuildContext(ctx, ec.Collation)) bf.SetCoercibility(ec.Coer) bf.SetRepertoire(ec.Repe) // note this function must be called after wrap cast function to the args diff --git a/pkg/expression/builtin_compare.go b/pkg/expression/builtin_compare.go index 74fa13d33ef0d..972e9ee50724b 100644 --- a/pkg/expression/builtin_compare.go +++ b/pkg/expression/builtin_compare.go @@ -703,7 +703,7 @@ func (b *builtinGreatestStringSig) evalString(ctx EvalContext, row chunk.Row) (m if isNull || err != nil { return maxv, isNull, err } - if types.CompareString(v, maxv, b.collation) > 0 { + if b.ctor.Compare(v, maxv) > 0 { maxv = v } } @@ -1047,7 +1047,7 @@ func (b *builtinLeastStringSig) evalString(ctx EvalContext, row chunk.Row) (minv if isNull || err != nil { return minv, isNull, err } - if types.CompareString(v, minv, b.collation) < 0 { + if b.ctor.Compare(v, minv) < 0 { minv = v } } @@ -1496,7 +1496,7 @@ func GetCmpFunction(ctx BuildContext, lhs, rhs Expression) CompareFunc { return CompareDecimal case types.ETString: coll, _ := CheckAndDeriveCollationFromExprs(ctx, "", types.ETInt, lhs, rhs) - return genCompareString(coll.Collation) + return genCompareStringWithCollator(getCollatorFromBuildContext(ctx, coll.Collation)) case types.ETDuration: return CompareDuration case types.ETDatetime, types.ETTimestamp: @@ -2267,7 +2267,7 @@ func (b *builtinLTStringSig) Clone() builtinFunc { } func (b *builtinLTStringSig) evalInt(ctx EvalContext, row chunk.Row) (val int64, isNull bool, err error) { - return resOfLT(CompareStringWithCollationInfo(ctx, b.args[0], b.args[1], row, row, b.collation)) + return resOfLT(compareStringWithCollator(ctx, b.args[0], b.args[1], row, row, b.ctor)) } type builtinLTDurationSig struct { @@ -2403,7 +2403,7 @@ func (b *builtinLEStringSig) Clone() builtinFunc { } func (b *builtinLEStringSig) evalInt(ctx EvalContext, row chunk.Row) (val int64, isNull bool, err error) { - return resOfLE(CompareStringWithCollationInfo(ctx, b.args[0], b.args[1], row, row, b.collation)) + return resOfLE(compareStringWithCollator(ctx, b.args[0], b.args[1], row, row, b.ctor)) } type builtinLEDurationSig struct { @@ -2539,7 +2539,7 @@ func (b *builtinGTStringSig) Clone() builtinFunc { } func (b *builtinGTStringSig) evalInt(ctx EvalContext, row chunk.Row) (val int64, isNull bool, err error) { - return resOfGT(CompareStringWithCollationInfo(ctx, b.args[0], b.args[1], row, row, b.collation)) + return resOfGT(compareStringWithCollator(ctx, b.args[0], b.args[1], row, row, b.ctor)) } type builtinGTDurationSig struct { @@ -2675,7 +2675,7 @@ func (b *builtinGEStringSig) Clone() builtinFunc { } func (b *builtinGEStringSig) evalInt(ctx EvalContext, row chunk.Row) (val int64, isNull bool, err error) { - return resOfGE(CompareStringWithCollationInfo(ctx, b.args[0], b.args[1], row, row, b.collation)) + return resOfGE(compareStringWithCollator(ctx, b.args[0], b.args[1], row, row, b.ctor)) } type builtinGEDurationSig struct { @@ -2811,7 +2811,7 @@ func (b *builtinEQStringSig) Clone() builtinFunc { } func (b *builtinEQStringSig) evalInt(ctx EvalContext, row chunk.Row) (val int64, isNull bool, err error) { - return resOfEQ(CompareStringWithCollationInfo(ctx, b.args[0], b.args[1], row, row, b.collation)) + return resOfEQ(compareStringWithCollator(ctx, b.args[0], b.args[1], row, row, b.ctor)) } type builtinEQDurationSig struct { @@ -2947,7 +2947,7 @@ func (b *builtinNEStringSig) Clone() builtinFunc { } func (b *builtinNEStringSig) evalInt(ctx EvalContext, row chunk.Row) (val int64, isNull bool, err error) { - return resOfNE(CompareStringWithCollationInfo(ctx, b.args[0], b.args[1], row, row, b.collation)) + return resOfNE(compareStringWithCollator(ctx, b.args[0], b.args[1], row, row, b.ctor)) } type builtinNEDurationSig struct { @@ -3151,7 +3151,7 @@ func (b *builtinNullEQStringSig) evalInt(ctx EvalContext, row chunk.Row) (val in res = 1 case isNull0 != isNull1: return res, false, nil - case types.CompareString(arg0, arg1, b.collation) == 0: + case b.ctor.Compare(arg0, arg1) == 0: res = 1 } return res, false, nil @@ -3410,13 +3410,21 @@ func CompareInt(sctx EvalContext, lhsArg, rhsArg Expression, lhsRow, rhsRow chun } func genCompareString(collation string) func(sctx EvalContext, lhsArg Expression, rhsArg Expression, lhsRow chunk.Row, rhsRow chunk.Row) (int64, bool, error) { + return genCompareStringWithCollator(collate.GetCollator(collation)) +} + +func genCompareStringWithCollator(collator collate.Collator) func(sctx EvalContext, lhsArg Expression, rhsArg Expression, lhsRow chunk.Row, rhsRow chunk.Row) (int64, bool, error) { return func(sctx EvalContext, lhsArg, rhsArg Expression, lhsRow, rhsRow chunk.Row) (int64, bool, error) { - return CompareStringWithCollationInfo(sctx, lhsArg, rhsArg, lhsRow, rhsRow, collation) + return compareStringWithCollator(sctx, lhsArg, rhsArg, lhsRow, rhsRow, collator) } } // CompareStringWithCollationInfo compares two strings with the specified collation information. func CompareStringWithCollationInfo(sctx EvalContext, lhsArg, rhsArg Expression, lhsRow, rhsRow chunk.Row, collation string) (int64, bool, error) { + return compareStringWithCollator(sctx, lhsArg, rhsArg, lhsRow, rhsRow, collate.GetCollator(collation)) +} + +func compareStringWithCollator(sctx EvalContext, lhsArg, rhsArg Expression, lhsRow, rhsRow chunk.Row, collator collate.Collator) (int64, bool, error) { arg0, isNull0, err := lhsArg.EvalString(sctx, lhsRow) if err != nil { return 0, true, err @@ -3430,7 +3438,7 @@ func CompareStringWithCollationInfo(sctx EvalContext, lhsArg, rhsArg Expression, if isNull0 || isNull1 { return compareNull(isNull0, isNull1), true, nil } - return int64(types.CompareString(arg0, arg1, collation)), false, nil + return int64(collator.Compare(arg0, arg1)), false, nil } // CompareReal compares two float-point values. diff --git a/pkg/expression/builtin_compare_vec.go b/pkg/expression/builtin_compare_vec.go index d2847e9d8c679..6633b94e136e4 100644 --- a/pkg/expression/builtin_compare_vec.go +++ b/pkg/expression/builtin_compare_vec.go @@ -267,7 +267,7 @@ func (b *builtinLeastStringSig) vecEvalString(ctx EvalContext, input *chunk.Chun } srcStr := src.GetString(i) argStr := arg.GetString(i) - if types.CompareString(srcStr, argStr, b.collation) < 0 { + if b.ctor.Compare(srcStr, argStr) < 0 { dst.AppendString(srcStr) } else { dst.AppendString(argStr) @@ -790,7 +790,7 @@ func (b *builtinGreatestStringSig) vecEvalString(ctx EvalContext, input *chunk.C } srcStr := src.GetString(i) argStr := arg.GetString(i) - if types.CompareString(srcStr, argStr, b.collation) > 0 { + if b.ctor.Compare(srcStr, argStr) > 0 { dst.AppendString(srcStr) } else { dst.AppendString(argStr) diff --git a/pkg/expression/builtin_compare_vec_generated.go b/pkg/expression/builtin_compare_vec_generated.go index 8f53e5c27ccbe..f246462d8bdbb 100644 --- a/pkg/expression/builtin_compare_vec_generated.go +++ b/pkg/expression/builtin_compare_vec_generated.go @@ -125,7 +125,7 @@ func (b *builtinLTStringSig) vecEvalInt(ctx EvalContext, input *chunk.Chunk, res if result.IsNull(i) { continue } - val := types.CompareString(buf0.GetString(i), buf1.GetString(i), b.collation) + val := b.ctor.Compare(buf0.GetString(i), buf1.GetString(i)) i64s[i] = boolToInt64(val < 0) } return nil @@ -349,7 +349,7 @@ func (b *builtinLEStringSig) vecEvalInt(ctx EvalContext, input *chunk.Chunk, res if result.IsNull(i) { continue } - val := types.CompareString(buf0.GetString(i), buf1.GetString(i), b.collation) + val := b.ctor.Compare(buf0.GetString(i), buf1.GetString(i)) i64s[i] = boolToInt64(val <= 0) } return nil @@ -573,7 +573,7 @@ func (b *builtinGTStringSig) vecEvalInt(ctx EvalContext, input *chunk.Chunk, res if result.IsNull(i) { continue } - val := types.CompareString(buf0.GetString(i), buf1.GetString(i), b.collation) + val := b.ctor.Compare(buf0.GetString(i), buf1.GetString(i)) i64s[i] = boolToInt64(val > 0) } return nil @@ -797,7 +797,7 @@ func (b *builtinGEStringSig) vecEvalInt(ctx EvalContext, input *chunk.Chunk, res if result.IsNull(i) { continue } - val := types.CompareString(buf0.GetString(i), buf1.GetString(i), b.collation) + val := b.ctor.Compare(buf0.GetString(i), buf1.GetString(i)) i64s[i] = boolToInt64(val >= 0) } return nil @@ -1021,7 +1021,7 @@ func (b *builtinEQStringSig) vecEvalInt(ctx EvalContext, input *chunk.Chunk, res if result.IsNull(i) { continue } - val := types.CompareString(buf0.GetString(i), buf1.GetString(i), b.collation) + val := b.ctor.Compare(buf0.GetString(i), buf1.GetString(i)) i64s[i] = boolToInt64(val == 0) } return nil @@ -1245,7 +1245,7 @@ func (b *builtinNEStringSig) vecEvalInt(ctx EvalContext, input *chunk.Chunk, res if result.IsNull(i) { continue } - val := types.CompareString(buf0.GetString(i), buf1.GetString(i), b.collation) + val := b.ctor.Compare(buf0.GetString(i), buf1.GetString(i)) i64s[i] = boolToInt64(val != 0) } return nil @@ -1480,7 +1480,7 @@ func (b *builtinNullEQStringSig) vecEvalInt(ctx EvalContext, input *chunk.Chunk, i64s[i] = 1 case isNull0 != isNull1: i64s[i] = 0 - case types.CompareString(buf0.GetString(i), buf1.GetString(i), b.collation) == 0: + case b.ctor.Compare(buf0.GetString(i), buf1.GetString(i)) == 0: i64s[i] = 1 } } diff --git a/pkg/expression/builtin_other.go b/pkg/expression/builtin_other.go index 35e017927f7bf..5f4246cd2773c 100644 --- a/pkg/expression/builtin_other.go +++ b/pkg/expression/builtin_other.go @@ -26,7 +26,6 @@ import ( "github.com/pingcap/tidb/pkg/parser/mysql" "github.com/pingcap/tidb/pkg/types" "github.com/pingcap/tidb/pkg/util/chunk" - "github.com/pingcap/tidb/pkg/util/collate" "github.com/pingcap/tidb/pkg/util/intest" "github.com/pingcap/tidb/pkg/util/set" "github.com/pingcap/tidb/pkg/util/stringutil" @@ -339,7 +338,6 @@ type builtinInStringSig struct { func (b *builtinInStringSig) buildHashMapForConstArgs(ctx BuildContext) error { b.nonConstArgsIdx = make([]int, 0) b.hashSet = set.NewStringSet() - collator := collate.GetCollator(b.collation) // Keep track of unique args count for in-place modification uniqueArgCount := 1 // Start with 1 for the first arg (value to check) @@ -363,7 +361,7 @@ func (b *builtinInStringSig) buildHashMapForConstArgs(ctx BuildContext) error { continue } - key := string(collator.Key(val)) // should do memory copy here + key := string(b.ctor.Key(val)) // should do memory copy here // Only keep this arg if value wasn't seen before if !b.hashSet.Exist(key) { b.hashSet.Insert(key) @@ -407,9 +405,8 @@ func (b *builtinInStringSig) evalInt(ctx EvalContext, row chunk.Row) (int64, boo } args := b.args[1:] - collator := collate.GetCollator(b.collation) if len(b.hashSet) != 0 { - if b.hashSet.Exist(string(collator.Key(arg0))) { + if b.hashSet.Exist(string(b.ctor.Key(arg0))) { return 1, false, nil } args = make([]Expression, 0, len(b.nonConstArgsIdx)) @@ -428,7 +425,7 @@ func (b *builtinInStringSig) evalInt(ctx EvalContext, row chunk.Row) (int64, boo hasNull = true continue } - if types.CompareString(arg0, evaledArg, b.collation) == 0 { + if b.ctor.Compare(arg0, evaledArg) == 0 { return 1, false, nil } } diff --git a/pkg/expression/builtin_other_vec_generated.go b/pkg/expression/builtin_other_vec_generated.go index 0294c85f3759b..beb54229818c0 100644 --- a/pkg/expression/builtin_other_vec_generated.go +++ b/pkg/expression/builtin_other_vec_generated.go @@ -22,7 +22,6 @@ import ( "github.com/pingcap/tidb/pkg/parser/mysql" "github.com/pingcap/tidb/pkg/types" "github.com/pingcap/tidb/pkg/util/chunk" - "github.com/pingcap/tidb/pkg/util/collate" ) func (b *builtinInIntSig) vecEvalInt(ctx EvalContext, input *chunk.Chunk, result *chunk.Column) error { @@ -160,14 +159,13 @@ func (b *builtinInStringSig) vecEvalInt(ctx EvalContext, input *chunk.Chunk, res var compareResult int args := b.args[1:] if len(b.hashSet) != 0 { - collator := collate.GetCollator(b.collation) for i := 0; i < n; i++ { if buf0.IsNull(i) { hasNull[i] = true continue } arg0 := buf0.GetString(i) - if _, ok := b.hashSet[string(collator.Key(arg0))]; ok { + if _, ok := b.hashSet[string(b.ctor.Key(arg0))]; ok { r64s[i] = 1 result.SetNull(i, false) } @@ -192,7 +190,7 @@ func (b *builtinInStringSig) vecEvalInt(ctx EvalContext, input *chunk.Chunk, res } arg0 := buf0.GetString(i) arg1 := buf1.GetString(i) - compareResult = types.CompareString(arg0, arg1, b.collation) + compareResult = b.ctor.Compare(arg0, arg1) if compareResult == 0 { result.SetNull(i, false) r64s[i] = 1 diff --git a/pkg/expression/builtin_string.go b/pkg/expression/builtin_string.go index f6299f988ff66..0dd18fbcfd8d3 100644 --- a/pkg/expression/builtin_string.go +++ b/pkg/expression/builtin_string.go @@ -1026,7 +1026,7 @@ func (b *builtinStrcmpSig) evalInt(ctx EvalContext, row chunk.Row) (int64, bool, if isNull || err != nil { return 0, isNull, err } - res := types.CompareString(left, right, b.collation) + res := b.ctor.Compare(left, right) return int64(res), false, nil } @@ -1593,7 +1593,7 @@ func (b *builtinLocate2ArgsUTF8Sig) evalInt(ctx EvalContext, row chunk.Row) (int return 1, false, nil } - return locateStringWithCollation(str, subStr, b.collation), false, nil + return locateStringWithCollator(str, subStr, b.ctor), false, nil } type builtinLocate3ArgsSig struct { @@ -1684,7 +1684,7 @@ func (b *builtinLocate3ArgsUTF8Sig) evalInt(ctx EvalContext, row chunk.Row) (int } slice := string([]rune(str)[pos:]) - idx := locateStringWithCollation(slice, subStr, b.collation) + idx := locateStringWithCollator(slice, subStr, b.ctor) if idx != 0 { return pos + idx, false, nil } @@ -4260,7 +4260,8 @@ func (c *weightStringFunctionClass) getFunction(ctx BuildContext, args []Express sig = &builtinWeightStringNullSig{bf} } else { maxAllowedPacket := ctx.GetEvalCtx().GetMaxAllowedPacket() - sig = &builtinWeightStringSig{bf, padding, length, maxAllowedPacket} + argCollator := getCollatorFromBuildContext(ctx, bf.args[0].GetType(ctx.GetEvalCtx()).GetCollate()) + sig = &builtinWeightStringSig{bf, padding, length, maxAllowedPacket, argCollator} } return sig, nil } @@ -4291,6 +4292,7 @@ type builtinWeightStringSig struct { padding weightStringPadding length int maxAllowedPacket uint64 + argCollator collate.Collator } func (b *builtinWeightStringSig) Clone() builtinFunc { @@ -4299,6 +4301,7 @@ func (b *builtinWeightStringSig) Clone() builtinFunc { newSig.padding = b.padding newSig.length = b.length newSig.maxAllowedPacket = b.maxAllowedPacket + newSig.argCollator = b.argCollator return newSig } @@ -4327,7 +4330,7 @@ func (b *builtinWeightStringSig) evalString(ctx EvalContext, row chunk.Row) (str } str += strings.Repeat(" ", b.length-lenRunes) } - ctor = collate.GetCollator(b.args[0].GetType(ctx).GetCollate()) + ctor = b.argCollator case weightStringPaddingAsBinary: lenStr := len(str) if b.length < lenStr { @@ -4343,7 +4346,7 @@ func (b *builtinWeightStringSig) evalString(ctx EvalContext, row chunk.Row) (str } ctor = collate.GetCollator(charset.CollationBin) case weightStringPaddingNone: - ctor = collate.GetCollator(b.args[0].GetType(ctx).GetCollate()) + ctor = b.argCollator default: return "", false, ErrIncorrectType.GenWithStackByArgs(ast.WeightString, string(b.padding)) } diff --git a/pkg/expression/builtin_string_vec.go b/pkg/expression/builtin_string_vec.go index 9de6e92372288..8f328b8505969 100644 --- a/pkg/expression/builtin_string_vec.go +++ b/pkg/expression/builtin_string_vec.go @@ -1294,7 +1294,7 @@ func (b *builtinStrcmpSig) vecEvalInt(ctx EvalContext, input *chunk.Chunk, resul if result.IsNull(i) { continue } - i64s[i] = int64(types.CompareString(leftBuf.GetString(i), rightBuf.GetString(i), b.collation)) + i64s[i] = int64(b.ctor.Compare(leftBuf.GetString(i), rightBuf.GetString(i))) } return nil } diff --git a/pkg/expression/collation_context.go b/pkg/expression/collation_context.go new file mode 100644 index 0000000000000..4c056da53d8bc --- /dev/null +++ b/pkg/expression/collation_context.go @@ -0,0 +1,49 @@ +// Copyright 2026 PingCAP, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package expression + +import "github.com/pingcap/tidb/pkg/util/collate" + +type collationBuildContext interface { + UseNewCollate() bool +} + +type buildContextWithCollation struct { + BuildContext + useNewCollate bool +} + +func (ctx *buildContextWithCollation) UseNewCollate() bool { + return ctx.useNewCollate +} + +// CtxWithUseNewCollate returns a build context with an explicit collation mode. +func CtxWithUseNewCollate(ctx BuildContext, useNewCollate bool) BuildContext { + return &buildContextWithCollation{ + BuildContext: ctx, + useNewCollate: useNewCollate, + } +} + +func useNewCollateFromBuildContext(ctx BuildContext) bool { + if c, ok := ctx.(collationBuildContext); ok { + return c.UseNewCollate() + } + return collate.NewCollationEnabled() +} + +func getCollatorFromBuildContext(ctx BuildContext, collation string) collate.Collator { + return collate.GetCollatorWithCollate(useNewCollateFromBuildContext(ctx), collation) +} diff --git a/pkg/expression/generator/compare_vec.go b/pkg/expression/generator/compare_vec.go index 37496dd2e7b60..c651162691de3 100644 --- a/pkg/expression/generator/compare_vec.go +++ b/pkg/expression/generator/compare_vec.go @@ -94,7 +94,7 @@ func (b *builtin{{ .compare.CompareName }}{{ .type.TypeName }}Sig) vecEvalInt(ct {{- else if eq .type.ETName "Real" }} val := cmp.Compare(arg0[i], arg1[i]) {{- else if eq .type.ETName "String" }} - val := types.CompareString(buf0.GetString(i), buf1.GetString(i), b.collation) + val := b.ctor.Compare(buf0.GetString(i), buf1.GetString(i)) {{- else if eq .type.ETName "Duration" }} val := cmp.Compare(arg0[i], arg1[i]) {{- else if eq .type.ETName "Datetime" }} @@ -151,7 +151,7 @@ func (b *builtin{{ .compare.CompareName }}{{ .type.TypeName }}Sig) vecEvalInt(ct {{- else if eq .type.ETName "Real" }} case cmp.Compare(arg0[i], arg1[i]) == 0: {{- else if eq .type.ETName "String" }} - case types.CompareString(buf0.GetString(i), buf1.GetString(i), b.collation) == 0: + case b.ctor.Compare(buf0.GetString(i), buf1.GetString(i)) == 0: {{- else if eq .type.ETName "Duration" }} case cmp.Compare(arg0[i], arg1[i]) == 0: {{- else if eq .type.ETName "Datetime" }} diff --git a/pkg/expression/generator/other_vec.go b/pkg/expression/generator/other_vec.go index 19f0338fa534a..aefcea4a5d23a 100644 --- a/pkg/expression/generator/other_vec.go +++ b/pkg/expression/generator/other_vec.go @@ -112,7 +112,7 @@ var builtinInTmpl = template.Must(template.New("builtinInTmpl").Parse(` {{- else if eq .Input.TypeName "JSON" -}} compareResult = types.CompareBinaryJSON(arg0, arg1) {{- else if eq .Input.TypeName "String" -}} - compareResult = types.CompareString(arg0, arg1, b.collation) + compareResult = b.ctor.Compare(arg0, arg1) {{- else if eq .Input.TypeNameInColumn "Float64" -}} compareResult = cmp.Compare(arg0, arg1) {{- else -}} @@ -153,9 +153,6 @@ func (b *{{.SigName}}) vecEvalInt(ctx EvalContext, input *chunk.Chunk, result *c args := b.args[1:] {{- if not $InputJSON}} if len(b.hashSet) != 0 { - {{- if $InputString }} - collator := collate.GetCollator(b.collation) - {{- end }} for i := 0; i < n; i++ { if buf0.IsNull(i) { hasNull[i] = true @@ -190,7 +187,7 @@ func (b *{{.SigName}}) vecEvalInt(ctx EvalContext, input *chunk.Chunk, result *c result.SetNull(i, false) } {{- else if $InputString }} - if _, ok := b.hashSet[string(collator.Key(arg0))]; ok { + if _, ok := b.hashSet[string(b.ctor.Key(arg0))]; ok { r64s[i] = 1 result.SetNull(i, false) } diff --git a/pkg/expression/util.go b/pkg/expression/util.go index dc537a03327de..bf8c4ede2eb06 100644 --- a/pkg/expression/util.go +++ b/pkg/expression/util.go @@ -818,7 +818,10 @@ func SubstituteCorCol2Constant(ctx BuildContext, expr Expression) (Expression, e } func locateStringWithCollation(str, substr, coll string) int64 { - collator := collate.GetCollator(coll) + return locateStringWithCollator(str, substr, collate.GetCollator(coll)) +} + +func locateStringWithCollator(str, substr string, collator collate.Collator) int64 { strKey := collator.KeyWithoutTrimRightSpace(str) subStrKey := collator.KeyWithoutTrimRightSpace(substr) diff --git a/pkg/planner/core/expression_rewriter.go b/pkg/planner/core/expression_rewriter.go index f380aa2bf0147..089359fba0319 100644 --- a/pkg/planner/core/expression_rewriter.go +++ b/pkg/planner/core/expression_rewriter.go @@ -118,6 +118,7 @@ func buildSimpleExpr(ctx expression.BuildContext, node ast.ExprNode, opts ...exp for _, opt := range opts { opt(&options) } + ctx = expression.CtxWithUseNewCollate(ctx, options.UseNewCollate) if options.InputSchema == nil && len(options.InputNames) > 0 { return nil, errors.New("InputSchema and InputNames should be specified at the same time") diff --git a/pkg/planner/core/expression_test.go b/pkg/planner/core/expression_test.go index 766e7c6f1b3e5..3a1cb05646f92 100644 --- a/pkg/planner/core/expression_test.go +++ b/pkg/planner/core/expression_test.go @@ -33,6 +33,7 @@ import ( "github.com/pingcap/tidb/pkg/testkit/testutil" "github.com/pingcap/tidb/pkg/types" "github.com/pingcap/tidb/pkg/util/chunk" + "github.com/pingcap/tidb/pkg/util/collate" "github.com/stretchr/testify/require" ) @@ -71,6 +72,114 @@ func buildExprAndEval(t *testing.T, ctx expression.BuildContext, exprNode any) t return val } +func TestBuildSimpleExprWithUseNewCollate(t *testing.T) { + origin := collate.NewCollationEnabled() + collate.SetNewCollationEnabledForTest(true) + defer collate.SetNewCollationEnabledForTest(origin) + + ctx := coretestsdk.MockContext() + defer func() { + domain.GetDomain(ctx).StatsHandle().Close() + }() + + tests := []struct { + name string + expr string + oldResult string + newResult string + }{ + { + name: "eq", + expr: "if(" + + "_utf8mb4'aaa' COLLATE utf8mb4_general_ci = " + + "_utf8mb4'AAA' COLLATE utf8mb4_general_ci, 'match', 'miss')", + oldResult: "miss", + newResult: "match", + }, + { + name: "null safe eq", + expr: "if(" + + "_utf8mb4'aaa' COLLATE utf8mb4_general_ci <=> " + + "_utf8mb4'AAA' COLLATE utf8mb4_general_ci, 'match', 'miss')", + oldResult: "miss", + newResult: "match", + }, + { + name: "in", + expr: "if(" + + "_utf8mb4'aaa' COLLATE utf8mb4_general_ci in " + + "(_utf8mb4'AAA' COLLATE utf8mb4_general_ci), 'match', 'miss')", + oldResult: "miss", + newResult: "match", + }, + { + name: "strcmp", + expr: "if(strcmp(" + + "_utf8mb4'aaa' COLLATE utf8mb4_general_ci, " + + "_utf8mb4'AAA' COLLATE utf8mb4_general_ci) = 0, 'match', 'miss')", + oldResult: "miss", + newResult: "match", + }, + { + name: "least", + expr: "least(" + + "_utf8mb4'aaa' COLLATE utf8mb4_general_ci, " + + "_utf8mb4'AAA' COLLATE utf8mb4_general_ci)", + oldResult: "AAA", + newResult: "aaa", + }, + { + name: "greatest", + expr: "greatest(" + + "_utf8mb4'AAA' COLLATE utf8mb4_general_ci, " + + "_utf8mb4'aaa' COLLATE utf8mb4_general_ci)", + oldResult: "aaa", + newResult: "AAA", + }, + { + name: "locate", + expr: "if(locate(" + + "_utf8mb4'A' COLLATE utf8mb4_general_ci, " + + "_utf8mb4'aaa' COLLATE utf8mb4_general_ci) = 1, 'match', 'miss')", + oldResult: "miss", + newResult: "match", + }, + { + name: "field", + expr: "if(field(" + + "_utf8mb4'aaa' COLLATE utf8mb4_general_ci, " + + "_utf8mb4'AAA' COLLATE utf8mb4_general_ci) = 1, 'match', 'miss')", + oldResult: "miss", + newResult: "match", + }, + { + name: "weight string", + expr: "if(weight_string(" + + "_utf8mb4'aaa' COLLATE utf8mb4_general_ci) = weight_string(" + + "_utf8mb4'AAA' COLLATE utf8mb4_general_ci), 'match', 'miss')", + oldResult: "miss", + newResult: "match", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + exprNode := parseExpr(t, tt.expr) + expr, err := buildExpr(t, ctx, exprNode, expression.WithUseNewCollate(false)) + require.NoError(t, err) + val, err := expr.Eval(ctx.GetExprCtx().GetEvalCtx(), chunk.Row{}) + require.NoError(t, err) + require.Equal(t, tt.oldResult, val.GetString()) + + expr, err = buildExpr(t, ctx, exprNode, expression.WithUseNewCollate(true)) + require.NoError(t, err) + val, err = expr.Eval(ctx.GetExprCtx().GetEvalCtx(), chunk.Row{}) + require.NoError(t, err) + require.Equal(t, tt.newResult, val.GetString()) + }) + } +} + type testCase struct { exprStr string resultStr string From 76fab39eb5f61797dff711d34bcc2052bc886256 Mon Sep 17 00:00:00 2001 From: Ruihao Chen Date: Mon, 6 Jul 2026 15:50:19 +0800 Subject: [PATCH 23/35] dxf: limit collation snapshot expression scope --- pkg/executor/importer/BUILD.bazel | 1 - pkg/executor/importer/import.go | 13 +-- pkg/expression/builtin.go | 6 +- pkg/expression/builtin_compare.go | 32 ++--- pkg/expression/builtin_compare_vec.go | 4 +- .../builtin_compare_vec_generated.go | 14 +-- pkg/expression/builtin_other.go | 9 +- pkg/expression/builtin_other_vec_generated.go | 6 +- pkg/expression/builtin_string.go | 15 +-- pkg/expression/builtin_string_vec.go | 2 +- pkg/expression/collation_context.go | 49 -------- pkg/expression/generator/compare_vec.go | 4 +- pkg/expression/generator/other_vec.go | 7 +- pkg/expression/util.go | 5 +- pkg/planner/core/expression_rewriter.go | 1 - pkg/planner/core/expression_test.go | 109 ------------------ 16 files changed, 55 insertions(+), 222 deletions(-) delete mode 100644 pkg/expression/collation_context.go diff --git a/pkg/executor/importer/BUILD.bazel b/pkg/executor/importer/BUILD.bazel index 8eaba81cfc195..005fe5fc0adc9 100644 --- a/pkg/executor/importer/BUILD.bazel +++ b/pkg/executor/importer/BUILD.bazel @@ -68,7 +68,6 @@ go_library( "//pkg/util", "//pkg/util/cdcutil", "//pkg/util/chunk", - "//pkg/util/collate", "//pkg/util/context", "//pkg/util/cpu", "//pkg/util/dbterror", diff --git a/pkg/executor/importer/import.go b/pkg/executor/importer/import.go index 33e33ea8c0eae..20f54fa5a28f0 100644 --- a/pkg/executor/importer/import.go +++ b/pkg/executor/importer/import.go @@ -65,7 +65,6 @@ import ( "github.com/pingcap/tidb/pkg/table" tidbutil "github.com/pingcap/tidb/pkg/util" "github.com/pingcap/tidb/pkg/util/chunk" - "github.com/pingcap/tidb/pkg/util/collate" contextutil "github.com/pingcap/tidb/pkg/util/context" "github.com/pingcap/tidb/pkg/util/cpu" "github.com/pingcap/tidb/pkg/util/dbterror" @@ -336,10 +335,10 @@ type Plan struct { ManualRecovery bool // the keyspace name when submitting this job, only for import-into Keyspace string - // UseNewCollate captures the submitting keyspace collation mode for key encoding. - // It is needed because an import task may execute in another keyspace with a - // different collation setting. Nil means old metadata and should fall back to - // the caller-provided default. + // UseNewCollate captures the collation mode used to build the target table + // snapshot. It is needed because an import task may execute in another + // keyspace with a different collation setting. Nil means old metadata and + // should fall back to the caller-provided default. UseNewCollate *bool `json:"use_new_collate,omitempty"` } @@ -364,7 +363,7 @@ func (p *Plan) GetUseNewCollateOrDefault(defaultVal bool) bool { return *p.UseNewCollate } -// SetUseNewCollate stores the submitting keyspace collation mode for key encoding. +// SetUseNewCollate stores the collation mode used to rebuild the target table snapshot. func (p *Plan) SetUseNewCollate(useNewCollate bool) { p.UseNewCollate = &useNewCollate } @@ -573,7 +572,7 @@ func NewImportPlan(ctx context.Context, userSctx sessionctx.Context, plan *plann User: userSctx.GetSessionVars().User.String(), Keyspace: userSctx.GetStore().GetKeyspace(), } - p.SetUseNewCollate(collate.NewCollationEnabled()) + p.SetUseNewCollate(tbl.UseNewCollate()) if err := p.initOptions(ctx, userSctx, plan.Options); err != nil { return nil, err } diff --git a/pkg/expression/builtin.go b/pkg/expression/builtin.go index 9edf9f12eb61c..c5e83893ed005 100644 --- a/pkg/expression/builtin.go +++ b/pkg/expression/builtin.go @@ -139,7 +139,7 @@ func newBaseBuiltinFunc(ctx BuildContext, funcName string, args []Expression, tp tp: tp, } bf.SetCharsetAndCollation(ec.Charset, ec.Collation) - bf.setCollator(getCollatorFromBuildContext(ctx, ec.Collation)) + bf.setCollator(collate.GetCollator(ec.Collation)) bf.SetCoercibility(ec.Coer) bf.SetRepertoire(ec.Repe) adjustNullFlagForReturnType(ctx.GetEvalCtx(), funcName, args, bf) @@ -230,7 +230,7 @@ func newBaseBuiltinFuncWithTp(ctx BuildContext, funcName string, args []Expressi tp: fieldType, } bf.SetCharsetAndCollation(ec.Charset, ec.Collation) - bf.setCollator(getCollatorFromBuildContext(ctx, ec.Collation)) + bf.setCollator(collate.GetCollator(ec.Collation)) bf.SetCoercibility(ec.Coer) bf.SetRepertoire(ec.Repe) // note this function must be called after wrap cast function to the args @@ -290,7 +290,7 @@ func newBaseBuiltinFuncWithFieldTypes(ctx BuildContext, funcName string, args [] tp: fieldType, } bf.SetCharsetAndCollation(ec.Charset, ec.Collation) - bf.setCollator(getCollatorFromBuildContext(ctx, ec.Collation)) + bf.setCollator(collate.GetCollator(ec.Collation)) bf.SetCoercibility(ec.Coer) bf.SetRepertoire(ec.Repe) // note this function must be called after wrap cast function to the args diff --git a/pkg/expression/builtin_compare.go b/pkg/expression/builtin_compare.go index 972e9ee50724b..74fa13d33ef0d 100644 --- a/pkg/expression/builtin_compare.go +++ b/pkg/expression/builtin_compare.go @@ -703,7 +703,7 @@ func (b *builtinGreatestStringSig) evalString(ctx EvalContext, row chunk.Row) (m if isNull || err != nil { return maxv, isNull, err } - if b.ctor.Compare(v, maxv) > 0 { + if types.CompareString(v, maxv, b.collation) > 0 { maxv = v } } @@ -1047,7 +1047,7 @@ func (b *builtinLeastStringSig) evalString(ctx EvalContext, row chunk.Row) (minv if isNull || err != nil { return minv, isNull, err } - if b.ctor.Compare(v, minv) < 0 { + if types.CompareString(v, minv, b.collation) < 0 { minv = v } } @@ -1496,7 +1496,7 @@ func GetCmpFunction(ctx BuildContext, lhs, rhs Expression) CompareFunc { return CompareDecimal case types.ETString: coll, _ := CheckAndDeriveCollationFromExprs(ctx, "", types.ETInt, lhs, rhs) - return genCompareStringWithCollator(getCollatorFromBuildContext(ctx, coll.Collation)) + return genCompareString(coll.Collation) case types.ETDuration: return CompareDuration case types.ETDatetime, types.ETTimestamp: @@ -2267,7 +2267,7 @@ func (b *builtinLTStringSig) Clone() builtinFunc { } func (b *builtinLTStringSig) evalInt(ctx EvalContext, row chunk.Row) (val int64, isNull bool, err error) { - return resOfLT(compareStringWithCollator(ctx, b.args[0], b.args[1], row, row, b.ctor)) + return resOfLT(CompareStringWithCollationInfo(ctx, b.args[0], b.args[1], row, row, b.collation)) } type builtinLTDurationSig struct { @@ -2403,7 +2403,7 @@ func (b *builtinLEStringSig) Clone() builtinFunc { } func (b *builtinLEStringSig) evalInt(ctx EvalContext, row chunk.Row) (val int64, isNull bool, err error) { - return resOfLE(compareStringWithCollator(ctx, b.args[0], b.args[1], row, row, b.ctor)) + return resOfLE(CompareStringWithCollationInfo(ctx, b.args[0], b.args[1], row, row, b.collation)) } type builtinLEDurationSig struct { @@ -2539,7 +2539,7 @@ func (b *builtinGTStringSig) Clone() builtinFunc { } func (b *builtinGTStringSig) evalInt(ctx EvalContext, row chunk.Row) (val int64, isNull bool, err error) { - return resOfGT(compareStringWithCollator(ctx, b.args[0], b.args[1], row, row, b.ctor)) + return resOfGT(CompareStringWithCollationInfo(ctx, b.args[0], b.args[1], row, row, b.collation)) } type builtinGTDurationSig struct { @@ -2675,7 +2675,7 @@ func (b *builtinGEStringSig) Clone() builtinFunc { } func (b *builtinGEStringSig) evalInt(ctx EvalContext, row chunk.Row) (val int64, isNull bool, err error) { - return resOfGE(compareStringWithCollator(ctx, b.args[0], b.args[1], row, row, b.ctor)) + return resOfGE(CompareStringWithCollationInfo(ctx, b.args[0], b.args[1], row, row, b.collation)) } type builtinGEDurationSig struct { @@ -2811,7 +2811,7 @@ func (b *builtinEQStringSig) Clone() builtinFunc { } func (b *builtinEQStringSig) evalInt(ctx EvalContext, row chunk.Row) (val int64, isNull bool, err error) { - return resOfEQ(compareStringWithCollator(ctx, b.args[0], b.args[1], row, row, b.ctor)) + return resOfEQ(CompareStringWithCollationInfo(ctx, b.args[0], b.args[1], row, row, b.collation)) } type builtinEQDurationSig struct { @@ -2947,7 +2947,7 @@ func (b *builtinNEStringSig) Clone() builtinFunc { } func (b *builtinNEStringSig) evalInt(ctx EvalContext, row chunk.Row) (val int64, isNull bool, err error) { - return resOfNE(compareStringWithCollator(ctx, b.args[0], b.args[1], row, row, b.ctor)) + return resOfNE(CompareStringWithCollationInfo(ctx, b.args[0], b.args[1], row, row, b.collation)) } type builtinNEDurationSig struct { @@ -3151,7 +3151,7 @@ func (b *builtinNullEQStringSig) evalInt(ctx EvalContext, row chunk.Row) (val in res = 1 case isNull0 != isNull1: return res, false, nil - case b.ctor.Compare(arg0, arg1) == 0: + case types.CompareString(arg0, arg1, b.collation) == 0: res = 1 } return res, false, nil @@ -3410,21 +3410,13 @@ func CompareInt(sctx EvalContext, lhsArg, rhsArg Expression, lhsRow, rhsRow chun } func genCompareString(collation string) func(sctx EvalContext, lhsArg Expression, rhsArg Expression, lhsRow chunk.Row, rhsRow chunk.Row) (int64, bool, error) { - return genCompareStringWithCollator(collate.GetCollator(collation)) -} - -func genCompareStringWithCollator(collator collate.Collator) func(sctx EvalContext, lhsArg Expression, rhsArg Expression, lhsRow chunk.Row, rhsRow chunk.Row) (int64, bool, error) { return func(sctx EvalContext, lhsArg, rhsArg Expression, lhsRow, rhsRow chunk.Row) (int64, bool, error) { - return compareStringWithCollator(sctx, lhsArg, rhsArg, lhsRow, rhsRow, collator) + return CompareStringWithCollationInfo(sctx, lhsArg, rhsArg, lhsRow, rhsRow, collation) } } // CompareStringWithCollationInfo compares two strings with the specified collation information. func CompareStringWithCollationInfo(sctx EvalContext, lhsArg, rhsArg Expression, lhsRow, rhsRow chunk.Row, collation string) (int64, bool, error) { - return compareStringWithCollator(sctx, lhsArg, rhsArg, lhsRow, rhsRow, collate.GetCollator(collation)) -} - -func compareStringWithCollator(sctx EvalContext, lhsArg, rhsArg Expression, lhsRow, rhsRow chunk.Row, collator collate.Collator) (int64, bool, error) { arg0, isNull0, err := lhsArg.EvalString(sctx, lhsRow) if err != nil { return 0, true, err @@ -3438,7 +3430,7 @@ func compareStringWithCollator(sctx EvalContext, lhsArg, rhsArg Expression, lhsR if isNull0 || isNull1 { return compareNull(isNull0, isNull1), true, nil } - return int64(collator.Compare(arg0, arg1)), false, nil + return int64(types.CompareString(arg0, arg1, collation)), false, nil } // CompareReal compares two float-point values. diff --git a/pkg/expression/builtin_compare_vec.go b/pkg/expression/builtin_compare_vec.go index 6633b94e136e4..d2847e9d8c679 100644 --- a/pkg/expression/builtin_compare_vec.go +++ b/pkg/expression/builtin_compare_vec.go @@ -267,7 +267,7 @@ func (b *builtinLeastStringSig) vecEvalString(ctx EvalContext, input *chunk.Chun } srcStr := src.GetString(i) argStr := arg.GetString(i) - if b.ctor.Compare(srcStr, argStr) < 0 { + if types.CompareString(srcStr, argStr, b.collation) < 0 { dst.AppendString(srcStr) } else { dst.AppendString(argStr) @@ -790,7 +790,7 @@ func (b *builtinGreatestStringSig) vecEvalString(ctx EvalContext, input *chunk.C } srcStr := src.GetString(i) argStr := arg.GetString(i) - if b.ctor.Compare(srcStr, argStr) > 0 { + if types.CompareString(srcStr, argStr, b.collation) > 0 { dst.AppendString(srcStr) } else { dst.AppendString(argStr) diff --git a/pkg/expression/builtin_compare_vec_generated.go b/pkg/expression/builtin_compare_vec_generated.go index f246462d8bdbb..8f53e5c27ccbe 100644 --- a/pkg/expression/builtin_compare_vec_generated.go +++ b/pkg/expression/builtin_compare_vec_generated.go @@ -125,7 +125,7 @@ func (b *builtinLTStringSig) vecEvalInt(ctx EvalContext, input *chunk.Chunk, res if result.IsNull(i) { continue } - val := b.ctor.Compare(buf0.GetString(i), buf1.GetString(i)) + val := types.CompareString(buf0.GetString(i), buf1.GetString(i), b.collation) i64s[i] = boolToInt64(val < 0) } return nil @@ -349,7 +349,7 @@ func (b *builtinLEStringSig) vecEvalInt(ctx EvalContext, input *chunk.Chunk, res if result.IsNull(i) { continue } - val := b.ctor.Compare(buf0.GetString(i), buf1.GetString(i)) + val := types.CompareString(buf0.GetString(i), buf1.GetString(i), b.collation) i64s[i] = boolToInt64(val <= 0) } return nil @@ -573,7 +573,7 @@ func (b *builtinGTStringSig) vecEvalInt(ctx EvalContext, input *chunk.Chunk, res if result.IsNull(i) { continue } - val := b.ctor.Compare(buf0.GetString(i), buf1.GetString(i)) + val := types.CompareString(buf0.GetString(i), buf1.GetString(i), b.collation) i64s[i] = boolToInt64(val > 0) } return nil @@ -797,7 +797,7 @@ func (b *builtinGEStringSig) vecEvalInt(ctx EvalContext, input *chunk.Chunk, res if result.IsNull(i) { continue } - val := b.ctor.Compare(buf0.GetString(i), buf1.GetString(i)) + val := types.CompareString(buf0.GetString(i), buf1.GetString(i), b.collation) i64s[i] = boolToInt64(val >= 0) } return nil @@ -1021,7 +1021,7 @@ func (b *builtinEQStringSig) vecEvalInt(ctx EvalContext, input *chunk.Chunk, res if result.IsNull(i) { continue } - val := b.ctor.Compare(buf0.GetString(i), buf1.GetString(i)) + val := types.CompareString(buf0.GetString(i), buf1.GetString(i), b.collation) i64s[i] = boolToInt64(val == 0) } return nil @@ -1245,7 +1245,7 @@ func (b *builtinNEStringSig) vecEvalInt(ctx EvalContext, input *chunk.Chunk, res if result.IsNull(i) { continue } - val := b.ctor.Compare(buf0.GetString(i), buf1.GetString(i)) + val := types.CompareString(buf0.GetString(i), buf1.GetString(i), b.collation) i64s[i] = boolToInt64(val != 0) } return nil @@ -1480,7 +1480,7 @@ func (b *builtinNullEQStringSig) vecEvalInt(ctx EvalContext, input *chunk.Chunk, i64s[i] = 1 case isNull0 != isNull1: i64s[i] = 0 - case b.ctor.Compare(buf0.GetString(i), buf1.GetString(i)) == 0: + case types.CompareString(buf0.GetString(i), buf1.GetString(i), b.collation) == 0: i64s[i] = 1 } } diff --git a/pkg/expression/builtin_other.go b/pkg/expression/builtin_other.go index 5f4246cd2773c..35e017927f7bf 100644 --- a/pkg/expression/builtin_other.go +++ b/pkg/expression/builtin_other.go @@ -26,6 +26,7 @@ import ( "github.com/pingcap/tidb/pkg/parser/mysql" "github.com/pingcap/tidb/pkg/types" "github.com/pingcap/tidb/pkg/util/chunk" + "github.com/pingcap/tidb/pkg/util/collate" "github.com/pingcap/tidb/pkg/util/intest" "github.com/pingcap/tidb/pkg/util/set" "github.com/pingcap/tidb/pkg/util/stringutil" @@ -338,6 +339,7 @@ type builtinInStringSig struct { func (b *builtinInStringSig) buildHashMapForConstArgs(ctx BuildContext) error { b.nonConstArgsIdx = make([]int, 0) b.hashSet = set.NewStringSet() + collator := collate.GetCollator(b.collation) // Keep track of unique args count for in-place modification uniqueArgCount := 1 // Start with 1 for the first arg (value to check) @@ -361,7 +363,7 @@ func (b *builtinInStringSig) buildHashMapForConstArgs(ctx BuildContext) error { continue } - key := string(b.ctor.Key(val)) // should do memory copy here + key := string(collator.Key(val)) // should do memory copy here // Only keep this arg if value wasn't seen before if !b.hashSet.Exist(key) { b.hashSet.Insert(key) @@ -405,8 +407,9 @@ func (b *builtinInStringSig) evalInt(ctx EvalContext, row chunk.Row) (int64, boo } args := b.args[1:] + collator := collate.GetCollator(b.collation) if len(b.hashSet) != 0 { - if b.hashSet.Exist(string(b.ctor.Key(arg0))) { + if b.hashSet.Exist(string(collator.Key(arg0))) { return 1, false, nil } args = make([]Expression, 0, len(b.nonConstArgsIdx)) @@ -425,7 +428,7 @@ func (b *builtinInStringSig) evalInt(ctx EvalContext, row chunk.Row) (int64, boo hasNull = true continue } - if b.ctor.Compare(arg0, evaledArg) == 0 { + if types.CompareString(arg0, evaledArg, b.collation) == 0 { return 1, false, nil } } diff --git a/pkg/expression/builtin_other_vec_generated.go b/pkg/expression/builtin_other_vec_generated.go index beb54229818c0..0294c85f3759b 100644 --- a/pkg/expression/builtin_other_vec_generated.go +++ b/pkg/expression/builtin_other_vec_generated.go @@ -22,6 +22,7 @@ import ( "github.com/pingcap/tidb/pkg/parser/mysql" "github.com/pingcap/tidb/pkg/types" "github.com/pingcap/tidb/pkg/util/chunk" + "github.com/pingcap/tidb/pkg/util/collate" ) func (b *builtinInIntSig) vecEvalInt(ctx EvalContext, input *chunk.Chunk, result *chunk.Column) error { @@ -159,13 +160,14 @@ func (b *builtinInStringSig) vecEvalInt(ctx EvalContext, input *chunk.Chunk, res var compareResult int args := b.args[1:] if len(b.hashSet) != 0 { + collator := collate.GetCollator(b.collation) for i := 0; i < n; i++ { if buf0.IsNull(i) { hasNull[i] = true continue } arg0 := buf0.GetString(i) - if _, ok := b.hashSet[string(b.ctor.Key(arg0))]; ok { + if _, ok := b.hashSet[string(collator.Key(arg0))]; ok { r64s[i] = 1 result.SetNull(i, false) } @@ -190,7 +192,7 @@ func (b *builtinInStringSig) vecEvalInt(ctx EvalContext, input *chunk.Chunk, res } arg0 := buf0.GetString(i) arg1 := buf1.GetString(i) - compareResult = b.ctor.Compare(arg0, arg1) + compareResult = types.CompareString(arg0, arg1, b.collation) if compareResult == 0 { result.SetNull(i, false) r64s[i] = 1 diff --git a/pkg/expression/builtin_string.go b/pkg/expression/builtin_string.go index 0dd18fbcfd8d3..f6299f988ff66 100644 --- a/pkg/expression/builtin_string.go +++ b/pkg/expression/builtin_string.go @@ -1026,7 +1026,7 @@ func (b *builtinStrcmpSig) evalInt(ctx EvalContext, row chunk.Row) (int64, bool, if isNull || err != nil { return 0, isNull, err } - res := b.ctor.Compare(left, right) + res := types.CompareString(left, right, b.collation) return int64(res), false, nil } @@ -1593,7 +1593,7 @@ func (b *builtinLocate2ArgsUTF8Sig) evalInt(ctx EvalContext, row chunk.Row) (int return 1, false, nil } - return locateStringWithCollator(str, subStr, b.ctor), false, nil + return locateStringWithCollation(str, subStr, b.collation), false, nil } type builtinLocate3ArgsSig struct { @@ -1684,7 +1684,7 @@ func (b *builtinLocate3ArgsUTF8Sig) evalInt(ctx EvalContext, row chunk.Row) (int } slice := string([]rune(str)[pos:]) - idx := locateStringWithCollator(slice, subStr, b.ctor) + idx := locateStringWithCollation(slice, subStr, b.collation) if idx != 0 { return pos + idx, false, nil } @@ -4260,8 +4260,7 @@ func (c *weightStringFunctionClass) getFunction(ctx BuildContext, args []Express sig = &builtinWeightStringNullSig{bf} } else { maxAllowedPacket := ctx.GetEvalCtx().GetMaxAllowedPacket() - argCollator := getCollatorFromBuildContext(ctx, bf.args[0].GetType(ctx.GetEvalCtx()).GetCollate()) - sig = &builtinWeightStringSig{bf, padding, length, maxAllowedPacket, argCollator} + sig = &builtinWeightStringSig{bf, padding, length, maxAllowedPacket} } return sig, nil } @@ -4292,7 +4291,6 @@ type builtinWeightStringSig struct { padding weightStringPadding length int maxAllowedPacket uint64 - argCollator collate.Collator } func (b *builtinWeightStringSig) Clone() builtinFunc { @@ -4301,7 +4299,6 @@ func (b *builtinWeightStringSig) Clone() builtinFunc { newSig.padding = b.padding newSig.length = b.length newSig.maxAllowedPacket = b.maxAllowedPacket - newSig.argCollator = b.argCollator return newSig } @@ -4330,7 +4327,7 @@ func (b *builtinWeightStringSig) evalString(ctx EvalContext, row chunk.Row) (str } str += strings.Repeat(" ", b.length-lenRunes) } - ctor = b.argCollator + ctor = collate.GetCollator(b.args[0].GetType(ctx).GetCollate()) case weightStringPaddingAsBinary: lenStr := len(str) if b.length < lenStr { @@ -4346,7 +4343,7 @@ func (b *builtinWeightStringSig) evalString(ctx EvalContext, row chunk.Row) (str } ctor = collate.GetCollator(charset.CollationBin) case weightStringPaddingNone: - ctor = b.argCollator + ctor = collate.GetCollator(b.args[0].GetType(ctx).GetCollate()) default: return "", false, ErrIncorrectType.GenWithStackByArgs(ast.WeightString, string(b.padding)) } diff --git a/pkg/expression/builtin_string_vec.go b/pkg/expression/builtin_string_vec.go index 8f328b8505969..9de6e92372288 100644 --- a/pkg/expression/builtin_string_vec.go +++ b/pkg/expression/builtin_string_vec.go @@ -1294,7 +1294,7 @@ func (b *builtinStrcmpSig) vecEvalInt(ctx EvalContext, input *chunk.Chunk, resul if result.IsNull(i) { continue } - i64s[i] = int64(b.ctor.Compare(leftBuf.GetString(i), rightBuf.GetString(i))) + i64s[i] = int64(types.CompareString(leftBuf.GetString(i), rightBuf.GetString(i), b.collation)) } return nil } diff --git a/pkg/expression/collation_context.go b/pkg/expression/collation_context.go deleted file mode 100644 index 4c056da53d8bc..0000000000000 --- a/pkg/expression/collation_context.go +++ /dev/null @@ -1,49 +0,0 @@ -// Copyright 2026 PingCAP, Inc. -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -package expression - -import "github.com/pingcap/tidb/pkg/util/collate" - -type collationBuildContext interface { - UseNewCollate() bool -} - -type buildContextWithCollation struct { - BuildContext - useNewCollate bool -} - -func (ctx *buildContextWithCollation) UseNewCollate() bool { - return ctx.useNewCollate -} - -// CtxWithUseNewCollate returns a build context with an explicit collation mode. -func CtxWithUseNewCollate(ctx BuildContext, useNewCollate bool) BuildContext { - return &buildContextWithCollation{ - BuildContext: ctx, - useNewCollate: useNewCollate, - } -} - -func useNewCollateFromBuildContext(ctx BuildContext) bool { - if c, ok := ctx.(collationBuildContext); ok { - return c.UseNewCollate() - } - return collate.NewCollationEnabled() -} - -func getCollatorFromBuildContext(ctx BuildContext, collation string) collate.Collator { - return collate.GetCollatorWithCollate(useNewCollateFromBuildContext(ctx), collation) -} diff --git a/pkg/expression/generator/compare_vec.go b/pkg/expression/generator/compare_vec.go index c651162691de3..37496dd2e7b60 100644 --- a/pkg/expression/generator/compare_vec.go +++ b/pkg/expression/generator/compare_vec.go @@ -94,7 +94,7 @@ func (b *builtin{{ .compare.CompareName }}{{ .type.TypeName }}Sig) vecEvalInt(ct {{- else if eq .type.ETName "Real" }} val := cmp.Compare(arg0[i], arg1[i]) {{- else if eq .type.ETName "String" }} - val := b.ctor.Compare(buf0.GetString(i), buf1.GetString(i)) + val := types.CompareString(buf0.GetString(i), buf1.GetString(i), b.collation) {{- else if eq .type.ETName "Duration" }} val := cmp.Compare(arg0[i], arg1[i]) {{- else if eq .type.ETName "Datetime" }} @@ -151,7 +151,7 @@ func (b *builtin{{ .compare.CompareName }}{{ .type.TypeName }}Sig) vecEvalInt(ct {{- else if eq .type.ETName "Real" }} case cmp.Compare(arg0[i], arg1[i]) == 0: {{- else if eq .type.ETName "String" }} - case b.ctor.Compare(buf0.GetString(i), buf1.GetString(i)) == 0: + case types.CompareString(buf0.GetString(i), buf1.GetString(i), b.collation) == 0: {{- else if eq .type.ETName "Duration" }} case cmp.Compare(arg0[i], arg1[i]) == 0: {{- else if eq .type.ETName "Datetime" }} diff --git a/pkg/expression/generator/other_vec.go b/pkg/expression/generator/other_vec.go index aefcea4a5d23a..19f0338fa534a 100644 --- a/pkg/expression/generator/other_vec.go +++ b/pkg/expression/generator/other_vec.go @@ -112,7 +112,7 @@ var builtinInTmpl = template.Must(template.New("builtinInTmpl").Parse(` {{- else if eq .Input.TypeName "JSON" -}} compareResult = types.CompareBinaryJSON(arg0, arg1) {{- else if eq .Input.TypeName "String" -}} - compareResult = b.ctor.Compare(arg0, arg1) + compareResult = types.CompareString(arg0, arg1, b.collation) {{- else if eq .Input.TypeNameInColumn "Float64" -}} compareResult = cmp.Compare(arg0, arg1) {{- else -}} @@ -153,6 +153,9 @@ func (b *{{.SigName}}) vecEvalInt(ctx EvalContext, input *chunk.Chunk, result *c args := b.args[1:] {{- if not $InputJSON}} if len(b.hashSet) != 0 { + {{- if $InputString }} + collator := collate.GetCollator(b.collation) + {{- end }} for i := 0; i < n; i++ { if buf0.IsNull(i) { hasNull[i] = true @@ -187,7 +190,7 @@ func (b *{{.SigName}}) vecEvalInt(ctx EvalContext, input *chunk.Chunk, result *c result.SetNull(i, false) } {{- else if $InputString }} - if _, ok := b.hashSet[string(b.ctor.Key(arg0))]; ok { + if _, ok := b.hashSet[string(collator.Key(arg0))]; ok { r64s[i] = 1 result.SetNull(i, false) } diff --git a/pkg/expression/util.go b/pkg/expression/util.go index bf8c4ede2eb06..dc537a03327de 100644 --- a/pkg/expression/util.go +++ b/pkg/expression/util.go @@ -818,10 +818,7 @@ func SubstituteCorCol2Constant(ctx BuildContext, expr Expression) (Expression, e } func locateStringWithCollation(str, substr, coll string) int64 { - return locateStringWithCollator(str, substr, collate.GetCollator(coll)) -} - -func locateStringWithCollator(str, substr string, collator collate.Collator) int64 { + collator := collate.GetCollator(coll) strKey := collator.KeyWithoutTrimRightSpace(str) subStrKey := collator.KeyWithoutTrimRightSpace(substr) diff --git a/pkg/planner/core/expression_rewriter.go b/pkg/planner/core/expression_rewriter.go index 089359fba0319..f380aa2bf0147 100644 --- a/pkg/planner/core/expression_rewriter.go +++ b/pkg/planner/core/expression_rewriter.go @@ -118,7 +118,6 @@ func buildSimpleExpr(ctx expression.BuildContext, node ast.ExprNode, opts ...exp for _, opt := range opts { opt(&options) } - ctx = expression.CtxWithUseNewCollate(ctx, options.UseNewCollate) if options.InputSchema == nil && len(options.InputNames) > 0 { return nil, errors.New("InputSchema and InputNames should be specified at the same time") diff --git a/pkg/planner/core/expression_test.go b/pkg/planner/core/expression_test.go index 3a1cb05646f92..766e7c6f1b3e5 100644 --- a/pkg/planner/core/expression_test.go +++ b/pkg/planner/core/expression_test.go @@ -33,7 +33,6 @@ import ( "github.com/pingcap/tidb/pkg/testkit/testutil" "github.com/pingcap/tidb/pkg/types" "github.com/pingcap/tidb/pkg/util/chunk" - "github.com/pingcap/tidb/pkg/util/collate" "github.com/stretchr/testify/require" ) @@ -72,114 +71,6 @@ func buildExprAndEval(t *testing.T, ctx expression.BuildContext, exprNode any) t return val } -func TestBuildSimpleExprWithUseNewCollate(t *testing.T) { - origin := collate.NewCollationEnabled() - collate.SetNewCollationEnabledForTest(true) - defer collate.SetNewCollationEnabledForTest(origin) - - ctx := coretestsdk.MockContext() - defer func() { - domain.GetDomain(ctx).StatsHandle().Close() - }() - - tests := []struct { - name string - expr string - oldResult string - newResult string - }{ - { - name: "eq", - expr: "if(" + - "_utf8mb4'aaa' COLLATE utf8mb4_general_ci = " + - "_utf8mb4'AAA' COLLATE utf8mb4_general_ci, 'match', 'miss')", - oldResult: "miss", - newResult: "match", - }, - { - name: "null safe eq", - expr: "if(" + - "_utf8mb4'aaa' COLLATE utf8mb4_general_ci <=> " + - "_utf8mb4'AAA' COLLATE utf8mb4_general_ci, 'match', 'miss')", - oldResult: "miss", - newResult: "match", - }, - { - name: "in", - expr: "if(" + - "_utf8mb4'aaa' COLLATE utf8mb4_general_ci in " + - "(_utf8mb4'AAA' COLLATE utf8mb4_general_ci), 'match', 'miss')", - oldResult: "miss", - newResult: "match", - }, - { - name: "strcmp", - expr: "if(strcmp(" + - "_utf8mb4'aaa' COLLATE utf8mb4_general_ci, " + - "_utf8mb4'AAA' COLLATE utf8mb4_general_ci) = 0, 'match', 'miss')", - oldResult: "miss", - newResult: "match", - }, - { - name: "least", - expr: "least(" + - "_utf8mb4'aaa' COLLATE utf8mb4_general_ci, " + - "_utf8mb4'AAA' COLLATE utf8mb4_general_ci)", - oldResult: "AAA", - newResult: "aaa", - }, - { - name: "greatest", - expr: "greatest(" + - "_utf8mb4'AAA' COLLATE utf8mb4_general_ci, " + - "_utf8mb4'aaa' COLLATE utf8mb4_general_ci)", - oldResult: "aaa", - newResult: "AAA", - }, - { - name: "locate", - expr: "if(locate(" + - "_utf8mb4'A' COLLATE utf8mb4_general_ci, " + - "_utf8mb4'aaa' COLLATE utf8mb4_general_ci) = 1, 'match', 'miss')", - oldResult: "miss", - newResult: "match", - }, - { - name: "field", - expr: "if(field(" + - "_utf8mb4'aaa' COLLATE utf8mb4_general_ci, " + - "_utf8mb4'AAA' COLLATE utf8mb4_general_ci) = 1, 'match', 'miss')", - oldResult: "miss", - newResult: "match", - }, - { - name: "weight string", - expr: "if(weight_string(" + - "_utf8mb4'aaa' COLLATE utf8mb4_general_ci) = weight_string(" + - "_utf8mb4'AAA' COLLATE utf8mb4_general_ci), 'match', 'miss')", - oldResult: "miss", - newResult: "match", - }, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - exprNode := parseExpr(t, tt.expr) - expr, err := buildExpr(t, ctx, exprNode, expression.WithUseNewCollate(false)) - require.NoError(t, err) - val, err := expr.Eval(ctx.GetExprCtx().GetEvalCtx(), chunk.Row{}) - require.NoError(t, err) - require.Equal(t, tt.oldResult, val.GetString()) - - expr, err = buildExpr(t, ctx, exprNode, expression.WithUseNewCollate(true)) - require.NoError(t, err) - val, err = expr.Eval(ctx.GetExprCtx().GetEvalCtx(), chunk.Row{}) - require.NoError(t, err) - require.Equal(t, tt.newResult, val.GetString()) - }) - } -} - type testCase struct { exprStr string resultStr string From 56df426953f7859017b52579ad2ac70d724f2146 Mon Sep 17 00:00:00 2001 From: Ruihao Chen Date: Mon, 6 Jul 2026 16:08:26 +0800 Subject: [PATCH 24/35] dxf: clarify collation snapshot contract --- pkg/executor/importer/import.go | 16 +++++++++------- pkg/executor/importer/import_test.go | 20 ++++++++++++++++++++ pkg/infoschema/tables.go | 6 ++++-- pkg/meta/model/job_test.go | 24 ++++++++++++++++++++++++ pkg/meta/model/reorg.go | 16 +++++++++------- pkg/table/table.go | 7 ++++++- pkg/table/tables/tables.go | 2 +- pkg/table/tables/tables_test.go | 23 +++++++++++++++++++++++ 8 files changed, 96 insertions(+), 18 deletions(-) diff --git a/pkg/executor/importer/import.go b/pkg/executor/importer/import.go index 20f54fa5a28f0..4257e3919d164 100644 --- a/pkg/executor/importer/import.go +++ b/pkg/executor/importer/import.go @@ -335,10 +335,11 @@ type Plan struct { ManualRecovery bool // the keyspace name when submitting this job, only for import-into Keyspace string - // UseNewCollate captures the collation mode used to build the target table - // snapshot. It is needed because an import task may execute in another - // keyspace with a different collation setting. Nil means old metadata and - // should fall back to the caller-provided default. + // UseNewCollate captures whether the new collation implementation was enabled + // when this import plan's target table snapshot was created. Import execution + // may happen in another keyspace, so key and expression encoding must use this + // captured value instead of the executor process default. Nil means old metadata + // and should fall back to the caller-provided default. UseNewCollate *bool `json:"use_new_collate,omitempty"` } @@ -354,8 +355,8 @@ func (p *Plan) GetOnDupKeyMode() OnDupKeyMode { return p.OnDupKey } -// GetUseNewCollateOrDefault returns the captured collation mode, or defaultVal -// for import metadata generated before the field existed. +// GetUseNewCollateOrDefault returns the captured new-collation mode, or +// defaultVal for import metadata generated before the field existed. func (p *Plan) GetUseNewCollateOrDefault(defaultVal bool) bool { if p.UseNewCollate == nil { return defaultVal @@ -363,7 +364,8 @@ func (p *Plan) GetUseNewCollateOrDefault(defaultVal bool) bool { return *p.UseNewCollate } -// SetUseNewCollate stores the collation mode used to rebuild the target table snapshot. +// SetUseNewCollate stores the new-collation mode captured from the target table +// snapshot. func (p *Plan) SetUseNewCollate(useNewCollate bool) { p.UseNewCollate = &useNewCollate } diff --git a/pkg/executor/importer/import_test.go b/pkg/executor/importer/import_test.go index b7cfa813fe337..50b57d62d18f5 100644 --- a/pkg/executor/importer/import_test.go +++ b/pkg/executor/importer/import_test.go @@ -105,6 +105,26 @@ func TestInitDefaultOptions(t *testing.T) { require.Equal(t, 5, plan.ThreadCnt) } +func TestPlanUseNewCollate(t *testing.T) { + plan := &Plan{} + require.True(t, plan.GetUseNewCollateOrDefault(true)) + require.False(t, plan.GetUseNewCollateOrDefault(false)) + + plan.SetUseNewCollate(false) + require.False(t, plan.GetUseNewCollateOrDefault(true)) + + data, err := json.Marshal(plan) + require.NoError(t, err) + require.Contains(t, string(data), `"use_new_collate":false`) + + var decoded Plan + require.NoError(t, json.Unmarshal(data, &decoded)) + require.False(t, decoded.GetUseNewCollateOrDefault(true)) + + decoded.SetUseNewCollate(true) + require.True(t, decoded.GetUseNewCollateOrDefault(false)) +} + // for negative case see TestImportIntoOptionsNegativeCase func TestInitOptionsPositiveCase(t *testing.T) { sctx := mock.NewContext() diff --git a/pkg/infoschema/tables.go b/pkg/infoschema/tables.go index 017f0f6178d0a..1c80ee1a4b6a3 100644 --- a/pkg/infoschema/tables.go +++ b/pkg/infoschema/tables.go @@ -2648,7 +2648,8 @@ func (it *infoschemaTable) Meta() *model.TableInfo { return it.meta } -// UseNewCollate implements table.Table UseNewCollate interface. +// UseNewCollate implements table.Table UseNewCollate interface. Info schema +// tables are not persisted user tables, so they use the current process default. func (it *infoschemaTable) UseNewCollate() bool { return collate.NewCollationEnabled() } @@ -2751,7 +2752,8 @@ func (vt *VirtualTable) Meta() *model.TableInfo { return nil } -// UseNewCollate implements table.Table UseNewCollate interface. +// UseNewCollate implements table.Table UseNewCollate interface. Virtual tables +// are not persisted user tables, so they use the current process default. func (vt *VirtualTable) UseNewCollate() bool { return collate.NewCollationEnabled() } diff --git a/pkg/meta/model/job_test.go b/pkg/meta/model/job_test.go index 5739b332e3007..3cab8a8153026 100644 --- a/pkg/meta/model/job_test.go +++ b/pkg/meta/model/job_test.go @@ -117,6 +117,30 @@ func TestJobCodec(t *testing.T) { require.Equal(t, int64(3), job.GetRowCount()) } +func TestDDLReorgMetaUseNewCollate(t *testing.T) { + var nilMeta *DDLReorgMeta + require.True(t, nilMeta.GetUseNewCollateOrDefault(true)) + require.False(t, nilMeta.GetUseNewCollateOrDefault(false)) + + meta := &DDLReorgMeta{} + require.True(t, meta.GetUseNewCollateOrDefault(true)) + require.False(t, meta.GetUseNewCollateOrDefault(false)) + + meta.SetUseNewCollate(false) + require.False(t, meta.GetUseNewCollateOrDefault(true)) + + data, err := json.Marshal(meta) + require.NoError(t, err) + require.Contains(t, string(data), `"use_new_collate":false`) + + var decoded DDLReorgMeta + require.NoError(t, json.Unmarshal(data, &decoded)) + require.False(t, decoded.GetUseNewCollateOrDefault(true)) + + decoded.SetUseNewCollate(true) + require.True(t, decoded.GetUseNewCollateOrDefault(false)) +} + func TestLocation(t *testing.T) { // test offset = 0 loc := &TimeZoneLocation{} diff --git a/pkg/meta/model/reorg.go b/pkg/meta/model/reorg.go index 34147d350616f..1a1d030dc5428 100644 --- a/pkg/meta/model/reorg.go +++ b/pkg/meta/model/reorg.go @@ -95,10 +95,11 @@ type DDLReorgMeta struct { MaxNodeCount int `json:"max_node_count"` AnalyzeState int8 `json:"analyze_state"` Stage ReorgStage `json:"stage"` - // UseNewCollate captures the submitting keyspace collation mode for key encoding. - // It is needed because a reorg task may execute in another keyspace with a - // different collation setting. Nil means old metadata and should fall back to - // the caller-provided default. + // UseNewCollate captures whether the new collation implementation was enabled + // when this reorg task's persisted table snapshot was created. Reorg execution + // may happen in another keyspace, so key and expression encoding must use this + // captured value instead of the executor process default. Nil means old metadata + // and should fall back to the caller-provided default. UseNewCollate *bool `json:"use_new_collate,omitempty"` // These two variables are used to control the concurrency and batch size of the reorganization process. // They can be adjusted dynamically through `admin alter ddl jobs` command. @@ -156,8 +157,8 @@ func (dm *DDLReorgMeta) SetMaxWriteSpeed(maxWriteSpeed int) { dm.MaxWriteSpeed.Store(int64(maxWriteSpeed)) } -// GetUseNewCollateOrDefault returns the captured collation mode, or defaultVal -// for reorg metadata generated before the field existed. +// GetUseNewCollateOrDefault returns the captured new-collation mode, or +// defaultVal for reorg metadata generated before the field existed. func (dm *DDLReorgMeta) GetUseNewCollateOrDefault(defaultVal bool) bool { if dm == nil || dm.UseNewCollate == nil { return defaultVal @@ -165,7 +166,8 @@ func (dm *DDLReorgMeta) GetUseNewCollateOrDefault(defaultVal bool) bool { return *dm.UseNewCollate } -// SetUseNewCollate stores the submitting keyspace collation mode for key encoding. +// SetUseNewCollate stores the new-collation mode captured from the persisted +// table snapshot. func (dm *DDLReorgMeta) SetUseNewCollate(useNewCollate bool) { dm.UseNewCollate = &useNewCollate } diff --git a/pkg/table/table.go b/pkg/table/table.go index 130bb18c1d119..7fd8fc9d11ca7 100644 --- a/pkg/table/table.go +++ b/pkg/table/table.go @@ -459,7 +459,12 @@ type Table interface { // Meta returns TableInfo. Meta() *model.TableInfo - // UseNewCollate returns the collation mode used for table key encoding. + // UseNewCollate returns whether the new collation implementation should be used + // when encoding persisted keys or expressions for this table. + // + // For persisted user tables, this value must come from the table snapshot used + // to build the task. Virtual or non-persisted table implementations may use the + // current process default. UseNewCollate() bool // Type returns the type of table diff --git a/pkg/table/tables/tables.go b/pkg/table/tables/tables.go index ba1cb844c443b..69ea903197bb3 100644 --- a/pkg/table/tables/tables.go +++ b/pkg/table/tables/tables.go @@ -169,7 +169,7 @@ func TableFromMeta(allocs autoid.Allocators, tblInfo *model.TableInfo) (table.Ta } // TableFromMetaWithCollate creates a Table instance from model.TableInfo with a -// fixed collation mode. +// fixed new-collation mode for persisted key and expression encoding. func TableFromMetaWithCollate(useNewCollate bool, allocs autoid.Allocators, tblInfo *model.TableInfo) (table.Table, error) { if tblInfo.State == model.StateNone { return nil, table.ErrTableStateCantNone.GenWithStackByArgs(tblInfo.Name) diff --git a/pkg/table/tables/tables_test.go b/pkg/table/tables/tables_test.go index 3e4fd0eca2041..41543f3a913bc 100644 --- a/pkg/table/tables/tables_test.go +++ b/pkg/table/tables/tables_test.go @@ -408,6 +408,29 @@ func TestTableFromMeta(t *testing.T) { require.Error(t, err) } +func TestTableFromMetaWithCollateUsesFixedMode(t *testing.T) { + tblInfo := &model.TableInfo{ + ID: 1, + Name: ast.NewCIStr("t"), + State: model.StatePublic, + Columns: []*model.ColumnInfo{ + { + ID: 1, + Name: ast.NewCIStr("a"), + Offset: 0, + State: model.StatePublic, + FieldType: *types.NewFieldType(mysql.TypeVarchar), + }, + }, + } + + for _, useNewCollate := range []bool{false, true} { + tbl, err := tables.TableFromMetaWithCollate(useNewCollate, autoid.NewAllocators(false), tblInfo) + require.NoError(t, err) + require.Equal(t, useNewCollate, tbl.UseNewCollate()) + } +} + func TestHiddenColumn(t *testing.T) { store, dom := testkit.CreateMockStoreAndDomain(t) tk := testkit.NewTestKit(t, store) From 3e58e05b746829ccd67c3345d18ebd7c409be2d7 Mon Sep 17 00:00:00 2001 From: Ruihao Chen Date: Mon, 6 Jul 2026 16:51:01 +0800 Subject: [PATCH 25/35] bazel: update test shard counts --- pkg/executor/importer/BUILD.bazel | 2 +- pkg/table/tables/BUILD.bazel | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/pkg/executor/importer/BUILD.bazel b/pkg/executor/importer/BUILD.bazel index 005fe5fc0adc9..ee86cb36c5eb6 100644 --- a/pkg/executor/importer/BUILD.bazel +++ b/pkg/executor/importer/BUILD.bazel @@ -120,7 +120,7 @@ go_test( embed = [":importer"], flaky = True, race = "on", - shard_count = 41, + shard_count = 42, deps = [ "//br/pkg/mock", "//br/pkg/streamhelper", diff --git a/pkg/table/tables/BUILD.bazel b/pkg/table/tables/BUILD.bazel index eb6dac9730867..9775792b76383 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 = 43, + shard_count = 44, deps = [ "//pkg/ddl", "//pkg/domain", From 0adebb63e28ea47f1bc6c078fa80540be4dc4361 Mon Sep 17 00:00:00 2001 From: Ruihao Chen Date: Mon, 6 Jul 2026 17:38:21 +0800 Subject: [PATCH 26/35] tests: cover DXF tasks with mismatched collation modes --- pkg/ddl/backfilling_dist_executor.go | 2 + tests/realtikvtest/addindextest1/BUILD.bazel | 3 +- .../addindextest1/cross_ks_test.go | 185 +++++++ .../realtikvtest/importintotest3/BUILD.bazel | 4 +- .../importintotest3/cross_ks_test.go | 486 ++++++++++++++++++ tests/realtikvtest/testkit.go | 38 +- 6 files changed, 714 insertions(+), 4 deletions(-) diff --git a/pkg/ddl/backfilling_dist_executor.go b/pkg/ddl/backfilling_dist_executor.go index 3f0436ff56c23..763299c3ea0b8 100644 --- a/pkg/ddl/backfilling_dist_executor.go +++ b/pkg/ddl/backfilling_dist_executor.go @@ -19,6 +19,7 @@ import ( "encoding/json" "github.com/pingcap/errors" + "github.com/pingcap/failpoint" "github.com/pingcap/tidb/pkg/ddl/logutil" sess "github.com/pingcap/tidb/pkg/ddl/session" "github.com/pingcap/tidb/pkg/dxf/framework/proto" @@ -139,6 +140,7 @@ func (s *backfillDistExecutor) newBackfillStepExecutor( sessPool := sess.NewSessionPool(s.TaskRuntime.SysSessionPool()) // TODO This is using DDL ctx which is never cancelled except when shutdown. // we should move this operation out of GetStepExecutor, and put into Init. + failpoint.InjectCall("beforeGetUserTableForBackfillStep", jobMeta) tblIface, err := getUserTableFromTaskStore(ddlObj.ctx, store, jobMeta) if err != nil { return nil, err diff --git a/tests/realtikvtest/addindextest1/BUILD.bazel b/tests/realtikvtest/addindextest1/BUILD.bazel index f0673be31315e..6a82811448c7d 100644 --- a/tests/realtikvtest/addindextest1/BUILD.bazel +++ b/tests/realtikvtest/addindextest1/BUILD.bazel @@ -9,7 +9,7 @@ go_test( "main_test.go", ], flaky = True, - shard_count = 14, + shard_count = 15, deps = [ "//pkg/config", "//pkg/config/kerneltype", @@ -30,6 +30,7 @@ go_test( "//pkg/store", "//pkg/testkit", "//pkg/testkit/testfailpoint", + "//pkg/util/collate", "//pkg/util/dbterror", "//tests/realtikvtest", "@com_github_pingcap_errors//:errors", diff --git a/tests/realtikvtest/addindextest1/cross_ks_test.go b/tests/realtikvtest/addindextest1/cross_ks_test.go index f03e897f20e97..467dae882843e 100644 --- a/tests/realtikvtest/addindextest1/cross_ks_test.go +++ b/tests/realtikvtest/addindextest1/cross_ks_test.go @@ -17,15 +17,18 @@ package addindextest import ( "fmt" "strconv" + "sync/atomic" "testing" "github.com/pingcap/tidb/pkg/config/kerneltype" "github.com/pingcap/tidb/pkg/ddl" "github.com/pingcap/tidb/pkg/ddl/schemaver" "github.com/pingcap/tidb/pkg/keyspace" + "github.com/pingcap/tidb/pkg/meta/model" kvstore "github.com/pingcap/tidb/pkg/store" "github.com/pingcap/tidb/pkg/testkit" "github.com/pingcap/tidb/pkg/testkit/testfailpoint" + "github.com/pingcap/tidb/pkg/util/collate" "github.com/pingcap/tidb/tests/realtikvtest" "github.com/stretchr/testify/require" ) @@ -60,6 +63,179 @@ func TestAddIndexOnSystemTable(t *testing.T) { tk.MustQuery(taskQuerySQL).Check(testkit.Rows("0")) } +func TestAddIndexOnUserKeyspaceWithDifferentNewCollation(t *testing.T) { + if kerneltype.IsClassic() { + t.Skip("This test is only for nextgen kernel, skip it in classic kernel") + } + + originNewCollationEnabled := collate.NewCollationEnabled() + t.Cleanup(func() { + collate.SetNewCollationEnabledForTest(originNewCollationEnabled) + }) + + const userKeyspace = "keyspacecollateadd" + runtimes := realtikvtest.PrepareForCrossKSTestWithNewCollation(t, map[string]bool{ + keyspace.System: true, + userKeyspace: false, + }, userKeyspace) + userStore := runtimes[userKeyspace].Store + tk := testkit.NewTestKit(t, userStore) + tk.MustQuery(`select variable_value from mysql.tidb where variable_name = 'new_collation_enabled'`). + Check(testkit.Rows("False")) + require.False(t, collate.NewCollationEnabled()) + + var ( + backfillInitCnt atomic.Int64 + metaMatched atomic.Bool + ) + metaMatched.Store(true) + testfailpoint.EnableCall( + t, + "github.com/pingcap/tidb/pkg/ddl/beforeGetUserTableForBackfillStep", + func(job *model.Job) { + if job == nil || job.ReorgMeta == nil || job.ReorgMeta.GetUseNewCollateOrDefault(true) { + metaMatched.Store(false) + } + collate.SetNewCollationEnabledForTest(true) + backfillInitCnt.Add(1) + }, + ) + + tk.PrepareDB("crossks_collate") + require.True(t, metaMatched.Load()) + + cases := []struct { + name string + table string + setupSQL []string + addIndexSQL []string + indexes []string + dmlSQL []string + }{ + { + name: "clustered varchar primary key and secondary varchar index", + table: "t_varchar_pk", + setupSQL: []string{ + "drop table if exists t_varchar_pk", + `create table t_varchar_pk ( + id varchar(32) collate utf8mb4_general_ci, + fk varchar(32) collate utf8mb4_general_ci, + primary key (id) clustered + )`, + "insert into t_varchar_pk values ('aaa', 'abc'), ('bbb', 'bbc'), ('ccc', 'cbc')", + }, + addIndexSQL: []string{ + "alter table t_varchar_pk add index idx_fk(fk)", + }, + indexes: []string{"idx_fk"}, + dmlSQL: []string{ + "insert into t_varchar_pk values ('ddd', 'dbc')", + "update t_varchar_pk set fk = 'updated' where id = 'ddd'", + "delete from t_varchar_pk where id = 'ddd'", + }, + }, + { + name: "composite clustered primary key with varchar part and secondary int index", + table: "t_composite_varchar_pk", + setupSQL: []string{ + "drop table if exists t_composite_varchar_pk", + `create table t_composite_varchar_pk ( + id1 varchar(32) collate utf8mb4_general_ci, + id2 int, + fk int, + primary key (id1, id2) clustered + )`, + "insert into t_composite_varchar_pk values ('ax', 1, 10), ('by', 2, 20), ('cz', 3, 30)", + }, + addIndexSQL: []string{ + "alter table t_composite_varchar_pk add index idx_fk(fk)", + }, + indexes: []string{"idx_fk"}, + dmlSQL: []string{ + "insert into t_composite_varchar_pk values ('dw', 4, 40)", + "update t_composite_varchar_pk set fk = 41 where id1 = 'dw' and id2 = 4", + "delete from t_composite_varchar_pk where id1 = 'dw' and id2 = 4", + }, + }, + { + name: "generated columns with string transformations", + table: "t_add_generated_column_index", + setupSQL: []string{ + "drop table if exists t_add_generated_column_index", + `create table t_add_generated_column_index ( + id varchar(32) collate utf8mb4_general_ci, + raw varchar(32) collate utf8mb4_general_ci, + g_lower varchar(32) generated always as (lower(raw)) virtual, + g_upper varchar(32) generated always as (upper(raw)) virtual, + g_concat varchar(80) generated always as (concat(id, ':', raw)) virtual, + g_substr varchar(32) generated always as (substr(raw, 1, 2)) virtual, + primary key (id) clustered + )`, + "insert into t_add_generated_column_index(id, raw) values ('aaa', 'abc'), ('bbb', 'bbc'), ('ccc', 'cbc')", + }, + addIndexSQL: []string{ + "alter table t_add_generated_column_index add index idx_g_lower(g_lower)", + "alter table t_add_generated_column_index add index idx_g_upper(g_upper)", + "alter table t_add_generated_column_index add index idx_g_concat(g_concat)", + "alter table t_add_generated_column_index add index idx_g_substr(g_substr)", + }, + indexes: []string{"idx_g_lower", "idx_g_upper", "idx_g_concat", "idx_g_substr"}, + dmlSQL: []string{ + "insert into t_add_generated_column_index(id, raw) values ('ddd', 'dbc')", + "update t_add_generated_column_index set raw = 'updated' where id = 'ddd'", + "delete from t_add_generated_column_index where id = 'ddd'", + }, + }, + { + name: "expression indexes with string transformations", + table: "t_add_expression_index", + setupSQL: []string{ + "drop table if exists t_add_expression_index", + `create table t_add_expression_index ( + id varchar(32) collate utf8mb4_general_ci, + raw varchar(32) collate utf8mb4_general_ci, + primary key (id) clustered + )`, + "insert into t_add_expression_index values ('aaa', 'abc'), ('bbb', 'bbc'), ('ccc', 'cbc')", + }, + addIndexSQL: []string{ + "alter table t_add_expression_index add index idx_lower ((lower(raw)))", + "alter table t_add_expression_index add index idx_upper ((upper(raw)))", + "alter table t_add_expression_index add index idx_concat ((concat(id, ':', raw)))", + "alter table t_add_expression_index add index idx_substr ((substr(raw, 1, 2)))", + }, + indexes: []string{"idx_lower", "idx_upper", "idx_concat", "idx_substr"}, + dmlSQL: []string{ + "insert into t_add_expression_index values ('ddd', 'dbc')", + "update t_add_expression_index set raw = 'updated' where id = 'ddd'", + "delete from t_add_expression_index where id = 'ddd'", + }, + }, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + collate.SetNewCollationEnabledForTest(false) + for _, sql := range tc.setupSQL { + tk.MustExec(sql) + } + for _, sql := range tc.addIndexSQL { + before := backfillInitCnt.Load() + collate.SetNewCollationEnabledForTest(false) + tk.MustExec(sql) + require.Greater(t, backfillInitCnt.Load(), before) + require.True(t, metaMatched.Load()) + collate.SetNewCollationEnabledForTest(false) + } + checkTableAndIndexes(tk, tc.table, tc.indexes, "3") + for _, sql := range tc.dmlSQL { + tk.MustExec(sql) + } + checkTableAndIndexes(tk, tc.table, tc.indexes, "3") + }) + } +} + func TestCrossKSInfoSchemaSync(t *testing.T) { if kerneltype.IsClassic() { t.Skip("This test is only for nextgen kernel, skip it in classic kernel") @@ -151,3 +327,12 @@ func TestCrossKSInfoSchemaSync(t *testing.T) { require.EqualValues(t, 2, sum.AssumedServerCount) }) } + +func checkTableAndIndexes(tk *testkit.TestKit, tableName string, indexes []string, expectedCount string) { + tk.MustExec("admin check table " + tableName) + tk.MustQuery("select count(*) from " + tableName).Check(testkit.Rows(expectedCount)) + for _, indexName := range indexes { + tk.MustQuery(fmt.Sprintf("select count(*) from %s force index(%s)", tableName, indexName)). + Check(testkit.Rows(expectedCount)) + } +} diff --git a/tests/realtikvtest/importintotest3/BUILD.bazel b/tests/realtikvtest/importintotest3/BUILD.bazel index 1eddd669eecda..c2262524cc068 100644 --- a/tests/realtikvtest/importintotest3/BUILD.bazel +++ b/tests/realtikvtest/importintotest3/BUILD.bazel @@ -13,7 +13,7 @@ go_test( ], flaky = True, race = "on", - shard_count = 2, + shard_count = 3, deps = [ "//pkg/config", "//pkg/config/deploymode", @@ -24,6 +24,7 @@ go_test( "//pkg/dxf/framework/testutil", "//pkg/dxf/importinto", "//pkg/executor/importer", + "//pkg/keyspace", "//pkg/kv", "//pkg/lightning/mydump", "//pkg/objstore", @@ -33,6 +34,7 @@ go_test( "//pkg/store/driver/error", "//pkg/testkit", "//pkg/testkit/testfailpoint", + "//pkg/util/collate", "//tests/realtikvtest", "@com_github_fsouza_fake_gcs_server//fakestorage", "@com_github_golang_snappy//:snappy", diff --git a/tests/realtikvtest/importintotest3/cross_ks_test.go b/tests/realtikvtest/importintotest3/cross_ks_test.go index 0f32ef77ad522..5623b1a1a2793 100644 --- a/tests/realtikvtest/importintotest3/cross_ks_test.go +++ b/tests/realtikvtest/importintotest3/cross_ks_test.go @@ -20,18 +20,23 @@ import ( "fmt" "strconv" "strings" + "sync/atomic" "testing" "time" "github.com/pingcap/tidb/pkg/config/kerneltype" + "github.com/pingcap/tidb/pkg/dxf/framework/proto" "github.com/pingcap/tidb/pkg/dxf/framework/taskexecutor/execute" "github.com/pingcap/tidb/pkg/dxf/importinto" "github.com/pingcap/tidb/pkg/executor/importer" + "github.com/pingcap/tidb/pkg/keyspace" "github.com/pingcap/tidb/pkg/objstore" plannercore "github.com/pingcap/tidb/pkg/planner/core" "github.com/pingcap/tidb/pkg/sessionctx/vardef" kvstore "github.com/pingcap/tidb/pkg/store" "github.com/pingcap/tidb/pkg/testkit" + "github.com/pingcap/tidb/pkg/testkit/testfailpoint" + "github.com/pingcap/tidb/pkg/util/collate" "github.com/pingcap/tidb/tests/realtikvtest" "github.com/stretchr/testify/require" ) @@ -128,3 +133,484 @@ func TestOnUserKeyspace(t *testing.T) { require.NoError(t, json.Unmarshal([]byte(rs[0][0].(string)), summary)) require.EqualValues(t, rowCount, summary.ImportedRows) } + +func TestImportIntoOnUserKeyspaceWithDifferentNewCollation(t *testing.T) { + if kerneltype.IsClassic() { + t.Skip("only runs in nextgen kernel") + } + + originNewCollationEnabled := collate.NewCollationEnabled() + t.Cleanup(func() { + collate.SetNewCollationEnabledForTest(originNewCollationEnabled) + }) + + const userKeyspace = "keyspacecollateimport" + runtimes := realtikvtest.PrepareForCrossKSTestWithNewCollation(t, map[string]bool{ + keyspace.System: true, + userKeyspace: false, + }, userKeyspace) + userStore := runtimes[userKeyspace].Store + userTK := testkit.NewTestKit(t, userStore) + userTK.MustQuery(`select variable_value from mysql.tidb where variable_name = 'new_collation_enabled'`). + Check(testkit.Rows("False")) + require.False(t, collate.NewCollationEnabled()) + + ctx := context.Background() + s3Args := "access-key=minioadmin&secret-access-key=minioadmin&endpoint=http%3a%2f%2f0.0.0.0%3a9000" + objStore, err := objstore.NewFromURL(ctx, fmt.Sprintf("s3://next-gen-test/collate-data?%s", s3Args)) + require.NoError(t, err) + t.Cleanup(func() { + objStore.Close() + }) + + var ( + prepareCnt atomic.Int64 + metaMatched atomic.Bool + ) + metaMatched.Store(true) + testfailpoint.EnableCall( + t, + "github.com/pingcap/tidb/pkg/dxf/importinto/afterPrepare", + func(task *proto.Task) { + var taskMeta importinto.TaskMeta + if err := json.Unmarshal(task.Meta, &taskMeta); err != nil || + taskMeta.Plan.GetUseNewCollateOrDefault(true) { + metaMatched.Store(false) + } + collate.SetNewCollationEnabledForTest(true) + prepareCnt.Add(1) + }, + ) + + require.True(t, metaMatched.Load()) + prepareAndUseDB("cross_ks_collate", userTK) + + cases := []struct { + name string + table string + setupSQL []string + fileName string + fileData string + importSQL string + indexes []string + postDMLSQL []string + }{ + { + name: "clustered varchar primary key and secondary varchar index", + table: "trigger_varchar_pk_varchar_idx", + setupSQL: []string{ + "drop table if exists trigger_varchar_pk_varchar_idx", + `create table trigger_varchar_pk_varchar_idx ( + id varchar(32) collate utf8mb4_general_ci, + fk varchar(32) collate utf8mb4_general_ci, + primary key (id) clustered, + key idx_fk (fk) + )`, + }, + fileName: "trigger_varchar_pk_varchar_idx.csv", + fileData: "x,aaa,abc\nx,bbb,bbc\nx,ccc,cbc\n", + importSQL: "import into trigger_varchar_pk_varchar_idx(@1,id,fk) from '%s'", + indexes: []string{"idx_fk"}, + postDMLSQL: []string{ + "insert into trigger_varchar_pk_varchar_idx values ('ddd', 'dbc')", + "update trigger_varchar_pk_varchar_idx set fk = 'updated' where id = 'ddd'", + "delete from trigger_varchar_pk_varchar_idx where id = 'ddd'", + }, + }, + { + name: "clustered varchar primary key and secondary int index", + table: "trigger_varchar_pk_int_idx", + setupSQL: []string{ + "drop table if exists trigger_varchar_pk_int_idx", + `create table trigger_varchar_pk_int_idx ( + id varchar(32) collate utf8mb4_general_ci, + fk int, + primary key (id) clustered, + key idx_fk (fk) + )`, + }, + fileName: "trigger_varchar_pk_int_idx.csv", + fileData: "10,aaa,x\n20,bbb,x\n30,ccc,x\n", + importSQL: "import into trigger_varchar_pk_int_idx(fk,id,@3) from '%s'", + indexes: []string{"idx_fk"}, + postDMLSQL: []string{ + "insert into trigger_varchar_pk_int_idx values ('ddd', 40)", + "update trigger_varchar_pk_int_idx set fk = 41 where id = 'ddd'", + "delete from trigger_varchar_pk_int_idx where id = 'ddd'", + }, + }, + { + name: "composite clustered primary key with varchar part and secondary int index", + table: "trigger_composite_varchar_int_pk_int_idx", + setupSQL: []string{ + "drop table if exists trigger_composite_varchar_int_pk_int_idx", + `create table trigger_composite_varchar_int_pk_int_idx ( + id1 varchar(32) collate utf8mb4_general_ci, + id2 int, + fk int, + primary key (id1, id2) clustered, + key idx_fk (fk) + )`, + }, + fileName: "trigger_composite_varchar_int_pk_int_idx.csv", + fileData: "1,10,aaa\n2,20,bbb\n3,30,ccc\n", + importSQL: "import into trigger_composite_varchar_int_pk_int_idx(id2,fk,id1) from '%s'", + indexes: []string{"idx_fk"}, + postDMLSQL: []string{ + "insert into trigger_composite_varchar_int_pk_int_idx values ('ddd', 4, 40)", + "update trigger_composite_varchar_int_pk_int_idx set fk = 41 where id1 = 'ddd' and id2 = 4", + "delete from trigger_composite_varchar_int_pk_int_idx where id1 = 'ddd' and id2 = 4", + }, + }, + { + name: "composite clustered int primary key and secondary varchar index", + table: "trigger_composite_int_int_pk_varchar_idx", + setupSQL: []string{ + "drop table if exists trigger_composite_int_int_pk_varchar_idx", + `create table trigger_composite_int_int_pk_varchar_idx ( + id1 int, + id2 int, + fk varchar(32) collate utf8mb4_general_ci, + primary key (id1, id2) clustered, + key idx_fk (fk) + )`, + }, + fileName: "trigger_composite_int_int_pk_varchar_idx.csv", + fileData: "1,10,aaa\n2,20,bbb\n3,30,ccc\n", + importSQL: "import into trigger_composite_int_int_pk_varchar_idx(id1,id2,fk) from '%s'", + indexes: []string{"idx_fk"}, + postDMLSQL: []string{ + "insert into trigger_composite_int_int_pk_varchar_idx values (4, 40, 'ddd')", + "update trigger_composite_int_int_pk_varchar_idx set fk = 'updated' where id1 = 4 and id2 = 40", + "delete from trigger_composite_int_int_pk_varchar_idx where id1 = 4 and id2 = 40", + }, + }, + { + name: "composite char primary key and secondary varchar index", + table: "trigger_composite_char_char_pk_varchar_idx", + setupSQL: []string{ + "drop table if exists trigger_composite_char_char_pk_varchar_idx", + `create table trigger_composite_char_char_pk_varchar_idx ( + id1 char(32) collate utf8mb4_general_ci, + id2 char(32) collate utf8mb4_general_ci, + fk varchar(32) collate utf8mb4_general_ci, + primary key (id1, id2) clustered, + key idx_fk (fk) + )`, + }, + fileName: "trigger_composite_char_char_pk_varchar_idx.csv", + fileData: "aaa,ax,ay\nbbb,bx,by\nccc,cx,cy\n", + importSQL: "import into trigger_composite_char_char_pk_varchar_idx(fk,id1,id2) from '%s'", + indexes: []string{"idx_fk"}, + postDMLSQL: []string{ + "insert into trigger_composite_char_char_pk_varchar_idx values ('dx', 'dy', 'ddd')", + "update trigger_composite_char_char_pk_varchar_idx set fk = 'updated' where id1 = 'dx' and id2 = 'dy'", + "delete from trigger_composite_char_char_pk_varchar_idx where id1 = 'dx' and id2 = 'dy'", + }, + }, + { + name: "clustered varchar primary key and prefix secondary varchar index", + table: "trigger_varchar_pk_prefix_varchar_idx", + setupSQL: []string{ + "drop table if exists trigger_varchar_pk_prefix_varchar_idx", + `create table trigger_varchar_pk_prefix_varchar_idx ( + id varchar(32) collate utf8mb4_general_ci, + fk varchar(32) collate utf8mb4_general_ci, + primary key (id) clustered, + key idx_fk_prefix (fk(2)) + )`, + }, + fileName: "trigger_varchar_pk_prefix_varchar_idx.csv", + fileData: "x,aaa,abc\nx,bbb,bbc\nx,ccc,cbc\n", + importSQL: "import into trigger_varchar_pk_prefix_varchar_idx(@1,id,fk) from '%s'", + indexes: []string{"idx_fk_prefix"}, + postDMLSQL: []string{ + "insert into trigger_varchar_pk_prefix_varchar_idx values ('ddd', 'dbc')", + "update trigger_varchar_pk_prefix_varchar_idx set fk = 'updated' where id = 'ddd'", + "delete from trigger_varchar_pk_prefix_varchar_idx where id = 'ddd'", + }, + }, + { + name: "clustered varchar primary key and secondary varchar index with extra payload", + table: "trigger_record_ok_index_bad_extra_payload", + setupSQL: []string{ + "drop table if exists trigger_record_ok_index_bad_extra_payload", + `create table trigger_record_ok_index_bad_extra_payload ( + id varchar(32) collate utf8mb4_general_ci, + fk varchar(32) collate utf8mb4_general_ci, + payload varchar(32) collate utf8mb4_general_ci default 'payload', + primary key (id) clustered, + key idx_fk (fk) + )`, + }, + fileName: "trigger_record_ok_index_bad_extra_payload.csv", + fileData: "x,aaa,abc\nx,bbb,bbc\nx,ccc,cbc\n", + importSQL: "import into trigger_record_ok_index_bad_extra_payload(@1,id,fk) from '%s'", + indexes: []string{"idx_fk"}, + postDMLSQL: []string{ + "insert into trigger_record_ok_index_bad_extra_payload(id, fk) values ('ddd', 'dbc')", + "update trigger_record_ok_index_bad_extra_payload set fk = 'updated' where id = 'ddd'", + "delete from trigger_record_ok_index_bad_extra_payload where id = 'ddd'", + }, + }, + { + name: "generated columns with string transformations", + table: "t_import_generated", + setupSQL: []string{ + "drop table if exists t_import_generated", + `create table t_import_generated ( + id varchar(32) collate utf8mb4_general_ci, + raw varchar(32) collate utf8mb4_general_ci, + g_lower varchar(32) generated always as (lower(raw)) stored, + g_upper varchar(32) generated always as (upper(raw)) stored, + g_concat varchar(80) generated always as (concat(id, ':', raw)) stored, + g_substr varchar(32) generated always as (substr(raw, 1, 2)) stored, + primary key (id) clustered, + key idx_lower (g_lower), + key idx_upper (g_upper), + key idx_concat (g_concat), + key idx_substr (g_substr) + )`, + }, + fileName: "t_import_generated.csv", + fileData: "x,aaa,abc\nx,bbb,bbc\nx,ccc,cbc\n", + importSQL: "import into t_import_generated(@1,id,raw) from '%s'", + indexes: []string{"idx_lower", "idx_upper", "idx_concat", "idx_substr"}, + postDMLSQL: []string{ + "insert into t_import_generated(id, raw) values ('ddd', 'dbc')", + "update t_import_generated set raw = 'updated' where id = 'ddd'", + "delete from t_import_generated where id = 'ddd'", + }, + }, + { + name: "assignment expressions with string transformations", + table: "t_import_assignment", + setupSQL: []string{ + "drop table if exists t_import_assignment", + `create table t_import_assignment ( + id varchar(32) collate utf8mb4_general_ci, + raw varchar(32) collate utf8mb4_general_ci, + a_lower varchar(32) collate utf8mb4_general_ci, + a_upper varchar(32) collate utf8mb4_general_ci, + a_concat varchar(80) collate utf8mb4_general_ci, + a_substr varchar(32) collate utf8mb4_general_ci, + primary key (id) clustered, + key idx_lower (a_lower), + key idx_upper (a_upper), + key idx_concat (a_concat), + key idx_substr (a_substr) + )`, + }, + fileName: "t_import_assignment.csv", + fileData: "x,aaa,abc\nx,bbb,bbc\nx,ccc,cbc\n", + importSQL: `import into t_import_assignment(@1,@2,@3) + set id=@2, + raw=@3, + a_lower=lower(@3), + a_upper=upper(@3), + a_concat=concat(@2, ':', @3), + a_substr=substr(@3, 1, 2) + from '%s'`, + indexes: []string{"idx_lower", "idx_upper", "idx_concat", "idx_substr"}, + postDMLSQL: []string{ + `insert into t_import_assignment + values ('ddd', 'dbc', lower('dbc'), upper('dbc'), concat('ddd', ':', 'dbc'), substr('dbc', 1, 2))`, + `update t_import_assignment + set raw = 'updated', + a_lower = lower('updated'), + a_upper = upper('updated'), + a_concat = concat(id, ':', 'updated'), + a_substr = substr('updated', 1, 2) + where id = 'ddd'`, + "delete from t_import_assignment where id = 'ddd'", + }, + }, + { + name: "control int primary key and varchar secondary index", + table: "ok_int_pk_varchar_idx", + setupSQL: []string{ + "drop table if exists ok_int_pk_varchar_idx", + `create table ok_int_pk_varchar_idx ( + id int, + fk varchar(32) collate utf8mb4_general_ci, + primary key (id) clustered, + key idx_fk (fk) + )`, + }, + fileName: "ok_int_pk_varchar_idx.csv", + fileData: "1,abc\n2,bbc\n3,cbc\n", + importSQL: "import into ok_int_pk_varchar_idx(id,fk) from '%s'", + indexes: []string{"idx_fk"}, + postDMLSQL: []string{ + "insert into ok_int_pk_varchar_idx values (4, 'dbc')", + "update ok_int_pk_varchar_idx set fk = 'updated' where id = 4", + "delete from ok_int_pk_varchar_idx where id = 4", + }, + }, + { + name: "control varchar primary key without secondary index", + table: "ok_varchar_pk_no_secondary", + setupSQL: []string{ + "drop table if exists ok_varchar_pk_no_secondary", + `create table ok_varchar_pk_no_secondary ( + id varchar(32) collate utf8mb4_general_ci, + v int, + primary key (id) clustered + )`, + }, + fileName: "ok_varchar_pk_no_secondary.csv", + fileData: "aaa,1\nbbb,2\nccc,3\n", + importSQL: "import into ok_varchar_pk_no_secondary(id,v) from '%s'", + postDMLSQL: []string{ + "insert into ok_varchar_pk_no_secondary values ('ddd', 4)", + "update ok_varchar_pk_no_secondary set v = 5 where id = 'ddd'", + "delete from ok_varchar_pk_no_secondary where id = 'ddd'", + }, + }, + { + name: "control nonclustered varchar primary key and varchar secondary index", + table: "ok_nonclustered_varchar_pk_varchar_idx", + setupSQL: []string{ + "drop table if exists ok_nonclustered_varchar_pk_varchar_idx", + `create table ok_nonclustered_varchar_pk_varchar_idx ( + id varchar(32) collate utf8mb4_general_ci, + fk varchar(32) collate utf8mb4_general_ci, + primary key (id) nonclustered, + key idx_fk (fk) + )`, + }, + fileName: "ok_nonclustered_varchar_pk_varchar_idx.csv", + fileData: "aaa,abc\nbbb,bbc\nccc,cbc\n", + importSQL: "import into ok_nonclustered_varchar_pk_varchar_idx(id,fk) from '%s'", + indexes: []string{"idx_fk"}, + postDMLSQL: []string{ + "insert into ok_nonclustered_varchar_pk_varchar_idx values ('ddd', 'dbc')", + "update ok_nonclustered_varchar_pk_varchar_idx set fk = 'updated' where id = 'ddd'", + "delete from ok_nonclustered_varchar_pk_varchar_idx where id = 'ddd'", + }, + }, + { + name: "control char primary key and char secondary index", + table: "ok_char_pk_char_idx", + setupSQL: []string{ + "drop table if exists ok_char_pk_char_idx", + `create table ok_char_pk_char_idx ( + id char(32) collate utf8mb4_general_ci, + fk char(32) collate utf8mb4_general_ci, + primary key (id) clustered, + key idx_fk (fk) + )`, + }, + fileName: "ok_char_pk_char_idx.csv", + fileData: "aaa,abc\nbbb,bbc\nccc,cbc\n", + importSQL: "import into ok_char_pk_char_idx(id,fk) from '%s'", + indexes: []string{"idx_fk"}, + postDMLSQL: []string{ + "insert into ok_char_pk_char_idx values ('ddd', 'dbc')", + "update ok_char_pk_char_idx set fk = 'updated' where id = 'ddd'", + "delete from ok_char_pk_char_idx where id = 'ddd'", + }, + }, + { + name: "control composite int primary key and int secondary index", + table: "ok_composite_int_int_pk_int_idx", + setupSQL: []string{ + "drop table if exists ok_composite_int_int_pk_int_idx", + `create table ok_composite_int_int_pk_int_idx ( + id1 int, + id2 int, + fk int, + primary key (id1, id2) clustered, + key idx_fk (fk) + )`, + }, + fileName: "ok_composite_int_int_pk_int_idx.csv", + fileData: "1,10,100\n2,20,200\n3,30,300\n", + importSQL: "import into ok_composite_int_int_pk_int_idx(id1,id2,fk) from '%s'", + indexes: []string{"idx_fk"}, + postDMLSQL: []string{ + "insert into ok_composite_int_int_pk_int_idx values (4, 40, 400)", + "update ok_composite_int_int_pk_int_idx set fk = 401 where id1 = 4 and id2 = 40", + "delete from ok_composite_int_int_pk_int_idx where id1 = 4 and id2 = 40", + }, + }, + { + name: "control composite varbinary int primary key and int secondary index", + table: "ok_composite_varbinary_int_pk_int_idx", + setupSQL: []string{ + "drop table if exists ok_composite_varbinary_int_pk_int_idx", + `create table ok_composite_varbinary_int_pk_int_idx ( + id1 varbinary(32), + id2 int, + fk int, + primary key (id1, id2) clustered, + key idx_fk (fk) + )`, + }, + fileName: "ok_composite_varbinary_int_pk_int_idx.csv", + fileData: "aaa,1,10\nbbb,2,20\nccc,3,30\n", + importSQL: "import into ok_composite_varbinary_int_pk_int_idx(id1,id2,fk) from '%s'", + indexes: []string{"idx_fk"}, + postDMLSQL: []string{ + "insert into ok_composite_varbinary_int_pk_int_idx values ('ddd', 4, 40)", + "update ok_composite_varbinary_int_pk_int_idx set fk = 41 where id1 = 'ddd' and id2 = 4", + "delete from ok_composite_varbinary_int_pk_int_idx where id1 = 'ddd' and id2 = 4", + }, + }, + { + name: "control composite char primary key and int secondary index", + table: "ok_composite_char_char_pk_int_idx", + setupSQL: []string{ + "drop table if exists ok_composite_char_char_pk_int_idx", + `create table ok_composite_char_char_pk_int_idx ( + id1 char(32) collate utf8mb4_general_ci, + id2 char(32) collate utf8mb4_general_ci, + fk int, + primary key (id1, id2) clustered, + key idx_fk (fk) + )`, + }, + fileName: "ok_composite_char_char_pk_int_idx.csv", + fileData: "aaa,ax,10\nbbb,bx,20\nccc,cx,30\n", + importSQL: "import into ok_composite_char_char_pk_int_idx(id1,id2,fk) from '%s'", + indexes: []string{"idx_fk"}, + postDMLSQL: []string{ + "insert into ok_composite_char_char_pk_int_idx values ('ddd', 'dx', 40)", + "update ok_composite_char_char_pk_int_idx set fk = 41 where id1 = 'ddd' and id2 = 'dx'", + "delete from ok_composite_char_char_pk_int_idx where id1 = 'ddd' and id2 = 'dx'", + }, + }, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + collate.SetNewCollationEnabledForTest(false) + for _, sql := range tc.setupSQL { + userTK.MustExec(sql) + } + require.NoError(t, objStore.WriteFile(ctx, tc.fileName, []byte(tc.fileData))) + before := prepareCnt.Load() + fileURL := fmt.Sprintf("s3://next-gen-test/collate-data/%s?%s", tc.fileName, s3Args) + result := userTK.MustQuery(fmt.Sprintf(tc.importSQL, fileURL)).Rows() + require.Len(t, result, 1) + require.Greater(t, prepareCnt.Load(), before) + require.True(t, metaMatched.Load()) + + collate.SetNewCollationEnabledForTest(false) + checkImportTableAndIndexes(userTK, tc.table, tc.indexes, "3") + for _, sql := range tc.postDMLSQL { + userTK.MustExec(sql) + } + checkImportTableAndIndexes(userTK, tc.table, tc.indexes, "3") + }) + } +} + +func checkImportTableAndIndexes(tk *testkit.TestKit, tableName string, indexes []string, expectedCount string) { + tk.MustExec("admin check table " + tableName) + tk.MustQuery("select count(*) from " + tableName).Check(testkit.Rows(expectedCount)) + for _, indexName := range indexes { + tk.MustQuery(fmt.Sprintf("select count(*) from %s force index(%s)", tableName, indexName)). + Check(testkit.Rows(expectedCount)) + } +} diff --git a/tests/realtikvtest/testkit.go b/tests/realtikvtest/testkit.go index a764f9a98a0fd..2842a65a8fe4a 100644 --- a/tests/realtikvtest/testkit.go +++ b/tests/realtikvtest/testkit.go @@ -106,6 +106,9 @@ func RunTestMain(m *testing.M) { type realtikvStoreOption struct { retainData bool keyspace string + // newCollationsEnabledOnFirstBootstrap is nil unless the test wants to + // bootstrap a keyspace with a specific persisted new-collation setting. + newCollationsEnabledOnFirstBootstrap *bool // only used when keyspace is not SYSTEM, in that case, the SYSTEM store will // be closed together with its domain, if we close it before domain of SYSTEM, // some routine might report errors, and we don't want close twice as the storage @@ -136,6 +139,14 @@ func WithKeyspaceName(name string) RealTiKVStoreOption { } } +// WithNewCollationsEnabledOnFirstBootstrap bootstraps a real TiKV test store +// with the requested persisted new-collation setting. +func WithNewCollationsEnabledOnFirstBootstrap(enabled bool) RealTiKVStoreOption { + return func(opt *realtikvStoreOption) { + opt.newCollationsEnabledOnFirstBootstrap = &enabled + } +} + // WithKeepSystemStore allows the store to keep the SYSTEM keyspace store func WithKeepSystemStore(keep bool) RealTiKVStoreOption { return func(opt *realtikvStoreOption) { @@ -165,6 +176,16 @@ type KSRuntime struct { // PrepareForCrossKSTest prepares the environment for cross keyspace tests. func PrepareForCrossKSTest(t *testing.T, userKSs ...string) map[string]*KSRuntime { + return PrepareForCrossKSTestWithNewCollation(t, nil, userKSs...) +} + +// PrepareForCrossKSTestWithNewCollation prepares cross-keyspace runtimes with +// optional per-keyspace persisted new-collation settings. +func PrepareForCrossKSTestWithNewCollation( + t *testing.T, + newCollationEnabled map[string]bool, + userKSs ...string, +) map[string]*KSRuntime { if !kerneltype.IsNextGen() { t.Fail() } @@ -179,8 +200,18 @@ func PrepareForCrossKSTest(t *testing.T, userKSs ...string) map[string]*KSRuntim ksList := append([]string{keyspace.System}, userKSs...) for _, ks := range ksList { - store, dom := CreateMockStoreAndDomainAndSetup(t, WithKeyspaceName(ks), - WithKeepSystemStore(true), WithKeepSelfStore(true), WithAllocPort(true)) + opts := []RealTiKVStoreOption{ + WithKeyspaceName(ks), + WithKeepSystemStore(true), + WithKeepSelfStore(true), + WithAllocPort(true), + } + if newCollationEnabled != nil { + if enabled, ok := newCollationEnabled[ks]; ok { + opts = append(opts, WithNewCollationsEnabledOnFirstBootstrap(enabled)) + } + } + store, dom := CreateMockStoreAndDomainAndSetup(t, opts...) res[ks] = &KSRuntime{ Store: store, Dom: dom, @@ -239,6 +270,9 @@ func CreateMockStoreAndDomainAndSetup(t *testing.T, opts ...RealTiKVStoreOption) conf.TxnLocalLatches.Enabled = false conf.KeyspaceName = ks conf.Store = config.StoreTypeTiKV + if option.newCollationsEnabledOnFirstBootstrap != nil { + conf.NewCollationsEnabledOnFirstBootstrap = *option.newCollationsEnabledOnFirstBootstrap + } if option.allocPort { conf.Port = uint(mockPortAlloc.Add(1)) } From 67428e420c34c4fa007e2807e1bae890a284793b Mon Sep 17 00:00:00 2001 From: Ruihao Chen Date: Mon, 6 Jul 2026 17:46:52 +0800 Subject: [PATCH 27/35] tests: switch collation before DXF task visibility --- .../addindextest1/cross_ks_test.go | 26 +++++++++++++------ 1 file changed, 18 insertions(+), 8 deletions(-) diff --git a/tests/realtikvtest/addindextest1/cross_ks_test.go b/tests/realtikvtest/addindextest1/cross_ks_test.go index 467dae882843e..f554de21d8652 100644 --- a/tests/realtikvtest/addindextest1/cross_ks_test.go +++ b/tests/realtikvtest/addindextest1/cross_ks_test.go @@ -23,6 +23,7 @@ import ( "github.com/pingcap/tidb/pkg/config/kerneltype" "github.com/pingcap/tidb/pkg/ddl" "github.com/pingcap/tidb/pkg/ddl/schemaver" + "github.com/pingcap/tidb/pkg/dxf/framework/proto" "github.com/pingcap/tidb/pkg/keyspace" "github.com/pingcap/tidb/pkg/meta/model" kvstore "github.com/pingcap/tidb/pkg/store" @@ -85,24 +86,32 @@ func TestAddIndexOnUserKeyspaceWithDifferentNewCollation(t *testing.T) { require.False(t, collate.NewCollationEnabled()) var ( - backfillInitCnt atomic.Int64 - metaMatched atomic.Bool + backfillInitCnt atomic.Int64 + unexpectedUseNewCollateMeta atomic.Bool + unexpectedExecutorGlobal atomic.Bool + ) + testfailpoint.EnableCall( + t, + "github.com/pingcap/tidb/pkg/dxf/framework/storage/beforeSubmitTask", + func(*int, *proto.ExtraParams) { + collate.SetNewCollationEnabledForTest(true) + }, ) - metaMatched.Store(true) testfailpoint.EnableCall( t, "github.com/pingcap/tidb/pkg/ddl/beforeGetUserTableForBackfillStep", func(job *model.Job) { - if job == nil || job.ReorgMeta == nil || job.ReorgMeta.GetUseNewCollateOrDefault(true) { - metaMatched.Store(false) + if job.ReorgMeta.GetUseNewCollateOrDefault(true) { + unexpectedUseNewCollateMeta.Store(true) + } + if !collate.NewCollationEnabled() { + unexpectedExecutorGlobal.Store(true) } - collate.SetNewCollationEnabledForTest(true) backfillInitCnt.Add(1) }, ) tk.PrepareDB("crossks_collate") - require.True(t, metaMatched.Load()) cases := []struct { name string @@ -224,7 +233,8 @@ func TestAddIndexOnUserKeyspaceWithDifferentNewCollation(t *testing.T) { collate.SetNewCollationEnabledForTest(false) tk.MustExec(sql) require.Greater(t, backfillInitCnt.Load(), before) - require.True(t, metaMatched.Load()) + require.False(t, unexpectedUseNewCollateMeta.Load()) + require.False(t, unexpectedExecutorGlobal.Load()) collate.SetNewCollationEnabledForTest(false) } checkTableAndIndexes(tk, tc.table, tc.indexes, "3") From 94bf00f60ebbd0e57ed4a23da8a52cc87f820d30 Mon Sep 17 00:00:00 2001 From: Ruihao Chen Date: Mon, 6 Jul 2026 17:48:55 +0800 Subject: [PATCH 28/35] tests: assert ADD INDEX collation state directly --- .../realtikvtest/addindextest1/cross_ks_test.go | 16 +++------------- 1 file changed, 3 insertions(+), 13 deletions(-) diff --git a/tests/realtikvtest/addindextest1/cross_ks_test.go b/tests/realtikvtest/addindextest1/cross_ks_test.go index f554de21d8652..b006bae815b29 100644 --- a/tests/realtikvtest/addindextest1/cross_ks_test.go +++ b/tests/realtikvtest/addindextest1/cross_ks_test.go @@ -85,11 +85,7 @@ func TestAddIndexOnUserKeyspaceWithDifferentNewCollation(t *testing.T) { Check(testkit.Rows("False")) require.False(t, collate.NewCollationEnabled()) - var ( - backfillInitCnt atomic.Int64 - unexpectedUseNewCollateMeta atomic.Bool - unexpectedExecutorGlobal atomic.Bool - ) + var backfillInitCnt atomic.Int64 testfailpoint.EnableCall( t, "github.com/pingcap/tidb/pkg/dxf/framework/storage/beforeSubmitTask", @@ -101,12 +97,8 @@ func TestAddIndexOnUserKeyspaceWithDifferentNewCollation(t *testing.T) { t, "github.com/pingcap/tidb/pkg/ddl/beforeGetUserTableForBackfillStep", func(job *model.Job) { - if job.ReorgMeta.GetUseNewCollateOrDefault(true) { - unexpectedUseNewCollateMeta.Store(true) - } - if !collate.NewCollationEnabled() { - unexpectedExecutorGlobal.Store(true) - } + require.False(t, job.ReorgMeta.GetUseNewCollateOrDefault(true)) + require.True(t, collate.NewCollationEnabled()) backfillInitCnt.Add(1) }, ) @@ -233,8 +225,6 @@ func TestAddIndexOnUserKeyspaceWithDifferentNewCollation(t *testing.T) { collate.SetNewCollationEnabledForTest(false) tk.MustExec(sql) require.Greater(t, backfillInitCnt.Load(), before) - require.False(t, unexpectedUseNewCollateMeta.Load()) - require.False(t, unexpectedExecutorGlobal.Load()) collate.SetNewCollationEnabledForTest(false) } checkTableAndIndexes(tk, tc.table, tc.indexes, "3") From 69efbf03c1a440f7d335f1523bec918d89a7c162 Mon Sep 17 00:00:00 2001 From: Ruihao Chen Date: Mon, 6 Jul 2026 17:52:16 +0800 Subject: [PATCH 29/35] tests: assert IMPORT collation state directly --- .../importintotest3/cross_ks_test.go | 21 +++++++++---------- 1 file changed, 10 insertions(+), 11 deletions(-) diff --git a/tests/realtikvtest/importintotest3/cross_ks_test.go b/tests/realtikvtest/importintotest3/cross_ks_test.go index 5623b1a1a2793..2ca48d16e90f2 100644 --- a/tests/realtikvtest/importintotest3/cross_ks_test.go +++ b/tests/realtikvtest/importintotest3/cross_ks_test.go @@ -163,26 +163,26 @@ func TestImportIntoOnUserKeyspaceWithDifferentNewCollation(t *testing.T) { objStore.Close() }) - var ( - prepareCnt atomic.Int64 - metaMatched atomic.Bool + var prepareCnt atomic.Int64 + testfailpoint.EnableCall( + t, + "github.com/pingcap/tidb/pkg/dxf/framework/storage/beforeSubmitTask", + func(*int, *proto.ExtraParams) { + collate.SetNewCollationEnabledForTest(true) + }, ) - metaMatched.Store(true) testfailpoint.EnableCall( t, "github.com/pingcap/tidb/pkg/dxf/importinto/afterPrepare", func(task *proto.Task) { var taskMeta importinto.TaskMeta - if err := json.Unmarshal(task.Meta, &taskMeta); err != nil || - taskMeta.Plan.GetUseNewCollateOrDefault(true) { - metaMatched.Store(false) - } - collate.SetNewCollationEnabledForTest(true) + require.NoError(t, json.Unmarshal(task.Meta, &taskMeta)) + require.False(t, taskMeta.Plan.GetUseNewCollateOrDefault(true)) + require.True(t, collate.NewCollationEnabled()) prepareCnt.Add(1) }, ) - require.True(t, metaMatched.Load()) prepareAndUseDB("cross_ks_collate", userTK) cases := []struct { @@ -594,7 +594,6 @@ func TestImportIntoOnUserKeyspaceWithDifferentNewCollation(t *testing.T) { result := userTK.MustQuery(fmt.Sprintf(tc.importSQL, fileURL)).Rows() require.Len(t, result, 1) require.Greater(t, prepareCnt.Load(), before) - require.True(t, metaMatched.Load()) collate.SetNewCollationEnabledForTest(false) checkImportTableAndIndexes(userTK, tc.table, tc.indexes, "3") From 7e9637b59e324d5c4d954054b0d1130cf3894c87 Mon Sep 17 00:00:00 2001 From: Ruihao Chen Date: Mon, 6 Jul 2026 17:55:22 +0800 Subject: [PATCH 30/35] Update pkg/meta/model/reorg.go Co-authored-by: D3Hunter --- pkg/meta/model/reorg.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pkg/meta/model/reorg.go b/pkg/meta/model/reorg.go index 1a1d030dc5428..644a97474fa1b 100644 --- a/pkg/meta/model/reorg.go +++ b/pkg/meta/model/reorg.go @@ -160,7 +160,7 @@ func (dm *DDLReorgMeta) SetMaxWriteSpeed(maxWriteSpeed int) { // GetUseNewCollateOrDefault returns the captured new-collation mode, or // defaultVal for reorg metadata generated before the field existed. func (dm *DDLReorgMeta) GetUseNewCollateOrDefault(defaultVal bool) bool { - if dm == nil || dm.UseNewCollate == nil { + if dm.UseNewCollate == nil { return defaultVal } return *dm.UseNewCollate From f95c567a33851faa5355f66bef9f323146957957 Mon Sep 17 00:00:00 2001 From: Ruihao Chen Date: Mon, 6 Jul 2026 18:02:47 +0800 Subject: [PATCH 31/35] model: align collation metadata test with API contract --- pkg/meta/model/job_test.go | 4 ---- 1 file changed, 4 deletions(-) diff --git a/pkg/meta/model/job_test.go b/pkg/meta/model/job_test.go index 3cab8a8153026..7cbe9d5f328b4 100644 --- a/pkg/meta/model/job_test.go +++ b/pkg/meta/model/job_test.go @@ -118,10 +118,6 @@ func TestJobCodec(t *testing.T) { } func TestDDLReorgMetaUseNewCollate(t *testing.T) { - var nilMeta *DDLReorgMeta - require.True(t, nilMeta.GetUseNewCollateOrDefault(true)) - require.False(t, nilMeta.GetUseNewCollateOrDefault(false)) - meta := &DDLReorgMeta{} require.True(t, meta.GetUseNewCollateOrDefault(true)) require.False(t, meta.GetUseNewCollateOrDefault(false)) From 2357d7ad6b4150b95bd52973ed718ba69b14c1cd Mon Sep 17 00:00:00 2001 From: Ruihao Chen Date: Mon, 6 Jul 2026 18:14:20 +0800 Subject: [PATCH 32/35] *: keep collation metadata setters internal --- pkg/executor/importer/import.go | 6 +++--- pkg/executor/importer/import_test.go | 4 ++-- pkg/meta/model/job_test.go | 4 ++-- pkg/meta/model/reorg.go | 4 ++-- 4 files changed, 9 insertions(+), 9 deletions(-) diff --git a/pkg/executor/importer/import.go b/pkg/executor/importer/import.go index 4257e3919d164..2b28a32577cb8 100644 --- a/pkg/executor/importer/import.go +++ b/pkg/executor/importer/import.go @@ -364,9 +364,9 @@ func (p *Plan) GetUseNewCollateOrDefault(defaultVal bool) bool { return *p.UseNewCollate } -// SetUseNewCollate stores the new-collation mode captured from the target table +// setUseNewCollate stores the new-collation mode captured from the target table // snapshot. -func (p *Plan) SetUseNewCollate(useNewCollate bool) { +func (p *Plan) setUseNewCollate(useNewCollate bool) { p.UseNewCollate = &useNewCollate } @@ -574,7 +574,7 @@ func NewImportPlan(ctx context.Context, userSctx sessionctx.Context, plan *plann User: userSctx.GetSessionVars().User.String(), Keyspace: userSctx.GetStore().GetKeyspace(), } - p.SetUseNewCollate(tbl.UseNewCollate()) + p.setUseNewCollate(tbl.UseNewCollate()) if err := p.initOptions(ctx, userSctx, plan.Options); err != nil { return nil, err } diff --git a/pkg/executor/importer/import_test.go b/pkg/executor/importer/import_test.go index 50b57d62d18f5..73406210467dc 100644 --- a/pkg/executor/importer/import_test.go +++ b/pkg/executor/importer/import_test.go @@ -110,7 +110,7 @@ func TestPlanUseNewCollate(t *testing.T) { require.True(t, plan.GetUseNewCollateOrDefault(true)) require.False(t, plan.GetUseNewCollateOrDefault(false)) - plan.SetUseNewCollate(false) + plan.setUseNewCollate(false) require.False(t, plan.GetUseNewCollateOrDefault(true)) data, err := json.Marshal(plan) @@ -121,7 +121,7 @@ func TestPlanUseNewCollate(t *testing.T) { require.NoError(t, json.Unmarshal(data, &decoded)) require.False(t, decoded.GetUseNewCollateOrDefault(true)) - decoded.SetUseNewCollate(true) + decoded.setUseNewCollate(true) require.True(t, decoded.GetUseNewCollateOrDefault(false)) } diff --git a/pkg/meta/model/job_test.go b/pkg/meta/model/job_test.go index 7cbe9d5f328b4..95d8f666e506b 100644 --- a/pkg/meta/model/job_test.go +++ b/pkg/meta/model/job_test.go @@ -122,7 +122,7 @@ func TestDDLReorgMetaUseNewCollate(t *testing.T) { require.True(t, meta.GetUseNewCollateOrDefault(true)) require.False(t, meta.GetUseNewCollateOrDefault(false)) - meta.SetUseNewCollate(false) + meta.setUseNewCollate(false) require.False(t, meta.GetUseNewCollateOrDefault(true)) data, err := json.Marshal(meta) @@ -133,7 +133,7 @@ func TestDDLReorgMetaUseNewCollate(t *testing.T) { require.NoError(t, json.Unmarshal(data, &decoded)) require.False(t, decoded.GetUseNewCollateOrDefault(true)) - decoded.SetUseNewCollate(true) + decoded.setUseNewCollate(true) require.True(t, decoded.GetUseNewCollateOrDefault(false)) } diff --git a/pkg/meta/model/reorg.go b/pkg/meta/model/reorg.go index 644a97474fa1b..76d7e34cf729b 100644 --- a/pkg/meta/model/reorg.go +++ b/pkg/meta/model/reorg.go @@ -166,9 +166,9 @@ func (dm *DDLReorgMeta) GetUseNewCollateOrDefault(defaultVal bool) bool { return *dm.UseNewCollate } -// SetUseNewCollate stores the new-collation mode captured from the persisted +// setUseNewCollate stores the new-collation mode captured from the persisted // table snapshot. -func (dm *DDLReorgMeta) SetUseNewCollate(useNewCollate bool) { +func (dm *DDLReorgMeta) setUseNewCollate(useNewCollate bool) { dm.UseNewCollate = &useNewCollate } From a739efc7fe5f8d03ed21218bf3f778ea89219453 Mon Sep 17 00:00:00 2001 From: Ruihao Chen Date: Mon, 6 Jul 2026 20:22:54 +0800 Subject: [PATCH 33/35] tests: cover DXF collation snapshot on user keyspace --- pkg/ddl/backfilling_dist_scheduler.go | 8 +++++- .../addindextest1/cross_ks_test.go | 25 +++++++++++-------- .../importintotest3/cross_ks_test.go | 15 ++++++----- 3 files changed, 31 insertions(+), 17 deletions(-) diff --git a/pkg/ddl/backfilling_dist_scheduler.go b/pkg/ddl/backfilling_dist_scheduler.go index 9a9ac1486fe28..7016592b2b268 100644 --- a/pkg/ddl/backfilling_dist_scheduler.go +++ b/pkg/ddl/backfilling_dist_scheduler.go @@ -202,8 +202,14 @@ func getUserTableFromTaskStore( return nil, err } // we don't touch table data during add-index, a fake Allocators is enough. + defaultUseNewCollate := collate.NewCollationEnabled() + failpoint.Inject("overrideDefaultUseNewCollateForBackfillStep", func(val failpoint.Value) { + defaultUseNewCollate = val.(bool) + }) + useNewCollate := job.ReorgMeta.GetUseNewCollateOrDefault(defaultUseNewCollate) + failpoint.InjectCall("afterResolveUserTableNewCollateForBackfillStep", job, defaultUseNewCollate, useNewCollate) tbl, err := tables.TableFromMetaWithCollate( - job.ReorgMeta.GetUseNewCollateOrDefault(collate.NewCollationEnabled()), + useNewCollate, autoid.NewAllocators(tblInfo.SepAutoInc()), tblInfo, ) diff --git a/tests/realtikvtest/addindextest1/cross_ks_test.go b/tests/realtikvtest/addindextest1/cross_ks_test.go index b006bae815b29..17789ebf2edda 100644 --- a/tests/realtikvtest/addindextest1/cross_ks_test.go +++ b/tests/realtikvtest/addindextest1/cross_ks_test.go @@ -20,10 +20,10 @@ import ( "sync/atomic" "testing" + "github.com/pingcap/tidb/pkg/config" "github.com/pingcap/tidb/pkg/config/kerneltype" "github.com/pingcap/tidb/pkg/ddl" "github.com/pingcap/tidb/pkg/ddl/schemaver" - "github.com/pingcap/tidb/pkg/dxf/framework/proto" "github.com/pingcap/tidb/pkg/keyspace" "github.com/pingcap/tidb/pkg/meta/model" kvstore "github.com/pingcap/tidb/pkg/store" @@ -68,13 +68,18 @@ func TestAddIndexOnUserKeyspaceWithDifferentNewCollation(t *testing.T) { if kerneltype.IsClassic() { t.Skip("This test is only for nextgen kernel, skip it in classic kernel") } + restore := config.RestoreFunc() + t.Cleanup(restore) + config.UpdateGlobal(func(conf *config.Config) { + conf.Experimental.AllowsExpressionIndex = true + }) originNewCollationEnabled := collate.NewCollationEnabled() t.Cleanup(func() { collate.SetNewCollationEnabledForTest(originNewCollationEnabled) }) - const userKeyspace = "keyspacecollateadd" + const userKeyspace = "keyspace2" runtimes := realtikvtest.PrepareForCrossKSTestWithNewCollation(t, map[string]bool{ keyspace.System: true, userKeyspace: false, @@ -86,19 +91,19 @@ func TestAddIndexOnUserKeyspaceWithDifferentNewCollation(t *testing.T) { require.False(t, collate.NewCollationEnabled()) var backfillInitCnt atomic.Int64 - testfailpoint.EnableCall( + testfailpoint.Enable( t, - "github.com/pingcap/tidb/pkg/dxf/framework/storage/beforeSubmitTask", - func(*int, *proto.ExtraParams) { - collate.SetNewCollationEnabledForTest(true) - }, + "github.com/pingcap/tidb/pkg/ddl/overrideDefaultUseNewCollateForBackfillStep", + "return(true)", ) testfailpoint.EnableCall( t, - "github.com/pingcap/tidb/pkg/ddl/beforeGetUserTableForBackfillStep", - func(job *model.Job) { + "github.com/pingcap/tidb/pkg/ddl/afterResolveUserTableNewCollateForBackfillStep", + func(job *model.Job, defaultUseNewCollate bool, useNewCollate bool) { require.False(t, job.ReorgMeta.GetUseNewCollateOrDefault(true)) - require.True(t, collate.NewCollationEnabled()) + require.True(t, defaultUseNewCollate) + require.False(t, useNewCollate) + require.False(t, collate.NewCollationEnabled()) backfillInitCnt.Add(1) }, ) diff --git a/tests/realtikvtest/importintotest3/cross_ks_test.go b/tests/realtikvtest/importintotest3/cross_ks_test.go index 2ca48d16e90f2..fcc20d0986584 100644 --- a/tests/realtikvtest/importintotest3/cross_ks_test.go +++ b/tests/realtikvtest/importintotest3/cross_ks_test.go @@ -144,7 +144,7 @@ func TestImportIntoOnUserKeyspaceWithDifferentNewCollation(t *testing.T) { collate.SetNewCollationEnabledForTest(originNewCollationEnabled) }) - const userKeyspace = "keyspacecollateimport" + const userKeyspace = "keyspace2" runtimes := realtikvtest.PrepareForCrossKSTestWithNewCollation(t, map[string]bool{ keyspace.System: true, userKeyspace: false, @@ -163,7 +163,7 @@ func TestImportIntoOnUserKeyspaceWithDifferentNewCollation(t *testing.T) { objStore.Close() }) - var prepareCnt atomic.Int64 + var taskRefreshCnt atomic.Int64 testfailpoint.EnableCall( t, "github.com/pingcap/tidb/pkg/dxf/framework/storage/beforeSubmitTask", @@ -173,13 +173,16 @@ func TestImportIntoOnUserKeyspaceWithDifferentNewCollation(t *testing.T) { ) testfailpoint.EnableCall( t, - "github.com/pingcap/tidb/pkg/dxf/importinto/afterPrepare", + "github.com/pingcap/tidb/pkg/dxf/framework/scheduler/afterRefreshTask", func(task *proto.Task) { + if task == nil || task.Type != proto.ImportInto || !strings.HasPrefix(task.Key, userKeyspace+"/ImportInto/") { + return + } var taskMeta importinto.TaskMeta require.NoError(t, json.Unmarshal(task.Meta, &taskMeta)) require.False(t, taskMeta.Plan.GetUseNewCollateOrDefault(true)) require.True(t, collate.NewCollationEnabled()) - prepareCnt.Add(1) + taskRefreshCnt.Add(1) }, ) @@ -589,11 +592,11 @@ func TestImportIntoOnUserKeyspaceWithDifferentNewCollation(t *testing.T) { userTK.MustExec(sql) } require.NoError(t, objStore.WriteFile(ctx, tc.fileName, []byte(tc.fileData))) - before := prepareCnt.Load() + before := taskRefreshCnt.Load() fileURL := fmt.Sprintf("s3://next-gen-test/collate-data/%s?%s", tc.fileName, s3Args) result := userTK.MustQuery(fmt.Sprintf(tc.importSQL, fileURL)).Rows() require.Len(t, result, 1) - require.Greater(t, prepareCnt.Load(), before) + require.Greater(t, taskRefreshCnt.Load(), before) collate.SetNewCollationEnabledForTest(false) checkImportTableAndIndexes(userTK, tc.table, tc.indexes, "3") From 38623f2c1a9eb13ff07360b5b8a2264eb600acb6 Mon Sep 17 00:00:00 2001 From: Ruihao Chen Date: Mon, 6 Jul 2026 21:08:34 +0800 Subject: [PATCH 34/35] tests: isolate DXF collation keyspace --- tests/realtikvtest/addindextest1/cross_ks_test.go | 2 +- tests/realtikvtest/configs/next-gen/pd.toml | 2 +- tests/realtikvtest/importintotest3/cross_ks_test.go | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/realtikvtest/addindextest1/cross_ks_test.go b/tests/realtikvtest/addindextest1/cross_ks_test.go index 17789ebf2edda..c16574fcb7ac9 100644 --- a/tests/realtikvtest/addindextest1/cross_ks_test.go +++ b/tests/realtikvtest/addindextest1/cross_ks_test.go @@ -79,7 +79,7 @@ func TestAddIndexOnUserKeyspaceWithDifferentNewCollation(t *testing.T) { collate.SetNewCollationEnabledForTest(originNewCollationEnabled) }) - const userKeyspace = "keyspace2" + const userKeyspace = "keyspacecollate" runtimes := realtikvtest.PrepareForCrossKSTestWithNewCollation(t, map[string]bool{ keyspace.System: true, userKeyspace: false, diff --git a/tests/realtikvtest/configs/next-gen/pd.toml b/tests/realtikvtest/configs/next-gen/pd.toml index 578cfbf2aaab4..07569ded421b1 100644 --- a/tests/realtikvtest/configs/next-gen/pd.toml +++ b/tests/realtikvtest/configs/next-gen/pd.toml @@ -1,3 +1,3 @@ # list all the keyspace names here, each keyspace represents a tenant. [keyspace] -pre-alloc = ["keyspace1", "keyspace2"] +pre-alloc = ["keyspace1", "keyspace2", "keyspacecollate"] diff --git a/tests/realtikvtest/importintotest3/cross_ks_test.go b/tests/realtikvtest/importintotest3/cross_ks_test.go index fcc20d0986584..9ff440eac0ca8 100644 --- a/tests/realtikvtest/importintotest3/cross_ks_test.go +++ b/tests/realtikvtest/importintotest3/cross_ks_test.go @@ -144,7 +144,7 @@ func TestImportIntoOnUserKeyspaceWithDifferentNewCollation(t *testing.T) { collate.SetNewCollationEnabledForTest(originNewCollationEnabled) }) - const userKeyspace = "keyspace2" + const userKeyspace = "keyspacecollate" runtimes := realtikvtest.PrepareForCrossKSTestWithNewCollation(t, map[string]bool{ keyspace.System: true, userKeyspace: false, From 51bbaa90fccca3b39870075053959489d528c015 Mon Sep 17 00:00:00 2001 From: Ruihao Chen Date: Mon, 6 Jul 2026 21:43:48 +0800 Subject: [PATCH 35/35] tests: avoid racy import collation assertion --- tests/realtikvtest/importintotest3/cross_ks_test.go | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/tests/realtikvtest/importintotest3/cross_ks_test.go b/tests/realtikvtest/importintotest3/cross_ks_test.go index 9ff440eac0ca8..75a287d3a5ca7 100644 --- a/tests/realtikvtest/importintotest3/cross_ks_test.go +++ b/tests/realtikvtest/importintotest3/cross_ks_test.go @@ -163,12 +163,14 @@ func TestImportIntoOnUserKeyspaceWithDifferentNewCollation(t *testing.T) { objStore.Close() }) + var taskSubmitCnt atomic.Int64 var taskRefreshCnt atomic.Int64 testfailpoint.EnableCall( t, "github.com/pingcap/tidb/pkg/dxf/framework/storage/beforeSubmitTask", func(*int, *proto.ExtraParams) { collate.SetNewCollationEnabledForTest(true) + taskSubmitCnt.Add(1) }, ) testfailpoint.EnableCall( @@ -180,8 +182,8 @@ func TestImportIntoOnUserKeyspaceWithDifferentNewCollation(t *testing.T) { } var taskMeta importinto.TaskMeta require.NoError(t, json.Unmarshal(task.Meta, &taskMeta)) - require.False(t, taskMeta.Plan.GetUseNewCollateOrDefault(true)) - require.True(t, collate.NewCollationEnabled()) + require.NotNil(t, taskMeta.Plan.UseNewCollate) + require.False(t, *taskMeta.Plan.UseNewCollate) taskRefreshCnt.Add(1) }, ) @@ -592,10 +594,12 @@ func TestImportIntoOnUserKeyspaceWithDifferentNewCollation(t *testing.T) { userTK.MustExec(sql) } require.NoError(t, objStore.WriteFile(ctx, tc.fileName, []byte(tc.fileData))) + beforeSubmit := taskSubmitCnt.Load() before := taskRefreshCnt.Load() fileURL := fmt.Sprintf("s3://next-gen-test/collate-data/%s?%s", tc.fileName, s3Args) result := userTK.MustQuery(fmt.Sprintf(tc.importSQL, fileURL)).Rows() require.Len(t, result, 1) + require.Greater(t, taskSubmitCnt.Load(), beforeSubmit) require.Greater(t, taskRefreshCnt.Load(), before) collate.SetNewCollationEnabledForTest(false)