From 48c7aba00046d4e564f46f1829f3ba04c1eebbc5 Mon Sep 17 00:00:00 2001 From: D3Hunter Date: Mon, 6 Jul 2026 11:36:24 +0800 Subject: [PATCH 1/3] This is an automated cherry-pick of #69566 Signed-off-by: ti-chi-bot --- pkg/ddl/column.go | 5 +- pkg/ddl/index.go | 8 +- pkg/executor/admin.go | 4 +- pkg/executor/test/executor/BUILD.bazel | 2 + pkg/executor/test/executor/executor_test.go | 4 +- pkg/executor/write.go | 2 + pkg/expression/expression.go | 10 + pkg/meta/model/table.go | 13 + pkg/planner/core/expression_rewriter.go | 21 +- pkg/server/handler/tests/BUILD.bazel | 1 + pkg/server/handler/tests/http_handler_test.go | 3 +- pkg/session/BUILD.bazel | 1 + pkg/session/global_init.go | 79 ++++++ pkg/session/session.go | 25 +- pkg/store/mockstore/BUILD.bazel | 1 + pkg/store/mockstore/cluster_test.go | 3 +- .../unistore/cophandler/cop_handler_test.go | 3 +- pkg/table/tables/index.go | 105 +++++++- pkg/table/tables/mutation_checker.go | 13 +- pkg/table/tables/mutation_checker_test.go | 212 +++++++++------ pkg/table/tables/partition.go | 2 +- pkg/table/tables/tables.go | 46 ++-- pkg/table/tables/testutil/BUILD.bazel | 19 ++ pkg/table/tables/testutil/indexcheck.go | 75 ++++++ pkg/table/tblctx/BUILD.bazel | 4 + pkg/table/tblctx/buffers.go | 8 +- pkg/table/tblctx/buffers_test.go | 12 +- pkg/tablecodec/tablecodec.go | 55 ++-- pkg/tablecodec/tablecodec_test.go | 250 +++++++++++++++++- pkg/testkit/mockstore.go | 14 +- pkg/types/etc.go | 10 +- pkg/util/codec/codec.go | 53 +++- pkg/util/codec/codec_test.go | 2 +- pkg/util/codec/collation_test.go | 37 +++ pkg/util/collate/collate.go | 20 +- pkg/util/rowDecoder/BUILD.bazel | 1 + pkg/util/rowDecoder/decoder_test.go | 7 +- pkg/util/rowcodec/bench_test.go | 4 +- pkg/util/rowcodec/rowcodec_test.go | 4 +- tests/realtikvtest/addindextest2/BUILD.bazel | 2 + .../addindextest2/global_sort_test.go | 4 +- 41 files changed, 951 insertions(+), 193 deletions(-) create mode 100644 pkg/session/global_init.go create mode 100644 pkg/table/tables/testutil/BUILD.bazel create mode 100644 pkg/table/tables/testutil/indexcheck.go diff --git a/pkg/ddl/column.go b/pkg/ddl/column.go index 93c1e40b6f3ad..fa9ea20d91d6b 100644 --- a/pkg/ddl/column.go +++ b/pkg/ddl/column.go @@ -44,6 +44,8 @@ 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" @@ -839,7 +841,8 @@ func (w *updateColumnWorker) getRowRecord(handle kv.Handle, recordKey []byte, ra if w.checksumNeeded { checksum = rowcodec.RawChecksum{Handle: handle} } - newRowVal, err := tablecodec.EncodeRow(sysTZ, newRow, newColumnIDs, nil, nil, checksum, rd) + enc := codec.NewEncoder(collate.NewCollationEnabled()) + newRowVal, err := tablecodec.EncodeRow(enc, sysTZ, newRow, newColumnIDs, nil, nil, checksum, rd) err = ec.HandleError(err) if err != nil { return errors.Trace(err) diff --git a/pkg/ddl/index.go b/pkg/ddl/index.go index 23f35e5e5bde5..0787dfe2c25f0 100644 --- a/pkg/ddl/index.go +++ b/pkg/ddl/index.go @@ -74,6 +74,7 @@ import ( "github.com/pingcap/tidb/pkg/util" "github.com/pingcap/tidb/pkg/util/backoff" "github.com/pingcap/tidb/pkg/util/chunk" + "github.com/pingcap/tidb/pkg/util/collate" "github.com/pingcap/tidb/pkg/util/dbterror" "github.com/pingcap/tidb/pkg/util/generatedexpr" "github.com/pingcap/tidb/pkg/util/intest" @@ -2092,7 +2093,7 @@ func (w *baseIndexWorker) getIndexRecord(idxInfo *model.IndexInfo, handle kv.Han idxVal[j] = idxColumnVal } - rsData := tables.TryGetHandleRestoredDataWrapper(w.table.Meta(), nil, w.rowMap, idxInfo) + rsData := tables.TryGetHandleRestoredDataWrapper(collate.NewCollationEnabled(), w.table.Meta(), nil, w.rowMap, idxInfo) idxRecord := &indexRecord{handle: handle, key: recordKey, vals: idxVal, rsData: rsData} return idxRecord, nil } @@ -2332,11 +2333,12 @@ func writeChunk( }() needRestoreForIndexes := make([]bool, len(indexes)) restore, pkNeedRestore := false, false + useNewCollate := collate.NewCollationEnabled() if c.PrimaryKeyInfo != nil && c.TableInfo.IsCommonHandle && c.TableInfo.CommonHandleVersion != 0 { - pkNeedRestore = tables.NeedRestoredData(c.PrimaryKeyInfo.Columns, c.TableInfo.Columns) + pkNeedRestore = tables.NeedRestoredData(useNewCollate, c.PrimaryKeyInfo.Columns, c.TableInfo.Columns) } for i, index := range indexes { - needRestore := pkNeedRestore || tables.NeedRestoredData(index.Meta().Columns, c.TableInfo.Columns) + needRestore := pkNeedRestore || tables.NeedRestoredData(useNewCollate, index.Meta().Columns, c.TableInfo.Columns) needRestoreForIndexes[i] = needRestore restore = restore || needRestore } diff --git a/pkg/executor/admin.go b/pkg/executor/admin.go index e67dd9a37c41e..365ddee7fe81d 100644 --- a/pkg/executor/admin.go +++ b/pkg/executor/admin.go @@ -35,6 +35,7 @@ import ( "github.com/pingcap/tidb/pkg/types" "github.com/pingcap/tidb/pkg/util" "github.com/pingcap/tidb/pkg/util/chunk" + "github.com/pingcap/tidb/pkg/util/collate" "github.com/pingcap/tidb/pkg/util/logutil" "github.com/pingcap/tidb/pkg/util/ranger" "github.com/pingcap/tidb/pkg/util/timeutil" @@ -363,6 +364,7 @@ func (e *RecoverIndexExec) fetchRecoverRows(ctx context.Context, srcResult dists idxValLen := len(e.index.Meta().Columns) result.scanRowCount = 0 + useNewCollate := collate.NewCollationEnabled() for { err := srcResult.Next(ctx, e.srcChunk) if err != nil { @@ -389,7 +391,7 @@ func (e *RecoverIndexExec) fetchRecoverRows(ctx context.Context, srcResult dists return nil, err } e.idxValsBufs[result.scanRowCount] = idxVals - rsData := tables.TryGetHandleRestoredDataWrapper(e.table.Meta(), plannercore.GetCommonHandleDatum(e.handleCols, row), nil, e.index.Meta()) + rsData := tables.TryGetHandleRestoredDataWrapper(useNewCollate, e.table.Meta(), plannercore.GetCommonHandleDatum(e.handleCols, row), nil, e.index.Meta()) e.recoverRows = append(e.recoverRows, recoverRows{handle: handle, idxVals: idxVals, rsData: rsData, skip: true}) result.scanRowCount++ result.currentHandle = handle diff --git a/pkg/executor/test/executor/BUILD.bazel b/pkg/executor/test/executor/BUILD.bazel index ec9501669c539..1f6af34ac385f 100644 --- a/pkg/executor/test/executor/BUILD.bazel +++ b/pkg/executor/test/executor/BUILD.bazel @@ -47,6 +47,8 @@ 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 45bc938245abc..ee25b25ea3f36 100644 --- a/pkg/executor/test/executor/executor_test.go +++ b/pkg/executor/test/executor/executor_test.go @@ -67,6 +67,8 @@ 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" @@ -251,7 +253,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(sc.TimeZone(), row, colIDs, nil, nil, nil, &rd) + value, err := tablecodec.EncodeRow(codec.NewEncoder(collate.NewCollationEnabled()), 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 3b45b9bbe559c..603dd965061de 100644 --- a/pkg/executor/write.go +++ b/pkg/executor/write.go @@ -35,6 +35,7 @@ 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/dbterror/plannererrors" "github.com/pingcap/tidb/pkg/util/memory" @@ -396,6 +397,7 @@ func addUnchangedKeysForLockByRow( return count, err } unchangedUniqueKey, _, err := tablecodec.GenIndexKey( + codec.NewEncoder(collate.NewCollationEnabled()), stmtCtx.TimeZone(), idx.TableMeta(), meta, diff --git a/pkg/expression/expression.go b/pkg/expression/expression.go index a648ba836556b..1361119727ed4 100644 --- a/pkg/expression/expression.go +++ b/pkg/expression/expression.go @@ -63,6 +63,8 @@ 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 @@ -100,6 +102,14 @@ 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 + } +} + // BuildSimpleExpr builds a simple expression from an ast node. // This function is used to build some "simple" expressions with limited context. // The below expressions are not supported: diff --git a/pkg/meta/model/table.go b/pkg/meta/model/table.go index d5386ce8c9dc2..9109bf5ddc69b 100644 --- a/pkg/meta/model/table.go +++ b/pkg/meta/model/table.go @@ -581,6 +581,19 @@ func FindFKInfoByName(fks []*FKInfo, name string) *FKInfo { return nil } +<<<<<<< HEAD +======= +// GetIdxChangingFieldType gets the field type of index column. +// Since both old/new type may coexist in one column during modify column, +// we need to get the correct type for index column. +func GetIdxChangingFieldType(idxCol *IndexColumn, col *ColumnInfo) *types.FieldType { + if idxCol.UseChangingType && col.ChangingFieldType != nil { + return col.ChangingFieldType + } + return &col.FieldType +} + +>>>>>>> ab1e19714d6 (codec, table: make new collation setting explicit in encoding (#69566)) // TableNameInfo provides meta data describing a table name info. type TableNameInfo struct { ID int64 `json:"id"` diff --git a/pkg/planner/core/expression_rewriter.go b/pkg/planner/core/expression_rewriter.go index 4a22b58f44659..534070d843e10 100644 --- a/pkg/planner/core/expression_rewriter.go +++ b/pkg/planner/core/expression_rewriter.go @@ -112,7 +112,9 @@ func buildSimpleExpr(ctx expression.BuildContext, node ast.ExprNode, opts ...exp return nil, errors.New("expression node should be present") } - var options expression.BuildOptions + options := expression.BuildOptions{ + UseNewCollate: collate.NewCollationEnabled(), + } for _, opt := range opts { opt(&options) } @@ -150,6 +152,7 @@ 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 { @@ -258,7 +261,8 @@ 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}, + planCtx: &exprRewriterPlanCtx{plan: p, builder: b, curClause: b.curClause, rollExpand: b.currentBlockExpand}, + useNewCollate: collate.NewCollationEnabled(), } b.rewriterPool = append(b.rewriterPool, rewriter) return @@ -373,7 +377,14 @@ type expressionRewriter struct { disableFoldCounter int tryFoldCounter int +<<<<<<< HEAD planCtx *exprRewriterPlanCtx +======= + astNodeStack []ast.Node + + planCtx *exprRewriterPlanCtx + useNewCollate bool +>>>>>>> ab1e19714d6 (codec, table: make new collation setting explicit in encoding (#69566)) } func (er *expressionRewriter) ctxStackLen() int { @@ -1652,7 +1663,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 collate.NewCollationEnabled() { + if er.useNewCollate { 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 { @@ -2068,7 +2079,7 @@ func (er *expressionRewriter) castCollationForIn(colLen int, elemCnt int, stkLen if colLen != 1 { return } - if !collate.NewCollationEnabled() { + if !er.useNewCollate { // 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. @@ -2172,7 +2183,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 && !collate.NewCollationEnabled() { + if patExpression, ok := er.ctxStack[l-1].(*expression.Constant); ok && !er.useNewCollate { patString, isNull, err := patExpression.EvalString(er.sctx.GetEvalCtx(), chunk.Row{}) if err != nil { er.err = err diff --git a/pkg/server/handler/tests/BUILD.bazel b/pkg/server/handler/tests/BUILD.bazel index 25baad0bc1ea0..adeb023d653de 100644 --- a/pkg/server/handler/tests/BUILD.bazel +++ b/pkg/server/handler/tests/BUILD.bazel @@ -54,6 +54,7 @@ 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 37d184ae87065..372b9666817e9 100644 --- a/pkg/server/handler/tests/http_handler_test.go +++ b/pkg/server/handler/tests/http_handler_test.go @@ -74,6 +74,7 @@ import ( "github.com/pingcap/tidb/pkg/testkit/external" "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" @@ -652,7 +653,7 @@ func TestDecodeColumnValue(t *testing.T) { } rd := rowcodec.Encoder{Enable: true} sc := stmtctx.NewStmtCtxWithTimeZone(time.UTC) - bs, err := tablecodec.EncodeRow(sc.TimeZone(), row, colIDs, nil, nil, nil, &rd) + bs, err := tablecodec.EncodeRow(codec.NewEncoder(collate.NewCollationEnabled()), 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/session/BUILD.bazel b/pkg/session/BUILD.bazel index bbfd75fdd1111..555b30bd3ae54 100644 --- a/pkg/session/BUILD.bazel +++ b/pkg/session/BUILD.bazel @@ -6,6 +6,7 @@ go_library( "advisory_locks.go", "bootstrap.go", "contextimpl.go", + "global_init.go", "mock_bootstrap.go", "nontransactional.go", "session.go", diff --git a/pkg/session/global_init.go b/pkg/session/global_init.go new file mode 100644 index 0000000000000..350bee9ac2233 --- /dev/null +++ b/pkg/session/global_init.go @@ -0,0 +1,79 @@ +// Copyright 2026 PingCAP, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package session + +import ( + "context" + + "github.com/pingcap/tidb/pkg/infoschema" + "github.com/pingcap/tidb/pkg/kv" + "github.com/pingcap/tidb/pkg/meta/metadef" + "github.com/pingcap/tidb/pkg/meta/model" + "github.com/pingcap/tidb/pkg/parser/mysql" + "github.com/pingcap/tidb/pkg/util/collate" + "github.com/pingcap/tidb/pkg/util/timeutil" +) + +type systemDBFilter struct{} + +func (systemDBFilter) SkipLoadDiff(*model.SchemaDiff, infoschema.InfoSchema) bool { + // when initialize global var, we don't start the domain, so this method + // will NOT be called, ok to return false here + return false +} + +func (systemDBFilter) SkipLoadSchema(dbInfo *model.DBInfo) bool { + return !metadef.IsSystemDB(dbInfo.Name.L) +} + +// we have to use a separate Domain to init the global variables, and then +// close it. as we are using a captured new_collate setting inside TableCommon +// now, otherwise, the captured new_collate setting in TableCommon is still +// cached in the Domain schema cache and is invalid. +func initGlobalVarFromSystemDB(ctx context.Context, store kv.Storage) error { + // the session used to load those global vars depends on the global vars + // themselves, which is a cyclic dependency. but luckily, mysql.tidb is + // created with DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_bin, it's safe to + // decode below rows without the correct TidbNewCollationEnabled initialized. + // + // Note: collate.newCollationEnabled is set to 1 in init(), so if the + // new_collate=false during first bootstrap, they mismatch. + dom, err := domap.GetOrCreateWithFilter(store, systemDBFilter{}) + if err != nil { + return err + } + defer dom.Close() + + sess, err := createSessionWithOpt(store, dom, dom.GetSchemaValidator(), dom.InfoCache(), nil) + if err != nil { + return err + } + + // get system tz from mysql.tidb + tz, err := sess.getTableValue(ctx, mysql.TiDBTable, tidbSystemTZ) + if err != nil { + return err + } + timeutil.SetSystemTZ(tz) + + // get the flag from `mysql`.`tidb` which indicating if new collations are enabled. + newCollationEnabled, err := loadCollationParameter(ctx, sess) + if err != nil { + return err + } + collate.SetNewCollationEnabledForTest(newCollationEnabled) + + return nil +} diff --git a/pkg/session/session.go b/pkg/session/session.go index a3fbb287c056e..387bfd9814431 100644 --- a/pkg/session/session.go +++ b/pkg/session/session.go @@ -133,7 +133,6 @@ import ( "github.com/pingcap/tidb/pkg/util/sqlescape" "github.com/pingcap/tidb/pkg/util/sqlexec" "github.com/pingcap/tidb/pkg/util/syncutil" - "github.com/pingcap/tidb/pkg/util/timeutil" "github.com/pingcap/tidb/pkg/util/topsql" topsqlstate "github.com/pingcap/tidb/pkg/util/topsql/state" "github.com/pingcap/tidb/pkg/util/topsql/stmtstats" @@ -3805,6 +3804,15 @@ func bootstrapSessionImpl(ctx context.Context, store kv.Storage, createSessionsI return nil, err } } + skipInitGlobalVarFromSystemDB := false + failpoint.Inject("skipInitGlobalVarFromSystemDB", func(val failpoint.Value) { + skipInitGlobalVarFromSystemDB = val.(bool) + }) + if !skipInitGlobalVarFromSystemDB { + if err = initGlobalVarFromSystemDB(ctx, store); err != nil { + return nil, err + } + } // initiate disttask framework components which need a store scheduler.RegisterSchedulerFactory( @@ -3828,27 +3836,12 @@ func bootstrapSessionImpl(ctx context.Context, store kv.Storage, createSessionsI if concurrency < 0 { // it is only for test, in the production, negative value is illegal. concurrency = 0 } - ses, err := createSessionsImpl(store, 10) if err != nil { return nil, err } ses[0].GetSessionVars().InRestrictedSQL = true - // get system tz from mysql.tidb - tz, err := ses[0].getTableValue(ctx, mysql.TiDBTable, tidbSystemTZ) - if err != nil { - return nil, err - } - timeutil.SetSystemTZ(tz) - - // get the flag from `mysql`.`tidb` which indicating if new collations are enabled. - newCollationEnabled, err := loadCollationParameter(ctx, ses[0]) - if err != nil { - return nil, err - } - collate.SetNewCollationEnabledForTest(newCollationEnabled) - // only start the domain after we have initialized some global variables. dom := domain.GetDomain(ses[0]) err = dom.Start(ddl.Normal) diff --git a/pkg/store/mockstore/BUILD.bazel b/pkg/store/mockstore/BUILD.bazel index ff8df81b8a30f..014aaea0b9078 100644 --- a/pkg/store/mockstore/BUILD.bazel +++ b/pkg/store/mockstore/BUILD.bazel @@ -52,6 +52,7 @@ 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 e5685cb1e13d8..8670adbd92006 100644 --- a/pkg/store/mockstore/cluster_test.go +++ b/pkg/store/mockstore/cluster_test.go @@ -28,6 +28,7 @@ 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" @@ -58,7 +59,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(sc.TimeZone(), []types.Datum{colValue}, []int64{colID}, nil, nil, nil, &rd) + rowValue, err1 := tablecodec.EncodeRow(codec.NewEncoder(collate.NewCollationEnabled()), 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 79a3fa67d9f04..cc86a0cd305e7 100644 --- a/pkg/store/mockstore/unistore/cophandler/cop_handler_test.go +++ b/pkg/store/mockstore/unistore/cophandler/cop_handler_test.go @@ -116,10 +116,11 @@ 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(stmtCtx.TimeZone(), datum, colIds, nil, nil, nil, encoder) + rowEncodedData, err := tablecodec.EncodeRow(codecEncoder, stmtCtx.TimeZone(), datum, colIds, nil, nil, nil, encoder) if err != nil { return nil, err } diff --git a/pkg/table/tables/index.go b/pkg/table/tables/index.go index 04f5286cce793..ab8d911189616 100644 --- a/pkg/table/tables/index.go +++ b/pkg/table/tables/index.go @@ -29,11 +29,32 @@ import ( "github.com/pingcap/tidb/pkg/tablecodec" "github.com/pingcap/tidb/pkg/types" "github.com/pingcap/tidb/pkg/util" +<<<<<<< HEAD +======= + "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" +>>>>>>> ab1e19714d6 (codec, table: make new collation setting explicit in encoding (#69566)) "github.com/pingcap/tidb/pkg/util/intest" "github.com/pingcap/tidb/pkg/util/rowcodec" "github.com/pingcap/tidb/pkg/util/tracing" ) +<<<<<<< HEAD +======= +var indexConditionECtx exprctx.BuildContext + +// indexPartialCondition is a data structure to help implement the partial index. +type indexPartialCondition struct { + conditionExpr expression.Expression + // conditionEvalBufferPool stores many eval buffer to avoid allocating chunk + // for evaluating partial index condition for each time. + // It's only initialized if the `partialConditionExpr` is not nil. + conditionEvalBufferPool sync.Pool +} + +>>>>>>> ab1e19714d6 (codec, table: make new collation setting explicit in encoding (#69566)) // index is the data structure for index data in the KV store. type index struct { idxInfo *model.IndexInfo @@ -44,13 +65,22 @@ type index struct { // the collation global variable is initialized *after* `NewIndex()`. initNeedRestoreData sync.Once needRestoredData bool +<<<<<<< HEAD +======= + encoder codec.Encoder + indexPartialCondition +>>>>>>> ab1e19714d6 (codec, table: make new collation setting explicit in encoding (#69566)) } // NeedRestoredData checks whether the index columns needs restored data. -func NeedRestoredData(idxCols []*model.IndexColumn, colInfos []*model.ColumnInfo) bool { +func NeedRestoredData(useNewCollate bool, idxCols []*model.IndexColumn, colInfos []*model.ColumnInfo) bool { for _, idxCol := range idxCols { col := colInfos[idxCol.Offset] +<<<<<<< HEAD if types.NeedRestoredData(&col.FieldType) { +======= + if types.NeedRestoredDataWithCollate(model.GetIdxChangingFieldType(idxCol, col), useNewCollate) { +>>>>>>> ab1e19714d6 (codec, table: make new collation setting explicit in encoding (#69566)) return true } } @@ -58,13 +88,64 @@ func NeedRestoredData(idxCols []*model.IndexColumn, colInfos []*model.ColumnInfo } // NewIndex builds a new Index object. +<<<<<<< HEAD func NewIndex(physicalID int64, tblInfo *model.TableInfo, indexInfo *model.IndexInfo) table.Index { +======= +func NewIndex(physicalID int64, tblInfo *model.TableInfo, indexInfo *model.IndexInfo) (table.Index, error) { + return NewIndexWithCollate(collate.NewCollationEnabled(), physicalID, tblInfo, indexInfo) +} + +// 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) { +>>>>>>> ab1e19714d6 (codec, table: make new collation setting explicit in encoding (#69566)) index := &index{ idxInfo: indexInfo, tblInfo: tblInfo, phyTblID: physicalID, + encoder: codec.NewEncoder(useNewCollate), } +<<<<<<< HEAD return index +======= + + 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) + } + + 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)) + } + + evalBuffer := chunk.MutRowFromTypes(evalBufferTypes) + return &evalBuffer + }, + } + } + return index, nil +>>>>>>> ab1e19714d6 (codec, table: make new collation setting explicit in encoding (#69566)) } // Meta returns index info. @@ -89,7 +170,17 @@ func (c *index) GenIndexKey(ec errctx.Context, loc *time.Location, indexedValues idxTblID = c.tblInfo.ID } } +<<<<<<< HEAD key, distinct, err = tablecodec.GenIndexKey(loc, c.tblInfo, c.idxInfo, idxTblID, indexedValues, h, buf) +======= + + if err = c.castIndexValuesToChangingTypes(indexedValues); err != nil { + return + } + + key, distinct, err = tablecodec.GenIndexKey(c.encoder, loc, c.tblInfo, c.idxInfo, + idxTblID, indexedValues, fullHandle, buf) +>>>>>>> ab1e19714d6 (codec, table: make new collation setting explicit in encoding (#69566)) err = ec.HandleError(err) return } @@ -98,9 +189,19 @@ func (c *index) GenIndexKey(ec errctx.Context, loc *time.Location, indexedValues func (c *index) GenIndexValue(ec errctx.Context, loc *time.Location, distinct bool, indexedValues []types.Datum, h kv.Handle, restoredData []types.Datum, buf []byte) ([]byte, error) { c.initNeedRestoreData.Do(func() { - c.needRestoredData = NeedRestoredData(c.idxInfo.Columns, c.tblInfo.Columns) + c.needRestoredData = NeedRestoredData(c.encoder.UseNewCollate(), c.idxInfo.Columns, c.tblInfo.Columns) }) +<<<<<<< HEAD idx, err := tablecodec.GenIndexValuePortal(loc, c.tblInfo, c.idxInfo, c.needRestoredData, distinct, false, indexedValues, h, c.phyTblID, restoredData, buf) +======= + + if err := c.castIndexValuesToChangingTypes(indexedValues); err != nil { + return nil, errors.Trace(err) + } + + idx, err := tablecodec.GenIndexValuePortal(c.encoder.UseNewCollate(), loc, c.tblInfo, + c.idxInfo, c.needRestoredData, distinct, untouched, indexedValues, h, c.phyTblID, restoredData, buf) +>>>>>>> ab1e19714d6 (codec, table: make new collation setting explicit in encoding (#69566)) err = ec.HandleError(err) return idx, err } diff --git a/pkg/table/tables/mutation_checker.go b/pkg/table/tables/mutation_checker.go index b4a65422b5623..521ebe6131185 100644 --- a/pkg/table/tables/mutation_checker.go +++ b/pkg/table/tables/mutation_checker.go @@ -218,6 +218,7 @@ func checkIndexKeys( indexIDToRowColInfos map[int64][]rowcodec.ColInfo, extraIndexesLayout table.IndexesLayout, ) error { + useNewCollate := t.encoder.UseNewCollate() var indexData []types.Datum for _, m := range indexMutations { var value []byte @@ -251,11 +252,12 @@ func checkIndexKeys( } // when we cannot decode the key to get the original value - if len(value) == 0 && NeedRestoredData(indexInfo.Columns, t.Meta().Columns) { + if len(value) == 0 && NeedRestoredData(useNewCollate, indexInfo.Columns, t.Meta().Columns) { continue } - decodedIndexValues, err := tablecodec.DecodeIndexKV( + decodedIndexValues, err := tablecodec.DecodeIndexKVWithCollate( + useNewCollate, m.key, value, len(indexInfo.Columns), tablecodec.HandleNotNeeded, rowColInfos, ) if err != nil { @@ -281,9 +283,9 @@ func checkIndexKeys( // When it is in add index new backfill state. if len(value) == 0 || isTmpIdxValAndDeleted { - err = compareIndexData(tc, t.Columns, indexData, rowToRemove, indexInfo, t.Meta(), extraIndexesLayout.GetIndexLayout(idxID)) + err = compareIndexData(useNewCollate, tc, t.Columns, indexData, rowToRemove, indexInfo, t.Meta(), extraIndexesLayout.GetIndexLayout(idxID)) } else { - err = compareIndexData(tc, t.Columns, indexData, rowToInsert, indexInfo, t.Meta(), extraIndexesLayout.GetIndexLayout(idxID)) + err = compareIndexData(useNewCollate, tc, t.Columns, indexData, rowToInsert, indexInfo, t.Meta(), extraIndexesLayout.GetIndexLayout(idxID)) } if err != nil { return errors.Trace(err) @@ -366,6 +368,7 @@ func collectTableMutationsFromBufferStage(t *TableCommon, memBuffer kv.MemBuffer // compareIndexData compares the decoded index data with the input data. // Returns error if the index data is not a subset of the input data. func compareIndexData( + useNewCollate bool, tc types.Context, cols []*table.Column, indexData, input []types.Datum, indexInfo *model.IndexInfo, tableInfo *model.TableInfo, extraIndexLayout table.IndexRowLayoutOption, @@ -389,7 +392,7 @@ func compareIndexData( ) comparison, err := CompareIndexAndVal(tc, expectedDatum, decodedMutationDatum, - collate.GetCollator(decodedMutationDatum.Collation()), + collate.GetCollatorWithCollate(useNewCollate, decodedMutationDatum.Collation()), cols[offsetInTable].ColumnInfo.FieldType.IsArray() && expectedDatum.Kind() == types.KindMysqlJSON) if err != nil { return errors.Trace(err) diff --git a/pkg/table/tables/mutation_checker_test.go b/pkg/table/tables/mutation_checker_test.go index 525360c75ad7f..e3b0e3076301f 100644 --- a/pkg/table/tables/mutation_checker_test.go +++ b/pkg/table/tables/mutation_checker_test.go @@ -19,7 +19,9 @@ import ( "testing" "time" + "github.com/pingcap/tidb/pkg/errctx" "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/parser/mysql" @@ -81,7 +83,7 @@ func TestCompareIndexData(t *testing.T) { } indexInfo := &model.IndexInfo{Name: ast.NewCIStr("i0"), Columns: indexCols} - err := compareIndexData(tc, cols, data.indexData, data.inputData, indexInfo, &model.TableInfo{Name: ast.NewCIStr("t")}, nil) + err := compareIndexData(collate.NewCollationEnabled(), tc, cols, data.indexData, data.inputData, indexInfo, &model.TableInfo{Name: ast.NewCIStr("t")}, nil) require.Equal(t, data.correct, err == nil, "case id = %v", caseID) } } @@ -92,7 +94,8 @@ func TestCheckRowInsertionConsistency(t *testing.T) { // mocked data mockRowKey233 := tablecodec.EncodeRowKeyWithHandle(1, kv.IntHandle(233)) - mockValue233, err := tablecodec.EncodeRow(sessVars.StmtCtx.TimeZone(), []types.Datum{types.NewIntDatum(233)}, []int64{101}, nil, nil, nil, &rd) + enc := codec.NewEncoder(collate.NewCollationEnabled()) + mockValue233, err := tablecodec.EncodeRow(enc, 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}} @@ -186,15 +189,18 @@ func TestCheckIndexKeysAndCheckHandleConsistency(t *testing.T) { indexInfos := []*model.IndexInfo{ { ID: 1, + Name: ast.NewCIStr("idx_unique"), State: model.StatePublic, Primary: false, Unique: true, Columns: []*model.IndexColumn{ { + Name: ast.NewCIStr("c2"), Offset: 1, Length: types.UnspecifiedLength, }, { + Name: ast.NewCIStr("c1"), Offset: 0, Length: types.UnspecifiedLength, }, @@ -202,15 +208,18 @@ func TestCheckIndexKeysAndCheckHandleConsistency(t *testing.T) { }, { ID: 2, + Name: ast.NewCIStr("idx_non_unique"), State: model.StatePublic, Primary: false, Unique: false, Columns: []*model.IndexColumn{ { + Name: ast.NewCIStr("c2"), Offset: 1, Length: types.UnspecifiedLength, }, { + Name: ast.NewCIStr("c1"), Offset: 0, Length: types.UnspecifiedLength, }, @@ -219,13 +228,13 @@ func TestCheckIndexKeysAndCheckHandleConsistency(t *testing.T) { } columnInfoSets := [][]*model.ColumnInfo{ { - {ID: 1, Offset: 0, FieldType: *types.NewFieldType(mysql.TypeString)}, - {ID: 2, Offset: 1, FieldType: *types.NewFieldType(mysql.TypeDatetime)}, + {ID: 1, Name: ast.NewCIStr("c1"), Offset: 0, State: model.StatePublic, FieldType: *types.NewFieldType(mysql.TypeString)}, + {ID: 2, Name: ast.NewCIStr("c2"), Offset: 1, State: model.StatePublic, FieldType: *types.NewFieldType(mysql.TypeDatetime)}, }, { - {ID: 1, Offset: 0, FieldType: *types.NewFieldTypeWithCollation(mysql.TypeString, "utf8_unicode_ci", + {ID: 1, Name: ast.NewCIStr("c1"), Offset: 0, State: model.StatePublic, FieldType: *types.NewFieldTypeWithCollation(mysql.TypeString, "utf8_unicode_ci", types.UnspecifiedLength)}, - {ID: 2, Offset: 1, FieldType: *types.NewFieldType(mysql.TypeDatetime)}, + {ID: 2, Name: ast.NewCIStr("c2"), Offset: 1, State: model.StatePublic, FieldType: *types.NewFieldType(mysql.TypeDatetime)}, }, } tc := types.DefaultStmtNoWarningContext @@ -249,92 +258,143 @@ func TestCheckIndexKeysAndCheckHandleConsistency(t *testing.T) { setter := func(maps map[int64]columnMaps) {} // test - collate.SetNewCollationEnabledForTest(true) - defer collate.SetNewCollationEnabledForTest(false) - for _, isCommonHandle := range []bool{true, false} { - for _, lc := range locations { - for _, columnInfos := range columnInfoSets { - tc = tc.WithLocation(lc) - tableInfo := model.TableInfo{ - ID: 1, - Name: ast.NewCIStr("t"), - Columns: columnInfos, - Indices: indexInfos, - PKIsHandle: false, - IsCommonHandle: isCommonHandle, - } - table := MockTableFromMeta(&tableInfo).(*TableCommon) - var handle, corruptedHandle kv.Handle - if isCommonHandle { - encoded, err := codec.EncodeKey(tc.Location(), nil, rowToInsert[0]) - require.Nil(t, err) - corrupted := make([]byte, len(encoded)) - copy(corrupted, encoded) - corrupted[len(corrupted)-1] ^= 1 - handle, err = kv.NewCommonHandle(encoded) - require.Nil(t, err) - corruptedHandle, err = kv.NewCommonHandle(corrupted) - require.Nil(t, err) - } else { - handle = kv.IntHandle(1) - corruptedHandle = kv.IntHandle(2) - } + originNewCollation := collate.NewCollationEnabled() + defer collate.SetNewCollationEnabledForTest(originNewCollation) + for _, useNewCollate := range []bool{true, false} { + collate.SetNewCollationEnabledForTest(useNewCollate) + for _, isCommonHandle := range []bool{true, false} { + for _, lc := range locations { + for _, columnInfos := range columnInfoSets { + tc = tc.WithLocation(lc) + tableInfo := model.TableInfo{ + ID: 1, + Name: ast.NewCIStr("t"), + State: model.StatePublic, + Columns: columnInfos, + Indices: indexInfos, + PKIsHandle: false, + IsCommonHandle: isCommonHandle, + } + tableFromMeta, err := TableFromMeta(autoid.NewAllocators(false), &tableInfo) + require.NoError(t, err) + table := tableFromMeta.(*TableCommon) + require.Equal(t, useNewCollate, table.encoder.UseNewCollate()) + var handle, corruptedHandle kv.Handle + if isCommonHandle { + encoded, err := codec.EncodeKey(tc.Location(), nil, rowToInsert[0]) + require.Nil(t, err) + corrupted := make([]byte, len(encoded)) + copy(corrupted, encoded) + corrupted[len(corrupted)-1] ^= 1 + handle, err = kv.NewCommonHandle(encoded) + require.Nil(t, err) + corruptedHandle, err = kv.NewCommonHandle(corrupted) + require.Nil(t, err) + } else { + handle = kv.IntHandle(1) + corruptedHandle = kv.IntHandle(2) + } - for i, indexInfo := range indexInfos { - index := table.indices[i] - maps := getOrBuildColumnMaps(getter, setter, table) + for i, indexInfo := range indexInfos { + index := table.indices[i] + maps := getOrBuildColumnMaps(getter, setter, table) - // test checkIndexKeys - insertionKey, insertionValue, err := buildIndexKeyValue(index, rowToInsert, tc.Location(), tableInfo, - indexInfo, table, handle) - require.Nil(t, err) - deletionKey, _, err := buildIndexKeyValue(index, rowToRemove, tc.Location(), tableInfo, indexInfo, table, - handle) - require.Nil(t, err) - indexMutations := []mutation{ - {key: insertionKey, value: insertionValue, indexID: indexInfo.ID}, - {key: deletionKey, indexID: indexInfo.ID}, - } - err = checkIndexKeys( - tc, table, rowToInsert, rowToRemove, indexMutations, maps.IndexIDToInfo, - maps.IndexIDToRowColInfos, nil, - ) - require.Nil(t, err) + // test checkIndexKeys + insertionKey, insertionValue, err := buildIndexKeyValue(index, rowToInsert, tc, indexInfo, table, handle) + require.Nil(t, err) + requireIndexKVDecodeMatchesRow(t, tc, table, indexInfo, rowToInsert, insertionKey, insertionValue, + maps.IndexIDToRowColInfos[indexInfo.ID]) + deletionKey, _, err := buildIndexKeyValue(index, rowToRemove, tc, indexInfo, table, handle) + require.Nil(t, err) + indexMutations := []mutation{ + {key: insertionKey, value: insertionValue, indexID: indexInfo.ID}, + {key: deletionKey, indexID: indexInfo.ID}, + } + err = checkIndexKeys( + tc, table, rowToInsert, rowToRemove, indexMutations, maps.IndexIDToInfo, + maps.IndexIDToRowColInfos, nil, + ) + require.Nil(t, err) - // test checkHandleConsistency - rowKey := tablecodec.EncodeRowKeyWithHandle(table.tableID, handle) - corruptedRowKey := tablecodec.EncodeRowKeyWithHandle(table.tableID, corruptedHandle) - 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} - err = checkHandleConsistency(rowMutation, indexMutations, maps.IndexIDToInfo, &tableInfo) - require.Nil(t, err) - err = checkHandleConsistency(corruptedRowMutation, indexMutations, maps.IndexIDToInfo, &tableInfo) - require.NotNil(t, err) + // 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) + require.Nil(t, err) + rowMutation := mutation{key: rowKey, value: rowValue} + corruptedRowMutation := mutation{key: corruptedRowKey, value: rowValue} + err = checkHandleConsistency(rowMutation, indexMutations, maps.IndexIDToInfo, &tableInfo) + require.Nil(t, err) + err = checkHandleConsistency(corruptedRowMutation, indexMutations, maps.IndexIDToInfo, &tableInfo) + require.NotNil(t, err) + } } } } } } -func buildIndexKeyValue(index table.Index, rowToInsert []types.Datum, loc *time.Location, - tableInfo model.TableInfo, indexInfo *model.IndexInfo, table *TableCommon, handle kv.Handle) ([]byte, []byte, error) { +func requireIndexKVDecodeMatchesRow( + t *testing.T, + tc types.Context, + table *TableCommon, + indexInfo *model.IndexInfo, + row []types.Datum, + key []byte, + value []byte, + rowColInfos []rowcodec.ColInfo, +) { + t.Helper() + + decodedIndexValues, err := tablecodec.DecodeIndexKV( + key, value, len(indexInfo.Columns), tablecodec.HandleNotNeeded, rowColInfos, + ) + require.NoError(t, err) + requireDecodedIndexValuesMatchRow(t, tc, table, indexInfo, row, decodedIndexValues) + + preAlloc := make([][]byte, len(indexInfo.Columns), len(indexInfo.Columns)+len(rowColInfos)) + decodedIndexValuesEx, err := tablecodec.DecodeIndexKVEx( + key, value, len(indexInfo.Columns), tablecodec.HandleNotNeeded, rowColInfos, nil, preAlloc, + ) + require.NoError(t, err) + require.Equal(t, decodedIndexValues, decodedIndexValuesEx) + requireDecodedIndexValuesMatchRow(t, tc, table, indexInfo, row, decodedIndexValuesEx) +} + +func requireDecodedIndexValuesMatchRow( + t *testing.T, + tc types.Context, + table *TableCommon, + indexInfo *model.IndexInfo, + row []types.Datum, + decodedIndexValues [][]byte, +) { + t.Helper() + + require.Len(t, decodedIndexValues, len(indexInfo.Columns)) + indexData := make([]types.Datum, 0, len(decodedIndexValues)) + for i, value := range decodedIndexValues { + fieldType := table.Columns[indexInfo.Columns[i].Offset].FieldType.ArrayType() + datum, err := tablecodec.DecodeColumnValue(value, fieldType, tc.Location()) + require.NoError(t, err) + indexData = append(indexData, datum) + } + require.NoError(t, compareIndexData(table.encoder.UseNewCollate(), tc, table.Columns, indexData, row, indexInfo, table.Meta(), nil)) +} + +func buildIndexKeyValue(index table.Index, rowToInsert []types.Datum, tc types.Context, + indexInfo *model.IndexInfo, table *TableCommon, handle kv.Handle) ([]byte, []byte, error) { indexedValues, err := index.FetchValues(rowToInsert, nil) if err != nil { return nil, nil, err } - key, distinct, err := tablecodec.GenIndexKey( - loc, &tableInfo, indexInfo, 1, indexedValues, handle, nil, - ) + key, distinct, err := index.GenIndexKey(errctx.StrictNoWarningContext, tc.Location(), indexedValues, handle, nil) if err != nil { return nil, nil, err } - rsData := TryGetHandleRestoredDataWrapper(table.meta, rowToInsert, nil, indexInfo) - value, err := tablecodec.GenIndexValuePortal( - loc, &tableInfo, indexInfo, NeedRestoredData(indexInfo.Columns, tableInfo.Columns), - distinct, false, indexedValues, handle, 0, rsData, nil, - ) + useNewCollate := table.encoder.UseNewCollate() + rsData := TryGetHandleRestoredDataWrapper(useNewCollate, table.meta, rowToInsert, nil, indexInfo) + value, err := index.GenIndexValue(errctx.StrictNoWarningContext, tc.Location(), distinct, false, indexedValues, handle, rsData, nil) if err != nil { return nil, nil, err } diff --git a/pkg/table/tables/partition.go b/pkg/table/tables/partition.go index ced458842926d..e4f8971fe537a 100644 --- a/pkg/table/tables/partition.go +++ b/pkg/table/tables/partition.go @@ -1676,7 +1676,7 @@ func GetReorganizedPartitionedTable(t table.Table) (table.PartitionedTable, erro return nil, err } var tc TableCommon - initTableCommon(&tc, tblInfo, tblInfo.ID, t.Cols(), t.Allocators(nil), constraints) + tc.initTableCommon(tblInfo, tblInfo.ID, t.Cols(), t.Allocators(nil), constraints) // and rebuild the partitioning structure return newPartitionedTable(&tc, tblInfo) diff --git a/pkg/table/tables/tables.go b/pkg/table/tables/tables.go index 04b5665620b6c..d650e7293e864 100644 --- a/pkg/table/tables/tables.go +++ b/pkg/table/tables/tables.go @@ -81,7 +81,9 @@ type TableCommon struct { // recordPrefix and indexPrefix are generated using physicalTableID. recordPrefix kv.Key indexPrefix kv.Key - + // encoder keeps the collation setting captured when the table is initialized. + // All row and index writes through this table must use this fixed setting. + encoder codec.Encoder // skipAssert is used for partitions that are in WriteOnly/DeleteOnly state. skipAssert bool } @@ -140,7 +142,7 @@ func MockTableFromMeta(tblInfo *model.TableInfo) table.Table { return nil } var t TableCommon - initTableCommon(&t, tblInfo, tblInfo.ID, columns, autoid.NewAllocators(false), constraints) + t.initTableCommon(tblInfo, tblInfo.ID, columns, autoid.NewAllocators(false), constraints) if tblInfo.TableCacheStatusType != model.TableCacheStatusDisable { ret, err := newCachedTable(&t) if err != nil { @@ -214,7 +216,7 @@ func TableFromMeta(allocs autoid.Allocators, tblInfo *model.TableInfo) (table.Ta return nil, err } var t TableCommon - initTableCommon(&t, tblInfo, tblInfo.ID, columns, allocs, constraints) + t.initTableCommon(tblInfo, tblInfo.ID, columns, allocs, constraints) if tblInfo.GetPartitionInfo() == nil { if err := initTableIndices(&t); err != nil { return nil, err @@ -240,7 +242,7 @@ func buildGeneratedExpr(tblInfo *model.TableInfo, genExpr string) (ast.ExprNode, } // initTableCommon initializes a TableCommon struct. -func initTableCommon(t *TableCommon, tblInfo *model.TableInfo, physicalTableID int64, cols []*table.Column, allocs autoid.Allocators, constraints []*table.Constraint) { +func (t *TableCommon) initTableCommon(tblInfo *model.TableInfo, physicalTableID int64, cols []*table.Column, allocs autoid.Allocators, constraints []*table.Constraint) { t.tableID = tblInfo.ID t.physicalTableID = physicalTableID t.allocs = allocs @@ -249,6 +251,7 @@ func initTableCommon(t *TableCommon, tblInfo *model.TableInfo, physicalTableID i t.Constraints = constraints t.recordPrefix = tablecodec.GenTableRecordPrefix(physicalTableID) t.indexPrefix = tablecodec.GenTableIndexPrefix(physicalTableID) + t.encoder = codec.NewEncoder(collate.NewCollationEnabled()) if tblInfo.IsSequence() { t.sequence = &sequenceCommon{meta: tblInfo.Sequence} } @@ -264,7 +267,14 @@ func initTableIndices(t *TableCommon) error { } // Use partition ID for index, because TableCommon may be table or partition. +<<<<<<< HEAD idx := NewIndex(t.physicalTableID, tblInfo, idxInfo) +======= + idx, err := NewIndexWithCollate(t.encoder.UseNewCollate(), t.physicalTableID, tblInfo, idxInfo) + if err != nil { + return err + } +>>>>>>> ab1e19714d6 (codec, table: make new collation setting explicit in encoding (#69566)) 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`. @@ -285,7 +295,7 @@ func asIndex(idx table.Index) *index { } func initTableCommonWithIndices(t *TableCommon, tblInfo *model.TableInfo, physicalTableID int64, cols []*table.Column, allocs autoid.Allocators, constraints []*table.Constraint) error { - initTableCommon(t, tblInfo, physicalTableID, cols, allocs, constraints) + t.initTableCommon(tblInfo, physicalTableID, cols, allocs, constraints) return initTableIndices(t) } @@ -492,7 +502,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(sctx.GetRowEncodingConfig(), tc.Location(), ec, memBuffer, key, h) + err = encodeRowBuffer.WriteMemBufferEncoded(t.encoder, sctx.GetRowEncodingConfig(), tc.Location(), ec, memBuffer, key, h) if err != nil { return err } @@ -723,7 +733,7 @@ func (t *TableCommon) addRecord(sctx table.MutateContext, txn kv.Transaction, r } tablecodec.TruncateIndexValues(tblInfo, pkIdx, pkDts) var handleBytes []byte - handleBytes, err = codec.EncodeKey(tc.Location(), nil, pkDts...) + handleBytes, err = t.encoder.EncodeKey(tc.Location(), nil, pkDts...) err = ec.HandleError(err) if err != nil { return @@ -855,7 +865,7 @@ func (t *TableCommon) addRecord(sctx table.MutateContext, txn kv.Transaction, r } } - err = encodeRowBuffer.WriteMemBufferEncoded(sctx.GetRowEncodingConfig(), tc.Location(), ec, memBuffer, key, recordID, flags...) + err = encodeRowBuffer.WriteMemBufferEncoded(t.encoder, sctx.GetRowEncodingConfig(), tc.Location(), ec, memBuffer, key, recordID, flags...) if err != nil { return nil, err } @@ -954,7 +964,7 @@ func (t *TableCommon) addIndices(sctx table.MutateContext, recordID kv.Handle, r } dupErr = kv.GenKeyExistsErr(colStrVals, fmt.Sprintf("%s.%s", v.TableMeta().Name.String(), v.Meta().Name.String())) } - rsData := TryGetHandleRestoredDataWrapper(t.meta, r, nil, v.Meta()) + rsData := TryGetHandleRestoredDataWrapper(t.encoder.UseNewCollate(), t.meta, r, nil, v.Meta()) if dupHandle, err := asIndex(v).create(sctx, txn, indexVals, recordID, rsData, false, opt); err != nil { if kv.ErrKeyExists.Equal(err) { return dupHandle, dupErr @@ -1221,7 +1231,7 @@ func (t *TableCommon) removeRowIndices(ctx table.MutateContext, txn kv.Transacti // buildIndexForRow implements table.Table BuildIndexForRow interface. func (t *TableCommon) buildIndexForRow(ctx table.MutateContext, h kv.Handle, vals []types.Datum, newData []types.Datum, idx *index, txn kv.Transaction, untouched bool, opt *table.CreateIdxOpt) error { - rsData := TryGetHandleRestoredDataWrapper(t.meta, newData, nil, idx.Meta()) + rsData := TryGetHandleRestoredDataWrapper(t.encoder.UseNewCollate(), t.meta, newData, nil, idx.Meta()) if _, err := idx.create(ctx, txn, vals, h, rsData, untouched, opt); err != nil { if kv.ErrKeyExists.Equal(err) { // Make error message consistent with MySQL. @@ -1428,14 +1438,14 @@ func (t *TableCommon) Type() table.Type { } func (t *TableCommon) canSkip(col *table.Column, value *types.Datum) bool { - return CanSkip(t.Meta(), col, value) + 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(info *model.TableInfo, col *table.Column, value *types.Datum) bool { +func CanSkip(useNewCollate bool, info *model.TableInfo, col *table.Column, value *types.Datum) bool { if col.IsPKHandleColumn(info) { return true } @@ -1446,7 +1456,7 @@ func CanSkip(info *model.TableInfo, col *table.Column, value *types.Datum) bool continue } canSkip := idxCol.Length == types.UnspecifiedLength - canSkip = canSkip && !types.NeedRestoredData(&col.FieldType) + canSkip = canSkip && !types.NeedRestoredDataWithCollate(&col.FieldType, useNewCollate) return canSkip } } @@ -1682,16 +1692,18 @@ func (t *TableCommon) GetSequenceCommon() *sequenceCommon { return t.sequence } -// TryGetHandleRestoredDataWrapper tries to get the restored data for handle if needed. The argument can be a slice or a map. -func TryGetHandleRestoredDataWrapper(tblInfo *model.TableInfo, row []types.Datum, rowMap map[int64]types.Datum, idx *model.IndexInfo) []types.Datum { - if !collate.NewCollationEnabled() || !tblInfo.IsCommonHandle || tblInfo.CommonHandleVersion == 0 { +// TryGetHandleRestoredDataWrapper tries to get the restored data for handle if +// needed. The argument can be a slice or a map. +func TryGetHandleRestoredDataWrapper(useNewCollate bool, tblInfo *model.TableInfo, + row []types.Datum, rowMap map[int64]types.Datum, idx *model.IndexInfo) []types.Datum { + if !useNewCollate || !tblInfo.IsCommonHandle || tblInfo.CommonHandleVersion == 0 { return nil } rsData := make([]types.Datum, 0, 4) pkIdx := FindPrimaryIndex(tblInfo) for _, pkIdxCol := range pkIdx.Columns { pkCol := tblInfo.Columns[pkIdxCol.Offset] - if !types.NeedRestoredData(&pkCol.FieldType) { + if !types.NeedRestoredDataWithCollate(&pkCol.FieldType, useNewCollate) { continue } var datum types.Datum diff --git a/pkg/table/tables/testutil/BUILD.bazel b/pkg/table/tables/testutil/BUILD.bazel new file mode 100644 index 0000000000000..6fb8235d08bb4 --- /dev/null +++ b/pkg/table/tables/testutil/BUILD.bazel @@ -0,0 +1,19 @@ +load("@io_bazel_rules_go//go:def.bzl", "go_library") + +go_library( + name = "testutil", + srcs = ["indexcheck.go"], + importpath = "github.com/pingcap/tidb/pkg/table/tables/testutil", + visibility = ["//visibility:public"], + deps = [ + "//pkg/domain", + "//pkg/parser/ast", + "//pkg/sessiontxn", + "//pkg/table", + "//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 new file mode 100644 index 0000000000000..bffb6cfff5793 --- /dev/null +++ b/pkg/table/tables/testutil/indexcheck.go @@ -0,0 +1,75 @@ +// Copyright 2025 PingCAP, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package testutil + +import ( + "context" + "testing" + "time" + + "github.com/pingcap/tidb/pkg/domain" + "github.com/pingcap/tidb/pkg/parser/ast" + "github.com/pingcap/tidb/pkg/sessiontxn" + "github.com/pingcap/tidb/pkg/table" + "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" +) + +// CheckIndexKVCount checks the number of index key-value pairs in the specified index of the specified table. +func CheckIndexKVCount(t *testing.T, tk *testkit.TestKit, dom *domain.Domain, tableName string, indexName string, expected int) { + tbl, err := dom.InfoSchema().TableByName(context.Background(), ast.NewCIStr("test"), ast.NewCIStr(tableName)) + require.NoError(t, err) + require.NotNil(t, tbl) + var idx table.Index + for _, i := range tbl.Indices() { + if i.Meta().Name.L == indexName { + idx = i + break + } + } + + enc := codec.NewEncoder(collate.NewCollationEnabled()) + minimumKey, _, err := tablecodec.GenIndexKey(enc, time.Local, tbl.Meta(), idx.Meta(), tbl.Meta().ID, nil, nil, nil) + require.NoError(t, err) + + tk.MustExec("BEGIN") + defer tk.MustExec("COMMIT") + txnManager := sessiontxn.GetTxnManager(tk.Session()) + snapshot, err := txnManager.GetSnapshotWithStmtReadTS() + require.NoError(t, err) + + iter, err := snapshot.Iter(minimumKey, nil) + require.NoError(t, err) + defer iter.Close() + count := 0 + for iter.Valid() { + key := iter.Key() + if !tablecodec.IsIndexKey(key) { + break + } + _, idxID, _, err := tablecodec.DecodeIndexKey(key) + require.NoError(t, err) + if idxID != idx.Meta().ID { + break + } + count++ + err = iter.Next() + require.NoError(t, err) + } + require.Equal(t, expected, count) +} diff --git a/pkg/table/tblctx/BUILD.bazel b/pkg/table/tblctx/BUILD.bazel index f7ea643d76115..3814046cea22b 100644 --- a/pkg/table/tblctx/BUILD.bazel +++ b/pkg/table/tblctx/BUILD.bazel @@ -20,6 +20,8 @@ go_library( "//pkg/tablecodec", "//pkg/types", "//pkg/util/chunk", + "//pkg/util/codec", + "//pkg/util/collate", "//pkg/util/intest", "//pkg/util/rowcodec", "//pkg/util/tableutil", @@ -40,6 +42,8 @@ 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 5e9e8d93d0bb1..2a51a7d8b4a9f 100644 --- a/pkg/table/tblctx/buffers.go +++ b/pkg/table/tblctx/buffers.go @@ -23,6 +23,8 @@ 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" ) @@ -51,6 +53,7 @@ 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 { @@ -69,7 +72,7 @@ func (b *EncodeRowBuffer) WriteMemBufferEncoded( stmtBufs.AddRowValues = ensureCapacityAndReset(stmtBufs.AddRowValues, len(b.row)*2) encoded, err := tablecodec.EncodeRow( - loc, b.row, b.colIDs, stmtBufs.RowValBuf, stmtBufs.AddRowValues, checksum, cfg.RowEncoder, + enc, loc, b.row, b.colIDs, stmtBufs.RowValBuf, stmtBufs.AddRowValues, checksum, cfg.RowEncoder, ) if err = ec.HandleError(err); err != nil { return err @@ -85,7 +88,8 @@ 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) { - value, err := tablecodec.EncodeOldRow(loc, b.row, b.colIDs, nil, nil) + enc := codec.NewEncoder(collate.NewCollationEnabled()) + value, err := tablecodec.EncodeOldRow(enc, 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 264efe7619cec..f3ef0fe6ed3ef 100644 --- a/pkg/table/tblctx/buffers_test.go +++ b/pkg/table/tblctx/buffers_test.go @@ -25,6 +25,8 @@ 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" @@ -78,6 +80,7 @@ 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 @@ -110,7 +113,7 @@ func TestEncodeRow(t *testing.T) { } expectedVal, err := tablecodec.EncodeRow( - c.loc, []types.Datum{d1, d2, d3}, []int64{1, 2, 3}, nil, nil, checksum, + enc, c.loc, []types.Datum{d1, d2, d3}, []int64{1, 2, 3}, nil, nil, checksum, &rowcodec.Encoder{Enable: !c.oldFormat}, ) require.NoError(t, err) @@ -124,7 +127,7 @@ func TestEncodeRow(t *testing.T) { Return(nil).Once() } err = buffer.WriteMemBufferEncoded( - cfg, c.loc, errctx.StrictNoWarningContext, + enc, cfg, c.loc, errctx.StrictNoWarningContext, memBuffer, kv.Key("key1"), kv.IntHandle(1), c.flags..., ) require.NoError(t, err) @@ -134,7 +137,7 @@ func TestEncodeRow(t *testing.T) { // test encode val for binlog expectedVal, err = - tablecodec.EncodeOldRow(c.loc, []types.Datum{d1, d2, d3}, []int64{1, 2, 3}, nil, nil) + tablecodec.EncodeOldRow(enc, 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) @@ -164,7 +167,8 @@ 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)) - require.NoError(t, buffer.WriteMemBufferEncoded(RowEncodingConfig{ + enc := codec.NewEncoder(collate.NewCollationEnabled()) + require.NoError(t, buffer.WriteMemBufferEncoded(enc, 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 aea8eca5a1c29..5b5fa2ea6a6a4 100644 --- a/pkg/tablecodec/tablecodec.go +++ b/pkg/tablecodec/tablecodec.go @@ -359,7 +359,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(loc *time.Location, row []types.Datum, colIDs []int64, valBuf []byte, values []types.Datum, checksum rowcodec.Checksum, e *rowcodec.Encoder) ([]byte, error) { +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) { if len(row) != len(colIDs) { return nil, errors.Errorf("EncodeRow error: data and columnID count not match %d vs %d", len(row), len(colIDs)) } @@ -367,14 +367,14 @@ func EncodeRow(loc *time.Location, row []types.Datum, colIDs []int64, valBuf []b valBuf = valBuf[:0] return e.Encode(loc, colIDs, row, checksum, valBuf) } - return EncodeOldRow(loc, row, colIDs, valBuf, values) + return EncodeOldRow(enc, 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(loc *time.Location, row []types.Datum, colIDs []int64, valBuf []byte, values []types.Datum) ([]byte, error) { +func EncodeOldRow(enc codec.Encoder, 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)) } @@ -394,7 +394,7 @@ func EncodeOldRow(loc *time.Location, row []types.Datum, colIDs []int64, valBuf // We could not set nil value into kv. return append(valBuf, codec.NilFlag), nil } - return codec.EncodeValue(loc, valBuf, values...) + return enc.EncodeValue(loc, valBuf, values...) } func flatten(loc *time.Location, data types.Datum, ret *types.Datum) error { @@ -827,7 +827,7 @@ func reEncodeHandleTo(handle kv.Handle, unsigned bool, buf []byte, result [][]by } // reEncodeHandleConsiderNewCollation encodes the handle as a Datum so it can be properly decoded later. -func reEncodeHandleConsiderNewCollation(handle kv.Handle, columns []rowcodec.ColInfo, restoreData []byte) ([][]byte, error) { +func reEncodeHandleConsiderNewCollation(useNewCollate bool, handle kv.Handle, columns []rowcodec.ColInfo, restoreData []byte) ([][]byte, error) { handleColLen := handle.NumCols() cHandleBytes := make([][]byte, 0, handleColLen) for i := range handleColLen { @@ -842,7 +842,7 @@ func reEncodeHandleConsiderNewCollation(handle kv.Handle, columns []rowcodec.Col for idx > 0 && columns[idx-1].ID < 0 { idx-- } - return decodeRestoredValuesV5(columns[:idx], cHandleBytes, restoreData) + return decodeRestoredValuesV5(useNewCollate, columns[:idx], cHandleBytes, restoreData) } func decodeRestoredValues(columns []rowcodec.ColInfo, restoredVal []byte) ([][]byte, error) { @@ -864,9 +864,9 @@ func decodeRestoredValues(columns []rowcodec.ColInfo, restoredVal []byte) ([][]b // 1. If the index is a composed index, only the non-binary string column's value need to write to value, not all. // 2. If a string column's collation is _bin, then we only write the number of the truncated spaces to value. // 3. If a string column is char, not varchar, then we use the sortKey directly. -func decodeRestoredValuesV5(columns []rowcodec.ColInfo, results [][]byte, restoredVal []byte) ([][]byte, error) { +func decodeRestoredValuesV5(useNewCollate bool, columns []rowcodec.ColInfo, results [][]byte, restoredVal []byte) ([][]byte, error) { colIDOffsets := buildColumnIDOffsets(columns) - colInfosNeedRestore := buildRestoredColumn(columns) + colInfosNeedRestore := buildRestoredColumn(useNewCollate, columns) rd := rowcodec.NewByteDecoder(colInfosNeedRestore, nil, nil, nil) newResults, err := rd.DecodeToBytesNoHandle(colIDOffsets, restoredVal) if err != nil { @@ -911,10 +911,10 @@ func buildColumnIDOffsets(allCols []rowcodec.ColInfo) map[int64]int { return colIDOffsets } -func buildRestoredColumn(allCols []rowcodec.ColInfo) []rowcodec.ColInfo { +func buildRestoredColumn(useNewCollate bool, allCols []rowcodec.ColInfo) []rowcodec.ColInfo { restoredColumns := make([]rowcodec.ColInfo, 0, len(allCols)) for i, col := range allCols { - if !types.NeedRestoredData(col.Ft) { + if !types.NeedRestoredDataWithCollate(col.Ft, useNewCollate) { continue } copyColInfo := rowcodec.ColInfo{ @@ -982,7 +982,7 @@ func DecodeIndexKVEx(key, value []byte, colsLen int, hdStatus HandleStatus, colu return decodeIndexKvOldCollation(key, value, hdStatus, buf, preAlloc) } if getIndexVersion(value) == 1 { - return decodeIndexKvForClusteredIndexVersion1(key, value, colsLen, hdStatus, columns) + return decodeIndexKvForClusteredIndexVersion1(collate.NewCollationEnabled(), key, value, colsLen, hdStatus, columns) } return decodeIndexKvGeneral(key, value, colsLen, hdStatus, columns) } @@ -992,12 +992,17 @@ func DecodeIndexKVEx(key, value []byte, colsLen int, hdStatus HandleStatus, colu // `colsLen` is expected to be index columns count. // `columns` is expected to be index columns + handle columns(if hdStatus is not HandleNotNeeded). func DecodeIndexKV(key, value []byte, colsLen int, hdStatus HandleStatus, columns []rowcodec.ColInfo) ([][]byte, error) { + return DecodeIndexKVWithCollate(collate.NewCollationEnabled(), key, value, colsLen, hdStatus, columns) +} + +// DecodeIndexKVWithCollate is similar to DecodeIndexKV but with explicit useNewCollate param. +func DecodeIndexKVWithCollate(useNewCollate bool, key, value []byte, colsLen int, hdStatus HandleStatus, columns []rowcodec.ColInfo) ([][]byte, error) { if len(value) <= MaxOldEncodeValueLen { preAlloc := make([][]byte, colsLen, colsLen+len(columns)) return decodeIndexKvOldCollation(key, value, hdStatus, nil, preAlloc) } if getIndexVersion(value) == 1 { - return decodeIndexKvForClusteredIndexVersion1(key, value, colsLen, hdStatus, columns) + return decodeIndexKvForClusteredIndexVersion1(useNewCollate, key, value, colsLen, hdStatus, columns) } return decodeIndexKvGeneral(key, value, colsLen, hdStatus, columns) } @@ -1199,7 +1204,7 @@ func GetIndexKeyBuf(buf []byte, defaultCap int) []byte { } // GenIndexKey generates index key using input physical table id -func GenIndexKey(loc *time.Location, tblInfo *model.TableInfo, idxInfo *model.IndexInfo, +func GenIndexKey(enc codec.Encoder, loc *time.Location, tblInfo *model.TableInfo, idxInfo *model.IndexInfo, phyTblID int64, indexedValues []types.Datum, h kv.Handle, buf []byte) (key []byte, distinct bool, err error) { if idxInfo.Unique { // See https://dev.mysql.com/doc/refman/5.7/en/create-index.html @@ -1220,7 +1225,7 @@ func GenIndexKey(loc *time.Location, tblInfo *model.TableInfo, idxInfo *model.In key = GetIndexKeyBuf(buf, RecordRowKeyLen+len(indexedValues)*9+9) key = appendTableIndexPrefix(key, phyTblID) key = codec.EncodeInt(key, idxInfo.ID) - key, err = codec.EncodeKey(loc, key, indexedValues...) + key, err = enc.EncodeKey(loc, key, indexedValues...) if err != nil { return nil, false, err } @@ -1550,18 +1555,18 @@ func TempIndexValueIsUntouched(b []byte) bool { // | In v5.0, restored data contains only non-binary data(except for char and _bin). In the above example, the restored data contains only the value of b. // | Besides, if the collation of b is _bin, then restored data is an integer indicate the spaces are truncated. Then we use sortKey // | and the restored data together to restore original data. -func GenIndexValuePortal(loc *time.Location, tblInfo *model.TableInfo, idxInfo *model.IndexInfo, +func GenIndexValuePortal(useNewCollate bool, loc *time.Location, tblInfo *model.TableInfo, idxInfo *model.IndexInfo, needRestoredData bool, distinct bool, untouched bool, indexedValues []types.Datum, h kv.Handle, partitionID int64, restoredData []types.Datum, buf []byte) ([]byte, error) { if tblInfo.IsCommonHandle && tblInfo.CommonHandleVersion == 1 { - return GenIndexValueForClusteredIndexVersion1(loc, tblInfo, idxInfo, needRestoredData, distinct, untouched, indexedValues, h, partitionID, restoredData, buf) + return GenIndexValueForClusteredIndexVersion1(useNewCollate, loc, tblInfo, idxInfo, needRestoredData, distinct, untouched, indexedValues, h, partitionID, restoredData, buf) } return genIndexValueVersion0(loc, tblInfo, idxInfo, needRestoredData, distinct, untouched, indexedValues, h, partitionID, buf) } // TryGetCommonPkColumnRestoredIds get the IDs of primary key columns which need restored data if the table has common handle. // Caller need to make sure the table has common handle. -func TryGetCommonPkColumnRestoredIds(tbl *model.TableInfo) []int64 { +func TryGetCommonPkColumnRestoredIds(useNewCollate bool, tbl *model.TableInfo) []int64 { var pkColIDs []int64 var pkIdx *model.IndexInfo for _, idx := range tbl.Indices { @@ -1574,7 +1579,7 @@ func TryGetCommonPkColumnRestoredIds(tbl *model.TableInfo) []int64 { return pkColIDs } for _, idxCol := range pkIdx.Columns { - if types.NeedRestoredData(&tbl.Columns[idxCol.Offset].FieldType) { + if types.NeedRestoredDataWithCollate(&tbl.Columns[idxCol.Offset].FieldType, useNewCollate) { pkColIDs = append(pkColIDs, tbl.Columns[idxCol.Offset].ID) } } @@ -1582,7 +1587,7 @@ func TryGetCommonPkColumnRestoredIds(tbl *model.TableInfo) []int64 { } // GenIndexValueForClusteredIndexVersion1 generates the index value for the clustered index with version 1(New in v5.0.0). -func GenIndexValueForClusteredIndexVersion1(loc *time.Location, tblInfo *model.TableInfo, idxInfo *model.IndexInfo, +func GenIndexValueForClusteredIndexVersion1(useNewCollate bool, loc *time.Location, tblInfo *model.TableInfo, idxInfo *model.IndexInfo, idxValNeedRestoredData bool, distinct bool, untouched bool, indexedValues []types.Datum, h kv.Handle, partitionID int64, handleRestoredData []types.Datum, buf []byte) ([]byte, error) { var idxVal []byte @@ -1613,7 +1618,11 @@ func GenIndexValueForClusteredIndexVersion1(loc *time.Location, tblInfo *model.T if mysql.HasPriKeyFlag(col.GetFlag()) { continue } +<<<<<<< HEAD if types.NeedRestoredData(&col.FieldType) { +======= + if types.NeedRestoredDataWithCollate(model.GetIdxChangingFieldType(idxCol, col), useNewCollate) { +>>>>>>> ab1e19714d6 (codec, table: make new collation setting explicit in encoding (#69566)) colIds = append(colIds, col.ID) if collate.IsBinCollation(col.GetCollate()) { allRestoredData = append(allRestoredData, types.NewUintDatum(uint64(stringutil.GetTailSpaceCount(indexedValues[i].GetString())))) @@ -1624,7 +1633,7 @@ func GenIndexValueForClusteredIndexVersion1(loc *time.Location, tblInfo *model.T } if len(handleRestoredData) > 0 { - pkColIDs := TryGetCommonPkColumnRestoredIds(tblInfo) + pkColIDs := TryGetCommonPkColumnRestoredIds(useNewCollate, tblInfo) colIds = append(colIds, pkColIDs...) allRestoredData = append(allRestoredData, handleRestoredData...) } @@ -1859,7 +1868,7 @@ func splitIndexValueForClusteredIndexVersion1(value []byte) (segs IndexValueSegm return } -func decodeIndexKvForClusteredIndexVersion1(key, value []byte, colsLen int, hdStatus HandleStatus, columns []rowcodec.ColInfo) ([][]byte, error) { +func decodeIndexKvForClusteredIndexVersion1(useNewCollate bool, key, value []byte, colsLen int, hdStatus HandleStatus, columns []rowcodec.ColInfo) ([][]byte, error) { var resultValues [][]byte var keySuffix []byte var handle kv.Handle @@ -1870,7 +1879,7 @@ func decodeIndexKvForClusteredIndexVersion1(key, value []byte, colsLen int, hdSt return nil, err } if segs.RestoredValues != nil { - resultValues, err = decodeRestoredValuesV5(columns[:colsLen], resultValues, segs.RestoredValues) + resultValues, err = decodeRestoredValuesV5(useNewCollate, columns[:colsLen], resultValues, segs.RestoredValues) if err != nil { return nil, err } @@ -1888,7 +1897,7 @@ func decodeIndexKvForClusteredIndexVersion1(key, value []byte, colsLen int, hdSt if err != nil { return nil, err } - handleBytes, err := reEncodeHandleConsiderNewCollation(handle, columns[colsLen:], segs.RestoredValues) + handleBytes, err := reEncodeHandleConsiderNewCollation(useNewCollate, handle, columns[colsLen:], segs.RestoredValues) if err != nil { return nil, err } diff --git a/pkg/tablecodec/tablecodec_test.go b/pkg/tablecodec/tablecodec_test.go index 9c72c3054ceb9..fa75652979bed 100644 --- a/pkg/tablecodec/tablecodec_test.go +++ b/pkg/tablecodec/tablecodec_test.go @@ -35,6 +35,10 @@ import ( "github.com/tikv/client-go/v2/tikv" ) +func defaultCodecEncoder() codec.Encoder { + return codec.NewEncoder(collate.NewCollationEnabled()) +} + // TestTableCodec tests some functions in package tablecodec // TODO: add more tests. func TestTableCodec(t *testing.T) { @@ -101,7 +105,7 @@ func TestRowCodec(t *testing.T) { } rd := rowcodec.Encoder{Enable: true} sc := stmtctx.NewStmtCtxWithTimeZone(time.Local) - bs, err := EncodeRow(sc.TimeZone(), row, colIDs, nil, nil, nil, &rd) + bs, err := EncodeRow(defaultCodecEncoder(), sc.TimeZone(), row, colIDs, nil, nil, nil, &rd) require.NoError(t, err) require.NotNil(t, bs) @@ -156,7 +160,7 @@ func TestRowCodec(t *testing.T) { } // Make sure empty row return not nil value. - bs, err = EncodeOldRow(sc.TimeZone(), []types.Datum{}, []int64{}, nil, nil) + bs, err = EncodeOldRow(defaultCodecEncoder(), sc.TimeZone(), []types.Datum{}, []int64{}, nil, nil) require.NoError(t, err) require.Len(t, bs, 1) @@ -170,7 +174,7 @@ func TestDecodeColumnValue(t *testing.T) { // test timestamp d := types.NewTimeDatum(types.NewTime(types.FromGoTime(time.Now()), mysql.TypeTimestamp, types.DefaultFsp)) - bs, err := EncodeOldRow(sc.TimeZone(), []types.Datum{d}, []int64{1}, nil, nil) + bs, err := EncodeOldRow(defaultCodecEncoder(), sc.TimeZone(), []types.Datum{d}, []int64{1}, nil, nil) require.NoError(t, err) require.NotNil(t, bs) _, bs, err = codec.CutOne(bs) // ignore colID @@ -186,7 +190,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(sc.TimeZone(), []types.Datum{d}, []int64{1}, nil, nil) + bs, err = EncodeOldRow(defaultCodecEncoder(), sc.TimeZone(), []types.Datum{d}, []int64{1}, nil, nil) require.NoError(t, err) require.NotNil(t, bs) _, bs, err = codec.CutOne(bs) // ignore colID @@ -201,7 +205,7 @@ func TestDecodeColumnValue(t *testing.T) { // test bit d = types.NewMysqlBitDatum(types.NewBinaryLiteralFromUint(3223600, 3)) - bs, err = EncodeOldRow(sc.TimeZone(), []types.Datum{d}, []int64{1}, nil, nil) + bs, err = EncodeOldRow(defaultCodecEncoder(), sc.TimeZone(), []types.Datum{d}, []int64{1}, nil, nil) require.NoError(t, err) require.NotNil(t, bs) _, bs, err = codec.CutOne(bs) // ignore colID @@ -216,7 +220,7 @@ func TestDecodeColumnValue(t *testing.T) { // test empty enum d = types.NewMysqlEnumDatum(types.Enum{}) - bs, err = EncodeOldRow(sc.TimeZone(), []types.Datum{d}, []int64{1}, nil, nil) + bs, err = EncodeOldRow(defaultCodecEncoder(), sc.TimeZone(), []types.Datum{d}, []int64{1}, nil, nil) require.NoError(t, err) require.NotNil(t, bs) _, bs, err = codec.CutOne(bs) // ignore colID @@ -276,7 +280,7 @@ func TestTimeCodec(t *testing.T) { } rd := rowcodec.Encoder{Enable: true} sc := stmtctx.NewStmtCtxWithTimeZone(time.UTC) - bs, err := EncodeRow(sc.TimeZone(), row, colIDs, nil, nil, nil, &rd) + bs, err := EncodeRow(defaultCodecEncoder(), sc.TimeZone(), row, colIDs, nil, nil, nil, &rd) require.NoError(t, err) require.NotNil(t, bs) @@ -324,7 +328,7 @@ func TestCutRow(t *testing.T) { for _, col := range cols { colIDs = append(colIDs, col.id) } - bs, err := EncodeOldRow(sc.TimeZone(), row, colIDs, nil, nil) + bs, err := EncodeOldRow(defaultCodecEncoder(), sc.TimeZone(), row, colIDs, nil, nil) require.NoError(t, err) require.NotNil(t, bs) @@ -744,3 +748,233 @@ func TestV2TableCodec(t *testing.T) { tbid = DecodeTableID(key) require.Equal(t, int64(0), tbid) } +<<<<<<< HEAD +======= + +// TestDecodeIndexHandleWithPartitionIDInKeyAndValue tests the scenario where +// a GlobalIndexVersionV1+ non-unique index has partition ID in both the key +// (new format) and the value (legacy global index format). This can produce +// a nested PartitionHandle if not handled correctly. +// See: https://github.com/pingcap/tidb/pull/65380#discussion_r2721786298 +func TestDecodeIndexHandleWithPartitionIDInKeyAndValue(t *testing.T) { + tableID := int64(100) + indexID := int64(1) + partitionID := int64(42) + handleID := int64(999) + colsLen := 1 + + // Build index key with one column value + sc := stmtctx.NewStmtCtxWithTimeZone(time.UTC) + indexedValues := []types.Datum{types.NewIntDatum(123)} + encodedCols, err := codec.EncodeKey(sc.TimeZone(), nil, indexedValues...) + require.NoError(t, err) + + // Build the key: table prefix + table ID + index ID + encoded columns + partition handle suffix + // For GlobalIndexVersionV1+ non-unique indexes, the key suffix is: + // PartitionIDFlag + partition_id (8 bytes) + IntHandleFlag + handle (8 bytes) + key := make([]byte, 0) + key = append(key, tablePrefix...) + key = codec.EncodeInt(key, tableID) + key = append(key, indexPrefixSep...) + key = codec.EncodeInt(key, indexID) + key = append(key, encodedCols...) + // Add partition handle suffix (GlobalIndexVersionV1+ format) + key = append(key, PartitionIDFlag) + key = codec.EncodeInt(key, partitionID) + key = append(key, codec.IntHandleFlag) + key = codec.EncodeInt(key, handleID) + + // Build index value with partition ID (global index value format) + // Format: TailLen | PartitionIDFlag | PartitionID | Padding + // We need len(value) >= 9 to trigger the partition ID check in DecodeIndexHandle + value := make([]byte, 0) + value = append(value, 0) // TailLen placeholder + value = append(value, PartitionIDFlag) + value = codec.EncodeInt(value, partitionID) + // Pad to make the value long enough (minimum 10 bytes for new encoding) + for len(value) < 10 { + value = append(value, 0) + } + value[0] = byte(len(value) - 1 - 1 - 8) // TailLen = total - 1(TailLen) - 1(PartitionIDFlag) - 8(PartitionID) + + // Decode the handle + handle, err := DecodeIndexHandle(key, value, colsLen) + require.NoError(t, err) + + // The handle should be a PartitionHandle + ph, ok := handle.(kv.PartitionHandle) + require.True(t, ok, "expected PartitionHandle, got %T", handle) + + // The correct behavior should be: + // - PartitionID should equal the expected partition ID (42) + // - Inner handle should be IntHandle(999), not another PartitionHandle + require.Equal(t, partitionID, ph.PartitionID, "partition ID mismatch") + + // Check that we do NOT have a nested PartitionHandle + // If the inner handle is also a PartitionHandle, that's the bug described in + // https://github.com/pingcap/tidb/pull/65380#discussion_r2721786298 + _, isNested := ph.Handle.(kv.PartitionHandle) + require.False(t, isNested, "DecodeIndexHandle should not create nested PartitionHandle; "+ + "when handle from key is already a PartitionHandle, skip value-based wrapping") + + // Verify the inner handle is the expected IntHandle + require.Equal(t, kv.IntHandle(handleID), ph.Handle, "inner handle should be IntHandle") +} + +// TestUniqueGlobalIndexKeyWithNullValues tests that for unique global indexes on +// non-clustered tables: +// - Non-NULL values do NOT have partition ID in the key (distinct = true) +// - NULL values DO have partition ID in the key (distinct = false) +// - Partition ID is always in the value for global indexes +// This is critical after EXCHANGE PARTITION where duplicate _tidb_rowid values can exist. +func TestUniqueGlobalIndexKeyWithNullValues(t *testing.T) { + tableID := int64(100) + partitionID := int64(42) + handleID := int64(999) + + // Build a simple TableInfo and IndexInfo for a unique global index + // on a non-clustered table (no clustered index) + tblInfo := &model.TableInfo{ + ID: tableID, + Name: ast.NewCIStr("test_table"), + Columns: []*model.ColumnInfo{ + { + ID: 1, + Name: ast.NewCIStr("a"), + Offset: 0, + FieldType: *types.NewFieldType(mysql.TypeLong), + }, + { + ID: 2, + Name: ast.NewCIStr("b"), + Offset: 1, + FieldType: *types.NewFieldType(mysql.TypeLong), + }, + }, + // Non-clustered table (PKIsHandle = false, IsCommonHandle = false) + PKIsHandle: false, + IsCommonHandle: false, + } + + idxInfo := &model.IndexInfo{ + ID: 1, + Name: ast.NewCIStr("idx_b"), + Columns: []*model.IndexColumn{ + { + Name: ast.NewCIStr("b"), + Offset: 1, + Length: types.UnspecifiedLength, + }, + }, + Unique: true, + Global: true, + GlobalIndexVersion: model.GlobalIndexVersionV1, + State: model.StatePublic, + } + + loc := time.UTC + + // For unique index with non-NULL values, distinct = true, + // so the handle is NOT encoded in the key at all. + indexedValues := []types.Datum{types.NewIntDatum(123)} + handle := kv.NewPartitionHandle(partitionID, kv.IntHandle(handleID)) + + key, distinct, err := GenIndexKey(defaultCodecEncoder(), loc, tblInfo, idxInfo, tableID, indexedValues, handle, nil) + require.NoError(t, err) + require.True(t, distinct, "unique index with non-NULL value should be distinct") + + // The key should NOT contain the partition ID flag since distinct = true + // means no handle (and thus no partition ID) is encoded in the key + require.NotContains(t, key, []byte{PartitionIDFlag}, + "unique index key with non-NULL value should NOT contain partition ID") + + // Verify key structure: tablePrefix + tableID + indexPrefixSep + indexID + encodedValues + // No handle suffix expected + require.True(t, len(key) > 0, "key should not be empty") + + // For unique index with NULL values, distinct = false, + // so the handle IS encoded in the key, including partition ID for V1+. + indexedValues = []types.Datum{types.NewDatum(nil)} // NULL value + handle = kv.NewPartitionHandle(partitionID, kv.IntHandle(handleID)) + + key, distinct, err = GenIndexKey(defaultCodecEncoder(), loc, tblInfo, idxInfo, tableID, indexedValues, handle, nil) + require.NoError(t, err) + require.False(t, distinct, "unique index with NULL value should NOT be distinct") + + // The key SHOULD contain the partition ID since distinct = false + // and GlobalIndexVersion >= V1 + containsPartitionIDFlag := false + for i := 0; i < len(key)-1; i++ { + if key[i] == PartitionIDFlag { + containsPartitionIDFlag = true + // Verify the partition ID is correctly encoded after the flag + if i+9 <= len(key) { + decodedPartID := codec.DecodeCmpUintToInt(binary.BigEndian.Uint64(key[i+1 : i+9])) + require.Equal(t, partitionID, decodedPartID, + "partition ID in key should match expected value") + } + break + } + } + require.True(t, containsPartitionIDFlag, + "unique index key with NULL value should contain partition ID flag") + + // For both distinct and non-distinct global indexes, partition ID + // should be encoded in the value. + indexedValues = []types.Datum{types.NewIntDatum(123)} + intHandle := kv.IntHandle(handleID) + + // Generate the index value + value, err := genIndexValueVersion0(loc, tblInfo, idxInfo, false, true, false, + indexedValues, intHandle, partitionID, nil) + require.NoError(t, err) + + // The value should contain the partition ID + containsPartitionIDFlag = false + for i := 0; i < len(value)-1; i++ { + if value[i] == PartitionIDFlag { + containsPartitionIDFlag = true + // Verify the partition ID is correctly encoded after the flag + if i+9 <= len(value) { + decodedPartID := codec.DecodeCmpUintToInt(binary.BigEndian.Uint64(value[i+1 : i+9])) + require.Equal(t, partitionID, decodedPartID, + "partition ID in value should match expected value") + } + break + } + } + require.True(t, containsPartitionIDFlag, + "global index value should contain partition ID flag") + + // Test that legacy (version 0) unique indexes do NOT have partition ID in key + // even with NULL values - this verifies backward compatibility + idxInfoV0 := &model.IndexInfo{ + ID: 1, + Name: ast.NewCIStr("idx_b_v0"), + Columns: []*model.IndexColumn{ + { + Name: ast.NewCIStr("b"), + Offset: 1, + Length: types.UnspecifiedLength, + }, + }, + Unique: true, + Global: true, + GlobalIndexVersion: 0, // Legacy version + State: model.StatePublic, + } + + indexedValues = []types.Datum{types.NewDatum(nil)} // NULL value + intHandle = kv.IntHandle(handleID) + + key, distinct, err = GenIndexKey(defaultCodecEncoder(), loc, tblInfo, idxInfoV0, tableID, indexedValues, intHandle, nil) + require.NoError(t, err) + require.False(t, distinct, "unique index with NULL value should NOT be distinct") + + // The key should NOT contain partition ID flag for version 0 + for i := 0; i < len(key)-1; i++ { + require.NotEqual(t, PartitionIDFlag, key[i], + "legacy (v0) global index key should NOT contain partition ID flag") + } +} +>>>>>>> ab1e19714d6 (codec, table: make new collation setting explicit in encoding (#69566)) diff --git a/pkg/testkit/mockstore.go b/pkg/testkit/mockstore.go index 8d632c6546235..38e838c4740f1 100644 --- a/pkg/testkit/mockstore.go +++ b/pkg/testkit/mockstore.go @@ -29,6 +29,7 @@ import ( "testing" "time" + "github.com/pingcap/failpoint" "github.com/pingcap/kvproto/pkg/keyspacepb" "github.com/pingcap/tidb/pkg/config" "github.com/pingcap/tidb/pkg/config/kerneltype" @@ -387,7 +388,18 @@ func bootstrap4DistExecution(t testing.TB, store kv.Storage, lease time.Duration session.DisableStats4Test() domain.DisablePlanReplayerBackgroundJob4Test() domain.DisableDumpHistoricalStats4Test() - dom, err := session.BootstrapSession4DistExecution(store) + + // Multi-domain tests intentionally keep multiple live domains for one mock + // store. Skip the temporary global-init domain so it cannot reuse and close + // an active test domain. + const skipInitGlobalVarFromSystemDB = "github.com/pingcap/tidb/pkg/session/skipInitGlobalVarFromSystemDB" + dom, err := func() (*domain.Domain, error) { + require.NoError(t, failpoint.Enable(skipInitGlobalVarFromSystemDB, `return(true)`)) + defer func() { + require.NoError(t, failpoint.Disable(skipInitGlobalVarFromSystemDB)) + }() + return session.BootstrapSession4DistExecution(store) + }() require.NoError(t, err) dom.SetStatsUpdating(true) diff --git a/pkg/types/etc.go b/pkg/types/etc.go index 8fcfb094516fd..aa3ad56010759 100644 --- a/pkg/types/etc.go +++ b/pkg/types/etc.go @@ -138,8 +138,14 @@ func IsNonBinaryStr(ft *FieldType) bool { // NeedRestoredData returns if a type needs restored data. // If the type is char and the collation is _bin, NeedRestoredData() returns false. func NeedRestoredData(ft *FieldType) bool { - if collate.NewCollationEnabled() && - IsNonBinaryStr(ft) && + return NeedRestoredDataWithCollate(ft, collate.NewCollationEnabled()) +} + +// NeedRestoredDataWithCollate reports restored-data needs under a caller-owned +// collation mode, so encode/decode paths can use the same setting captured by +// their table or index. +func NeedRestoredDataWithCollate(ft *FieldType, useNewCollate bool) bool { + if useNewCollate && IsNonBinaryStr(ft) && (!collate.IsBinCollation(ft.GetCollate()) || IsTypeVarchar(ft.GetType())) && ft.GetCollate() != "utf8mb4_0900_bin" { return true diff --git a/pkg/util/codec/codec.go b/pkg/util/codec/codec.go index 22e30d04023b7..9b39fa5a572f4 100644 --- a/pkg/util/codec/codec.go +++ b/pkg/util/codec/codec.go @@ -61,6 +61,21 @@ const ( sizeFloat64 = unsafe.Sizeof(float64(0)) ) +// Encoder encodes Datum values with a fixed new collation setting. +type Encoder struct { + useNewCollate bool +} + +// NewEncoder creates an Encoder with the given new collation setting. +func NewEncoder(useNewCollate bool) Encoder { + return Encoder{useNewCollate: useNewCollate} +} + +// UseNewCollate returns whether the encoder is using new collation. +func (enc Encoder) UseNewCollate() bool { + return enc.useNewCollate +} + func preRealloc(b []byte, vals []types.Datum, comparable1 bool) []byte { var size int for i := range vals { @@ -88,7 +103,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 encode(loc *time.Location, b []byte, vals []types.Datum, comparable1 bool) (_ []byte, err error) { +func (enc Encoder) encode(loc *time.Location, b []byte, vals []types.Datum, comparable1 bool) (_ []byte, err error) { b = preRealloc(b, vals, comparable1) for i, length := 0, len(vals); i < length; i++ { switch vals[i].Kind() { @@ -100,7 +115,7 @@ func encode(loc *time.Location, b []byte, vals []types.Datum, comparable1 bool) b = append(b, floatFlag) b = EncodeFloat(b, vals[i].GetFloat64()) case types.KindString: - b = encodeString(b, vals[i], comparable1) + b = enc.encodeString(b, vals[i], comparable1) case types.KindBytes: b = encodeBytes(b, vals[i].GetBytes(), comparable1) case types.KindMysqlTime: @@ -212,9 +227,9 @@ func EncodeMySQLTime(loc *time.Location, t types.Time, tp byte, b []byte) (_ []b return b, nil } -func encodeString(b []byte, val types.Datum, comparable1 bool) []byte { - if collate.NewCollationEnabled() && comparable1 { - return encodeBytes(b, collate.GetCollator(val.Collation()).ImmutableKey(val.GetString()), true) +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) } return encodeBytes(b, val.GetBytes(), comparable1) } @@ -301,13 +316,27 @@ func sizeInt(comparable1 bool) int { // slice. It guarantees the encoded value is in ascending order for comparison. // For decimal type, datum must set datum's length and frac. func EncodeKey(loc *time.Location, b []byte, v ...types.Datum) ([]byte, error) { - return encode(loc, b, v, true) + return NewEncoder(collate.NewCollationEnabled()).EncodeKey(loc, b, v...) +} + +// EncodeKey appends the encoded values to byte slice b using the encoder's +// 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) } // 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 encode(loc, b, v, false) + 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) } func encodeHashChunkRowIdx(typeCtx types.Context, row chunk.Row, tp *types.FieldType, idx int) (flag byte, b []byte, err error) { @@ -1620,6 +1649,14 @@ 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) @@ -1629,7 +1666,7 @@ func HashCode(b []byte, d types.Datum) []byte { b = append(b, floatFlag) b = EncodeFloat(b, d.GetFloat64()) case types.KindString: - b = encodeString(b, d, false) + b = enc.encodeString(b, d, false) case types.KindBytes: b = encodeBytes(b, d.GetBytes(), false) case types.KindMysqlTime: diff --git a/pkg/util/codec/codec_test.go b/pkg/util/codec/codec_test.go index 88f5c6d677ecf..94f47cfd647cc 100644 --- a/pkg/util/codec/codec_test.go +++ b/pkg/util/codec/codec_test.go @@ -807,7 +807,7 @@ func TestJSON(t *testing.T) { } buf := make([]byte, 0, 4096) - buf, err := encode(nil, buf, originalDatums, false) + buf, err := EncodeValue(nil, buf, originalDatums...) require.NoError(t, err) decodedDatums, err := Decode(buf, 2) diff --git a/pkg/util/codec/collation_test.go b/pkg/util/codec/collation_test.go index 4ac902b8e5e05..8238dfb987a1f 100644 --- a/pkg/util/codec/collation_test.go +++ b/pkg/util/codec/collation_test.go @@ -24,6 +24,7 @@ import ( "github.com/pingcap/tidb/pkg/parser/mysql" "github.com/pingcap/tidb/pkg/types" "github.com/pingcap/tidb/pkg/util/chunk" + "github.com/pingcap/tidb/pkg/util/collate" "github.com/stretchr/testify/require" ) @@ -43,6 +44,42 @@ func prepareCollationData() (int, *chunk.Chunk, *chunk.Chunk) { return 3, chk1, chk2 } +func TestEncoderNewCollationEnabled(t *testing.T) { + origin := collate.NewCollationEnabled() + defer collate.SetNewCollationEnabledForTest(origin) + + lower := types.NewCollationStringDatum("aaa", "utf8_general_ci") + upper := types.NewCollationStringDatum("AAA", "utf8_general_ci") + enabledEncoder := Encoder{useNewCollate: true} + disabledEncoder := Encoder{useNewCollate: false} + + collate.SetNewCollationEnabledForTest(true) + enabledLower, err := enabledEncoder.EncodeKey(time.Local, nil, lower) + require.NoError(t, err) + enabledUpper, err := enabledEncoder.EncodeKey(time.Local, nil, upper) + require.NoError(t, err) + require.Equal(t, enabledLower, enabledUpper) + + disabledLower, err := disabledEncoder.EncodeKey(time.Local, nil, lower) + require.NoError(t, err) + disabledUpper, err := disabledEncoder.EncodeKey(time.Local, nil, upper) + require.NoError(t, err) + require.NotEqual(t, disabledLower, disabledUpper) + + exportedEnabledLower, err := EncodeKey(time.Local, nil, lower) + require.NoError(t, err) + require.Equal(t, enabledLower, exportedEnabledLower) + + collate.SetNewCollationEnabledForTest(false) + 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) { tp := types.NewFieldType(mysql.TypeString) n, chk1, chk2 := prepareCollationData() diff --git a/pkg/util/collate/collate.go b/pkg/util/collate/collate.go index 968878d1d67a1..22bbd07c42ee5 100644 --- a/pkg/util/collate/collate.go +++ b/pkg/util/collate/collate.go @@ -113,7 +113,7 @@ func CompatibleCollate(collate1, collate2 string) bool { // the protocol definition. // When new collations are not enabled, collation id remains the same. func RewriteNewCollationIDIfNeeded(id int32) int32 { - if atomic.LoadInt32(&newCollationEnabled) == 1 { + if NewCollationEnabled() { if id >= 0 { return -id } @@ -124,7 +124,7 @@ func RewriteNewCollationIDIfNeeded(id int32) int32 { // RestoreCollationIDIfNeeded restores a collation id if the new collations are enabled. func RestoreCollationIDIfNeeded(id int32) int32 { - if atomic.LoadInt32(&newCollationEnabled) == 1 { + if NewCollationEnabled() { if id <= 0 { return -id } @@ -133,9 +133,15 @@ func RestoreCollationIDIfNeeded(id int32) int32 { return id } -// GetCollator get the collator according to collate, it will return the binary collator if the corresponding collator doesn't exist. +// GetCollator get the collator according to collate, it will return the binary +// collator if the corresponding collator doesn't exist. func GetCollator(collate string) Collator { - if atomic.LoadInt32(&newCollationEnabled) == 1 { + return GetCollatorWithCollate(NewCollationEnabled(), collate) +} + +// GetCollatorWithCollate is similar with GetCollator but allow explicit useNewCollate. +func GetCollatorWithCollate(useNewCollate bool, collate string) Collator { + if useNewCollate { ctor, ok := newCollatorMap[collate] if !ok { if collate != "" { @@ -170,7 +176,7 @@ func GetBinaryCollatorSlice(n int) []Collator { // GetCollatorByID get the collator according to id, it will return the binary collator if the corresponding collator doesn't exist. func GetCollatorByID(id int) Collator { - if atomic.LoadInt32(&newCollationEnabled) == 1 { + if NewCollationEnabled() { ctor, ok := newCollatorIDMap[id] if !ok { logutil.BgLogger().Warn( @@ -228,7 +234,7 @@ func GetCollationByName(name string) (coll *charset.Collation, err error) { if coll, err = charset.GetCollationByName(name); err != nil { return nil, errors.Trace(err) } - if atomic.LoadInt32(&newCollationEnabled) == 1 { + if NewCollationEnabled() { if _, ok := newCollatorIDMap[coll.ID]; !ok { return nil, ErrUnsupportedCollation.GenWithStackByArgs(name) } @@ -238,7 +244,7 @@ func GetCollationByName(name string) (coll *charset.Collation, err error) { // GetSupportedCollations gets information for all collations supported so far. func GetSupportedCollations() []*charset.Collation { - if atomic.LoadInt32(&newCollationEnabled) == 1 { + if NewCollationEnabled() { newSupportedCollations := make([]*charset.Collation, 0, len(newCollatorMap)) for name := range newCollatorMap { // utf8mb4_zh_pinyin_tidb_as_cs is under developing, should not be shown to user. diff --git a/pkg/util/rowDecoder/BUILD.bazel b/pkg/util/rowDecoder/BUILD.bazel index 9a6aa73fc6e02..b163f546dff15 100644 --- a/pkg/util/rowDecoder/BUILD.bazel +++ b/pkg/util/rowDecoder/BUILD.bazel @@ -41,6 +41,7 @@ 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 00373b8dd3f98..0db975740fb86 100644 --- a/pkg/util/rowDecoder/decoder_test.go +++ b/pkg/util/rowDecoder/decoder_test.go @@ -29,6 +29,7 @@ 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" @@ -108,12 +109,13 @@ 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(sc.TimeZone(), row.input, row.cols, nil, nil, nil, &rd) + bs, err := tablecodec.EncodeRow(codecEncoder, sc.TimeZone(), row.input, row.cols, nil, nil, nil, &rd) require.NoError(t, err) require.NotNil(t, bs) @@ -187,8 +189,9 @@ func TestClusterIndexRowDecoder(t *testing.T) { }, } rd := rowcodec.Encoder{Enable: true} + codecEncoder := codec.NewEncoder(collate.NewCollationEnabled()) for _, row := range testRows { - bs, err := tablecodec.EncodeRow(sc.TimeZone(), row.input, row.cols, nil, nil, nil, &rd) + bs, err := tablecodec.EncodeRow(codecEncoder, 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 2e342134200ac..35ada49123337 100644 --- a/pkg/util/rowcodec/bench_test.go +++ b/pkg/util/rowcodec/bench_test.go @@ -25,6 +25,8 @@ 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" ) @@ -67,7 +69,7 @@ func BenchmarkEncode(b *testing.B) { func BenchmarkEncodeFromOldRow(b *testing.B) { b.ReportAllocs() oldRow := types.MakeDatums(1, "abc", 1.1) - oldRowData, err := tablecodec.EncodeOldRow(nil, oldRow, []int64{1, 2, 3}, nil, nil) + oldRowData, err := tablecodec.EncodeOldRow(codec.NewEncoder(collate.NewCollationEnabled()), 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 5409210876a47..deb6182ff882e 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(sc.TimeZone(), types.MakeDatums(1, 2, 3, nil), colIDs, nil, nil) + oldRow, err := tablecodec.EncodeOldRow(codec.NewEncoder(collate.NewCollationEnabled()), 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(sc.TimeZone(), types.MakeDatums(1, 2, 3, nil), colIDs, nil, nil) + oldRow, err := tablecodec.EncodeOldRow(codec.NewEncoder(collate.NewCollationEnabled()), sc.TimeZone(), types.MakeDatums(1, 2, 3, nil), colIDs, nil, nil) require.NoError(t, err) var ( diff --git a/tests/realtikvtest/addindextest2/BUILD.bazel b/tests/realtikvtest/addindextest2/BUILD.bazel index aa76db172bc1e..2b4c16bc34cfb 100644 --- a/tests/realtikvtest/addindextest2/BUILD.bazel +++ b/tests/realtikvtest/addindextest2/BUILD.bazel @@ -34,6 +34,8 @@ go_test( "//pkg/testkit", "//pkg/testkit/testfailpoint", "//pkg/types", + "//pkg/util/codec", + "//pkg/util/collate", "//tests/realtikvtest", "//tests/realtikvtest/testutils", "@com_github_fsouza_fake_gcs_server//fakestorage", diff --git a/tests/realtikvtest/addindextest2/global_sort_test.go b/tests/realtikvtest/addindextest2/global_sort_test.go index 5d8244fce09b2..9d87d0c268ac8 100644 --- a/tests/realtikvtest/addindextest2/global_sort_test.go +++ b/tests/realtikvtest/addindextest2/global_sort_test.go @@ -50,6 +50,8 @@ import ( "github.com/pingcap/tidb/pkg/testkit" "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/tests/realtikvtest" "github.com/pingcap/tidb/tests/realtikvtest/testutils" "github.com/stretchr/testify/require" @@ -572,7 +574,7 @@ func TestIngestUseGivenTS(t *testing.T) { dts := []types.Datum{types.NewIntDatum(1)} sctx := tk.Session().GetSessionVars().StmtCtx - idxKey, _, err := tablecodec.GenIndexKey(sctx.TimeZone(), tblInfo, idxInfo, tblInfo.ID, dts, kv.IntHandle(1), nil) + idxKey, _, err := tablecodec.GenIndexKey(codec.NewEncoder(collate.NewCollationEnabled()), sctx.TimeZone(), tblInfo, idxInfo, tblInfo.ID, dts, kv.IntHandle(1), nil) require.NoError(t, err) tikvStore := dom.Store().(helper.Storage) newHelper := helper.NewHelper(tikvStore) From ce26a1a0e7b550e816ba24117e1aa9ac999b40d3 Mon Sep 17 00:00:00 2001 From: D3Hunter Date: Mon, 6 Jul 2026 19:37:29 +0800 Subject: [PATCH 2/3] resolve cherry-pick conflicts for #69566 --- pkg/domain/domain.go | 4 +- pkg/infoschema/issyncer/BUILD.bazel | 1 + pkg/infoschema/issyncer/filter.go | 37 ++++ pkg/infoschema/issyncer/loader.go | 23 ++- pkg/infoschema/issyncer/loader_test.go | 6 +- pkg/infoschema/issyncer/syncer.go | 3 +- pkg/infoschema/issyncer/syncer_test.go | 2 +- pkg/meta/model/table.go | 13 -- pkg/planner/core/expression_rewriter.go | 6 - pkg/session/BUILD.bazel | 1 + pkg/session/bootstrap_test.go | 2 +- pkg/session/tidb.go | 14 +- pkg/store/copr/copr_test/BUILD.bazel | 2 +- pkg/table/tables/index.go | 96 +-------- pkg/table/tables/mutation_checker_test.go | 2 +- pkg/table/tables/tables.go | 9 +- pkg/tablecodec/tablecodec.go | 6 +- pkg/tablecodec/tablecodec_test.go | 230 ---------------------- 18 files changed, 93 insertions(+), 364 deletions(-) create mode 100644 pkg/infoschema/issyncer/filter.go diff --git a/pkg/domain/domain.go b/pkg/domain/domain.go index 86597c1701c25..cf026fd0b54b8 100644 --- a/pkg/domain/domain.go +++ b/pkg/domain/domain.go @@ -542,7 +542,7 @@ const resourceIdleTimeout = 3 * time.Minute // resources in the ResourcePool wil // NewDomain creates a new domain. Should not create multiple domains for the same store. func NewDomain(store kv.Storage, schemaLease time.Duration, statsLease time.Duration, dumpFileGcLease time.Duration, factory pools.Factory) *Domain { - return NewDomainWithEtcdClient(store, schemaLease, statsLease, dumpFileGcLease, factory, nil, nil) + return NewDomainWithEtcdClient(store, schemaLease, statsLease, dumpFileGcLease, factory, nil, nil, nil) } // NewDomainWithEtcdClient creates a new domain with etcd client. Should not create multiple domains for the same store. @@ -554,6 +554,7 @@ func NewDomainWithEtcdClient( factory pools.Factory, crossKSSessFactoryGetter func(targetKS string, validator validatorapi.Validator) pools.Factory, etcdClient *clientv3.Client, + schemaFilter issyncer.Filter, ) *Domain { intest.Assert(schemaLease > 0, "schema lease should be a positive duration") do := &Domain{ @@ -597,6 +598,7 @@ func NewDomainWithEtcdClient( do.schemaLease, do.sysSessionPool, isvalidator.New(do.schemaLease), + schemaFilter, ) do.initDomainSysVars() diff --git a/pkg/infoschema/issyncer/BUILD.bazel b/pkg/infoschema/issyncer/BUILD.bazel index c264e7e8b9389..7f1d59b781d83 100644 --- a/pkg/infoschema/issyncer/BUILD.bazel +++ b/pkg/infoschema/issyncer/BUILD.bazel @@ -4,6 +4,7 @@ go_library( name = "issyncer", srcs = [ "deferfn.go", + "filter.go", "loader.go", "mdl_check.go", "syncer.go", diff --git a/pkg/infoschema/issyncer/filter.go b/pkg/infoschema/issyncer/filter.go new file mode 100644 index 0000000000000..d116fc666a715 --- /dev/null +++ b/pkg/infoschema/issyncer/filter.go @@ -0,0 +1,37 @@ +// Copyright 2026 PingCAP, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package issyncer + +import ( + "github.com/pingcap/tidb/pkg/infoschema" + "github.com/pingcap/tidb/pkg/meta/model" +) + +// Filter allows caller to customize which schema objects should be loaded +// when syncing the information schema. +// +// The semantics of the return values are: +// - true: skip the corresponding operation +// - false: continue the default loading logic +type Filter interface { + // SkipLoadDiff returns true when the given schema diff should be ignored. + // latestIS is the newest InfoSchema cached by the loader. It can be nil when + // the loader hasn't loaded any schema yet. + SkipLoadDiff(diff *model.SchemaDiff, latestIS infoschema.InfoSchema) bool + + // SkipLoadSchema returns true when the given DB should not be loaded during + // a full schema load. + SkipLoadSchema(dbInfo *model.DBInfo) bool +} diff --git a/pkg/infoschema/issyncer/loader.go b/pkg/infoschema/issyncer/loader.go index c852bf5006c56..85ce2a78a2cc0 100644 --- a/pkg/infoschema/issyncer/loader.go +++ b/pkg/infoschema/issyncer/loader.go @@ -85,6 +85,7 @@ type Loader struct { // loading system tables crossKS bool logger *zap.Logger + filter Filter // below fields are set when running background routines // Note: for cross keyspace loader, we don't set below fields as system tables @@ -98,13 +99,14 @@ type Loader struct { sysExecutorFactory func() (pools.Resource, error) } -func newLoader(store kv.Storage, infoCache *infoschema.InfoCache, deferFn *deferFn) *Loader { +func newLoader(store kv.Storage, infoCache *infoschema.InfoCache, deferFn *deferFn, filter Filter) *Loader { mode := LoadModeAuto return &Loader{ mode: mode, store: store, infoCache: infoCache, deferFn: deferFn, + filter: filter, logger: logutil.BgLogger().With(zap.Stringer("mode", mode)), } } @@ -289,6 +291,16 @@ func (l *Loader) LoadWithTS(startTS uint64, isSnapshot bool) (infoschema.InfoSch } func (l *Loader) skipLoadingDiff(diff *model.SchemaDiff) bool { + if l.filter != nil { + var latestIS infoschema.InfoSchema + if l.infoCache != nil { + latestIS = l.infoCache.GetLatest() + } + if l.filter.SkipLoadDiff(diff, latestIS) { + return true + } + } + if !l.crossKS { return false } @@ -417,6 +429,15 @@ func (l *Loader) fetchAllSchemasWithTables(m meta.Reader, schemaCacheSize uint64 if err != nil { return nil, err } + if l.filter != nil { + filteredSchemas := allSchemas[:0] + for _, dbInfo := range allSchemas { + if !l.filter.SkipLoadSchema(dbInfo) { + filteredSchemas = append(filteredSchemas, dbInfo) + } + } + allSchemas = filteredSchemas + } } if len(allSchemas) == 0 { return nil, nil diff --git a/pkg/infoschema/issyncer/loader_test.go b/pkg/infoschema/issyncer/loader_test.go index 716bc6d240122..8e06e3abfb226 100644 --- a/pkg/infoschema/issyncer/loader_test.go +++ b/pkg/infoschema/issyncer/loader_test.go @@ -33,7 +33,7 @@ import ( func TestLoadFromTS(t *testing.T) { store, err := mockstore.NewMockStore() require.NoError(t, err) - l := newLoader(store, infoschema.NewCache(nil, 1), nil) + l := newLoader(store, infoschema.NewCache(nil, 1), nil, nil) ver, err := store.CurrentVersion(tidbkv.GlobalTxnScope) require.NoError(t, err) is, hitCache, oldSchemaVersion, changes, err := l.LoadWithTS(ver.Ver, false) @@ -70,7 +70,7 @@ func TestLoadFromTS(t *testing.T) { })) ver, err = store.CurrentVersion(tidbkv.GlobalTxnScope) require.NoError(t, err) - l = newLoader(store, infoschema.NewCache(nil, 1), nil) + l = newLoader(store, infoschema.NewCache(nil, 1), nil, nil) is, hitCache, oldSchemaVersion, changes, err = l.LoadWithTS(ver.Ver, false) require.NoError(t, err) allSchemas = is.AllSchemas() @@ -217,7 +217,7 @@ func (testStoreWithKS) GetKeyspace() string { } func TestLoaderSkipLoadingDiff(t *testing.T) { - syncer := New(nil, nil, 0, nil, nil) + syncer := New(nil, nil, 0, nil, nil, nil) require.False(t, syncer.loader.skipLoadingDiff(&model.SchemaDiff{})) require.False(t, syncer.loader.skipLoadingDiff(&model.SchemaDiff{TableID: 100})) require.False(t, syncer.loader.skipLoadingDiff(&model.SchemaDiff{OldTableID: 100})) diff --git a/pkg/infoschema/issyncer/syncer.go b/pkg/infoschema/issyncer/syncer.go index 6084e09548462..364c0617734cd 100644 --- a/pkg/infoschema/issyncer/syncer.go +++ b/pkg/infoschema/issyncer/syncer.go @@ -91,9 +91,10 @@ func New( schemaLease time.Duration, sysSessionPool util.DestroyableSessionPool, isValidator validatorapi.Validator, + filter Filter, ) *Syncer { s := newSyncer(store, logutil.BgLogger(), schemaLease, sysSessionPool, isValidator) - s.loader = newLoader(store, infoCache, &s.deferFn) + s.loader = newLoader(store, infoCache, &s.deferFn, filter) return s } diff --git a/pkg/infoschema/issyncer/syncer_test.go b/pkg/infoschema/issyncer/syncer_test.go index 0fe00791c3a83..526e82fdad4cc 100644 --- a/pkg/infoschema/issyncer/syncer_test.go +++ b/pkg/infoschema/issyncer/syncer_test.go @@ -22,7 +22,7 @@ import ( ) func TestSyncerSkipMDLCheck(t *testing.T) { - syncer := New(nil, nil, 0, nil, nil) + syncer := New(nil, nil, 0, nil, nil, nil) require.False(t, syncer.skipMDLCheck(map[int64]struct{}{})) require.False(t, syncer.skipMDLCheck(map[int64]struct{}{123: {}})) require.False(t, syncer.skipMDLCheck(map[int64]struct{}{123: {}, 456: {}})) diff --git a/pkg/meta/model/table.go b/pkg/meta/model/table.go index 9109bf5ddc69b..d5386ce8c9dc2 100644 --- a/pkg/meta/model/table.go +++ b/pkg/meta/model/table.go @@ -581,19 +581,6 @@ func FindFKInfoByName(fks []*FKInfo, name string) *FKInfo { return nil } -<<<<<<< HEAD -======= -// GetIdxChangingFieldType gets the field type of index column. -// Since both old/new type may coexist in one column during modify column, -// we need to get the correct type for index column. -func GetIdxChangingFieldType(idxCol *IndexColumn, col *ColumnInfo) *types.FieldType { - if idxCol.UseChangingType && col.ChangingFieldType != nil { - return col.ChangingFieldType - } - return &col.FieldType -} - ->>>>>>> ab1e19714d6 (codec, table: make new collation setting explicit in encoding (#69566)) // TableNameInfo provides meta data describing a table name info. type TableNameInfo struct { ID int64 `json:"id"` diff --git a/pkg/planner/core/expression_rewriter.go b/pkg/planner/core/expression_rewriter.go index 534070d843e10..0eb563da1d486 100644 --- a/pkg/planner/core/expression_rewriter.go +++ b/pkg/planner/core/expression_rewriter.go @@ -377,14 +377,8 @@ type expressionRewriter struct { disableFoldCounter int tryFoldCounter int -<<<<<<< HEAD - planCtx *exprRewriterPlanCtx -======= - astNodeStack []ast.Node - planCtx *exprRewriterPlanCtx useNewCollate bool ->>>>>>> ab1e19714d6 (codec, table: make new collation setting explicit in encoding (#69566)) } func (er *expressionRewriter) ctxStackLen() int { diff --git a/pkg/session/BUILD.bazel b/pkg/session/BUILD.bazel index 555b30bd3ae54..cc128bba195fe 100644 --- a/pkg/session/BUILD.bazel +++ b/pkg/session/BUILD.bazel @@ -47,6 +47,7 @@ go_library( "//pkg/extension/extensionimpl", "//pkg/infoschema", "//pkg/infoschema/context", + "//pkg/infoschema/issyncer", "//pkg/infoschema/validatorapi", "//pkg/kv", "//pkg/meta", diff --git a/pkg/session/bootstrap_test.go b/pkg/session/bootstrap_test.go index d802647f4aecd..76ef361b08a2a 100644 --- a/pkg/session/bootstrap_test.go +++ b/pkg/session/bootstrap_test.go @@ -2626,7 +2626,7 @@ func makeStore(t *testing.T, keyspaceMeta *keyspacepb.KeyspaceMeta, isHasPrefix t.Cleanup(func() { ddl.CloseOwnerManager(mockStore) }) - dom, err := domap.getWithEtcdClient(mockStore, etcdClient) + dom, err := domap.getWithEtcdClient(mockStore, etcdClient, nil) require.NoError(t, err) defer dom.Close() diff --git a/pkg/session/tidb.go b/pkg/session/tidb.go index e856260b335e1..2aee51a3d7339 100644 --- a/pkg/session/tidb.go +++ b/pkg/session/tidb.go @@ -32,6 +32,7 @@ import ( "github.com/pingcap/tidb/pkg/errno" "github.com/pingcap/tidb/pkg/executor" "github.com/pingcap/tidb/pkg/infoschema" + "github.com/pingcap/tidb/pkg/infoschema/issyncer" "github.com/pingcap/tidb/pkg/infoschema/validatorapi" "github.com/pingcap/tidb/pkg/kv" "github.com/pingcap/tidb/pkg/parser" @@ -65,10 +66,18 @@ type domainMap struct { // TODO decouple domain create from it, it's more clear to create domain explicitly // before any usage of it. func (dm *domainMap) Get(store kv.Storage) (d *domain.Domain, err error) { - return dm.getWithEtcdClient(store, nil) + return dm.getWithEtcdClient(store, nil, nil) } -func (dm *domainMap) getWithEtcdClient(store kv.Storage, etcdClient *clientv3.Client) (d *domain.Domain, err error) { +// GetOrCreateWithFilter gets or creates the domain for store with a schema filter. +// +// Caveat: If there is already a domain opened with your `store`, the filter passed in will be ignored and +// the actual schema filter of the returned `Domain` is the one when the domain was created. +func (dm *domainMap) GetOrCreateWithFilter(store kv.Storage, filter issyncer.Filter) (d *domain.Domain, err error) { + return dm.getWithEtcdClient(store, nil, filter) +} + +func (dm *domainMap) getWithEtcdClient(store kv.Storage, etcdClient *clientv3.Client, schemaFilter issyncer.Filter) (d *domain.Domain, err error) { dm.mu.Lock() defer dm.mu.Unlock() @@ -102,6 +111,7 @@ func (dm *domainMap) getWithEtcdClient(store kv.Storage, etcdClient *clientv3.Cl return getCrossKSSessionFactory(store, targetKS, schemaValidator) }, etcdClient, + schemaFilter, ) var ddlInjector func(ddl.DDL, ddl.Executor, *infoschema.InfoCache) *schematracker.Checker diff --git a/pkg/store/copr/copr_test/BUILD.bazel b/pkg/store/copr/copr_test/BUILD.bazel index 3e92f01ca1dfd..452a2981877ee 100644 --- a/pkg/store/copr/copr_test/BUILD.bazel +++ b/pkg/store/copr/copr_test/BUILD.bazel @@ -8,7 +8,7 @@ go_test( "main_test.go", ], flaky = True, - shard_count = 5, + shard_count = 6, deps = [ "//pkg/config", "//pkg/config/kerneltype", diff --git a/pkg/table/tables/index.go b/pkg/table/tables/index.go index ab8d911189616..f4ef8cca195fb 100644 --- a/pkg/table/tables/index.go +++ b/pkg/table/tables/index.go @@ -29,32 +29,13 @@ import ( "github.com/pingcap/tidb/pkg/tablecodec" "github.com/pingcap/tidb/pkg/types" "github.com/pingcap/tidb/pkg/util" -<<<<<<< HEAD -======= - "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" ->>>>>>> ab1e19714d6 (codec, table: make new collation setting explicit in encoding (#69566)) "github.com/pingcap/tidb/pkg/util/intest" "github.com/pingcap/tidb/pkg/util/rowcodec" "github.com/pingcap/tidb/pkg/util/tracing" ) -<<<<<<< HEAD -======= -var indexConditionECtx exprctx.BuildContext - -// indexPartialCondition is a data structure to help implement the partial index. -type indexPartialCondition struct { - conditionExpr expression.Expression - // conditionEvalBufferPool stores many eval buffer to avoid allocating chunk - // for evaluating partial index condition for each time. - // It's only initialized if the `partialConditionExpr` is not nil. - conditionEvalBufferPool sync.Pool -} - ->>>>>>> ab1e19714d6 (codec, table: make new collation setting explicit in encoding (#69566)) // index is the data structure for index data in the KV store. type index struct { idxInfo *model.IndexInfo @@ -65,22 +46,14 @@ type index struct { // the collation global variable is initialized *after* `NewIndex()`. initNeedRestoreData sync.Once needRestoredData bool -<<<<<<< HEAD -======= encoder codec.Encoder - indexPartialCondition ->>>>>>> ab1e19714d6 (codec, table: make new collation setting explicit in encoding (#69566)) } // NeedRestoredData checks whether the index columns needs restored data. func NeedRestoredData(useNewCollate bool, idxCols []*model.IndexColumn, colInfos []*model.ColumnInfo) bool { for _, idxCol := range idxCols { col := colInfos[idxCol.Offset] -<<<<<<< HEAD - if types.NeedRestoredData(&col.FieldType) { -======= - if types.NeedRestoredDataWithCollate(model.GetIdxChangingFieldType(idxCol, col), useNewCollate) { ->>>>>>> ab1e19714d6 (codec, table: make new collation setting explicit in encoding (#69566)) + if types.NeedRestoredDataWithCollate(&col.FieldType, useNewCollate) { return true } } @@ -88,10 +61,7 @@ func NeedRestoredData(useNewCollate bool, idxCols []*model.IndexColumn, colInfos } // NewIndex builds a new Index object. -<<<<<<< HEAD func NewIndex(physicalID int64, tblInfo *model.TableInfo, indexInfo *model.IndexInfo) table.Index { -======= -func NewIndex(physicalID int64, tblInfo *model.TableInfo, indexInfo *model.IndexInfo) (table.Index, error) { return NewIndexWithCollate(collate.NewCollationEnabled(), physicalID, tblInfo, indexInfo) } @@ -101,51 +71,14 @@ func NewIndexWithCollate( physicalID int64, tblInfo *model.TableInfo, indexInfo *model.IndexInfo, -) (table.Index, error) { ->>>>>>> ab1e19714d6 (codec, table: make new collation setting explicit in encoding (#69566)) +) table.Index { index := &index{ idxInfo: indexInfo, tblInfo: tblInfo, phyTblID: physicalID, encoder: codec.NewEncoder(useNewCollate), } -<<<<<<< HEAD return index -======= - - 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) - } - - 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)) - } - - evalBuffer := chunk.MutRowFromTypes(evalBufferTypes) - return &evalBuffer - }, - } - } - return index, nil ->>>>>>> ab1e19714d6 (codec, table: make new collation setting explicit in encoding (#69566)) } // Meta returns index info. @@ -170,17 +103,8 @@ func (c *index) GenIndexKey(ec errctx.Context, loc *time.Location, indexedValues idxTblID = c.tblInfo.ID } } -<<<<<<< HEAD - key, distinct, err = tablecodec.GenIndexKey(loc, c.tblInfo, c.idxInfo, idxTblID, indexedValues, h, buf) -======= - - if err = c.castIndexValuesToChangingTypes(indexedValues); err != nil { - return - } - key, distinct, err = tablecodec.GenIndexKey(c.encoder, loc, c.tblInfo, c.idxInfo, - idxTblID, indexedValues, fullHandle, buf) ->>>>>>> ab1e19714d6 (codec, table: make new collation setting explicit in encoding (#69566)) + idxTblID, indexedValues, h, buf) err = ec.HandleError(err) return } @@ -191,17 +115,9 @@ func (c *index) GenIndexValue(ec errctx.Context, loc *time.Location, distinct bo c.initNeedRestoreData.Do(func() { c.needRestoredData = NeedRestoredData(c.encoder.UseNewCollate(), c.idxInfo.Columns, c.tblInfo.Columns) }) -<<<<<<< HEAD - idx, err := tablecodec.GenIndexValuePortal(loc, c.tblInfo, c.idxInfo, c.needRestoredData, distinct, false, indexedValues, h, c.phyTblID, restoredData, buf) -======= - - if err := c.castIndexValuesToChangingTypes(indexedValues); err != nil { - return nil, errors.Trace(err) - } idx, err := tablecodec.GenIndexValuePortal(c.encoder.UseNewCollate(), loc, c.tblInfo, - c.idxInfo, c.needRestoredData, distinct, untouched, indexedValues, h, c.phyTblID, restoredData, buf) ->>>>>>> ab1e19714d6 (codec, table: make new collation setting explicit in encoding (#69566)) + c.idxInfo, c.needRestoredData, distinct, false, indexedValues, h, c.phyTblID, restoredData, buf) err = ec.HandleError(err) return idx, err } @@ -346,9 +262,9 @@ func (c *index) create(sctx table.MutateContext, txn kv.Transaction, indexedValu // save the key buffer to reuse. writeBufs.IndexKeyBuf = key c.initNeedRestoreData.Do(func() { - c.needRestoredData = NeedRestoredData(c.idxInfo.Columns, c.tblInfo.Columns) + c.needRestoredData = NeedRestoredData(c.encoder.UseNewCollate(), c.idxInfo.Columns, c.tblInfo.Columns) }) - idxVal, err := tablecodec.GenIndexValuePortal(loc, c.tblInfo, c.idxInfo, + idxVal, err := tablecodec.GenIndexValuePortal(c.encoder.UseNewCollate(), loc, c.tblInfo, c.idxInfo, c.needRestoredData, distinct, untouched, value, h, c.phyTblID, handleRestoreData, nil) err = ec.HandleError(err) if err != nil { diff --git a/pkg/table/tables/mutation_checker_test.go b/pkg/table/tables/mutation_checker_test.go index e3b0e3076301f..361d8753da85a 100644 --- a/pkg/table/tables/mutation_checker_test.go +++ b/pkg/table/tables/mutation_checker_test.go @@ -394,7 +394,7 @@ func buildIndexKeyValue(index table.Index, rowToInsert []types.Datum, tc types.C } useNewCollate := table.encoder.UseNewCollate() rsData := TryGetHandleRestoredDataWrapper(useNewCollate, table.meta, rowToInsert, nil, indexInfo) - value, err := index.GenIndexValue(errctx.StrictNoWarningContext, tc.Location(), distinct, false, indexedValues, handle, rsData, nil) + value, err := index.GenIndexValue(errctx.StrictNoWarningContext, tc.Location(), distinct, indexedValues, handle, rsData, nil) if err != nil { return nil, nil, err } diff --git a/pkg/table/tables/tables.go b/pkg/table/tables/tables.go index d650e7293e864..e58eb87406419 100644 --- a/pkg/table/tables/tables.go +++ b/pkg/table/tables/tables.go @@ -267,14 +267,7 @@ func initTableIndices(t *TableCommon) error { } // Use partition ID for index, because TableCommon may be table or partition. -<<<<<<< HEAD - idx := NewIndex(t.physicalTableID, tblInfo, idxInfo) -======= - idx, err := NewIndexWithCollate(t.encoder.UseNewCollate(), t.physicalTableID, tblInfo, idxInfo) - if err != nil { - return err - } ->>>>>>> ab1e19714d6 (codec, table: make new collation setting explicit in encoding (#69566)) + idx := NewIndexWithCollate(t.encoder.UseNewCollate(), t.physicalTableID, tblInfo, idxInfo) 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`. diff --git a/pkg/tablecodec/tablecodec.go b/pkg/tablecodec/tablecodec.go index 5b5fa2ea6a6a4..9ed827c94490e 100644 --- a/pkg/tablecodec/tablecodec.go +++ b/pkg/tablecodec/tablecodec.go @@ -1618,11 +1618,7 @@ func GenIndexValueForClusteredIndexVersion1(useNewCollate bool, loc *time.Locati if mysql.HasPriKeyFlag(col.GetFlag()) { continue } -<<<<<<< HEAD - if types.NeedRestoredData(&col.FieldType) { -======= - if types.NeedRestoredDataWithCollate(model.GetIdxChangingFieldType(idxCol, col), useNewCollate) { ->>>>>>> ab1e19714d6 (codec, table: make new collation setting explicit in encoding (#69566)) + if types.NeedRestoredDataWithCollate(&col.FieldType, useNewCollate) { colIds = append(colIds, col.ID) if collate.IsBinCollation(col.GetCollate()) { allRestoredData = append(allRestoredData, types.NewUintDatum(uint64(stringutil.GetTailSpaceCount(indexedValues[i].GetString())))) diff --git a/pkg/tablecodec/tablecodec_test.go b/pkg/tablecodec/tablecodec_test.go index fa75652979bed..611c48c0e052a 100644 --- a/pkg/tablecodec/tablecodec_test.go +++ b/pkg/tablecodec/tablecodec_test.go @@ -748,233 +748,3 @@ func TestV2TableCodec(t *testing.T) { tbid = DecodeTableID(key) require.Equal(t, int64(0), tbid) } -<<<<<<< HEAD -======= - -// TestDecodeIndexHandleWithPartitionIDInKeyAndValue tests the scenario where -// a GlobalIndexVersionV1+ non-unique index has partition ID in both the key -// (new format) and the value (legacy global index format). This can produce -// a nested PartitionHandle if not handled correctly. -// See: https://github.com/pingcap/tidb/pull/65380#discussion_r2721786298 -func TestDecodeIndexHandleWithPartitionIDInKeyAndValue(t *testing.T) { - tableID := int64(100) - indexID := int64(1) - partitionID := int64(42) - handleID := int64(999) - colsLen := 1 - - // Build index key with one column value - sc := stmtctx.NewStmtCtxWithTimeZone(time.UTC) - indexedValues := []types.Datum{types.NewIntDatum(123)} - encodedCols, err := codec.EncodeKey(sc.TimeZone(), nil, indexedValues...) - require.NoError(t, err) - - // Build the key: table prefix + table ID + index ID + encoded columns + partition handle suffix - // For GlobalIndexVersionV1+ non-unique indexes, the key suffix is: - // PartitionIDFlag + partition_id (8 bytes) + IntHandleFlag + handle (8 bytes) - key := make([]byte, 0) - key = append(key, tablePrefix...) - key = codec.EncodeInt(key, tableID) - key = append(key, indexPrefixSep...) - key = codec.EncodeInt(key, indexID) - key = append(key, encodedCols...) - // Add partition handle suffix (GlobalIndexVersionV1+ format) - key = append(key, PartitionIDFlag) - key = codec.EncodeInt(key, partitionID) - key = append(key, codec.IntHandleFlag) - key = codec.EncodeInt(key, handleID) - - // Build index value with partition ID (global index value format) - // Format: TailLen | PartitionIDFlag | PartitionID | Padding - // We need len(value) >= 9 to trigger the partition ID check in DecodeIndexHandle - value := make([]byte, 0) - value = append(value, 0) // TailLen placeholder - value = append(value, PartitionIDFlag) - value = codec.EncodeInt(value, partitionID) - // Pad to make the value long enough (minimum 10 bytes for new encoding) - for len(value) < 10 { - value = append(value, 0) - } - value[0] = byte(len(value) - 1 - 1 - 8) // TailLen = total - 1(TailLen) - 1(PartitionIDFlag) - 8(PartitionID) - - // Decode the handle - handle, err := DecodeIndexHandle(key, value, colsLen) - require.NoError(t, err) - - // The handle should be a PartitionHandle - ph, ok := handle.(kv.PartitionHandle) - require.True(t, ok, "expected PartitionHandle, got %T", handle) - - // The correct behavior should be: - // - PartitionID should equal the expected partition ID (42) - // - Inner handle should be IntHandle(999), not another PartitionHandle - require.Equal(t, partitionID, ph.PartitionID, "partition ID mismatch") - - // Check that we do NOT have a nested PartitionHandle - // If the inner handle is also a PartitionHandle, that's the bug described in - // https://github.com/pingcap/tidb/pull/65380#discussion_r2721786298 - _, isNested := ph.Handle.(kv.PartitionHandle) - require.False(t, isNested, "DecodeIndexHandle should not create nested PartitionHandle; "+ - "when handle from key is already a PartitionHandle, skip value-based wrapping") - - // Verify the inner handle is the expected IntHandle - require.Equal(t, kv.IntHandle(handleID), ph.Handle, "inner handle should be IntHandle") -} - -// TestUniqueGlobalIndexKeyWithNullValues tests that for unique global indexes on -// non-clustered tables: -// - Non-NULL values do NOT have partition ID in the key (distinct = true) -// - NULL values DO have partition ID in the key (distinct = false) -// - Partition ID is always in the value for global indexes -// This is critical after EXCHANGE PARTITION where duplicate _tidb_rowid values can exist. -func TestUniqueGlobalIndexKeyWithNullValues(t *testing.T) { - tableID := int64(100) - partitionID := int64(42) - handleID := int64(999) - - // Build a simple TableInfo and IndexInfo for a unique global index - // on a non-clustered table (no clustered index) - tblInfo := &model.TableInfo{ - ID: tableID, - Name: ast.NewCIStr("test_table"), - Columns: []*model.ColumnInfo{ - { - ID: 1, - Name: ast.NewCIStr("a"), - Offset: 0, - FieldType: *types.NewFieldType(mysql.TypeLong), - }, - { - ID: 2, - Name: ast.NewCIStr("b"), - Offset: 1, - FieldType: *types.NewFieldType(mysql.TypeLong), - }, - }, - // Non-clustered table (PKIsHandle = false, IsCommonHandle = false) - PKIsHandle: false, - IsCommonHandle: false, - } - - idxInfo := &model.IndexInfo{ - ID: 1, - Name: ast.NewCIStr("idx_b"), - Columns: []*model.IndexColumn{ - { - Name: ast.NewCIStr("b"), - Offset: 1, - Length: types.UnspecifiedLength, - }, - }, - Unique: true, - Global: true, - GlobalIndexVersion: model.GlobalIndexVersionV1, - State: model.StatePublic, - } - - loc := time.UTC - - // For unique index with non-NULL values, distinct = true, - // so the handle is NOT encoded in the key at all. - indexedValues := []types.Datum{types.NewIntDatum(123)} - handle := kv.NewPartitionHandle(partitionID, kv.IntHandle(handleID)) - - key, distinct, err := GenIndexKey(defaultCodecEncoder(), loc, tblInfo, idxInfo, tableID, indexedValues, handle, nil) - require.NoError(t, err) - require.True(t, distinct, "unique index with non-NULL value should be distinct") - - // The key should NOT contain the partition ID flag since distinct = true - // means no handle (and thus no partition ID) is encoded in the key - require.NotContains(t, key, []byte{PartitionIDFlag}, - "unique index key with non-NULL value should NOT contain partition ID") - - // Verify key structure: tablePrefix + tableID + indexPrefixSep + indexID + encodedValues - // No handle suffix expected - require.True(t, len(key) > 0, "key should not be empty") - - // For unique index with NULL values, distinct = false, - // so the handle IS encoded in the key, including partition ID for V1+. - indexedValues = []types.Datum{types.NewDatum(nil)} // NULL value - handle = kv.NewPartitionHandle(partitionID, kv.IntHandle(handleID)) - - key, distinct, err = GenIndexKey(defaultCodecEncoder(), loc, tblInfo, idxInfo, tableID, indexedValues, handle, nil) - require.NoError(t, err) - require.False(t, distinct, "unique index with NULL value should NOT be distinct") - - // The key SHOULD contain the partition ID since distinct = false - // and GlobalIndexVersion >= V1 - containsPartitionIDFlag := false - for i := 0; i < len(key)-1; i++ { - if key[i] == PartitionIDFlag { - containsPartitionIDFlag = true - // Verify the partition ID is correctly encoded after the flag - if i+9 <= len(key) { - decodedPartID := codec.DecodeCmpUintToInt(binary.BigEndian.Uint64(key[i+1 : i+9])) - require.Equal(t, partitionID, decodedPartID, - "partition ID in key should match expected value") - } - break - } - } - require.True(t, containsPartitionIDFlag, - "unique index key with NULL value should contain partition ID flag") - - // For both distinct and non-distinct global indexes, partition ID - // should be encoded in the value. - indexedValues = []types.Datum{types.NewIntDatum(123)} - intHandle := kv.IntHandle(handleID) - - // Generate the index value - value, err := genIndexValueVersion0(loc, tblInfo, idxInfo, false, true, false, - indexedValues, intHandle, partitionID, nil) - require.NoError(t, err) - - // The value should contain the partition ID - containsPartitionIDFlag = false - for i := 0; i < len(value)-1; i++ { - if value[i] == PartitionIDFlag { - containsPartitionIDFlag = true - // Verify the partition ID is correctly encoded after the flag - if i+9 <= len(value) { - decodedPartID := codec.DecodeCmpUintToInt(binary.BigEndian.Uint64(value[i+1 : i+9])) - require.Equal(t, partitionID, decodedPartID, - "partition ID in value should match expected value") - } - break - } - } - require.True(t, containsPartitionIDFlag, - "global index value should contain partition ID flag") - - // Test that legacy (version 0) unique indexes do NOT have partition ID in key - // even with NULL values - this verifies backward compatibility - idxInfoV0 := &model.IndexInfo{ - ID: 1, - Name: ast.NewCIStr("idx_b_v0"), - Columns: []*model.IndexColumn{ - { - Name: ast.NewCIStr("b"), - Offset: 1, - Length: types.UnspecifiedLength, - }, - }, - Unique: true, - Global: true, - GlobalIndexVersion: 0, // Legacy version - State: model.StatePublic, - } - - indexedValues = []types.Datum{types.NewDatum(nil)} // NULL value - intHandle = kv.IntHandle(handleID) - - key, distinct, err = GenIndexKey(defaultCodecEncoder(), loc, tblInfo, idxInfoV0, tableID, indexedValues, intHandle, nil) - require.NoError(t, err) - require.False(t, distinct, "unique index with NULL value should NOT be distinct") - - // The key should NOT contain partition ID flag for version 0 - for i := 0; i < len(key)-1; i++ { - require.NotEqual(t, PartitionIDFlag, key[i], - "legacy (v0) global index key should NOT contain partition ID flag") - } -} ->>>>>>> ab1e19714d6 (codec, table: make new collation setting explicit in encoding (#69566)) From 35fe581cab7f34a8e3c85b3535b898a8b95fe046 Mon Sep 17 00:00:00 2001 From: D3Hunter Date: Mon, 6 Jul 2026 19:47:51 +0800 Subject: [PATCH 3/3] change --- pkg/infoschema/issyncer/loader.go | 20 ++++--- pkg/meta/meta.go | 13 +++++ pkg/meta/reader.go | 1 + pkg/table/tables/testutil/BUILD.bazel | 19 ------- pkg/table/tables/testutil/indexcheck.go | 75 ------------------------- 5 files changed, 25 insertions(+), 103 deletions(-) delete mode 100644 pkg/table/tables/testutil/BUILD.bazel delete mode 100644 pkg/table/tables/testutil/indexcheck.go diff --git a/pkg/infoschema/issyncer/loader.go b/pkg/infoschema/issyncer/loader.go index 85ce2a78a2cc0..6c4f1573bb9b7 100644 --- a/pkg/infoschema/issyncer/loader.go +++ b/pkg/infoschema/issyncer/loader.go @@ -424,20 +424,22 @@ func (l *Loader) fetchAllSchemasWithTables(m meta.Reader, schemaCacheSize uint64 return nil, errors.New("system database not found") } allSchemas = []*model.DBInfo{dbInfo} + } else if l.filter != nil { + allSchemas = make([]*model.DBInfo, 0, 6) + err := m.IterDatabases(func(dbInfo *model.DBInfo) error { + if !l.filter.SkipLoadSchema(dbInfo) { + allSchemas = append(allSchemas, dbInfo) + } + return nil + }) + if err != nil { + return nil, err + } } else { allSchemas, err = m.ListDatabases() if err != nil { return nil, err } - if l.filter != nil { - filteredSchemas := allSchemas[:0] - for _, dbInfo := range allSchemas { - if !l.filter.SkipLoadSchema(dbInfo) { - filteredSchemas = append(filteredSchemas, dbInfo) - } - } - allSchemas = filteredSchemas - } } if len(allSchemas) == 0 { return nil, nil diff --git a/pkg/meta/meta.go b/pkg/meta/meta.go index 189ef58dd4c43..4646e1d376ac8 100644 --- a/pkg/meta/meta.go +++ b/pkg/meta/meta.go @@ -1045,6 +1045,19 @@ func (m *Mutator) UpdateTable(dbID int64, tableInfo *model.TableInfo) error { return errors.Trace(err) } +// IterDatabases iterates all the Databases at once, stop iterate when fn returns an error. +func (m *Mutator) IterDatabases(fn func(info *model.DBInfo) error) error { + err := m.txn.HGetIter(mDBs, func(r structure.HashPair) error { + dbInfo := &model.DBInfo{} + err := json.Unmarshal(r.Value, dbInfo) + if err != nil { + return errors.Trace(err) + } + return fn(dbInfo) + }) + return errors.Trace(err) +} + // IterTables iterates all the table at once, in order to avoid oom. func (m *Mutator) IterTables(dbID int64, fn func(info *model.TableInfo) error) error { dbKey := m.dbKey(dbID) diff --git a/pkg/meta/reader.go b/pkg/meta/reader.go index 5441713045848..5017266a4743b 100644 --- a/pkg/meta/reader.go +++ b/pkg/meta/reader.go @@ -29,6 +29,7 @@ type Reader interface { GetTable(dbID int64, tableID int64) (*model.TableInfo, error) ListTables(ctx context.Context, dbID int64) ([]*model.TableInfo, error) ListSimpleTables(dbID int64) ([]*model.TableNameInfo, error) + IterDatabases(func(info *model.DBInfo) error) error IterTables(dbID int64, fn func(info *model.TableInfo) error) error GetAutoIDAccessors(dbID, tableID int64) AutoIDAccessors GetAllNameToIDAndTheMustLoadedTableInfo(dbID int64) (map[string]int64, []*model.TableInfo, error) diff --git a/pkg/table/tables/testutil/BUILD.bazel b/pkg/table/tables/testutil/BUILD.bazel deleted file mode 100644 index 6fb8235d08bb4..0000000000000 --- a/pkg/table/tables/testutil/BUILD.bazel +++ /dev/null @@ -1,19 +0,0 @@ -load("@io_bazel_rules_go//go:def.bzl", "go_library") - -go_library( - name = "testutil", - srcs = ["indexcheck.go"], - importpath = "github.com/pingcap/tidb/pkg/table/tables/testutil", - visibility = ["//visibility:public"], - deps = [ - "//pkg/domain", - "//pkg/parser/ast", - "//pkg/sessiontxn", - "//pkg/table", - "//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 deleted file mode 100644 index bffb6cfff5793..0000000000000 --- a/pkg/table/tables/testutil/indexcheck.go +++ /dev/null @@ -1,75 +0,0 @@ -// Copyright 2025 PingCAP, Inc. -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -package testutil - -import ( - "context" - "testing" - "time" - - "github.com/pingcap/tidb/pkg/domain" - "github.com/pingcap/tidb/pkg/parser/ast" - "github.com/pingcap/tidb/pkg/sessiontxn" - "github.com/pingcap/tidb/pkg/table" - "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" -) - -// CheckIndexKVCount checks the number of index key-value pairs in the specified index of the specified table. -func CheckIndexKVCount(t *testing.T, tk *testkit.TestKit, dom *domain.Domain, tableName string, indexName string, expected int) { - tbl, err := dom.InfoSchema().TableByName(context.Background(), ast.NewCIStr("test"), ast.NewCIStr(tableName)) - require.NoError(t, err) - require.NotNil(t, tbl) - var idx table.Index - for _, i := range tbl.Indices() { - if i.Meta().Name.L == indexName { - idx = i - break - } - } - - enc := codec.NewEncoder(collate.NewCollationEnabled()) - minimumKey, _, err := tablecodec.GenIndexKey(enc, time.Local, tbl.Meta(), idx.Meta(), tbl.Meta().ID, nil, nil, nil) - require.NoError(t, err) - - tk.MustExec("BEGIN") - defer tk.MustExec("COMMIT") - txnManager := sessiontxn.GetTxnManager(tk.Session()) - snapshot, err := txnManager.GetSnapshotWithStmtReadTS() - require.NoError(t, err) - - iter, err := snapshot.Iter(minimumKey, nil) - require.NoError(t, err) - defer iter.Close() - count := 0 - for iter.Valid() { - key := iter.Key() - if !tablecodec.IsIndexKey(key) { - break - } - _, idxID, _, err := tablecodec.DecodeIndexKey(key) - require.NoError(t, err) - if idxID != idx.Meta().ID { - break - } - count++ - err = iter.Next() - require.NoError(t, err) - } - require.Equal(t, expected, count) -}