diff --git a/pkg/ddl/backfilling_dist_scheduler.go b/pkg/ddl/backfilling_dist_scheduler.go index 7016592b2b268..76d48d2d98b20 100644 --- a/pkg/ddl/backfilling_dist_scheduler.go +++ b/pkg/ddl/backfilling_dist_scheduler.go @@ -208,15 +208,7 @@ func getUserTableFromTaskStore( }) useNewCollate := job.ReorgMeta.GetUseNewCollateOrDefault(defaultUseNewCollate) failpoint.InjectCall("afterResolveUserTableNewCollateForBackfillStep", job, defaultUseNewCollate, useNewCollate) - tbl, err := tables.TableFromMetaWithCollate( - useNewCollate, - autoid.NewAllocators(tblInfo.SepAutoInc()), - tblInfo, - ) - if err != nil { - return nil, err - } - return tbl, nil + return tables.TableFromMetaWithCollate(useNewCollate, autoid.NewAllocators(tblInfo.SepAutoInc()), tblInfo) } // GetNextStep implements scheduler.Extension interface. diff --git a/pkg/ddl/backfilling_operators.go b/pkg/ddl/backfilling_operators.go index 80d8ac9f285e8..8ccd064c40e17 100644 --- a/pkg/ddl/backfilling_operators.go +++ b/pkg/ddl/backfilling_operators.go @@ -46,7 +46,6 @@ import ( "github.com/pingcap/tidb/pkg/resourcemanager/util" "github.com/pingcap/tidb/pkg/sessionctx" "github.com/pingcap/tidb/pkg/table" - "github.com/pingcap/tidb/pkg/table/tables" "github.com/pingcap/tidb/pkg/tablecodec" "github.com/pingcap/tidb/pkg/util/chunk" contextutil "github.com/pingcap/tidb/pkg/util/context" @@ -109,13 +108,9 @@ func NewAddIndexIngestPipeline( concurrency int, collector execute.Collector, ) (*operator.AsyncPipeline, error) { - indexes := make([]table.Index, 0, len(idxInfos)) - for _, idxInfo := range idxInfos { - index, err := tables.NewIndexWithCollate(tbl.UseNewCollate(), tbl.GetPhysicalID(), tbl.Meta(), idxInfo) - if err != nil { - return nil, err - } - indexes = append(indexes, index) + indexes, err := indexesForBackfill(tbl, idxInfos) + if err != nil { + return nil, err } reqSrc := getDDLRequestSource(model.ActionAddIndex) copCtx, err := NewReorgCopContext(reorgMeta, tbl.Meta(), idxInfos, reqSrc) @@ -167,13 +162,9 @@ func NewWriteIndexToExternalStoragePipeline( collector execute.Collector, tikvCodec tikv.Codec, ) (*operator.AsyncPipeline, error) { - indexes := make([]table.Index, 0, len(idxInfos)) - for _, idxInfo := range idxInfos { - index, err := tables.NewIndexWithCollate(tbl.UseNewCollate(), tbl.GetPhysicalID(), tbl.Meta(), idxInfo) - if err != nil { - return nil, err - } - indexes = append(indexes, index) + indexes, err := indexesForBackfill(tbl, idxInfos) + if err != nil { + return nil, err } reqSrc := getDDLRequestSource(model.ActionAddIndex) copCtx, err := NewReorgCopContext(reorgMeta, tbl.Meta(), idxInfos, reqSrc) @@ -219,6 +210,23 @@ func NewWriteIndexToExternalStoragePipeline( ), nil } +func indexesForBackfill(tbl table.PhysicalTable, idxInfos []*model.IndexInfo) ([]table.Index, error) { + indexesByID := make(map[int64]table.Index, len(tbl.Indices())) + for _, idx := range tbl.Indices() { + indexesByID[idx.Meta().ID] = idx + } + + indexes := make([]table.Index, 0, len(idxInfos)) + for _, idxInfo := range idxInfos { + idx, ok := indexesByID[idxInfo.ID] + if !ok { + return nil, errors.Errorf("index ID %d not found in physical table %d", idxInfo.ID, tbl.GetPhysicalID()) + } + indexes = append(indexes, idx) + } + return indexes, nil +} + func createChunkPool(copCtx copr.CopContext, reorgMeta *model.DDLReorgMeta) *sync.Pool { return &sync.Pool{ New: func() any { @@ -915,7 +923,7 @@ func (w *indexIngestWorker) WriteChunk(rs *IndexRecordChunk) (count int, bytes i indexConditionCheckers = nil } cnt, kvBytes, err := writeChunk(w.ctx, w.writers, w.indexes, indexConditionCheckers, w.copCtx, - sc.TimeZone(), sc.ErrCtx(), vars.GetWriteStmtBufs(), rs.Chunk, w.tbl.Meta(), w.tbl.UseNewCollate()) + sc.TimeZone(), sc.ErrCtx(), vars.GetWriteStmtBufs(), rs.Chunk, w.tbl.Meta()) if err != nil || cnt == 0 { return 0, 0, err } diff --git a/pkg/ddl/backfilling_test.go b/pkg/ddl/backfilling_test.go index e0b6b84f284cf..565e1995beb75 100644 --- a/pkg/ddl/backfilling_test.go +++ b/pkg/ddl/backfilling_test.go @@ -26,9 +26,11 @@ import ( "github.com/pingcap/tidb/pkg/ddl/ingest" distsqlctx "github.com/pingcap/tidb/pkg/distsql/context" "github.com/pingcap/tidb/pkg/errctx" + "github.com/pingcap/tidb/pkg/expression" "github.com/pingcap/tidb/pkg/expression/exprstatic" "github.com/pingcap/tidb/pkg/kv" "github.com/pingcap/tidb/pkg/meta/model" + "github.com/pingcap/tidb/pkg/parser/ast" "github.com/pingcap/tidb/pkg/parser/mysql" "github.com/pingcap/tidb/pkg/resourcemanager/pool/workerpool" "github.com/pingcap/tidb/pkg/sessionctx" @@ -36,6 +38,8 @@ import ( "github.com/pingcap/tidb/pkg/sessionctx/variable" "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" contextutil "github.com/pingcap/tidb/pkg/util/context" "github.com/pingcap/tidb/pkg/util/deeptest" "github.com/pingcap/tidb/pkg/util/mock" @@ -71,6 +75,85 @@ func TestIndexInfoNotFoundIsNonRetryable(t *testing.T) { require.False(t, (&backfillDistExecutor{}).IsRetryableError(err)) } +func TestBuildIndexConditionCheckerUsesFixedCollation(t *testing.T) { + origin := collate.NewCollationEnabled() + collate.SetNewCollationEnabledForTest(true) + defer collate.SetNewCollationEnabledForTest(origin) + + originBuildSimpleExpr := expression.BuildSimpleExpr + defer func() { + expression.BuildSimpleExpr = originBuildSimpleExpr + }() + expression.BuildSimpleExpr = func(ctx expression.BuildContext, _ ast.ExprNode, opts ...expression.BuildOption) (expression.Expression, error) { + var options expression.BuildOptions + for _, opt := range opts { + opt(&options) + } + if options.InputSchema == nil { + return expression.NewOne(), nil + } + constantTp := types.NewFieldTypeWithCollation(mysql.TypeVarchar, "utf8mb4_general_ci", 16) + return expression.NewFunction( + ctx, + ast.EQ, + types.NewFieldType(mysql.TypeTiny), + options.InputSchema.Columns[0], + &expression.Constant{Value: types.NewDatum("A"), RetType: constantTp}, + ) + } + + colTp := types.NewFieldTypeWithCollation(mysql.TypeVarchar, "utf8mb4_general_ci", 16) + colInfo := &model.ColumnInfo{ + ID: 1, + Offset: 0, + Name: ast.NewCIStr("c0"), + FieldType: *colTp, + State: model.StatePublic, + } + idxInfo := &model.IndexInfo{ + ID: 1, + Name: ast.NewCIStr("idx"), + Columns: []*model.IndexColumn{{Name: colInfo.Name, Offset: colInfo.Offset}}, + State: model.StatePublic, + ConditionExprString: "c0 = 'A'", + } + tblInfo := &model.TableInfo{ + Name: ast.NewCIStr("t"), + Columns: []*model.ColumnInfo{colInfo}, + Indices: []*model.IndexInfo{idxInfo}, + } + + sctx := mock.NewContext() + exprCtx := sctx.ExprContext.IntoStatic() + copCtx, err := copr.NewCopContextSingleIndex( + exprCtx.Apply(exprstatic.WithNewCollationEnabled(false)), + sctx.GetSessionVars().StmtCtx.PushDownFlags(), + tblInfo, + idxInfo, + "", + ) + require.NoError(t, err) + checker, err := buildIndexConditionChecker(copCtx, tblInfo, idxInfo) + require.NoError(t, err) + matched, err := checker(chunk.MutRowFromValues("a").ToRow()) + require.NoError(t, err) + require.False(t, matched) + + copCtx, err = copr.NewCopContextSingleIndex( + exprCtx.Apply(exprstatic.WithNewCollationEnabled(true)), + sctx.GetSessionVars().StmtCtx.PushDownFlags(), + tblInfo, + idxInfo, + "", + ) + require.NoError(t, err) + checker, err = buildIndexConditionChecker(copCtx, tblInfo, idxInfo) + require.NoError(t, err) + matched, err = checker(chunk.MutRowFromValues("a").ToRow()) + require.NoError(t, err) + require.True(t, matched) +} + func TestPickBackfillType(t *testing.T) { ingest.LitDiskRoot = ingest.NewDiskRootImpl(t.TempDir()) ingest.LitMemRoot = ingest.NewMemRootImpl(math.MaxInt64) @@ -194,6 +277,7 @@ func assertStaticExprContextEqual(t *testing.T, sctx sessionctx.Context, exprCtx f.check(exprCtx) ignoreFields = append(ignoreFields, "$.exprCtxState."+f.field) } + ignoreFields = append(ignoreFields, "$.exprCtxState.newCollationEnabled") deeptest.AssertDeepClonedEqual(t, expected, exprCtx, deeptest.WithIgnorePath(ignoreFields)) // check EvalContext @@ -227,6 +311,10 @@ func newMockReorgSessCtx(store kv.Storage) sessionctx.Context { // compatible with newMockReorgSessCtx(nil).GetExprCtx() to make it safe to replace `mock.Context` usage. // After refactor, the TestReorgExprContext can be removed. func TestReorgExprContext(t *testing.T) { + origin := collate.NewCollationEnabled() + collate.SetNewCollationEnabledForTest(true) + defer collate.SetNewCollationEnabledForTest(origin) + // test default expr context store := &mockStorage{client: &mock.Client{}} sctx := newMockReorgSessCtx(store) @@ -238,26 +326,51 @@ func TestReorgExprContext(t *testing.T) { defaultTypeCtx := evalCtx.TypeCtx() defaultErrCtx := evalCtx.ErrCtx() + oldCollation := false + newCollation := true + // test expr context from DDLReorgMeta - for _, reorg := range []model.DDLReorgMeta{ + for _, testCase := range []struct { + reorg model.DDLReorgMeta + expectedUseNewCollate bool + }{ + { + reorg: model.DDLReorgMeta{ + SQLMode: mysql.ModeStrictTransTables | mysql.ModeAllowInvalidDates, + Location: &model.TimeZoneLocation{Name: "Asia/Tokyo"}, + ReorgTp: model.ReorgTypeIngest, + ResourceGroupName: "rg1", + UseNewCollate: &oldCollation, + }, + expectedUseNewCollate: false, + }, { - SQLMode: mysql.ModeStrictTransTables | mysql.ModeAllowInvalidDates, - Location: &model.TimeZoneLocation{Name: "Asia/Tokyo"}, - ReorgTp: model.ReorgTypeIngest, - ResourceGroupName: "rg1", + reorg: model.DDLReorgMeta{ + SQLMode: mysql.ModeAllowInvalidDates, + // should load location from system value when reorg.Location is nil + Location: nil, + ReorgTp: model.ReorgTypeTxnMerge, + ResourceGroupName: "rg2", + UseNewCollate: &newCollation, + }, + expectedUseNewCollate: true, }, { - SQLMode: mysql.ModeAllowInvalidDates, - // should load location from system value when reorg.Location is nil - Location: nil, - ReorgTp: model.ReorgTypeTxnMerge, - ResourceGroupName: "rg2", + reorg: model.DDLReorgMeta{ + SQLMode: mysql.ModeAllowInvalidDates, + Location: nil, + ReorgTp: model.ReorgTypeTxnMerge, + ResourceGroupName: "rg3", + }, + expectedUseNewCollate: true, }, } { + reorg := testCase.reorg sctx = newMockReorgSessCtx(store) require.NoError(t, initSessCtx(sctx, &reorg)) ctx, err := newReorgExprCtxWithReorgMeta(&reorg, sctx.GetSessionVars().StmtCtx.WarnHandler) require.NoError(t, err) + require.Equal(t, testCase.expectedUseNewCollate, ctx.NewCollationEnabled()) assertStaticExprContextEqual(t, sctx, ctx, ctx.GetStaticEvalCtx().GetWarnHandler()) evalCtx := ctx.GetEvalCtx() tc, ec := evalCtx.TypeCtx(), evalCtx.ErrCtx() diff --git a/pkg/ddl/backfilling_txn_executor.go b/pkg/ddl/backfilling_txn_executor.go index f2a72e3fc1382..46c2121ce6acf 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" @@ -143,7 +142,6 @@ func NewReorgCopContext( tblInfo, allIdxInfo, requestSource, - reorgMeta.GetUseNewCollateOrDefault(collate.NewCollationEnabled()), ) } diff --git a/pkg/ddl/column.go b/pkg/ddl/column.go index 84dd980ce75da..ee81ef083ee93 100644 --- a/pkg/ddl/column.go +++ b/pkg/ddl/column.go @@ -43,8 +43,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/codec" - "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/intest" @@ -812,8 +810,7 @@ func (w *updateColumnWorker) getRowRecord(handle kv.Handle, recordKey []byte, ra if w.checksumNeeded { checksum = rowcodec.RawChecksum{Handle: handle} } - enc := codec.NewEncoder(collate.NewCollationEnabled()) - newRowVal, err := tablecodec.EncodeRow(enc, sysTZ, newRow, newColumnIDs, nil, nil, checksum, rd) + newRowVal, err := tablecodec.EncodeRow(sysTZ, newRow, newColumnIDs, nil, nil, checksum, rd) err = ec.HandleError(err) if err != nil { return errors.Trace(err) diff --git a/pkg/ddl/copr/BUILD.bazel b/pkg/ddl/copr/BUILD.bazel index 0cf654c6aaed3..0abea7617d975 100644 --- a/pkg/ddl/copr/BUILD.bazel +++ b/pkg/ddl/copr/BUILD.bazel @@ -23,7 +23,7 @@ go_test( srcs = ["copr_ctx_test.go"], embed = [":copr"], flaky = True, - shard_count = 3, + shard_count = 4, deps = [ "//pkg/expression", "//pkg/expression/exprstatic", @@ -31,6 +31,7 @@ go_test( "//pkg/parser/ast", "//pkg/parser/mysql", "//pkg/types", + "//pkg/util/collate", "//pkg/util/mock", "@com_github_stretchr_testify//require", ], diff --git a/pkg/ddl/copr/copr_ctx.go b/pkg/ddl/copr/copr_ctx.go index 26cab165d69e4..b050b6fc1c9dd 100644 --- a/pkg/ddl/copr/copr_ctx.go +++ b/pkg/ddl/copr/copr_ctx.go @@ -46,7 +46,6 @@ type CopContextBase struct { ExprCtx exprctx.BuildContext PushDownFlags uint64 RequestSource string - UseNewCollate bool ColumnInfos []*model.ColumnInfo FieldTypes []*types.FieldType @@ -75,13 +74,13 @@ type CopContextMultiIndex struct { // NewCopContextBase creates a CopContextBase. // `idxCols` contains all the index columns and also the columns referenced by the index condition. +// The new-collation mode is carried by `exprCtx`. 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)) @@ -127,13 +126,12 @@ func NewCopContextBase( handleIDs = []int64{extra.ID} } - expColInfos, _, err := expression.ColumnInfos2ColumnsAndNamesWithCollate( + expColInfos, _, err := expression.ColumnInfos2ColumnsAndNames( exprCtx, ast.CIStr{}, // unused tblInfo.Name, colInfos, tblInfo, - useNewCollate, ) if err != nil { return nil, err @@ -147,7 +145,6 @@ func NewCopContextBase( ExprCtx: exprCtx, PushDownFlags: pushDownFlags, RequestSource: requestSource, - UseNewCollate: useNewCollate, ColumnInfos: colInfos, FieldTypes: fieldTps, ExprColumnInfos: expColInfos, @@ -157,14 +154,13 @@ func NewCopContextBase( }, nil } -// NewCopContext creates a CopContext with a fixed collation mode. +// NewCopContext creates a CopContext. The new-collation mode is carried by `exprCtx`. func NewCopContext( exprCtx exprctx.BuildContext, pushDownFlags uint64, tblInfo *model.TableInfo, allIdxInfo []*model.IndexInfo, requestSource string, - useNewCollate bool, ) (CopContext, error) { if len(allIdxInfo) == 1 { return NewCopContextSingleIndex( @@ -173,20 +169,18 @@ func NewCopContext( tblInfo, allIdxInfo[0], requestSource, - useNewCollate, ) } - return NewCopContextMultiIndex(exprCtx, pushDownFlags, tblInfo, allIdxInfo, requestSource, useNewCollate) + return NewCopContextMultiIndex(exprCtx, pushDownFlags, tblInfo, allIdxInfo, requestSource) } -// NewCopContextSingleIndex creates a CopContextSingleIndex with a fixed collation mode. +// NewCopContextSingleIndex creates a CopContextSingleIndex. 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) @@ -196,7 +190,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) if err != nil { return nil, err } @@ -246,14 +240,13 @@ func (c *CopContextSingleIndex) GetCondition() (expression.Expression, error) { return expr, nil } -// NewCopContextMultiIndex creates a CopContextMultiIndex with a fixed collation mode. +// NewCopContextMultiIndex creates a CopContextMultiIndex. 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 { @@ -271,7 +264,7 @@ func NewCopContextMultiIndex( } allIdxCols = tables.DedupIndexColumns(allIdxCols) - base, err := NewCopContextBase(exprCtx, pushDownFlags, tblInfo, allIdxCols, requestSource, useNewCollate) + base, err := NewCopContextBase(exprCtx, pushDownFlags, tblInfo, allIdxCols, requestSource) if err != nil { return nil, err } diff --git a/pkg/ddl/copr/copr_ctx_test.go b/pkg/ddl/copr/copr_ctx_test.go index 80b4a676f63c9..82b7130219c36 100644 --- a/pkg/ddl/copr/copr_ctx_test.go +++ b/pkg/ddl/copr/copr_ctx_test.go @@ -24,6 +24,7 @@ import ( "github.com/pingcap/tidb/pkg/parser/ast" "github.com/pingcap/tidb/pkg/parser/mysql" "github.com/pingcap/tidb/pkg/types" + "github.com/pingcap/tidb/pkg/util/collate" "github.com/pingcap/tidb/pkg/util/mock" "github.com/stretchr/testify/require" ) @@ -113,7 +114,6 @@ func TestNewCopContextSingleIndex(t *testing.T) { mockTableInfo, mockIdxInfo, "", - false, ) require.NoError(t, err) base := copCtx.GetBase() @@ -131,6 +131,71 @@ func TestNewCopContextSingleIndex(t *testing.T) { } } +func TestCopContextConditionUsesFixedCollation(t *testing.T) { + origin := collate.NewCollationEnabled() + collate.SetNewCollationEnabledForTest(true) + defer collate.SetNewCollationEnabledForTest(origin) + + colTp := types.NewFieldTypeWithCollation(mysql.TypeVarchar, "utf8mb4_general_ci", 16) + colInfo := &model.ColumnInfo{ + ID: 1, + Offset: 0, + Name: ast.NewCIStr("c0"), + FieldType: *colTp, + State: model.StatePublic, + } + generatedColInfo := &model.ColumnInfo{ + ID: 2, + Offset: 1, + Name: ast.NewCIStr("g0"), + FieldType: *colTp, + State: model.StatePublic, + GeneratedExprString: "lower(c0)", + GeneratedStored: false, + Dependences: map[string]struct{}{"c0": {}}, + } + + originBuildSimpleExpr := expression.BuildSimpleExpr + defer func() { + expression.BuildSimpleExpr = originBuildSimpleExpr + }() + var seenUseNewCollates []bool + expression.BuildSimpleExpr = func(ctx expression.BuildContext, expr ast.ExprNode, _ ...expression.BuildOption) (expression.Expression, error) { + seenUseNewCollates = append(seenUseNewCollates, ctx.NewCollationEnabled()) + return expression.NewOne(), nil + } + idxInfo := &model.IndexInfo{ + ID: 1, + Name: ast.NewCIStr("idx"), + Columns: []*model.IndexColumn{{Name: generatedColInfo.Name, Offset: generatedColInfo.Offset}}, + State: model.StatePublic, + ConditionExprString: "1", + } + tblInfo := &model.TableInfo{ + Name: ast.NewCIStr("t"), + Columns: []*model.ColumnInfo{colInfo, generatedColInfo}, + Indices: []*model.IndexInfo{idxInfo}, + } + + sctx := mock.NewContext() + exprCtx := sctx.ExprContext.IntoStatic().Apply(exprstatic.WithNewCollationEnabled(false)) + copCtx, err := NewCopContextSingleIndex( + exprCtx, + sctx.GetSessionVars().StmtCtx.PushDownFlags(), + tblInfo, + idxInfo, + "", + ) + require.NoError(t, err) + condition, err := copCtx.GetCondition() + require.NoError(t, err) + require.NotNil(t, condition) + require.NotEmpty(t, seenUseNewCollates) + for _, useNewCollate := range seenUseNewCollates { + require.False(t, useNewCollate) + } +} + func TestResolveIndicesForHandle(t *testing.T) { type args struct { cols []*expression.Column diff --git a/pkg/ddl/index.go b/pkg/ddl/index.go index 756582f542fb9..a043e67065de5 100644 --- a/pkg/ddl/index.go +++ b/pkg/ddl/index.go @@ -2710,11 +2710,11 @@ 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() ectx := c.ExprCtx.GetEvalCtx() + useNewCollate := c.ExprCtx.NewCollationEnabled() maxIdxColCnt := maxIndexColumnCount(indexes) idxDataBuf := make([]types.Datum, maxIdxColCnt) diff --git a/pkg/ddl/index_cop.go b/pkg/ddl/index_cop.go index 65adf1ee8e7c6..f647f55c60a19 100644 --- a/pkg/ddl/index_cop.go +++ b/pkg/ddl/index_cop.go @@ -36,6 +36,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/logutil" "github.com/pingcap/tidb/pkg/util/timeutil" "github.com/pingcap/tipb/go-tipb" @@ -174,6 +175,7 @@ func getRestoreData(useNewCollate bool, tblInfo *model.TableInfo, targetIdx, pkI func buildDAGPB(ctx context.Context, exprCtx exprctx.BuildContext, distSQLCtx *distsqlctx.DistSQLContext, pushDownFlags uint64, tblInfo *model.TableInfo, colInfos []*model.ColumnInfo, selectExpr expression.Expression) (*tipb.DAGRequest, bool, error) { conditionPushed := false + useNewCollate := exprCtx.NewCollationEnabled() dagReq := &tipb.DAGRequest{} dagReq.TimeZoneName, dagReq.TimeZoneOffset = timeutil.Zone(exprCtx.GetEvalCtx().Location()) @@ -187,7 +189,9 @@ func buildDAGPB(ctx context.Context, exprCtx exprctx.BuildContext, distSQLCtx *d } var selectionPB *tipb.Executor - if selectExpr != nil { + // Pushdown cannot preserve the reorg task's captured collation mode when it + // differs from the executor's global mode. Evaluate the condition in TiDB instead. + if selectExpr != nil && useNewCollate == collate.NewCollationEnabled() { selectionPB, err = constructSelectionPB(exprCtx, selectExpr, distSQLCtx, tblScanPB) } @@ -199,10 +203,18 @@ func buildDAGPB(ctx context.Context, exprCtx exprctx.BuildContext, distSQLCtx *d } else { if selectExpr != nil { selectExprStr := selectExpr.StringWithCtx(exprCtx.GetEvalCtx(), errors.RedactLogDisable) - logutil.Logger(ctx).Info("fail to push down the selection expression for index condition", - zap.String("table", tblInfo.Name.O), - zap.String("expr", selectExprStr), - zap.Error(err)) + if useNewCollate != collate.NewCollationEnabled() { + logutil.Logger(ctx).Info("skip pushing down the selection expression for index condition due to collation mode mismatch", + zap.String("table", tblInfo.Name.O), + zap.String("expr", selectExprStr), + zap.Bool("useNewCollate", useNewCollate), + zap.Bool("globalUseNewCollate", collate.NewCollationEnabled())) + } else { + logutil.Logger(ctx).Info("fail to push down the selection expression for index condition", + zap.String("table", tblInfo.Name.O), + zap.String("expr", selectExprStr), + zap.Error(err)) + } } dagReq.Executors = append(dagReq.Executors, tblScanPB) } diff --git a/pkg/ddl/reorg.go b/pkg/ddl/reorg.go index 2f7dad05b6526..2b5bd847a6a58 100644 --- a/pkg/ddl/reorg.go +++ b/pkg/ddl/reorg.go @@ -55,6 +55,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" contextutil "github.com/pingcap/tidb/pkg/util/context" "github.com/pingcap/tidb/pkg/util/dbterror" "github.com/pingcap/tidb/pkg/util/intest" @@ -126,7 +127,10 @@ func newReorgExprCtxWithReorgMeta(reorgMeta *model.DDLReorgMeta, warnHandler con exprstatic.WithErrLevelMap(reorgErrLevelsWithSQLMode(reorgMeta.SQLMode)), exprstatic.WithWarnHandler(warnHandler), ) - return ctx.Apply(exprstatic.WithEvalCtx(evalCtx)), nil + return ctx.Apply( + exprstatic.WithEvalCtx(evalCtx), + exprstatic.WithNewCollationEnabled(reorgMeta.GetUseNewCollateOrDefault(collate.NewCollationEnabled())), + ), nil } // reorgTableMutateContext implements table.MutateContext for reorganization. diff --git a/pkg/executor/importer/BUILD.bazel b/pkg/executor/importer/BUILD.bazel index ee86cb36c5eb6..0a2f752a84b94 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/executor/importer/import.go b/pkg/executor/importer/import.go index 2b28a32577cb8..cec1f17dfa6bb 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" @@ -336,10 +337,10 @@ type Plan struct { // the keyspace name when submitting this job, only for import-into Keyspace string // 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. + // in the submitting keyspace. 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"` } @@ -364,8 +365,8 @@ func (p *Plan) GetUseNewCollateOrDefault(defaultVal bool) bool { return *p.UseNewCollate } -// setUseNewCollate stores the new-collation mode captured from the target table -// snapshot. +// setUseNewCollate stores the new-collation mode captured from the submitting +// keyspace. func (p *Plan) setUseNewCollate(useNewCollate bool) { p.UseNewCollate = &useNewCollate } @@ -574,7 +575,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 } @@ -1946,7 +1947,6 @@ func createColAssignSimpleExprs( assignments []*ast.Assignment, ctx expression.BuildContext, mu *sync.Mutex, - useNewCollate bool, ) (_ []expression.Expression, _ []contextutil.SQLWarn, retErr error) { if mu != nil { mu.Lock() @@ -1955,7 +1955,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, expression.WithUseNewCollate(useNewCollate)) + newExpr, err := expression.BuildSimpleExpr(ctx, assign.Expr) // 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 { @@ -1975,7 +1975,6 @@ func (e *LoadDataController) CreateColAssignSimpleExprs(ctx expression.BuildCont e.ColumnAssignments, ctx, &e.colAssignMu, - e.Table.UseNewCollate(), ) } diff --git a/pkg/executor/importer/sampler.go b/pkg/executor/importer/sampler.go index 073ac642d30a7..f10f6ac0d2f09 100644 --- a/pkg/executor/importer/sampler.go +++ b/pkg/executor/importer/sampler.go @@ -213,7 +213,6 @@ func (s *kvSizeSampler) CreateColAssignSimpleExprs( s.cfg.ColumnAssignments, ctx, &s.colAssignMu, - s.table.UseNewCollate(), ) } diff --git a/pkg/executor/test/executor/BUILD.bazel b/pkg/executor/test/executor/BUILD.bazel index 28ff06bf07711..366993fafcfbd 100644 --- a/pkg/executor/test/executor/BUILD.bazel +++ b/pkg/executor/test/executor/BUILD.bazel @@ -47,8 +47,6 @@ go_test( "//pkg/testkit/testfailpoint", "//pkg/types", "//pkg/util", - "//pkg/util/codec", - "//pkg/util/collate", "//pkg/util/dbterror/exeerrors", "//pkg/util/dbterror/plannererrors", "//pkg/util/mock", diff --git a/pkg/executor/test/executor/executor_test.go b/pkg/executor/test/executor/executor_test.go index 4ed339c5eae78..e39fa70a871b2 100644 --- a/pkg/executor/test/executor/executor_test.go +++ b/pkg/executor/test/executor/executor_test.go @@ -67,8 +67,6 @@ import ( "github.com/pingcap/tidb/pkg/testkit/testfailpoint" "github.com/pingcap/tidb/pkg/types" "github.com/pingcap/tidb/pkg/util" - "github.com/pingcap/tidb/pkg/util/codec" - "github.com/pingcap/tidb/pkg/util/collate" "github.com/pingcap/tidb/pkg/util/dbterror/exeerrors" "github.com/pingcap/tidb/pkg/util/dbterror/plannererrors" "github.com/pingcap/tidb/pkg/util/mock" @@ -253,7 +251,7 @@ func setColValue(t *testing.T, txn kv.Transaction, key kv.Key, v types.Datum) { colIDs := []int64{2, 3} sc := stmtctx.NewStmtCtxWithTimeZone(time.Local) rd := rowcodec.Encoder{Enable: true} - value, err := tablecodec.EncodeRow(codec.NewEncoder(collate.NewCollationEnabled()), sc.TimeZone(), row, colIDs, nil, nil, nil, &rd) + value, err := tablecodec.EncodeRow(sc.TimeZone(), row, colIDs, nil, nil, nil, &rd) require.NoError(t, err) err = txn.Set(key, value) require.NoError(t, err) diff --git a/pkg/executor/write.go b/pkg/executor/write.go index 28a73bb6ba730..9445a400923b2 100644 --- a/pkg/executor/write.go +++ b/pkg/executor/write.go @@ -411,7 +411,7 @@ func addUnchangedKeysForLockByRow( } } unchangedUniqueKey, _, err := tablecodec.GenIndexKey( - codec.NewEncoder(collate.NewCollationEnabled()), + codec.NewEncoder(t.UseNewCollate()), stmtCtx.TimeZone(), idx.TableMeta(), meta, diff --git a/pkg/expression/builtin.go b/pkg/expression/builtin.go index c5e83893ed005..d307ba74dc6b1 100644 --- a/pkg/expression/builtin.go +++ b/pkg/expression/builtin.go @@ -55,7 +55,11 @@ type baseBuiltinFunc struct { pbCode tipb.ScalarFuncSig ctor collate.Collator - childrenVectorized bool + childrenVectorized bool + useNewCollate bool + // collatorOverridden prevents an explicit collator from following later + // collation metadata changes. + collatorOverridden bool childrenVectorizedOnce *sync.Once safeToShareAcrossSessionFlag uint32 // 0 not-initialized, 1 safe, 2 unsafe @@ -95,12 +99,20 @@ func (b *baseBuiltinFunc) setPbCode(c tipb.ScalarFuncSig) { func (b *baseBuiltinFunc) setCollator(ctor collate.Collator) { b.ctor = ctor + b.collatorOverridden = true } func (b *baseBuiltinFunc) collator() collate.Collator { return b.ctor } +func (b *baseBuiltinFunc) SetCharsetAndCollation(chs, coll string) { + b.collationInfo.SetCharsetAndCollation(chs, coll) + if !b.collatorOverridden { + b.ctor = collate.GetCollatorWithCollate(b.useNewCollate, coll) + } +} + func adjustNullFlagForReturnType(ctx EvalContext, funcName string, args []Expression, bf baseBuiltinFunc) { if functionSetForReturnTypeAlwaysNotNull.Exist(funcName) { bf.tp.AddFlag(mysql.NotNullFlag) @@ -134,12 +146,12 @@ func newBaseBuiltinFunc(ctx BuildContext, funcName string, args []Expression, tp bf := baseBuiltinFunc{ childrenVectorizedOnce: new(sync.Once), + useNewCollate: ctx.NewCollationEnabled(), args: args, tp: tp, } bf.SetCharsetAndCollation(ec.Charset, ec.Collation) - bf.setCollator(collate.GetCollator(ec.Collation)) bf.SetCoercibility(ec.Coer) bf.SetRepertoire(ec.Repe) adjustNullFlagForReturnType(ctx.GetEvalCtx(), funcName, args, bf) @@ -225,12 +237,12 @@ func newBaseBuiltinFuncWithTp(ctx BuildContext, funcName string, args []Expressi fieldType := newReturnFieldTypeForBaseBuiltinFunc(funcName, retType, ec) bf = baseBuiltinFunc{ childrenVectorizedOnce: new(sync.Once), + useNewCollate: ctx.NewCollationEnabled(), args: args, tp: fieldType, } bf.SetCharsetAndCollation(ec.Charset, 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 @@ -285,12 +297,12 @@ func newBaseBuiltinFuncWithFieldTypes(ctx BuildContext, funcName string, args [] fieldType := newReturnFieldTypeForBaseBuiltinFunc(funcName, retType, ec) bf = baseBuiltinFunc{ childrenVectorizedOnce: new(sync.Once), + useNewCollate: ctx.NewCollationEnabled(), args: args, tp: fieldType, } bf.SetCharsetAndCollation(ec.Charset, 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 @@ -303,12 +315,12 @@ func newBaseBuiltinFuncWithFieldTypes(ctx BuildContext, funcName string, args [] func newBaseBuiltinFuncWithFieldType(tp *types.FieldType, args []Expression) (baseBuiltinFunc, error) { bf := baseBuiltinFunc{ childrenVectorizedOnce: new(sync.Once), + useNewCollate: collate.NewCollationEnabled(), args: args, tp: tp, } bf.SetCharsetAndCollation(tp.GetCharset(), tp.GetCollate()) - bf.setCollator(collate.GetCollator(tp.GetCollate())) return bf, nil } @@ -437,6 +449,8 @@ func (b *baseBuiltinFunc) cloneFrom(from *baseBuiltinFunc) { } b.tp = from.tp b.pbCode = from.pbCode + b.useNewCollate = from.useNewCollate + b.collatorOverridden = from.collatorOverridden b.childrenVectorizedOnce = new(sync.Once) if from.ctor != nil { b.ctor = from.ctor.Clone() @@ -481,12 +495,12 @@ func newBaseBuiltinCastFunc4String(ctx BuildContext, funcName string, args []Exp if isExplicitCharset { bf = baseBuiltinFunc{ childrenVectorizedOnce: new(sync.Once), + useNewCollate: ctx.NewCollationEnabled(), args: args, tp: tp, } bf.SetCharsetAndCollation(tp.GetCharset(), tp.GetCollate()) - bf.setCollator(collate.GetCollator(tp.GetCollate())) bf.SetCoercibility(CoercibilityExplicit) bf.SetExplicitCharset(true) if tp.GetCharset() == charset.CharsetASCII { diff --git a/pkg/expression/builtin_compare.go b/pkg/expression/builtin_compare.go index 74fa13d33ef0d..7be742b274fd4 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.collator().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.collator().Compare(v, minv) < 0 { minv = v } } @@ -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.collator())) } 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.collator())) } 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.collator())) } 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.collator())) } 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.collator())) } 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.collator())) } 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.collator().Compare(arg0, arg1) == 0: res = 1 } return res, false, nil @@ -3417,6 +3417,10 @@ func genCompareString(collation string) func(sctx EvalContext, lhsArg Expression // 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 +3434,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_test.go b/pkg/expression/builtin_compare_test.go index d91c4544fe148..7835f20b82098 100644 --- a/pkg/expression/builtin_compare_test.go +++ b/pkg/expression/builtin_compare_test.go @@ -169,6 +169,14 @@ func TestCompare(t *testing.T) { args = bf.getArgs() require.Equal(t, mysql.TypeJSON, args[0].GetType(ctx).GetType()) require.Equal(t, mysql.TypeJSON, args[1].GetType(ctx).GetType()) + + // Some callers override a comparison function's derived collation after construction. + bf, err = funcs[ast.EQ].getFunction(ctx, primitiveValsToConstants(ctx, []any{"a", "A"})) + require.NoError(t, err) + bf.SetCharsetAndCollation("utf8mb4", "utf8mb4_unicode_ci") + result, err := evalBuiltinFunc(bf, ctx, chunk.Row{}) + require.NoError(t, err) + require.Equal(t, int64(1), result.GetInt64()) } func TestCoalesce(t *testing.T) { diff --git a/pkg/expression/builtin_ilike.go b/pkg/expression/builtin_ilike.go index d664f459ac6d3..22f9b113ba020 100644 --- a/pkg/expression/builtin_ilike.go +++ b/pkg/expression/builtin_ilike.go @@ -44,6 +44,7 @@ func (c *ilikeFunctionClass) getFunction(ctx BuildContext, args []Expression) (b if err != nil { return nil, err } + bf.setCollator(getCollator(ctx, collate.ConvertAndGetBinCollation(bf.collation))) bf.tp.SetFlen(1) sig := &builtinIlikeSig{baseBuiltinFunc: bf} sig.setPbCode(tipb.ScalarFuncSig_IlikeSig) @@ -96,7 +97,7 @@ func (b *builtinIlikeSig) evalInt(ctx EvalContext, row chunk.Row) (int64, bool, var pattern collate.WildcardPattern if b.args[1].ConstLevel() >= ConstOnlyInContext && b.args[2].ConstLevel() >= ConstOnlyInContext { pattern, err = b.patternCache.getOrInitCache(ctx, func() (collate.WildcardPattern, error) { - ret := collate.ConvertAndGetBinCollator(b.collation).Pattern() + ret := b.collator().Pattern() ret.Compile(patternStr, byte(escape)) return ret, nil }) @@ -106,7 +107,7 @@ func (b *builtinIlikeSig) evalInt(ctx EvalContext, row chunk.Row) (int64, bool, return 0, true, err } } else { - pattern = collate.ConvertAndGetBinCollator(b.collation).Pattern() + pattern = b.collator().Pattern() pattern.Compile(patternStr, byte(escape)) } return boolToInt64(pattern.DoMatch(valStr)), false, nil diff --git a/pkg/expression/builtin_ilike_test.go b/pkg/expression/builtin_ilike_test.go index 1a89b3e2eccc7..7ba84eb034d5b 100644 --- a/pkg/expression/builtin_ilike_test.go +++ b/pkg/expression/builtin_ilike_test.go @@ -83,7 +83,7 @@ func TestIlike(t *testing.T) { f, err := fc.getFunction(ctx, inputs) require.NoError(t, err, comment) f.SetCharsetAndCollation(charsetAndCollation[0], charsetAndCollation[1]) - f.setCollator(collate.GetCollator(charsetAndCollation[1])) + f.setCollator(collate.GetCollator(collate.ConvertAndGetBinCollation(charsetAndCollation[1]))) r, err := evalBuiltinFunc(f, ctx, chunk.Row{}) require.NoError(t, err, comment) testutil.DatumEqual(t, types.NewDatum(tt.generalMatch), r, comment) @@ -104,7 +104,7 @@ func TestIlike(t *testing.T) { f, err := fc.getFunction(ctx, inputs) require.NoError(t, err, comment) f.SetCharsetAndCollation(charsetAndCollation[0], charsetAndCollation[1]) - f.setCollator(collate.GetCollator(charsetAndCollation[1])) + f.setCollator(collate.GetCollator(collate.ConvertAndGetBinCollation(charsetAndCollation[1]))) r, err := evalBuiltinFunc(f, ctx, chunk.Row{}) require.NoError(t, err, comment) testutil.DatumEqual(t, types.NewDatum(tt.unicodeMatch), r, comment) diff --git a/pkg/expression/builtin_other.go b/pkg/expression/builtin_other.go index 35e017927f7bf..c988ced93ee4a 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,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) + collator := b.collator() // Keep track of unique args count for in-place modification uniqueArgCount := 1 // Start with 1 for the first arg (value to check) @@ -407,7 +406,7 @@ func (b *builtinInStringSig) evalInt(ctx EvalContext, row chunk.Row) (int64, boo } args := b.args[1:] - collator := collate.GetCollator(b.collation) + collator := b.collator() if len(b.hashSet) != 0 { if b.hashSet.Exist(string(collator.Key(arg0))) { return 1, false, nil @@ -428,7 +427,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.collator().Compare(arg0, evaledArg) == 0 { return 1, false, nil } } diff --git a/pkg/expression/builtin_string.go b/pkg/expression/builtin_string.go index f6299f988ff66..b4240873ea686 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.collator().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.collator()), false, nil } type builtinLocate3ArgsSig struct { @@ -1666,7 +1666,7 @@ func (b *builtinLocate3ArgsUTF8Sig) evalInt(ctx EvalContext, row chunk.Row) (int if isNull || err != nil { return 0, isNull, err } - if collate.IsCICollation(b.collation) { + if !b.collatorOverridden && collate.IsCICollation(b.collation) { subStr = strings.ToLower(subStr) str = strings.ToLower(str) } @@ -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.collator()) if idx != 0 { return pos + idx, false, nil } @@ -4254,6 +4254,7 @@ func (c *weightStringFunctionClass) getFunction(ctx BuildContext, args []Express if err != nil { return nil, err } + bf.setCollator(getCollator(ctx, bf.args[0].GetType(ctx.GetEvalCtx()).GetCollate())) types.SetBinChsClnFlag(bf.tp) var sig builtinFunc if padding == weightStringPaddingNull { @@ -4327,7 +4328,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.collator() case weightStringPaddingAsBinary: lenStr := len(str) if b.length < lenStr { @@ -4341,9 +4342,9 @@ func (b *builtinWeightStringSig) evalString(ctx EvalContext, row chunk.Row) (str } str += strings.Repeat("\x00", b.length-lenStr) } - ctor = collate.GetCollator(charset.CollationBin) + ctor = collate.GetBinaryCollator() case weightStringPaddingNone: - ctor = collate.GetCollator(b.args[0].GetType(ctx).GetCollate()) + ctor = b.collator() default: return "", false, ErrIncorrectType.GenWithStackByArgs(ast.WeightString, string(b.padding)) } diff --git a/pkg/expression/builtin_string_test.go b/pkg/expression/builtin_string_test.go index c9814878f2813..e42c6338f8ba6 100644 --- a/pkg/expression/builtin_string_test.go +++ b/pkg/expression/builtin_string_test.go @@ -1078,6 +1078,24 @@ func TestLocate(t *testing.T) { require.NotNil(t, f) require.Equalf(t, c["Want"][0], got, "[%d]: args: %v", i, c["Args"]) } + + args := primitiveValsToConstants(ctx, []any{"A", "a", 1}) + for _, arg := range args[:2] { + arg.GetType(ctx).SetCharset(charset.CharsetUTF8MB4) + arg.GetType(ctx).SetCollate("utf8mb4_general_ci") + } + f, err := instr.getFunction(ctx, args) + require.NoError(t, err) + require.IsType(t, &builtinLocate3ArgsUTF8Sig{}, f) + f.SetCharsetAndCollation(charset.CharsetUTF8MB4, "utf8mb4_general_ci") + got, err := evalBuiltinFunc(f, ctx, chunk.Row{}) + require.NoError(t, err) + require.Equal(t, int64(1), got.GetInt64()) + f.setCollator(getCollator(ctx, charset.CollationBin)) + got, err = evalBuiltinFunc(f, ctx, chunk.Row{}) + require.NoError(t, err) + require.Equal(t, int64(0), got.GetInt64()) + // 2. Test LOCATE with binary input tbl2 := []struct { Args []any diff --git a/pkg/expression/exprctx/context.go b/pkg/expression/exprctx/context.go index e752af29c6a74..54975c86af103 100644 --- a/pkg/expression/exprctx/context.go +++ b/pkg/expression/exprctx/context.go @@ -102,6 +102,8 @@ type BuildContext interface { GetCharsetInfo() (string, string) // GetDefaultCollationForUTF8MB4 returns the default collation of UTF8MB4. GetDefaultCollationForUTF8MB4() string + // NewCollationEnabled returns whether expression building should use new collation semantics. + NewCollationEnabled() bool // GetBlockEncryptionMode returns the variable `block_encryption_mode`. GetBlockEncryptionMode() string // GetSysdateIsNow returns a bool to determine whether Sysdate is an alias of Now function. diff --git a/pkg/expression/expression.go b/pkg/expression/expression.go index aa396e8d5a33a..cd28fdba3e415 100644 --- a/pkg/expression/expression.go +++ b/pkg/expression/expression.go @@ -65,8 +65,6 @@ type BuildOptions struct { AllowCastArray bool // TargetFieldType indicates to cast the expression to the target field type if it is not nil TargetFieldType *types.FieldType - // UseNewCollate whether to use new collate when building expression. - UseNewCollate bool } // BuildOption is a function to apply optional settings @@ -104,12 +102,8 @@ func WithCastExprTo(targetFt *types.FieldType) BuildOption { } } -// WithUseNewCollate fixes the collation mode used while building expressions -// that must stay consistent with table or index encoding created earlier. -func WithUseNewCollate(useNewCollate bool) BuildOption { - return func(options *BuildOptions) { - options.UseNewCollate = useNewCollate - } +func getCollator(ctx BuildContext, collation string) collate.Collator { + return collate.GetCollatorWithCollate(ctx.NewCollationEnabled(), collation) } // BuildSimpleExpr builds a simple expression from an ast node. @@ -1106,18 +1100,11 @@ 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( +func ColumnInfos2ColumnsAndNames( 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)) @@ -1161,8 +1148,7 @@ func ColumnInfos2ColumnsAndNamesWithCollate( } e, err := BuildSimpleExpr(ctx, expr, WithInputSchemaAndNames(mockSchema, names, tblInfo), - WithAllowCastArray(true), - WithUseNewCollate(useNewCollate)) + WithAllowCastArray(true)) if err != nil { return nil, nil, errors.Trace(err) } diff --git a/pkg/expression/exprstatic/BUILD.bazel b/pkg/expression/exprstatic/BUILD.bazel index b5d2e3789bde2..484660a6a4481 100644 --- a/pkg/expression/exprstatic/BUILD.bazel +++ b/pkg/expression/exprstatic/BUILD.bazel @@ -17,6 +17,7 @@ go_library( "//pkg/sessionctx/vardef", "//pkg/sessionctx/variable", "//pkg/types", + "//pkg/util/collate", "//pkg/util/context", "//pkg/util/intest", "//pkg/util/mathutil", diff --git a/pkg/expression/exprstatic/exprctx.go b/pkg/expression/exprstatic/exprctx.go index 60b9bb07a2fd2..12ccdb04378b4 100644 --- a/pkg/expression/exprstatic/exprctx.go +++ b/pkg/expression/exprstatic/exprctx.go @@ -22,6 +22,7 @@ import ( "github.com/pingcap/tidb/pkg/parser/mysql" "github.com/pingcap/tidb/pkg/sessionctx/vardef" "github.com/pingcap/tidb/pkg/sessionctx/variable" + "github.com/pingcap/tidb/pkg/util/collate" contextutil "github.com/pingcap/tidb/pkg/util/context" "github.com/pingcap/tidb/pkg/util/intest" "github.com/pingcap/tidb/pkg/util/mathutil" @@ -46,6 +47,7 @@ type exprCtxState struct { connectionID uint64 windowingUseHighPrecision bool groupConcatMaxLen uint64 + newCollationEnabled *bool } // ExprCtxOption is the option to create or update the `ExprContext` @@ -141,6 +143,13 @@ func WithGroupConcatMaxLen(maxLen uint64) ExprCtxOption { } } +// WithNewCollationEnabled fixes the new-collation mode for expression building. +func WithNewCollationEnabled(enabled bool) ExprCtxOption { + return func(s *exprCtxState) { + s.newCollationEnabled = &enabled + } +} + // ExprContext implements the `exprctx.ExprContext` interface. // The "static" means comparing with `ExprContext`, its internal state does not relay on the session or other // complex contexts that keeps immutable for most fields. @@ -223,6 +232,14 @@ func (ctx *ExprContext) GetDefaultCollationForUTF8MB4() string { return ctx.defaultCollationForUTF8MB4 } +// NewCollationEnabled implements the `ExprContext.NewCollationEnabled`. +func (ctx *ExprContext) NewCollationEnabled() bool { + if ctx.newCollationEnabled != nil { + return *ctx.newCollationEnabled + } + return collate.NewCollationEnabled() +} + // GetBlockEncryptionMode implements the `ExprContext.GetBlockEncryptionMode`. func (ctx *ExprContext) GetBlockEncryptionMode() string { return ctx.blockEncryptionMode diff --git a/pkg/expression/exprstatic/exprctx_test.go b/pkg/expression/exprstatic/exprctx_test.go index 2950c5af703cd..177747d2a6407 100644 --- a/pkg/expression/exprstatic/exprctx_test.go +++ b/pkg/expression/exprstatic/exprctx_test.go @@ -182,6 +182,7 @@ func TestMakeExprContextStatic(t *testing.T) { ignorePath := []string{ "$.exprCtxState.evalCtx**", + "$.exprCtxState.newCollationEnabled", } deeptest.AssertRecursivelyNotEqual(t, obj, NewExprContext(), deeptest.WithIgnorePath(ignorePath), @@ -296,6 +297,7 @@ func TestExprCtxLoadSystemVars(t *testing.T) { "$.planCacheTracker", "$.columnIDAllocator", "$.connectionID", + "$.newCollationEnabled", } // varsRelatedFields means the fields related to diff --git a/pkg/expression/sessionexpr/BUILD.bazel b/pkg/expression/sessionexpr/BUILD.bazel index a1dbf22fe313c..e1d179ca49f3a 100644 --- a/pkg/expression/sessionexpr/BUILD.bazel +++ b/pkg/expression/sessionexpr/BUILD.bazel @@ -20,6 +20,7 @@ go_library( "//pkg/sessionctx/variable", "//pkg/types", "//pkg/util", + "//pkg/util/collate", "//pkg/util/context", "//pkg/util/intest", "//pkg/util/logutil", diff --git a/pkg/expression/sessionexpr/sessionctx.go b/pkg/expression/sessionexpr/sessionctx.go index a3ec0c7a9d4b9..1cc445d4452b2 100644 --- a/pkg/expression/sessionexpr/sessionctx.go +++ b/pkg/expression/sessionexpr/sessionctx.go @@ -33,6 +33,7 @@ import ( "github.com/pingcap/tidb/pkg/sessionctx/variable" "github.com/pingcap/tidb/pkg/types" "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/intest" "github.com/pingcap/tidb/pkg/util/logutil" @@ -73,6 +74,11 @@ func (ctx *ExprContext) GetDefaultCollationForUTF8MB4() string { return ctx.sctx.GetSessionVars().DefaultCollationForUTF8MB4 } +// NewCollationEnabled returns whether expression building should use new collation semantics. +func (ctx *ExprContext) NewCollationEnabled() bool { + return collate.NewCollationEnabled() +} + // GetBlockEncryptionMode returns the variable block_encryption_mode func (ctx *ExprContext) GetBlockEncryptionMode() string { blockMode, ok := ctx.sctx.GetSessionVars().GetSystemVar(vardef.BlockEncryptionMode) diff --git a/pkg/expression/util.go b/pkg/expression/util.go index ea8f4d1aa448a..c227629b8a640 100644 --- a/pkg/expression/util.go +++ b/pkg/expression/util.go @@ -818,8 +818,7 @@ func SubstituteCorCol2Constant(ctx BuildContext, expr Expression) (Expression, e return expr, nil } -func locateStringWithCollation(str, substr, coll string) int64 { - collator := collate.GetCollator(coll) +func locateStringWithCollator(str, substr string, collator collate.Collator) int64 { strKey := collator.KeyWithoutTrimRightSpace(str) subStrKey := collator.KeyWithoutTrimRightSpace(substr) diff --git a/pkg/lightning/backend/kv/base.go b/pkg/lightning/backend/kv/base.go index b05c4dc784731..cb7cc14abc3d4 100644 --- a/pkg/lightning/backend/kv/base.go +++ b/pkg/lightning/backend/kv/base.go @@ -141,6 +141,7 @@ func NewBaseKVEncoder(config *encode.EncodingConfig) (*BaseKVEncoder, error) { if err != nil { return nil, err } + se.exprCtx.setNewCollationEnabled(config.Table.UseNewCollate()) var autoRandomColID int64 autoIDFn := func(id int64) int64 { return id } diff --git a/pkg/lightning/backend/kv/base_test.go b/pkg/lightning/backend/kv/base_test.go index bfd551a7fa540..a7525e099fd51 100644 --- a/pkg/lightning/backend/kv/base_test.go +++ b/pkg/lightning/backend/kv/base_test.go @@ -44,7 +44,7 @@ func TestLogKVConvertFailed(t *testing.T) { cols := []*model.ColumnInfo{c1} tblInfo := &model.TableInfo{ID: 1, Columns: cols, PKIsHandle: false, State: model.StatePublic} var tbl table.Table - tbl, err = tables.TableFromMeta(NewPanickingAllocators(tblInfo.SepAutoInc()), tblInfo) + tbl, err = tables.TableFromMetaWithCollate(false, NewPanickingAllocators(tblInfo.SepAutoInc()), tblInfo) require.NoError(t, err) var baseKVEncoder *BaseKVEncoder @@ -56,6 +56,13 @@ func TestLogKVConvertFailed(t *testing.T) { }, Logger: log.L(), }) + require.NoError(t, err) + require.False(t, baseKVEncoder.SessionCtx.GetExprCtx().NewCollationEnabled()) + tbl, err = tables.TableFromMetaWithCollate(true, NewPanickingAllocators(tblInfo.SepAutoInc()), tblInfo) + require.NoError(t, err) + newCollationEncoder, err := NewBaseKVEncoder(&encode.EncodingConfig{Table: tbl, Logger: log.L()}) + require.NoError(t, err) + require.True(t, newCollationEncoder.SessionCtx.GetExprCtx().NewCollationEnabled()) var newString strings.Builder for range 100000 { newString.WriteString("test_test_test_test_") diff --git a/pkg/lightning/backend/kv/context.go b/pkg/lightning/backend/kv/context.go index 134209447d049..3f2eb700c8495 100644 --- a/pkg/lightning/backend/kv/context.go +++ b/pkg/lightning/backend/kv/context.go @@ -96,6 +96,10 @@ func newLitExprContext(sqlMode mysql.SQLMode, sysVars map[string]string, timesta }, nil } +func (ctx *litExprContext) setNewCollationEnabled(enabled bool) { + ctx.ExprContext = ctx.ExprContext.Apply(exprstatic.WithNewCollationEnabled(enabled)) +} + // setUserVarVal sets the value of a user variable. func (ctx *litExprContext) setUserVarVal(name string, dt types.Datum) { ctx.userVars.SetUserVarVal(name, dt) diff --git a/pkg/lightning/backend/kv/kv2sql.go b/pkg/lightning/backend/kv/kv2sql.go index fa589c8fb2829..c1b78745066d0 100644 --- a/pkg/lightning/backend/kv/kv2sql.go +++ b/pkg/lightning/backend/kv/kv2sql.go @@ -133,6 +133,7 @@ func NewTableKVDecoder( if err != nil { return nil, err } + se.exprCtx.setNewCollationEnabled(tbl.UseNewCollate()) genCols, err := CollectGeneratedColumns(se, tbl) if err != nil { diff --git a/pkg/lightning/backend/kv/sql2kv.go b/pkg/lightning/backend/kv/sql2kv.go index 140c8eadccd30..4493d44bd5189 100644 --- a/pkg/lightning/backend/kv/sql2kv.go +++ b/pkg/lightning/backend/kv/sql2kv.go @@ -114,7 +114,6 @@ func CollectGeneratedColumns(se *Session, tbl table.Table) ([]GeneratedCol, erro col.GeneratedExpr.Internal(), expression.WithInputSchemaAndNames(schema, names, meta), expression.WithAllowCastArray(true), - 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 76d7e34cf729b..2f6b59077fe7c 100644 --- a/pkg/meta/model/reorg.go +++ b/pkg/meta/model/reorg.go @@ -96,10 +96,10 @@ type DDLReorgMeta struct { AnalyzeState int8 `json:"analyze_state"` Stage ReorgStage `json:"stage"` // 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. + // in the submitting keyspace. 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. @@ -166,8 +166,8 @@ func (dm *DDLReorgMeta) GetUseNewCollateOrDefault(defaultVal bool) bool { return *dm.UseNewCollate } -// setUseNewCollate stores the new-collation mode captured from the persisted -// table snapshot. +// setUseNewCollate stores the new-collation mode captured from the submitting +// keyspace. func (dm *DDLReorgMeta) setUseNewCollate(useNewCollate bool) { dm.UseNewCollate = &useNewCollate } diff --git a/pkg/planner/core/expression_rewriter.go b/pkg/planner/core/expression_rewriter.go index f380aa2bf0147..2caa19302534e 100644 --- a/pkg/planner/core/expression_rewriter.go +++ b/pkg/planner/core/expression_rewriter.go @@ -112,9 +112,7 @@ func buildSimpleExpr(ctx expression.BuildContext, node ast.ExprNode, opts ...exp return nil, errors.New("expression node should be present") } - options := expression.BuildOptions{ - UseNewCollate: collate.NewCollationEnabled(), - } + var options expression.BuildOptions for _, opt := range opts { opt(&options) } @@ -152,7 +150,6 @@ func buildSimpleExpr(ctx expression.BuildContext, node ast.ExprNode, opts ...exp sourceTable: options.SourceTable, allowBuildCastArray: options.AllowCastArray, asScalar: true, - useNewCollate: options.UseNewCollate, } if tbl := options.SourceTable; tbl != nil && rewriter.schema == nil { @@ -261,8 +258,7 @@ func (b *PlanBuilder) getExpressionRewriter(ctx context.Context, p base.LogicalP if len(b.rewriterPool) < b.rewriterCounter { rewriter = &expressionRewriter{ sctx: b.ctx.GetExprCtx(), ctx: ctx, - planCtx: &exprRewriterPlanCtx{plan: p, builder: b, curClause: b.curClause, rollExpand: b.currentBlockExpand}, - useNewCollate: collate.NewCollationEnabled(), + planCtx: &exprRewriterPlanCtx{plan: p, builder: b, curClause: b.curClause, rollExpand: b.currentBlockExpand}, } b.rewriterPool = append(b.rewriterPool, rewriter) return @@ -379,8 +375,7 @@ type expressionRewriter struct { astNodeStack []ast.Node - planCtx *exprRewriterPlanCtx - useNewCollate bool + planCtx *exprRewriterPlanCtx } func (er *expressionRewriter) ctxStackLen() int { @@ -1852,7 +1847,7 @@ func (er *expressionRewriter) Leave(originInNode ast.Node) (retNode ast.Node, ok }, types.EmptyName) case *ast.SetCollationExpr: arg := er.ctxStack[len(er.ctxStack)-1] - if er.useNewCollate { + if er.sctx.NewCollationEnabled() { var collInfo *charset.Collation // TODO(bb7133): use charset.ValidCharsetAndCollation when its bug is fixed. if collInfo, er.err = collate.GetCollationByName(v.Collate); er.err != nil { @@ -2280,7 +2275,7 @@ func (er *expressionRewriter) castCollationForIn(colLen int, elemCnt int, stkLen if colLen != 1 { return } - if !er.useNewCollate { + if !er.sctx.NewCollationEnabled() { // See https://github.com/pingcap/tidb/issues/52772 // This function will apply CoercibilityExplicit to the casted expression, but some checks(during ColumnSubstituteImpl) is missed when the new // collation is disabled, then lead to panic. @@ -2398,7 +2393,7 @@ func (er *expressionRewriter) patternLikeOrIlikeToExpression(v *ast.PatternLikeO fieldType := &types.FieldType{} isPatternExactMatch := false // Treat predicate 'like' or 'ilike' the same way as predicate '=' when it is an exact match and new collation is not enabled. - if patExpression, ok := er.ctxStack[l-1].(*expression.Constant); ok && !er.useNewCollate { + if patExpression, ok := er.ctxStack[l-1].(*expression.Constant); ok && !er.sctx.NewCollationEnabled() { patString, isNull, err := patExpression.EvalString(er.sctx.GetEvalCtx(), chunk.Row{}) if err != nil { er.err = err diff --git a/pkg/planner/core/expression_test.go b/pkg/planner/core/expression_test.go index 766e7c6f1b3e5..6d7d57c4eda0d 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" ) @@ -488,6 +489,74 @@ func TestBuildExpression(t *testing.T) { require.NoError(t, err) require.Equal(t, int64(10), val) + origin := collate.NewCollationEnabled() + collate.SetNewCollationEnabledForTest(true) + defer collate.SetNewCollationEnabledForTest(origin) + collationSensitiveTbl := &model.TableInfo{ + Name: ast.NewCIStr("tc"), + Columns: []*model.ColumnInfo{ + { + Name: ast.NewCIStr("c0"), + Offset: 0, + State: model.StatePublic, + FieldType: *types.NewFieldTypeWithCollation(mysql.TypeVarchar, "utf8mb4_general_ci", 16), + }, + { + Name: ast.NewCIStr("c1"), + Offset: 1, + State: model.StatePublic, + FieldType: *types.NewFieldTypeWithCollation(mysql.TypeVarchar, "utf8mb4_general_ci", 16), + }, + }, + } + collationSensitiveRow := chunk.MutRowFromValues("a", "A").ToRow() + oldCollationCtx := ctx.Apply(exprstatic.WithNewCollationEnabled(false)) + newCollationCtx := ctx.Apply(exprstatic.WithNewCollationEnabled(true)) + for _, test := range []struct { + expr string + oldExpected any + newExpected any + }{ + {"c0 = 'A'", int64(0), int64(1)}, + {"c0 <=> c1", int64(0), int64(1)}, + {"c0 != c1", int64(1), int64(0)}, + {"c1 < c0", int64(1), int64(0)}, + {"c0 <= c1", int64(0), int64(1)}, + {"c0 > c1", int64(1), int64(0)}, + {"c1 >= c0", int64(0), int64(1)}, + {"c0 between c1 and c1", int64(0), int64(1)}, + {"c0 in (c1)", int64(0), int64(1)}, + {"c0 like c1", int64(0), int64(1)}, + {"c0 ilike c1", int64(1), int64(1)}, + {"c0 regexp c1", int64(1), int64(1)}, + {"if(c0 = c1, 'same', 'different')", "different", "same"}, + {"nullif(c0, c1)", "a", nil}, + {"case when c0 = c1 then 'same' else 'different' end", "different", "same"}, + {"case c0 when c1 then 'same' else 'different' end", "different", "same"}, + {"strcmp(c0, c1)", int64(1), int64(0)}, + {"field(c0, c1)", int64(0), int64(1)}, + {"find_in_set(c0, c1)", int64(0), int64(1)}, + {"greatest(c1, c0)", "a", "A"}, + {"least(c0, c1)", "A", "a"}, + {"locate(c1, c0)", int64(0), int64(1)}, + {"position(c1 in c0)", int64(0), int64(1)}, + {"locate(c1, c0, 1)", int64(1), int64(1)}, + {"instr(c0, c1)", int64(1), int64(1)}, + {"weight_string(c0)", "a", "\x00A"}, + } { + expr, err = buildExpr(t, oldCollationCtx, test.expr, expression.WithTableInfo("", collationSensitiveTbl)) + require.NoError(t, err) + result, err := expr.Eval(evalCtx, collationSensitiveRow) + require.NoError(t, err) + require.Equal(t, test.oldExpected, result.GetValue(), test.expr) + + expr, err = buildExpr(t, newCollationCtx, test.expr, expression.WithTableInfo("", collationSensitiveTbl)) + require.NoError(t, err) + result, err = expr.Eval(evalCtx, collationSensitiveRow) + require.NoError(t, err) + require.Equal(t, test.newExpected, result.GetValue(), test.expr) + } + // build expression without enough columns _, err = buildExpr(t, ctx, "1+a") require.EqualError(t, err, "[planner:1054]Unknown column 'a' in 'expression'") diff --git a/pkg/planner/core/rule/rule_partition_processor.go b/pkg/planner/core/rule/rule_partition_processor.go index 33dff10418c26..aa10aeca23244 100644 --- a/pkg/planner/core/rule/rule_partition_processor.go +++ b/pkg/planner/core/rule/rule_partition_processor.go @@ -278,7 +278,9 @@ func (s *PartitionProcessor) getUsedKeyPartitions(ctx base.PlanContext, pi := tbl.Meta().Partition partExpr := tbl.(base.PartitionTable).PartitionExpr() partCols, colLen := partExpr.GetPartColumnsForKeyPartition(columns) - pe := &tables.ForKeyPruning{KeyPartCols: partCols} + // Copy the existing pruning state to preserve its captured collation mode. + pe := *partExpr.ForKeyPruning + pe.KeyPartCols = partCols detachedResult, err := ranger.DetachCondAndBuildRangeForPartition(ctx.GetRangerCtx(), conds, partCols, colLen, ctx.GetSessionVars().RangeMaxSize) if err != nil { return nil, err diff --git a/pkg/server/handler/tests/BUILD.bazel b/pkg/server/handler/tests/BUILD.bazel index 4dc6e5de6d81b..f505f12b2098b 100644 --- a/pkg/server/handler/tests/BUILD.bazel +++ b/pkg/server/handler/tests/BUILD.bazel @@ -58,7 +58,6 @@ go_test( "//pkg/testkit/testsetup", "//pkg/types", "//pkg/util/codec", - "//pkg/util/collate", "//pkg/util/deadlockhistory", "//pkg/util/rowcodec", "//pkg/util/topsql/state", diff --git a/pkg/server/handler/tests/http_handler_test.go b/pkg/server/handler/tests/http_handler_test.go index c83eb753b543a..146b0684f91c1 100644 --- a/pkg/server/handler/tests/http_handler_test.go +++ b/pkg/server/handler/tests/http_handler_test.go @@ -79,7 +79,6 @@ import ( "github.com/pingcap/tidb/pkg/testkit/testfailpoint" "github.com/pingcap/tidb/pkg/types" "github.com/pingcap/tidb/pkg/util/codec" - "github.com/pingcap/tidb/pkg/util/collate" "github.com/pingcap/tidb/pkg/util/rowcodec" "github.com/stretchr/testify/require" "github.com/tikv/client-go/v2/tikv" @@ -656,7 +655,7 @@ func TestDecodeColumnValue(t *testing.T) { } rd := rowcodec.Encoder{Enable: true} sc := stmtctx.NewStmtCtxWithTimeZone(time.UTC) - bs, err := tablecodec.EncodeRow(codec.NewEncoder(collate.NewCollationEnabled()), sc.TimeZone(), row, colIDs, nil, nil, nil, &rd) + bs, err := tablecodec.EncodeRow(sc.TimeZone(), row, colIDs, nil, nil, nil, &rd) require.NoError(t, err) require.NotNil(t, bs) bin := base64.StdEncoding.EncodeToString(bs) diff --git a/pkg/store/mockstore/BUILD.bazel b/pkg/store/mockstore/BUILD.bazel index 014aaea0b9078..ff8df81b8a30f 100644 --- a/pkg/store/mockstore/BUILD.bazel +++ b/pkg/store/mockstore/BUILD.bazel @@ -52,7 +52,6 @@ go_test( "//pkg/testkit/testsetup", "//pkg/types", "//pkg/util/codec", - "//pkg/util/collate", "//pkg/util/rowcodec", "@com_github_pingcap_kvproto//pkg/kvrpcpb", "@com_github_stretchr_testify//require", diff --git a/pkg/store/mockstore/cluster_test.go b/pkg/store/mockstore/cluster_test.go index 8670adbd92006..e5685cb1e13d8 100644 --- a/pkg/store/mockstore/cluster_test.go +++ b/pkg/store/mockstore/cluster_test.go @@ -28,7 +28,6 @@ import ( "github.com/pingcap/tidb/pkg/tablecodec" "github.com/pingcap/tidb/pkg/types" "github.com/pingcap/tidb/pkg/util/codec" - "github.com/pingcap/tidb/pkg/util/collate" "github.com/pingcap/tidb/pkg/util/rowcodec" "github.com/stretchr/testify/require" "github.com/tikv/client-go/v2/testutils" @@ -59,7 +58,7 @@ func TestClusterSplit(t *testing.T) { colValue := types.NewStringDatum(strconv.Itoa(int(handle))) // TODO: Should use session's TimeZone instead of UTC. rd := rowcodec.Encoder{Enable: true} - rowValue, err1 := tablecodec.EncodeRow(codec.NewEncoder(collate.NewCollationEnabled()), sc.TimeZone(), []types.Datum{colValue}, []int64{colID}, nil, nil, nil, &rd) + rowValue, err1 := tablecodec.EncodeRow(sc.TimeZone(), []types.Datum{colValue}, []int64{colID}, nil, nil, nil, &rd) require.NoError(t, err1) txn.Set(rowKey, rowValue) diff --git a/pkg/store/mockstore/unistore/cophandler/cop_handler_test.go b/pkg/store/mockstore/unistore/cophandler/cop_handler_test.go index 849377052fabd..484e21cd722d4 100644 --- a/pkg/store/mockstore/unistore/cophandler/cop_handler_test.go +++ b/pkg/store/mockstore/unistore/cophandler/cop_handler_test.go @@ -119,11 +119,10 @@ func prepareTestTableData(keyNumber int, tableID int64) (*data, error) { rows := map[int64][]types.Datum{} encodedTestKVDatas := make([]*encodedTestKVData, keyNumber) encoder := &rowcodec.Encoder{Enable: true} - codecEncoder := codec.NewEncoder(collate.NewCollationEnabled()) for i := range keyNumber { datum := types.MakeDatums(i, "abc", 10.0) rows[int64(i)] = datum - rowEncodedData, err := tablecodec.EncodeRow(codecEncoder, stmtCtx.TimeZone(), datum, colIds, nil, nil, nil, encoder) + rowEncodedData, err := tablecodec.EncodeRow(stmtCtx.TimeZone(), datum, colIds, nil, nil, nil, encoder) if err != nil { return nil, err } diff --git a/pkg/table/BUILD.bazel b/pkg/table/BUILD.bazel index 28b868f92b05a..aac5f9db03f9c 100644 --- a/pkg/table/BUILD.bazel +++ b/pkg/table/BUILD.bazel @@ -58,6 +58,7 @@ go_test( "//pkg/errctx", "//pkg/errno", "//pkg/expression", + "//pkg/expression/exprstatic", "//pkg/meta/model", "//pkg/parser/ast", "//pkg/parser/charset", diff --git a/pkg/table/column.go b/pkg/table/column.go index d6d8e4d0db943..f06903f2d2e08 100644 --- a/pkg/table/column.go +++ b/pkg/table/column.go @@ -343,7 +343,21 @@ func CastColumnValueWithStrictMode(val types.Datum, tp *types.FieldType) (casted // CastColumnValue casts a value based on column type with expression BuildContext func CastColumnValue(ctx expression.BuildContext, val types.Datum, col *model.ColumnInfo, returnErr, forceIgnoreTruncate bool) (casted types.Datum, err error) { evalCtx := ctx.GetEvalCtx() - return castColumnValue(evalCtx.TypeCtx(), evalCtx.ErrCtx(), evalCtx.SQLMode(), val, &col.FieldType, col.Name.O, ctx.ConnectionID(), returnErr, forceIgnoreTruncate) + ft := &col.FieldType + useLegacyCollation := (ft.GetType() == mysql.TypeEnum || ft.GetType() == mysql.TypeSet) && !ctx.NewCollationEnabled() + if useLegacyCollation { + // Legacy ENUM/SET name matching is binary even when the field metadata has a non-binary collation. + ft = ft.Clone() + ft.SetCollate(charset.CollationBin) + } + casted, err = castColumnValue( + evalCtx.TypeCtx(), evalCtx.ErrCtx(), evalCtx.SQLMode(), val, ft, + col.Name.O, ctx.ConnectionID(), returnErr, forceIgnoreTruncate, + ) + if useLegacyCollation { + casted.SetCollation(col.GetCollate()) + } + return casted, err } // castColumnValue casts a value based on column type. diff --git a/pkg/table/column_test.go b/pkg/table/column_test.go index 987860ab19f6c..e2d3034234e7e 100644 --- a/pkg/table/column_test.go +++ b/pkg/table/column_test.go @@ -22,6 +22,7 @@ import ( "github.com/pingcap/tidb/pkg/errctx" "github.com/pingcap/tidb/pkg/expression" + "github.com/pingcap/tidb/pkg/expression/exprstatic" "github.com/pingcap/tidb/pkg/meta/model" "github.com/pingcap/tidb/pkg/parser/ast" "github.com/pingcap/tidb/pkg/parser/charset" @@ -331,6 +332,26 @@ func TestCastValue(t *testing.T) { val, err = CastValue(ctx, types.NewDatum([]byte{0xE5, 0xA5, 0xBD, 0x81}), &colInfoS, false, false) require.ErrorContains(t, err, "[table:1366]Incorrect string value '\\x81' for column ''") require.Equal(t, "utf8mb4_general_ci", val.Collation()) + + t.Run("fixed legacy collation for enum and set", func(t *testing.T) { + original := collate.NewCollationEnabled() + collate.SetNewCollationEnabledForTest(true) + t.Cleanup(func() { collate.SetNewCollationEnabledForTest(original) }) + + exprCtx := exprstatic.NewExprContext(exprstatic.WithNewCollationEnabled(false)) + for _, tp := range []byte{mysql.TypeEnum, mysql.TypeSet} { + colInfo := &model.ColumnInfo{FieldType: *types.NewFieldType(tp)} + colInfo.SetCharset(charset.CharsetUTF8MB4) + colInfo.SetCollate("utf8mb4_general_ci") + colInfo.SetElems([]string{"A", "a", "B"}) + + casted, err := CastColumnValue(exprCtx, types.NewStringDatum("a"), colInfo, false, false) + require.NoError(t, err) + require.Equal(t, "a", casted.GetString()) + require.Equal(t, uint64(2), casted.GetUint64()) + require.Equal(t, "utf8mb4_general_ci", casted.Collation()) + } + }) } func TestGetDefaultValue(t *testing.T) { diff --git a/pkg/table/tables/index.go b/pkg/table/tables/index.go index 5d433e5d656c8..296e38800b033 100644 --- a/pkg/table/tables/index.go +++ b/pkg/table/tables/index.go @@ -22,7 +22,6 @@ import ( "github.com/pingcap/errors" "github.com/pingcap/tidb/pkg/errctx" "github.com/pingcap/tidb/pkg/expression" - "github.com/pingcap/tidb/pkg/expression/exprctx" "github.com/pingcap/tidb/pkg/expression/exprstatic" "github.com/pingcap/tidb/pkg/kv" "github.com/pingcap/tidb/pkg/meta/model" @@ -44,7 +43,7 @@ import ( "go.uber.org/zap" ) -var indexConditionECtx exprctx.BuildContext +var indexConditionECtx *exprstatic.ExprContext // indexPartialCondition is a data structure to help implement the partial index. type indexPartialCondition struct { @@ -60,9 +59,8 @@ type index struct { idxInfo *model.IndexInfo tblInfo *model.TableInfo phyTblID int64 - // initNeedRestoreData is used to initialize `needRestoredData` in `index.Create()`. - // This routine cannot be done in `NewIndex()` because `needRestoreData` relies on `NewCollationEnabled()` and - // the collation global variable is initialized *after* `NewIndex()`. + // needRestoredData is initialized on first use. sync.Once makes the cached value + // safe for index instances shared by concurrent writers. initNeedRestoreData sync.Once needRestoredData bool encoder codec.Encoder @@ -82,55 +80,55 @@ func NeedRestoredData(useNewCollate bool, idxCols []*model.IndexColumn, colInfos // NewIndex builds a new Index object. func NewIndex(physicalID int64, tblInfo *model.TableInfo, indexInfo *model.IndexInfo) (table.Index, error) { - return NewIndexWithCollate(collate.NewCollationEnabled(), physicalID, tblInfo, indexInfo) + return newIndex(physicalID, tblInfo, indexInfo, collate.NewCollationEnabled()) } -// NewIndexWithCollate builds a new Index object with the specified collation setting. -func NewIndexWithCollate( - useNewCollate bool, - physicalID int64, - tblInfo *model.TableInfo, - indexInfo *model.IndexInfo, -) (table.Index, error) { - index := &index{ +func newIndex(physicalID int64, tblInfo *model.TableInfo, indexInfo *model.IndexInfo, useNewCollate bool) (*index, error) { + idx := &index{ idxInfo: indexInfo, tblInfo: tblInfo, phyTblID: physicalID, encoder: codec.NewEncoder(useNewCollate), } + if err := idx.initPartialCondition(); err != nil { + return nil, err + } + return idx, nil +} - conditionString := indexInfo.ConditionExprString - if len(conditionString) > 0 { - var err error - index.conditionExpr, err = expression.ParseSimpleExpr(indexConditionECtx, conditionString, - expression.WithTableInfo("", tblInfo), - expression.WithUseNewCollate(useNewCollate)) - if err != nil { - return nil, errors.Trace(err) - } - index.conditionEvalBufferPool = sync.Pool{ - New: func() any { - // For INSERT path, it'll only pass all writable columns. - // For UPDATE/DELETE path, it'll contain all columns. - // As the writable columns are always at the beginning of the `tblInfo.Columns`, it'll not affect - // the offsets of related columns in the expression. Therefore, it's fine to always record all - // columns here. - evalBufferTypes := make([]*types.FieldType, 0, len(tblInfo.Columns)+1) - for _, col := range tblInfo.Columns { - evalBufferTypes = append(evalBufferTypes, &col.FieldType) - } +func (c *index) initPartialCondition() error { + conditionString := c.idxInfo.ConditionExprString + if len(conditionString) == 0 { + return nil + } + ctx := indexConditionECtx.Apply(exprstatic.WithNewCollationEnabled(c.encoder.UseNewCollate())) + conditionExpr, err := expression.ParseSimpleExpr(ctx, conditionString, expression.WithTableInfo("", c.tblInfo)) + if err != nil { + return errors.Trace(err) + } + c.conditionExpr = conditionExpr + c.conditionEvalBufferPool = sync.Pool{ + New: func() any { + // For INSERT path, it'll only pass all writable columns. + // For UPDATE/DELETE path, it'll contain all columns. + // As the writable columns are always at the beginning of the `tblInfo.Columns`, it'll not affect + // the offsets of related columns in the expression. Therefore, it's fine to always record all + // columns here. + evalBufferTypes := make([]*types.FieldType, 0, len(c.tblInfo.Columns)+1) + for _, col := range c.tblInfo.Columns { + evalBufferTypes = append(evalBufferTypes, &col.FieldType) + } - if !tblInfo.HasClusteredIndex() { - // If the table doesn't have clustered index, we need to append an extra handle column. - evalBufferTypes = append(evalBufferTypes, types.NewFieldType(mysql.TypeLonglong)) - } + if !c.tblInfo.HasClusteredIndex() { + // If the table doesn't have clustered index, we need to append an extra handle column. + evalBufferTypes = append(evalBufferTypes, types.NewFieldType(mysql.TypeLonglong)) + } - evalBuffer := chunk.MutRowFromTypes(evalBufferTypes) - return &evalBuffer - }, - } + evalBuffer := chunk.MutRowFromTypes(evalBufferTypes) + return &evalBuffer + }, } - return index, nil + return nil } // Meta returns index info. diff --git a/pkg/table/tables/mutation_checker_test.go b/pkg/table/tables/mutation_checker_test.go index 7e83ba2fa251e..58aa36724a356 100644 --- a/pkg/table/tables/mutation_checker_test.go +++ b/pkg/table/tables/mutation_checker_test.go @@ -94,8 +94,7 @@ func TestCheckRowInsertionConsistency(t *testing.T) { // mocked data mockRowKey233 := tablecodec.EncodeRowKeyWithHandle(1, kv.IntHandle(233)) - enc := codec.NewEncoder(collate.NewCollationEnabled()) - mockValue233, err := tablecodec.EncodeRow(enc, sessVars.StmtCtx.TimeZone(), []types.Datum{types.NewIntDatum(233)}, []int64{101}, nil, nil, nil, &rd) + mockValue233, err := tablecodec.EncodeRow(sessVars.StmtCtx.TimeZone(), []types.Datum{types.NewIntDatum(233)}, []int64{101}, nil, nil, nil, &rd) require.Nil(t, err) fakeRowInsertion := mutation{key: []byte{1, 1}, value: []byte{1, 1, 1}} @@ -319,7 +318,7 @@ func TestCheckIndexKeysAndCheckHandleConsistency(t *testing.T) { // test checkHandleConsistency rowKey := tablecodec.EncodeRowKeyWithHandle(table.tableID, handle) corruptedRowKey := tablecodec.EncodeRowKeyWithHandle(table.tableID, corruptedHandle) - rowValue, err := tablecodec.EncodeRow(table.encoder, tc.Location(), rowToInsert, []int64{1, 2}, nil, nil, nil, &rd) + rowValue, err := tablecodec.EncodeRow(tc.Location(), rowToInsert, []int64{1, 2}, nil, nil, nil, &rd) require.Nil(t, err) rowMutation := mutation{key: rowKey, value: rowValue} corruptedRowMutation := mutation{key: corruptedRowKey, value: rowValue} diff --git a/pkg/table/tables/partition.go b/pkg/table/tables/partition.go index 0f2cd754e9a56..215d2db9ec531 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" @@ -115,7 +116,7 @@ func newPartitionedTable(tbl *TableCommon, tblInfo *model.TableInfo) (table.Part return nil, table.ErrUnknownPartition } ret := &partitionedTable{TableCommon: tbl.Copy()} - partitionExpr, err := newPartitionExpr(tblInfo, pi.Type, pi.Expr, pi.Columns, pi.Definitions) + partitionExpr, err := ret.newPartitionExpr(pi.Type, pi.Expr, pi.Columns, pi.Definitions) if err != nil { return nil, errors.Trace(err) } @@ -215,9 +216,9 @@ func newPartitionedTable(tbl *TableCommon, tblInfo *model.TableInfo) (table.Part if pi.DDLState == model.StateDeleteReorganization || pi.DDLState == model.StatePublic { // TODO: Explicitly explain the different DDL/New fields! if pi.NewTableID != 0 { - ret.reorgPartitionExpr, err = newPartitionExpr(tblInfo, pi.DDLType, pi.DDLExpr, pi.DDLColumns, pi.DroppingDefinitions) + ret.reorgPartitionExpr, err = ret.newPartitionExpr(pi.DDLType, pi.DDLExpr, pi.DDLColumns, pi.DroppingDefinitions) } else { - ret.reorgPartitionExpr, err = newPartitionExpr(tblInfo, pi.Type, pi.Expr, pi.Columns, pi.DroppingDefinitions) + ret.reorgPartitionExpr, err = ret.newPartitionExpr(pi.Type, pi.Expr, pi.Columns, pi.DroppingDefinitions) } if err != nil { return nil, errors.Trace(err) @@ -241,10 +242,10 @@ func newPartitionedTable(tbl *TableCommon, tblInfo *model.TableInfo) (table.Part if len(pi.AddingDefinitions) > 0 { if pi.NewTableID != 0 { // REMOVE PARTITIONING or PARTITION BY - ret.reorgPartitionExpr, err = newPartitionExpr(tblInfo, pi.DDLType, pi.DDLExpr, pi.DDLColumns, pi.AddingDefinitions) + ret.reorgPartitionExpr, err = ret.newPartitionExpr(pi.DDLType, pi.DDLExpr, pi.DDLColumns, pi.AddingDefinitions) } else { // REORGANIZE PARTITION - ret.reorgPartitionExpr, err = newPartitionExpr(tblInfo, pi.Type, pi.Expr, pi.Columns, pi.AddingDefinitions) + ret.reorgPartitionExpr, err = ret.newPartitionExpr(pi.Type, pi.Expr, pi.Columns, pi.AddingDefinitions) } if err != nil { return nil, errors.Trace(err) @@ -282,7 +283,7 @@ func initPartition(t *partitionedTable, def model.PartitionDefinition) (*partiti } // NewPartitionExprBuildCtx returns a context to build partition expression. -func NewPartitionExprBuildCtx() expression.BuildContext { +func NewPartitionExprBuildCtx() *exprstatic.ExprContext { return exprstatic.NewExprContext( exprstatic.WithEvalCtx(exprstatic.NewEvalContext( // Set a non-strict SQL mode and allow all date values if possible to make sure constant fold can work to @@ -302,8 +303,9 @@ func NewPartitionExprBuildCtx() expression.BuildContext { ) } -func newPartitionExpr(tblInfo *model.TableInfo, tp ast.PartitionType, expr string, partCols []ast.CIStr, defs []model.PartitionDefinition) (*PartitionExpr, error) { - ctx := NewPartitionExprBuildCtx() +func (t *partitionedTable) newPartitionExpr(tp ast.PartitionType, expr string, partCols []ast.CIStr, defs []model.PartitionDefinition) (*PartitionExpr, error) { + tblInfo := t.meta + ctx := NewPartitionExprBuildCtx().Apply(exprstatic.WithNewCollationEnabled(t.UseNewCollate())) dbName := ast.NewCIStr(ctx.GetEvalCtx().CurrentDB()) columns, names, err := expression.ColumnInfos2ColumnsAndNames(ctx, dbName, tblInfo.Name, tblInfo.Cols(), tblInfo) if err != nil { @@ -365,7 +367,7 @@ func (kp *ForKeyPruning) LocateKeyPartition(numParts uint64, r []types.Datum) (i if val.Kind() == types.KindNull { h.Write([]byte{0}) } else { - data, err := val.ToHashKey() + data, err := kp.datumToHashKey(&val) if err != nil { return 0, err } @@ -375,6 +377,19 @@ func (kp *ForKeyPruning) LocateKeyPartition(numParts uint64, r []types.Datum) (i return int(h.Sum32() % uint32(numParts)), nil } +func (kp *ForKeyPruning) datumToHashKey(d *types.Datum) ([]byte, error) { + switch d.Kind() { + case types.KindString, types.KindBytes: + return collate.GetCollatorWithCollate(kp.useNewCollate, d.Collation()).Key(d.GetString()), nil + default: + str, err := d.ToString() + if err != nil { + return nil, errors.Trace(err) + } + return collate.GetCollatorWithCollate(kp.useNewCollate, d.Collation()).Key(str), nil + } +} + func initEvalBufferType(t *partitionedTable) { hasExtraHandle := false numCols := len(t.WritableCols()) @@ -450,7 +465,8 @@ func parseSimpleExprWithNames(p *parser.Parser, ctx expression.BuildContext, exp // ForKeyPruning is used for key partition pruning. type ForKeyPruning struct { - KeyPartCols []*expression.Column + KeyPartCols []*expression.Column + useNewCollate bool } // ForListPruning is used for list partition pruning. @@ -508,6 +524,7 @@ func lessBtreeListColumnItem(a, b *btreeListColumnItem) bool { type ForListColumnPruning struct { ExprCol *expression.Column valueTp *types.FieldType + encoder codec.Encoder valueMap map[string]ListPartitionLocation sorted *btree.BTreeG[*btreeListColumnItem] @@ -735,7 +752,7 @@ func rangePartitionExprStrings(cols []ast.CIStr, expr string) []string { func generateKeyPartitionExpr(ctx expression.BuildContext, expr string, partCols []ast.CIStr, columns []*expression.Column, names types.NameSlice) (*PartitionExpr, error) { ret := &PartitionExpr{ - ForKeyPruning: &ForKeyPruning{}, + ForKeyPruning: &ForKeyPruning{useNewCollate: ctx.NewCollationEnabled()}, } _, partColumns, offset, err := extractPartitionExprColumns(ctx, expr, partCols, columns, names) if err != nil { @@ -954,6 +971,7 @@ func (lp *ForListPruning) buildListColumnsPruner(ctx expression.BuildContext, p := parser.New() colPrunes := make([]*ForListColumnPruning, 0, len(partCols)) lp.defaultPartitionIdx = -1 + useNewCollate := ctx.NewCollationEnabled() for colIdx := range partCols { colInfo := model.FindColumnInfo(tblInfo.Columns, partCols[colIdx].L) if colInfo == nil { @@ -971,6 +989,7 @@ func (lp *ForListPruning) buildListColumnsPruner(ctx expression.BuildContext, colIdx: colIdx, ExprCol: columns[idx], valueTp: &colInfo.FieldType, + encoder: codec.NewEncoder(useNewCollate), valueMap: make(map[string]ListPartitionLocation), sorted: btree.NewG[*btreeListColumnItem](btreeDegree, lessBtreeListColumnItem), } @@ -1241,7 +1260,7 @@ func (lp *ForListColumnPruning) genKey(tc types.Context, ec errctx.Context, v ty if err != nil { return nil, errors.Trace(err) } - valByte, err := codec.EncodeKey(tc.Location(), nil, v) + valByte, err := lp.encoder.EncodeKey(tc.Location(), nil, v) err = ec.HandleError(err) return valByte, err } diff --git a/pkg/table/tables/tables.go b/pkg/table/tables/tables.go index 69ea903197bb3..bfb4426cfc912 100644 --- a/pkg/table/tables/tables.go +++ b/pkg/table/tables/tables.go @@ -273,19 +273,10 @@ func (t *TableCommon) initTableIndices() error { } // Use partition ID for index, because TableCommon may be table or partition. - idx, err := NewIndexWithCollate(t.encoder.UseNewCollate(), t.physicalTableID, tblInfo, idxInfo) + idx, err := newIndex(t.physicalTableID, tblInfo, idxInfo, t.encoder.UseNewCollate()) if err != nil { return err } - intest.AssertFunc(func() bool { - // `TableCommon.indices` is type of `[]table.Index` to implement interface method `Table.Indices`. - // However, we have an assumption that the specific type of each element in it should always be `*index`. - // We have this assumption because some codes access the inner method of `*index`, - // and they use `asIndex` to cast `table.Index` to `*index`. - _, ok := idx.(*index) - intest.Assert(ok, "index should be type of `*index`") - return true - }) t.indices = append(t.indices, idx) } return nil @@ -523,7 +514,7 @@ func (t *TableCommon) updateRecord(sctx table.MutateContext, txn kv.Transaction, key := t.RecordKey(h) evalCtx := sctx.GetExprCtx().GetEvalCtx() tc, ec := evalCtx.TypeCtx(), evalCtx.ErrCtx() - err = encodeRowBuffer.WriteMemBufferEncoded(t.encoder, sctx.GetRowEncodingConfig(), tc.Location(), ec, memBuffer, key, h) + err = encodeRowBuffer.WriteMemBufferEncoded(sctx.GetRowEncodingConfig(), tc.Location(), ec, memBuffer, key, h) if err != nil { return err } @@ -928,7 +919,7 @@ func (t *TableCommon) addRecord(sctx table.MutateContext, txn kv.Transaction, r } } - err = encodeRowBuffer.WriteMemBufferEncoded(t.encoder, sctx.GetRowEncodingConfig(), tc.Location(), ec, memBuffer, key, recordID, flags...) + err = encodeRowBuffer.WriteMemBufferEncoded(sctx.GetRowEncodingConfig(), tc.Location(), ec, memBuffer, key, recordID, flags...) if err != nil { return nil, err } @@ -1529,19 +1520,15 @@ func (t *TableCommon) UseNewCollate() bool { return t.encoder.UseNewCollate() } +// canSkip reports whether a column can be omitted from the encoded row: it is +// represented by the handle, has an absent NULL default, or is virtual generated. func (t *TableCommon) canSkip(col *table.Column, value *types.Datum) bool { - return CanSkip(t.encoder.UseNewCollate(), t.Meta(), col, value) -} - -// CanSkip is for these cases, we can skip the columns in encoded row: -// 1. the column is included in primary key; -// 2. the column's default value is null, and the value equals to that but has no origin default; -// 3. the column is virtual generated. -func CanSkip(useNewCollate bool, info *model.TableInfo, col *table.Column, value *types.Datum) bool { + info := t.Meta() if col.IsPKHandleColumn(info) { return true } if col.IsCommonHandleColumn(info) { + useNewCollate := t.encoder.UseNewCollate() pkIdx := FindPrimaryIndex(info) for _, idxCol := range pkIdx.Columns { if info.Columns[idxCol.Offset].ID != col.ID { diff --git a/pkg/table/tables/tables_test.go b/pkg/table/tables/tables_test.go index 41543f3a913bc..7315ffb76cca1 100644 --- a/pkg/table/tables/tables_test.go +++ b/pkg/table/tables/tables_test.go @@ -40,6 +40,7 @@ import ( "github.com/pingcap/tidb/pkg/testkit" "github.com/pingcap/tidb/pkg/types" "github.com/pingcap/tidb/pkg/util" + "github.com/pingcap/tidb/pkg/util/collate" "github.com/stretchr/testify/require" ) @@ -422,12 +423,29 @@ func TestTableFromMetaWithCollateUsesFixedMode(t *testing.T) { FieldType: *types.NewFieldType(mysql.TypeVarchar), }, }, + Indices: []*model.IndexInfo{ + { + ID: 1, + Name: ast.NewCIStr("idx"), + State: model.StatePublic, + Columns: []*model.IndexColumn{{Name: ast.NewCIStr("a"), Offset: 0, Length: types.UnspecifiedLength}}, + ConditionExprString: "a = 'A'", + }, + }, } + tblInfo.Columns[0].SetCharset("utf8mb4") + tblInfo.Columns[0].SetCollate("utf8mb4_general_ci") + origin := collate.NewCollationEnabled() + defer collate.SetNewCollationEnabledForTest(origin) for _, useNewCollate := range []bool{false, true} { + collate.SetNewCollationEnabledForTest(!useNewCollate) tbl, err := tables.TableFromMetaWithCollate(useNewCollate, autoid.NewAllocators(false), tblInfo) require.NoError(t, err) require.Equal(t, useNewCollate, tbl.UseNewCollate()) + meet, err := tbl.Indices()[0].MeetPartialCondition(types.MakeDatums("a")) + require.NoError(t, err) + require.Equal(t, useNewCollate, meet) } } diff --git a/pkg/table/tables/test/partition/BUILD.bazel b/pkg/table/tables/test/partition/BUILD.bazel index 299e44f0d0f40..06004dc48dc76 100644 --- a/pkg/table/tables/test/partition/BUILD.bazel +++ b/pkg/table/tables/test/partition/BUILD.bazel @@ -8,10 +8,11 @@ go_test( "partition_test.go", ], flaky = True, - shard_count = 23, + shard_count = 24, deps = [ "//pkg/domain", "//pkg/kv", + "//pkg/meta/autoid", "//pkg/meta/model", "//pkg/parser/ast", "//pkg/sessiontxn", @@ -22,6 +23,7 @@ go_test( "//pkg/testkit/testsetup", "//pkg/types", "//pkg/util", + "//pkg/util/collate", "//pkg/util/logutil", "@com_github_pingcap_errors//:errors", "@com_github_stretchr_testify//require", diff --git a/pkg/table/tables/test/partition/partition_test.go b/pkg/table/tables/test/partition/partition_test.go index b8a6c227152cd..be95e94f40934 100644 --- a/pkg/table/tables/test/partition/partition_test.go +++ b/pkg/table/tables/test/partition/partition_test.go @@ -17,6 +17,7 @@ package partition import ( "context" "fmt" + "hash/crc32" "math/rand" "slices" "strconv" @@ -27,6 +28,7 @@ import ( "github.com/pingcap/errors" "github.com/pingcap/tidb/pkg/domain" "github.com/pingcap/tidb/pkg/kv" + "github.com/pingcap/tidb/pkg/meta/autoid" "github.com/pingcap/tidb/pkg/meta/model" "github.com/pingcap/tidb/pkg/parser/ast" "github.com/pingcap/tidb/pkg/sessiontxn" @@ -36,11 +38,76 @@ import ( "github.com/pingcap/tidb/pkg/testkit/testfailpoint" "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/logutil" "github.com/stretchr/testify/require" "go.uber.org/zap" ) +func TestPartitionTableUsesTableCollationSnapshot(t *testing.T) { + origin := collate.NewCollationEnabled() + collate.SetNewCollationEnabledForTest(false) + defer collate.SetNewCollationEnabledForTest(origin) + + store, dom := testkit.CreateMockStoreAndDomain(t) + tk := testkit.NewTestKit(t, store) + tk.MustExec("use test") + tk.MustExec(`create table t_collate_snapshot ( + a varchar(16) collate utf8mb4_general_ci + ) partition by list columns (a) ( + partition p0 values in ('a'), + partition p_default values in (default) + )`) + tk.MustExec(`create table t_key_collate_snapshot ( + a varchar(16) collate utf8mb4_general_ci + ) partition by key(a) partitions 8`) + + tbl, err := dom.InfoSchema().TableByName(context.Background(), ast.NewCIStr("test"), ast.NewCIStr("t_collate_snapshot")) + require.NoError(t, err) + tblInfo := tbl.Meta() + keyTbl, err := dom.InfoSchema().TableByName(context.Background(), ast.NewCIStr("test"), ast.NewCIStr("t_key_collate_snapshot")) + require.NoError(t, err) + keyTblInfo := keyTbl.Meta() + + collate.SetNewCollationEnabledForTest(true) + snapshotTbl, err := tables.TableFromMetaWithCollate(false, autoid.NewAllocators(tblInfo.SepAutoInc()), tblInfo) + require.NoError(t, err) + + pt := snapshotTbl.GetPartitionedTable() + require.NotNil(t, pt) + physicalTbl, err := pt.GetPartitionByRow(tk.Session().GetExprCtx().GetEvalCtx(), types.MakeDatums("A")) + require.NoError(t, err) + require.Equal(t, tblInfo.Partition.Definitions[1].ID, physicalTbl.GetPhysicalID()) + + keySnapshotTbl, err := tables.TableFromMetaWithCollate(false, autoid.NewAllocators(keyTblInfo.SepAutoInc()), keyTblInfo) + require.NoError(t, err) + keyPt := keySnapshotTbl.GetPartitionedTable() + require.NotNil(t, keyPt) + keyPartitionIdx := func(useNewCollate bool, val string) int { + h := crc32.NewIEEE() + h.Write(collate.GetCollatorWithCollate(useNewCollate, "utf8mb4_general_ci").Key(val)) + return int(h.Sum32() % uint32(len(keyTblInfo.Partition.Definitions))) + } + var ( + keyVal string + expectedKeyIdx int + ) + for _, val := range []string{"A", "B", "Aa", "aA", "abc", "ABC", "xYz"} { + oldIdx := keyPartitionIdx(false, val) + if oldIdx != keyPartitionIdx(true, val) { + keyVal = val + expectedKeyIdx = oldIdx + break + } + } + require.NotEmpty(t, keyVal) + keyDatum := types.NewStringDatum(keyVal) + keyDatum.SetCollation("utf8mb4_general_ci") + physicalTbl, err = keyPt.GetPartitionByRow(tk.Session().GetExprCtx().GetEvalCtx(), []types.Datum{keyDatum}) + require.NoError(t, err) + require.Equal(t, keyTblInfo.Partition.Definitions[expectedKeyIdx].ID, physicalTbl.GetPhysicalID()) +} + func TestPartitionAddRecord(t *testing.T) { createTable1 := `CREATE TABLE test.t1 (id int(11), index(id)) PARTITION BY RANGE ( id ) ( diff --git a/pkg/table/tables/testutil/BUILD.bazel b/pkg/table/tables/testutil/BUILD.bazel index 6fb8235d08bb4..e4f0c0eb1187a 100644 --- a/pkg/table/tables/testutil/BUILD.bazel +++ b/pkg/table/tables/testutil/BUILD.bazel @@ -13,7 +13,6 @@ go_library( "//pkg/tablecodec", "//pkg/testkit", "//pkg/util/codec", - "//pkg/util/collate", "@com_github_stretchr_testify//require", ], ) diff --git a/pkg/table/tables/testutil/indexcheck.go b/pkg/table/tables/testutil/indexcheck.go index bffb6cfff5793..4c9184573d947 100644 --- a/pkg/table/tables/testutil/indexcheck.go +++ b/pkg/table/tables/testutil/indexcheck.go @@ -26,7 +26,6 @@ import ( "github.com/pingcap/tidb/pkg/tablecodec" "github.com/pingcap/tidb/pkg/testkit" "github.com/pingcap/tidb/pkg/util/codec" - "github.com/pingcap/tidb/pkg/util/collate" "github.com/stretchr/testify/require" ) @@ -43,7 +42,7 @@ func CheckIndexKVCount(t *testing.T, tk *testkit.TestKit, dom *domain.Domain, ta } } - enc := codec.NewEncoder(collate.NewCollationEnabled()) + enc := codec.NewEncoder(tbl.UseNewCollate()) minimumKey, _, err := tablecodec.GenIndexKey(enc, time.Local, tbl.Meta(), idx.Meta(), tbl.Meta().ID, nil, nil, nil) require.NoError(t, err) diff --git a/pkg/table/tblctx/BUILD.bazel b/pkg/table/tblctx/BUILD.bazel index 3814046cea22b..f7ea643d76115 100644 --- a/pkg/table/tblctx/BUILD.bazel +++ b/pkg/table/tblctx/BUILD.bazel @@ -20,8 +20,6 @@ go_library( "//pkg/tablecodec", "//pkg/types", "//pkg/util/chunk", - "//pkg/util/codec", - "//pkg/util/collate", "//pkg/util/intest", "//pkg/util/rowcodec", "//pkg/util/tableutil", @@ -42,8 +40,6 @@ go_test( "//pkg/sessionctx/variable", "//pkg/tablecodec", "//pkg/types", - "//pkg/util/codec", - "//pkg/util/collate", "//pkg/util/rowcodec", "@com_github_stretchr_testify//mock", "@com_github_stretchr_testify//require", diff --git a/pkg/table/tblctx/buffers.go b/pkg/table/tblctx/buffers.go index 2a51a7d8b4a9f..5e9e8d93d0bb1 100644 --- a/pkg/table/tblctx/buffers.go +++ b/pkg/table/tblctx/buffers.go @@ -23,8 +23,6 @@ import ( "github.com/pingcap/tidb/pkg/tablecodec" "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/intest" "github.com/pingcap/tidb/pkg/util/rowcodec" ) @@ -53,7 +51,6 @@ func (b *EncodeRowBuffer) AddColVal(colID int64, val types.Datum) { // WriteMemBufferEncoded writes the encoded row to the memBuffer. func (b *EncodeRowBuffer) WriteMemBufferEncoded( - enc codec.Encoder, cfg RowEncodingConfig, loc *time.Location, ec errctx.Context, memBuffer kv.MemBuffer, key kv.Key, handle kv.Handle, flags ...kv.FlagsOp, ) error { @@ -72,7 +69,7 @@ func (b *EncodeRowBuffer) WriteMemBufferEncoded( stmtBufs.AddRowValues = ensureCapacityAndReset(stmtBufs.AddRowValues, len(b.row)*2) encoded, err := tablecodec.EncodeRow( - enc, loc, b.row, b.colIDs, stmtBufs.RowValBuf, stmtBufs.AddRowValues, checksum, cfg.RowEncoder, + loc, b.row, b.colIDs, stmtBufs.RowValBuf, stmtBufs.AddRowValues, checksum, cfg.RowEncoder, ) if err = ec.HandleError(err); err != nil { return err @@ -88,8 +85,7 @@ func (b *EncodeRowBuffer) WriteMemBufferEncoded( // EncodeBinlogRowData encodes the row data for binlog and returns the encoded row value. // The returned slice is not referenced in the buffer, so you can cache and modify them freely. func (b *EncodeRowBuffer) EncodeBinlogRowData(loc *time.Location, ec errctx.Context) ([]byte, error) { - enc := codec.NewEncoder(collate.NewCollationEnabled()) - value, err := tablecodec.EncodeOldRow(enc, loc, b.row, b.colIDs, nil, nil) + value, err := tablecodec.EncodeOldRow(loc, b.row, b.colIDs, nil, nil) err = ec.HandleError(err) if err != nil { return nil, err diff --git a/pkg/table/tblctx/buffers_test.go b/pkg/table/tblctx/buffers_test.go index f3ef0fe6ed3ef..ac9e7716d79a8 100644 --- a/pkg/table/tblctx/buffers_test.go +++ b/pkg/table/tblctx/buffers_test.go @@ -25,8 +25,6 @@ import ( "github.com/pingcap/tidb/pkg/sessionctx/variable" "github.com/pingcap/tidb/pkg/tablecodec" "github.com/pingcap/tidb/pkg/types" - "github.com/pingcap/tidb/pkg/util/codec" - "github.com/pingcap/tidb/pkg/util/collate" "github.com/pingcap/tidb/pkg/util/rowcodec" "github.com/stretchr/testify/mock" "github.com/stretchr/testify/require" @@ -80,8 +78,6 @@ func TestEncodeRow(t *testing.T) { buffer.AddColVal(3, d3) require.Equal(t, []int64{1, 2, 3}, buffer.colIDs) require.Equal(t, []types.Datum{d1, d2, d3}, buffer.row) - enc := codec.NewEncoder(collate.NewCollationEnabled()) - for _, c := range []struct { loc *time.Location rowLevelChecksum bool @@ -113,7 +109,7 @@ func TestEncodeRow(t *testing.T) { } expectedVal, err := tablecodec.EncodeRow( - enc, c.loc, []types.Datum{d1, d2, d3}, []int64{1, 2, 3}, nil, nil, checksum, + c.loc, []types.Datum{d1, d2, d3}, []int64{1, 2, 3}, nil, nil, checksum, &rowcodec.Encoder{Enable: !c.oldFormat}, ) require.NoError(t, err) @@ -127,7 +123,7 @@ func TestEncodeRow(t *testing.T) { Return(nil).Once() } err = buffer.WriteMemBufferEncoded( - enc, cfg, c.loc, errctx.StrictNoWarningContext, + cfg, c.loc, errctx.StrictNoWarningContext, memBuffer, kv.Key("key1"), kv.IntHandle(1), c.flags..., ) require.NoError(t, err) @@ -137,7 +133,7 @@ func TestEncodeRow(t *testing.T) { // test encode val for binlog expectedVal, err = - tablecodec.EncodeOldRow(enc, c.loc, []types.Datum{d1, d2, d3}, []int64{1, 2, 3}, nil, nil) + tablecodec.EncodeOldRow(c.loc, []types.Datum{d1, d2, d3}, []int64{1, 2, 3}, nil, nil) require.NoError(t, err) encoded, err := buffer.EncodeBinlogRowData(c.loc, errctx.StrictNoWarningContext) require.NoError(t, err) @@ -167,8 +163,7 @@ func TestEncodeBufferReserve(t *testing.T) { buffer.AddColVal(2, types.NewIntDatum(2)) require.Equal(t, 2, len(buffer.colIDs)) require.Equal(t, 2, len(buffer.row)) - enc := codec.NewEncoder(collate.NewCollationEnabled()) - require.NoError(t, buffer.WriteMemBufferEncoded(enc, RowEncodingConfig{ + require.NoError(t, buffer.WriteMemBufferEncoded(RowEncodingConfig{ RowEncoder: &rowcodec.Encoder{Enable: true}, }, time.UTC, errctx.StrictNoWarningContext, mb, kv.Key("key1"), kv.IntHandle(1))) encodedCap := cap(buffer.writeStmtBufs.RowValBuf) diff --git a/pkg/tablecodec/tablecodec.go b/pkg/tablecodec/tablecodec.go index 410c1dc1a6fdb..b1cbc803ad162 100644 --- a/pkg/tablecodec/tablecodec.go +++ b/pkg/tablecodec/tablecodec.go @@ -362,7 +362,7 @@ func EncodeValue(loc *time.Location, b []byte, raw types.Datum) ([]byte, error) // EncodeRow will allocate it. // This function may return both a valid encoded bytes and an error (actually `"pingcap/errors".ErrorGroup`). If the caller // expects to handle these errors according to `SQL_MODE` or other configuration, please refer to `pkg/errctx`. -func EncodeRow(enc codec.Encoder, loc *time.Location, row []types.Datum, colIDs []int64, valBuf []byte, values []types.Datum, checksum rowcodec.Checksum, e *rowcodec.Encoder) ([]byte, error) { +func EncodeRow(loc *time.Location, row []types.Datum, colIDs []int64, valBuf []byte, values []types.Datum, checksum rowcodec.Checksum, e *rowcodec.Encoder) ([]byte, error) { if len(row) != len(colIDs) { return nil, errors.Errorf("EncodeRow error: data and columnID count not match %d vs %d", len(row), len(colIDs)) } @@ -370,14 +370,14 @@ func EncodeRow(enc codec.Encoder, loc *time.Location, row []types.Datum, colIDs valBuf = valBuf[:0] return e.Encode(loc, colIDs, row, checksum, valBuf) } - return EncodeOldRow(enc, loc, row, colIDs, valBuf, values) + return EncodeOldRow(loc, row, colIDs, valBuf, values) } // EncodeOldRow encode row data and column ids into a slice of byte. // Row layout: colID1, value1, colID2, value2, ..... // valBuf and values pass by caller, for reducing EncodeOldRow allocates temporary bufs. If you pass valBuf and values as nil, // EncodeOldRow will allocate it. -func EncodeOldRow(enc codec.Encoder, loc *time.Location, row []types.Datum, colIDs []int64, valBuf []byte, values []types.Datum) ([]byte, error) { +func EncodeOldRow(loc *time.Location, row []types.Datum, colIDs []int64, valBuf []byte, values []types.Datum) ([]byte, error) { if len(row) != len(colIDs) { return nil, errors.Errorf("EncodeRow error: data and columnID count not match %d vs %d", len(row), len(colIDs)) } @@ -397,7 +397,7 @@ func EncodeOldRow(enc codec.Encoder, loc *time.Location, row []types.Datum, colI // We could not set nil value into kv. return append(valBuf, codec.NilFlag), nil } - return enc.EncodeValue(loc, valBuf, values...) + return codec.EncodeValue(loc, valBuf, values...) } func flatten(loc *time.Location, data types.Datum, ret *types.Datum) error { diff --git a/pkg/tablecodec/tablecodec_test.go b/pkg/tablecodec/tablecodec_test.go index c948b50b996fa..28677127871e4 100644 --- a/pkg/tablecodec/tablecodec_test.go +++ b/pkg/tablecodec/tablecodec_test.go @@ -107,7 +107,7 @@ func TestRowCodec(t *testing.T) { } rd := rowcodec.Encoder{Enable: true} sc := stmtctx.NewStmtCtxWithTimeZone(time.Local) - bs, err := EncodeRow(defaultCodecEncoder(), sc.TimeZone(), row, colIDs, nil, nil, nil, &rd) + bs, err := EncodeRow(sc.TimeZone(), row, colIDs, nil, nil, nil, &rd) require.NoError(t, err) require.NotNil(t, bs) @@ -162,7 +162,7 @@ func TestRowCodec(t *testing.T) { } // Make sure empty row return not nil value. - bs, err = EncodeOldRow(defaultCodecEncoder(), sc.TimeZone(), []types.Datum{}, []int64{}, nil, nil) + bs, err = EncodeOldRow(sc.TimeZone(), []types.Datum{}, []int64{}, nil, nil) require.NoError(t, err) require.Len(t, bs, 1) @@ -176,7 +176,7 @@ func TestDecodeColumnValue(t *testing.T) { // test timestamp d := types.NewTimeDatum(types.NewTime(types.FromGoTime(time.Now()), mysql.TypeTimestamp, types.DefaultFsp)) - bs, err := EncodeOldRow(defaultCodecEncoder(), sc.TimeZone(), []types.Datum{d}, []int64{1}, nil, nil) + bs, err := EncodeOldRow(sc.TimeZone(), []types.Datum{d}, []int64{1}, nil, nil) require.NoError(t, err) require.NotNil(t, bs) _, bs, err = codec.CutOne(bs) // ignore colID @@ -192,7 +192,7 @@ func TestDecodeColumnValue(t *testing.T) { elems := []string{"a", "b", "c", "d", "e"} e, _ := types.ParseSetValue(elems, uint64(1)) d = types.NewMysqlSetDatum(e, "") - bs, err = EncodeOldRow(defaultCodecEncoder(), sc.TimeZone(), []types.Datum{d}, []int64{1}, nil, nil) + bs, err = EncodeOldRow(sc.TimeZone(), []types.Datum{d}, []int64{1}, nil, nil) require.NoError(t, err) require.NotNil(t, bs) _, bs, err = codec.CutOne(bs) // ignore colID @@ -207,7 +207,7 @@ func TestDecodeColumnValue(t *testing.T) { // test bit d = types.NewMysqlBitDatum(types.NewBinaryLiteralFromUint(3223600, 3)) - bs, err = EncodeOldRow(defaultCodecEncoder(), sc.TimeZone(), []types.Datum{d}, []int64{1}, nil, nil) + bs, err = EncodeOldRow(sc.TimeZone(), []types.Datum{d}, []int64{1}, nil, nil) require.NoError(t, err) require.NotNil(t, bs) _, bs, err = codec.CutOne(bs) // ignore colID @@ -222,7 +222,7 @@ func TestDecodeColumnValue(t *testing.T) { // test empty enum d = types.NewMysqlEnumDatum(types.Enum{}) - bs, err = EncodeOldRow(defaultCodecEncoder(), sc.TimeZone(), []types.Datum{d}, []int64{1}, nil, nil) + bs, err = EncodeOldRow(sc.TimeZone(), []types.Datum{d}, []int64{1}, nil, nil) require.NoError(t, err) require.NotNil(t, bs) _, bs, err = codec.CutOne(bs) // ignore colID @@ -282,7 +282,7 @@ func TestTimeCodec(t *testing.T) { } rd := rowcodec.Encoder{Enable: true} sc := stmtctx.NewStmtCtxWithTimeZone(time.UTC) - bs, err := EncodeRow(defaultCodecEncoder(), sc.TimeZone(), row, colIDs, nil, nil, nil, &rd) + bs, err := EncodeRow(sc.TimeZone(), row, colIDs, nil, nil, nil, &rd) require.NoError(t, err) require.NotNil(t, bs) @@ -330,7 +330,7 @@ func TestCutRow(t *testing.T) { for _, col := range cols { colIDs = append(colIDs, col.id) } - bs, err := EncodeOldRow(defaultCodecEncoder(), sc.TimeZone(), row, colIDs, nil, nil) + bs, err := EncodeOldRow(sc.TimeZone(), row, colIDs, nil, nil) require.NoError(t, err) require.NotNil(t, bs) diff --git a/pkg/util/codec/codec.go b/pkg/util/codec/codec.go index 7dccb4c39535c..ff2c0e7aff136 100644 --- a/pkg/util/codec/codec.go +++ b/pkg/util/codec/codec.go @@ -64,7 +64,7 @@ const ( sizeFloat64 = unsafe.Sizeof(float64(0)) ) -// Encoder encodes Datum values with a fixed new collation setting. +// Encoder encodes comparable Datum keys with a fixed new collation setting. type Encoder struct { useNewCollate bool } @@ -106,7 +106,7 @@ func preRealloc(b []byte, vals []types.Datum, comparable1 bool) []byte { // encode will encode a datum and append it to a byte slice. If comparable1 is true, the encoded bytes can be sorted as it's original order. // If hash is true, the encoded bytes can be checked equal as it's original value. -func (enc Encoder) encode(loc *time.Location, b []byte, vals []types.Datum, comparable1 bool) (_ []byte, err error) { +func encode(loc *time.Location, b []byte, vals []types.Datum, comparable1, useNewCollate bool) (_ []byte, err error) { b = preRealloc(b, vals, comparable1) for i, length := 0, len(vals); i < length; i++ { switch vals[i].Kind() { @@ -118,7 +118,7 @@ func (enc Encoder) encode(loc *time.Location, b []byte, vals []types.Datum, comp b = append(b, floatFlag) b = EncodeFloat(b, vals[i].GetFloat64()) case types.KindString: - b = enc.encodeString(b, vals[i], comparable1) + b = encodeString(b, vals[i], comparable1, useNewCollate) case types.KindBytes: b = encodeBytes(b, vals[i].GetBytes(), comparable1) case types.KindMysqlTime: @@ -230,9 +230,9 @@ func EncodeMySQLTime(loc *time.Location, t types.Time, tp byte, b []byte) (_ []b return b, nil } -func (enc Encoder) encodeString(b []byte, val types.Datum, comparable1 bool) []byte { - if enc.useNewCollate && comparable1 { - return encodeBytes(b, collate.GetCollatorWithCollate(enc.useNewCollate, val.Collation()).ImmutableKey(val.GetString()), true) +func encodeString(b []byte, val types.Datum, comparable1, useNewCollate bool) []byte { + if useNewCollate && comparable1 { + return encodeBytes(b, collate.GetCollatorWithCollate(useNewCollate, val.Collation()).ImmutableKey(val.GetString()), true) } return encodeBytes(b, val.GetBytes(), comparable1) } @@ -326,20 +326,13 @@ func EncodeKey(loc *time.Location, b []byte, v ...types.Datum) ([]byte, error) { // fixed collation setting. It guarantees the encoded value is in ascending order // for comparison. For decimal type, datum must set datum's length and frac. func (enc Encoder) EncodeKey(loc *time.Location, b []byte, v ...types.Datum) ([]byte, error) { - return enc.encode(loc, b, v, true) + return encode(loc, b, v, true, enc.useNewCollate) } // EncodeValue appends the encoded values to byte slice b, returning the appended // slice. It does not guarantee the order for comparison. func EncodeValue(loc *time.Location, b []byte, v ...types.Datum) ([]byte, error) { - return NewEncoder(collate.NewCollationEnabled()).EncodeValue(loc, b, v...) -} - -// EncodeValue appends the encoded values to byte slice b using the encoder's -// fixed collation setting, returning the appended slice. It does not guarantee -// the order for comparison. -func (enc Encoder) EncodeValue(loc *time.Location, b []byte, v ...types.Datum) ([]byte, error) { - return enc.encode(loc, b, v, false) + return encode(loc, b, v, false, false) } // EncodeHashChunkRowIdx encodes value for further comparison @@ -1907,14 +1900,6 @@ func init() { // HashCode encodes a Datum into a unique byte slice. // It is mostly the same as EncodeValue, but it doesn't contain truncation or verification logic in order to make the encoding lossless. func HashCode(b []byte, d types.Datum) []byte { - return NewEncoder(collate.NewCollationEnabled()).HashCode(b, d) -} - -// HashCode encodes a Datum into a unique byte slice using the encoder's fixed -// collation setting. It is mostly the same as EncodeValue, but it doesn't -// contain truncation or verification logic in order to make the encoding -// lossless. -func (enc Encoder) HashCode(b []byte, d types.Datum) []byte { switch d.Kind() { case types.KindInt64: b = encodeSignedInt(b, d.GetInt64(), false) @@ -1924,7 +1909,7 @@ func (enc Encoder) HashCode(b []byte, d types.Datum) []byte { b = append(b, floatFlag) b = EncodeFloat(b, d.GetFloat64()) case types.KindString: - b = enc.encodeString(b, d, false) + b = encodeBytes(b, d.GetBytes(), false) case types.KindBytes: b = encodeBytes(b, d.GetBytes(), false) case types.KindMysqlTime: diff --git a/pkg/util/codec/collation_test.go b/pkg/util/codec/collation_test.go index 8238dfb987a1f..9b31c6fc5485c 100644 --- a/pkg/util/codec/collation_test.go +++ b/pkg/util/codec/collation_test.go @@ -74,10 +74,6 @@ func TestEncoderNewCollationEnabled(t *testing.T) { exportedDisabledLower, err := EncodeKey(time.Local, nil, lower) require.NoError(t, err) require.Equal(t, disabledLower, exportedDisabledLower) - - enabledHash := enabledEncoder.HashCode(nil, lower) - disabledHash := disabledEncoder.HashCode(nil, lower) - require.Equal(t, enabledHash, disabledHash) } func TestHashGroupKeyCollation(t *testing.T) { diff --git a/pkg/util/rowDecoder/BUILD.bazel b/pkg/util/rowDecoder/BUILD.bazel index b163f546dff15..9a6aa73fc6e02 100644 --- a/pkg/util/rowDecoder/BUILD.bazel +++ b/pkg/util/rowDecoder/BUILD.bazel @@ -41,7 +41,6 @@ go_test( "//pkg/testkit/testsetup", "//pkg/testkit/testutil", "//pkg/types", - "//pkg/util/codec", "//pkg/util/collate", "//pkg/util/mock", "//pkg/util/rowcodec", diff --git a/pkg/util/rowDecoder/decoder_test.go b/pkg/util/rowDecoder/decoder_test.go index 0db975740fb86..00373b8dd3f98 100644 --- a/pkg/util/rowDecoder/decoder_test.go +++ b/pkg/util/rowDecoder/decoder_test.go @@ -29,7 +29,6 @@ import ( "github.com/pingcap/tidb/pkg/tablecodec" "github.com/pingcap/tidb/pkg/testkit/testutil" "github.com/pingcap/tidb/pkg/types" - "github.com/pingcap/tidb/pkg/util/codec" "github.com/pingcap/tidb/pkg/util/collate" "github.com/pingcap/tidb/pkg/util/mock" decoder "github.com/pingcap/tidb/pkg/util/rowDecoder" @@ -109,13 +108,12 @@ func TestRowDecoder(t *testing.T) { }, } rd := rowcodec.Encoder{Enable: true} - codecEncoder := codec.NewEncoder(collate.NewCollationEnabled()) for i, row := range testRows { // test case for pk is unsigned. if i > 0 { c7.AddFlag(mysql.UnsignedFlag) } - bs, err := tablecodec.EncodeRow(codecEncoder, sc.TimeZone(), row.input, row.cols, nil, nil, nil, &rd) + bs, err := tablecodec.EncodeRow(sc.TimeZone(), row.input, row.cols, nil, nil, nil, &rd) require.NoError(t, err) require.NotNil(t, bs) @@ -189,9 +187,8 @@ func TestClusterIndexRowDecoder(t *testing.T) { }, } rd := rowcodec.Encoder{Enable: true} - codecEncoder := codec.NewEncoder(collate.NewCollationEnabled()) for _, row := range testRows { - bs, err := tablecodec.EncodeRow(codecEncoder, sc.TimeZone(), row.input, row.cols, nil, nil, nil, &rd) + bs, err := tablecodec.EncodeRow(sc.TimeZone(), row.input, row.cols, nil, nil, nil, &rd) require.NoError(t, err) require.NotNil(t, bs) diff --git a/pkg/util/rowcodec/bench_test.go b/pkg/util/rowcodec/bench_test.go index 3474fe3420eab..eea2c8373fbe4 100644 --- a/pkg/util/rowcodec/bench_test.go +++ b/pkg/util/rowcodec/bench_test.go @@ -25,8 +25,6 @@ import ( "github.com/pingcap/tidb/pkg/types" "github.com/pingcap/tidb/pkg/util/benchdaily" "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/rowcodec" ) @@ -69,7 +67,7 @@ func BenchmarkEncode(b *testing.B) { func BenchmarkEncodeFromOldRow(b *testing.B) { b.ReportAllocs() oldRow := types.MakeDatums(1, "abc", 1.1) - oldRowData, err := tablecodec.EncodeOldRow(codec.NewEncoder(collate.NewCollationEnabled()), nil, oldRow, []int64{1, 2, 3}, nil, nil) + oldRowData, err := tablecodec.EncodeOldRow(nil, oldRow, []int64{1, 2, 3}, nil, nil) if err != nil { b.Fatal(err) } diff --git a/pkg/util/rowcodec/rowcodec_test.go b/pkg/util/rowcodec/rowcodec_test.go index 3ea9206376719..20095f8bb4834 100644 --- a/pkg/util/rowcodec/rowcodec_test.go +++ b/pkg/util/rowcodec/rowcodec_test.go @@ -757,7 +757,7 @@ func TestCodecUtil(t *testing.T) { } tps[3] = types.NewFieldType(mysql.TypeNull) sc := stmtctx.NewStmtCtx() - oldRow, err := tablecodec.EncodeOldRow(codec.NewEncoder(collate.NewCollationEnabled()), sc.TimeZone(), types.MakeDatums(1, 2, 3, nil), colIDs, nil, nil) + oldRow, err := tablecodec.EncodeOldRow(sc.TimeZone(), types.MakeDatums(1, 2, 3, nil), colIDs, nil, nil) require.NoError(t, err) var ( @@ -807,7 +807,7 @@ func TestOldRowCodec(t *testing.T) { } tps[3] = types.NewFieldType(mysql.TypeNull) sc := stmtctx.NewStmtCtx() - oldRow, err := tablecodec.EncodeOldRow(codec.NewEncoder(collate.NewCollationEnabled()), sc.TimeZone(), types.MakeDatums(1, 2, 3, nil), colIDs, nil, nil) + oldRow, err := tablecodec.EncodeOldRow(sc.TimeZone(), types.MakeDatums(1, 2, 3, nil), colIDs, nil, nil) require.NoError(t, err) var (