This is an automated cherry-pick of pingcap#53268

Signed-off-by: ti-chi-bot <[email protected]>
ti-chi-bot · May 14, 2024 · e2e902c · e2e902c
1 parent da6dc8b
commit e2e902c
Show file tree

Hide file tree

Showing 12 changed files with 3,178 additions and 8 deletions.
diff --git a/build/BUILD.bazel b/build/BUILD.bazel
@@ -150,5 +150,13 @@ nogo(
                    "//build/linter/revive",
                ],
                "//conditions:default": [],
+<<<<<<< HEAD
+=======
+           }) +
+           select({
+               "//build:without_rbe": [
+               ],
+               "//conditions:default": [],
+>>>>>>> acdb6f58e3b (planner: UPDATE's select plan's output col IDs should be stable (#53268))
            }),
 )
diff --git a/expression/BUILD.bazel b/expression/BUILD.bazel
@@ -118,8 +118,11 @@ go_library(
         "@com_github_pingcap_tipb//go-tipb",
         "@com_github_pkg_errors//:errors",
         "@com_github_tikv_client_go_v2//oracle",
+<<<<<<< HEAD:expression/BUILD.bazel
         "@org_golang_x_exp//slices",
         "@org_golang_x_tools//container/intsets",
+=======
+>>>>>>> acdb6f58e3b (planner: UPDATE's select plan's output col IDs should be stable (#53268)):pkg/expression/BUILD.bazel
         "@org_uber_go_atomic//:atomic",
         "@org_uber_go_zap//:zap",
     ],

diff --git a/expression/util.go b/expression/util.go
@@ -25,6 +25,7 @@ import (
 
 	"github.com/pingcap/errors"
 	"github.com/pingcap/failpoint"
+<<<<<<< HEAD:expression/util.go
 	"github.com/pingcap/tidb/kv"
 	"github.com/pingcap/tidb/parser/ast"
 	"github.com/pingcap/tidb/parser/mysql"
@@ -37,8 +38,22 @@ import (
 	"github.com/pingcap/tidb/util/collate"
 	"github.com/pingcap/tidb/util/logutil"
 	"github.com/pingcap/tidb/util/sqlexec"
+=======
+	"github.com/pingcap/tidb/pkg/expression/contextopt"
+	"github.com/pingcap/tidb/pkg/kv"
+	"github.com/pingcap/tidb/pkg/parser/ast"
+	"github.com/pingcap/tidb/pkg/parser/mysql"
+	"github.com/pingcap/tidb/pkg/parser/opcode"
+	"github.com/pingcap/tidb/pkg/parser/terror"
+	"github.com/pingcap/tidb/pkg/sessionctx/variable"
+	"github.com/pingcap/tidb/pkg/types"
+	driver "github.com/pingcap/tidb/pkg/types/parser_driver"
+	"github.com/pingcap/tidb/pkg/util/chunk"
+	"github.com/pingcap/tidb/pkg/util/collate"
+	"github.com/pingcap/tidb/pkg/util/intset"
+	"github.com/pingcap/tidb/pkg/util/logutil"
+>>>>>>> acdb6f58e3b (planner: UPDATE's select plan's output col IDs should be stable (#53268)):pkg/expression/util.go
 	"go.uber.org/zap"
-	"golang.org/x/tools/container/intsets"
 )
 
 // cowExprRef is a copy-on-write slice ref util using in `ColumnSubstitute`
@@ -372,15 +387,15 @@ func ExtractColumnsAndCorColumnsFromExpressions(result []*Column, list []Express
 }
 
 // ExtractColumnSet extracts the different values of `UniqueId` for columns in expressions.
-func ExtractColumnSet(exprs ...Expression) *intsets.Sparse {
-	set := &intsets.Sparse{}
+func ExtractColumnSet(exprs ...Expression) intset.FastIntSet {
+	set := intset.NewFastIntSet()
 	for _, expr := range exprs {
-		extractColumnSet(expr, set)
+		extractColumnSet(expr, &set)
 	}
 	return set
 }
 
-func extractColumnSet(expr Expression, set *intsets.Sparse) {
+func extractColumnSet(expr Expression, set *intset.FastIntSet) {
 	switch v := expr.(type) {
 	case *Column:
 		set.Insert(int(v.UniqueID))

diff --git a/pkg/sessionctx/stmtctx/BUILD.bazel b/pkg/sessionctx/stmtctx/BUILD.bazel
@@ -0,0 +1,63 @@
+load("@io_bazel_rules_go//go:def.bzl", "go_library", "go_test")
+
+go_library(
+    name = "stmtctx",
+    srcs = ["stmtctx.go"],
+    importpath = "github.com/pingcap/tidb/pkg/sessionctx/stmtctx",
+    visibility = ["//visibility:public"],
+    deps = [
+        "//pkg/distsql/context",
+        "//pkg/domain/resourcegroup",
+        "//pkg/errctx",
+        "//pkg/parser",
+        "//pkg/parser/model",
+        "//pkg/parser/mysql",
+        "//pkg/parser/terror",
+        "//pkg/statistics/handle/usage/indexusage",
+        "//pkg/types",
+        "//pkg/util/context",
+        "//pkg/util/dbterror/plannererrors",
+        "//pkg/util/disk",
+        "//pkg/util/execdetails",
+        "//pkg/util/hint",
+        "//pkg/util/intest",
+        "//pkg/util/intset",
+        "//pkg/util/linter/constructor",
+        "//pkg/util/memory",
+        "//pkg/util/nocopy",
+        "//pkg/util/resourcegrouptag",
+        "//pkg/util/topsql/stmtstats",
+        "//pkg/util/tracing",
+        "@com_github_tikv_client_go_v2//tikvrpc",
+        "@org_golang_x_exp//maps",
+        "@org_golang_x_sync//singleflight",
+        "@org_uber_go_atomic//:atomic",
+    ],
+)
+
+go_test(
+    name = "stmtctx_test",
+    timeout = "short",
+    srcs = [
+        "main_test.go",
+        "stmtctx_test.go",
+    ],
+    embed = [":stmtctx"],
+    flaky = True,
+    shard_count = 12,
+    deps = [
+        "//pkg/errctx",
+        "//pkg/kv",
+        "//pkg/sessionctx/variable",
+        "//pkg/testkit",
+        "//pkg/testkit/testsetup",
+        "//pkg/types",
+        "//pkg/util/context",
+        "//pkg/util/execdetails",
+        "//pkg/util/hint",
+        "@com_github_pingcap_errors//:errors",
+        "@com_github_stretchr_testify//require",
+        "@com_github_tikv_client_go_v2//util",
+        "@org_uber_go_goleak//:goleak",
+    ],
+)
diff --git a/planner/core/logical_plan_builder.go b/planner/core/logical_plan_builder.go
@@ -6054,8 +6054,8 @@ func (b *PlanBuilder) buildUpdateLists(ctx context.Context, tableList []*ast.Tab
 			allAssignmentsAreConstant = false
 		}
 		p = np
-		if col, ok := newExpr.(*expression.Column); ok {
-			b.ctx.GetSessionVars().StmtCtx.ColRefFromUpdatePlan = append(b.ctx.GetSessionVars().StmtCtx.ColRefFromUpdatePlan, col.UniqueID)
+		if cols := expression.ExtractColumnSet(newExpr); cols.Len() > 0 {
+			b.ctx.GetSessionVars().StmtCtx.ColRefFromUpdatePlan.UnionWith(cols)
 		}
 		newList = append(newList, &expression.Assignment{Col: col, ColName: name.ColName, Expr: newExpr})
 		dbName := name.DBName.L

diff --git a/planner/core/rule_eliminate_projection.go b/planner/core/rule_eliminate_projection.go
@@ -72,13 +72,16 @@ func canProjectionBeEliminatedStrict(p *PhysicalProjection) bool {
 	if p.Schema().Len() != child.Schema().Len() {
 		return false
 	}
+<<<<<<< HEAD:planner/core/rule_eliminate_projection.go
 	for _, ref := range p.ctx.GetSessionVars().StmtCtx.ColRefFromUpdatePlan {
 		for _, one := range p.Schema().Columns {
 			if ref == one.UniqueID {
 				return false
 			}
 		}
 	}
+=======
+>>>>>>> acdb6f58e3b (planner: UPDATE's select plan's output col IDs should be stable (#53268)):pkg/planner/core/rule_eliminate_projection.go
 	for i, expr := range p.Exprs {
 		col, ok := expr.(*expression.Column)
 		if !ok || !col.Equal(nil, child.Schema().Columns[i]) {
@@ -132,6 +135,11 @@ func doPhysicalProjectionElimination(p PhysicalPlan) PhysicalPlan {
 	if childProj, ok := child.(*PhysicalProjection); ok {
 		childProj.SetSchema(p.Schema())
 	}
+	for i, col := range p.Schema().Columns {
+		if p.SCtx().GetSessionVars().StmtCtx.ColRefFromUpdatePlan.Has(int(col.UniqueID)) && !child.Schema().Columns[i].Equal(nil, col) {
+			return p
+		}
+	}
 	return child
 }
 

diff --git a/sessionctx/stmtctx/stmtctx.go b/sessionctx/stmtctx/stmtctx.go
@@ -26,6 +26,7 @@ import (
 	"sync/atomic"
 	"time"
 
+<<<<<<< HEAD:sessionctx/stmtctx/stmtctx.go
 	"github.com/pingcap/errors"
 	"github.com/pingcap/tidb/errno"
 	"github.com/pingcap/tidb/parser"
@@ -39,6 +40,30 @@ import (
 	"github.com/pingcap/tidb/util/resourcegrouptag"
 	"github.com/pingcap/tidb/util/topsql/stmtstats"
 	"github.com/pingcap/tidb/util/tracing"
+=======
+	distsqlctx "github.com/pingcap/tidb/pkg/distsql/context"
+	"github.com/pingcap/tidb/pkg/domain/resourcegroup"
+	"github.com/pingcap/tidb/pkg/errctx"
+	"github.com/pingcap/tidb/pkg/parser"
+	"github.com/pingcap/tidb/pkg/parser/model"
+	"github.com/pingcap/tidb/pkg/parser/mysql"
+	"github.com/pingcap/tidb/pkg/parser/terror"
+	"github.com/pingcap/tidb/pkg/statistics/handle/usage/indexusage"
+	"github.com/pingcap/tidb/pkg/types"
+	contextutil "github.com/pingcap/tidb/pkg/util/context"
+	"github.com/pingcap/tidb/pkg/util/dbterror/plannererrors"
+	"github.com/pingcap/tidb/pkg/util/disk"
+	"github.com/pingcap/tidb/pkg/util/execdetails"
+	"github.com/pingcap/tidb/pkg/util/hint"
+	"github.com/pingcap/tidb/pkg/util/intest"
+	"github.com/pingcap/tidb/pkg/util/intset"
+	"github.com/pingcap/tidb/pkg/util/linter/constructor"
+	"github.com/pingcap/tidb/pkg/util/memory"
+	"github.com/pingcap/tidb/pkg/util/nocopy"
+	"github.com/pingcap/tidb/pkg/util/resourcegrouptag"
+	"github.com/pingcap/tidb/pkg/util/topsql/stmtstats"
+	"github.com/pingcap/tidb/pkg/util/tracing"
+>>>>>>> acdb6f58e3b (planner: UPDATE's select plan's output col IDs should be stable (#53268)):pkg/sessionctx/stmtctx/stmtctx.go
 	"github.com/tikv/client-go/v2/tikvrpc"
 	"github.com/tikv/client-go/v2/util"
 	atomic2 "go.uber.org/atomic"
@@ -371,7 +396,7 @@ type StatementContext struct {
 	// UseDynamicPruneMode indicates whether use UseDynamicPruneMode in query stmt
 	UseDynamicPruneMode bool
 	// ColRefFromPlan mark the column ref used by assignment in update statement.
-	ColRefFromUpdatePlan []int64
+	ColRefFromUpdatePlan intset.FastIntSet
 
 	// RangeFallback indicates that building complete ranges exceeds the memory limit so it falls back to less accurate ranges such as full range.
 	RangeFallback bool