Skip to content

Commit

Permalink
planner: UPDATE's select plan's output col IDs should be stable (#53268…
Browse files Browse the repository at this point in the history 
…) (#53275)

close #53236
  • Loading branch information
@ti-chi-bot
ti-chi-bot authored Nov 11, 2024
1 parent 3929a07 commit 6e21798
Show file tree
Hide file tree
Showing 12 changed files with 1,149 additions and 32 deletions.
2 changes: 1 addition & 1 deletion expression/BUILD.bazel
View file
Original file line number Diff line number Diff line change
Expand Up @@ -99,6 +99,7 @@ go_library(
"//util/encrypt",
"//util/generatedexpr",
"//util/hack",
"//util/intset",
"//util/logutil",
"//util/mathutil",
"//util/mock",
Expand All @@ -120,7 +121,6 @@ go_library(
"@com_github_pkg_errors//:errors",
"@com_github_tikv_client_go_v2//oracle",
"@org_golang_x_exp//slices",
"@org_golang_x_tools//container/intsets",
"@org_uber_go_atomic//:atomic",
"@org_uber_go_zap//:zap",
],
Expand Down
10 changes: 5 additions & 5 deletions expression/util.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -36,10 +36,10 @@ import (
driver "github.com/pingcap/tidb/types/parser_driver"
"github.com/pingcap/tidb/util/chunk"
"github.com/pingcap/tidb/util/collate"
"github.com/pingcap/tidb/util/intset"
"github.com/pingcap/tidb/util/logutil"
"github.com/pingcap/tidb/util/sqlexec"
"go.uber.org/zap"
"golang.org/x/tools/container/intsets"
)

// cowExprRef is a copy-on-write slice ref util using in `ColumnSubstitute`
Expand Down Expand Up @@ -373,15 +373,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))
Expand Down
2 changes: 1 addition & 1 deletion planner/core/issuetest/BUILD.bazel
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,6 @@ go_test(
srcs = ["planner_issue_test.go"],
flaky = True,
race = "on",
shard_count = 15,
shard_count = 16,
deps = ["//testkit"],
)
56 changes: 41 additions & 15 deletions planner/core/issuetest/planner_issue_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -244,21 +244,17 @@ func TestIssue50614(t *testing.T) {
testkit.Rows(
"Update N/A root N/A",
"└─Projection 0.00 root test.tt.a, test.tt.b, test.tt.c, test.tt.d, test.tt.e, Column#18, Column#19, Column#20, Column#21",
" └─Projection 0.00 root test.tt.a, test.tt.b, test.tt.c, test.tt.d, test.tt.e, Column#18, Column#19, Column#20, Column#21",
" └─IndexJoin 0.00 root inner join, inner:TableReader, outer key:Column#20, Column#21, inner key:test.tt.c, test.tt.d, equal cond:eq(Column#20, test.tt.c), eq(Column#21, test.tt.d), other cond:or(or(and(eq(Column#20, 11), eq(test.tt.d, 111)), and(eq(Column#20, 22), eq(test.tt.d, 222))), or(and(eq(Column#20, 33), eq(test.tt.d, 333)), and(eq(Column#20, 44), eq(test.tt.d, 444)))), or(or(and(eq(test.tt.c, 11), eq(Column#21, 111)), and(eq(test.tt.c, 22), eq(Column#21, 222))), or(and(eq(test.tt.c, 33), eq(Column#21, 333)), and(eq(test.tt.c, 44), eq(Column#21, 444))))",
" ├─Union(Build) 0.00 root ",
" │ ├─Projection 0.00 root Column#6, Column#7, Column#8, Column#9",
" │ │ └─Projection 0.00 root 1->Column#6, 2->Column#7, 3->Column#8, 4->Column#9",
" │ │ └─TableDual 0.00 root rows:0",
" │ ├─Projection 0.00 root Column#10, Column#11, Column#12, Column#13",
" │ │ └─Projection 0.00 root 2->Column#10, 3->Column#11, 4->Column#12, 5->Column#13",
" │ │ └─TableDual 0.00 root rows:0",
" │ └─Projection 0.00 root Column#14, Column#15, Column#16, Column#17",
" │ └─Projection 0.00 root 3->Column#14, 4->Column#15, 5->Column#16, 6->Column#17",
" │ └─TableDual 0.00 root rows:0",
" └─TableReader(Probe) 0.00 root data:Selection",
" └─Selection 0.00 cop[tikv] or(or(and(eq(test.tt.c, 11), eq(test.tt.d, 111)), and(eq(test.tt.c, 22), eq(test.tt.d, 222))), or(and(eq(test.tt.c, 33), eq(test.tt.d, 333)), and(eq(test.tt.c, 44), eq(test.tt.d, 444)))), or(or(eq(test.tt.c, 11), eq(test.tt.c, 22)), or(eq(test.tt.c, 33), eq(test.tt.c, 44))), or(or(eq(test.tt.d, 111), eq(test.tt.d, 222)), or(eq(test.tt.d, 333), eq(test.tt.d, 444)))",
" └─TableRangeScan 0.00 cop[tikv] table:tt range: decided by [eq(test.tt.c, Column#20) eq(test.tt.d, Column#21)], keep order:false, stats:pseudo",
" └─IndexJoin 0.00 root inner join, inner:TableReader, outer key:Column#20, Column#21, inner key:test.tt.c, test.tt.d, equal cond:eq(Column#20, test.tt.c), eq(Column#21, test.tt.d), other cond:or(or(and(eq(Column#20, 11), eq(test.tt.d, 111)), and(eq(Column#20, 22), eq(test.tt.d, 222))), or(and(eq(Column#20, 33), eq(test.tt.d, 333)), and(eq(Column#20, 44), eq(test.tt.d, 444)))), or(or(and(eq(test.tt.c, 11), eq(Column#21, 111)), and(eq(test.tt.c, 22), eq(Column#21, 222))), or(and(eq(test.tt.c, 33), eq(Column#21, 333)), and(eq(test.tt.c, 44), eq(Column#21, 444))))",
" ├─Union(Build) 0.00 root ",
" │ ├─Projection 0.00 root 1->Column#18, 2->Column#19, 3->Column#20, 4->Column#21",
" │ │ └─TableDual 0.00 root rows:0",
" │ ├─Projection 0.00 root 2->Column#18, 3->Column#19, 4->Column#20, 5->Column#21",
" │ │ └─TableDual 0.00 root rows:0",
" │ └─Projection 0.00 root 3->Column#18, 4->Column#19, 5->Column#20, 6->Column#21",
" │ └─TableDual 0.00 root rows:0",
" └─TableReader(Probe) 0.00 root data:Selection",
" └─Selection 0.00 cop[tikv] or(or(and(eq(test.tt.c, 11), eq(test.tt.d, 111)), and(eq(test.tt.c, 22), eq(test.tt.d, 222))), or(and(eq(test.tt.c, 33), eq(test.tt.d, 333)), and(eq(test.tt.c, 44), eq(test.tt.d, 444)))), or(or(eq(test.tt.c, 11), eq(test.tt.c, 22)), or(eq(test.tt.c, 33), eq(test.tt.c, 44))), or(or(eq(test.tt.d, 111), eq(test.tt.d, 222)), or(eq(test.tt.d, 333), eq(test.tt.d, 444)))",
" └─TableRangeScan 0.00 cop[tikv] table:tt range: decided by [eq(test.tt.c, Column#20) eq(test.tt.d, Column#21)], keep order:false, stats:pseudo",
),
)
}
Expand Down Expand Up @@ -340,6 +336,36 @@ WHERE
tk.MustQuery(`show warnings`).Check(testkit.Rows())
}

// https://github.com/pingcap/tidb/issues/53236
func TestIssue53236(t *testing.T) {
store := testkit.CreateMockStore(t)
tk := testkit.NewTestKit(t, store)

tk.MustExec("use test;")
tk.MustExec("create table t1(id int primary key, a varchar(128));")
tk.MustExec("create table t2(id int primary key, b varchar(128), c varchar(128));")
tk.MustExec(`UPDATE
t1
SET
t1.a = IFNULL(
(
SELECT
t2.c
FROM
t2
WHERE
t2.b = t1.a
ORDER BY
t2.b DESC,
t2.c DESC
LIMIT
1
), ''
)
WHERE
t1.id = 1;`)
}

func TestIssue52687(t *testing.T) {
store := testkit.CreateMockStore(t)
tk := testkit.NewTestKit(t, store)
Expand Down
4 changes: 2 additions & 2 deletions planner/core/logical_plan_builder.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6085,8 +6085,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
Expand Down
12 changes: 5 additions & 7 deletions planner/core/rule_eliminate_projection.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -72,13 +72,6 @@ func canProjectionBeEliminatedStrict(p *PhysicalProjection) bool {
if p.Schema().Len() != child.Schema().Len() {
return false
}
for _, ref := range p.ctx.GetSessionVars().StmtCtx.ColRefFromUpdatePlan {
for _, one := range p.Schema().Columns {
if ref == one.UniqueID {
return false
}
}
}
for i, expr := range p.Exprs {
col, ok := expr.(*expression.Column)
if !ok || !col.Equal(nil, child.Schema().Columns[i]) {
Expand Down Expand Up @@ -136,6 +129,11 @@ func doPhysicalProjectionElimination(p PhysicalPlan) PhysicalPlan {
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
}

Expand Down
1 change: 1 addition & 0 deletions sessionctx/stmtctx/BUILD.bazel
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@ go_library(
"//parser/terror",
"//util/disk",
"//util/execdetails",
"//util/intset",
"//util/memory",
"//util/resourcegrouptag",
"//util/topsql/stmtstats",
Expand Down
3 changes: 2 additions & 1 deletion sessionctx/stmtctx/stmtctx.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,7 @@ import (
"github.com/pingcap/tidb/parser/terror"
"github.com/pingcap/tidb/util/disk"
"github.com/pingcap/tidb/util/execdetails"
"github.com/pingcap/tidb/util/intset"
"github.com/pingcap/tidb/util/memory"
"github.com/pingcap/tidb/util/resourcegrouptag"
"github.com/pingcap/tidb/util/topsql/stmtstats"
Expand Down Expand Up @@ -374,7 +375,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
Expand Down
26 changes: 26 additions & 0 deletions util/intset/BUILD.bazel
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,26 @@
load("@io_bazel_rules_go//go:def.bzl", "go_library", "go_test")

go_library(
name = "intset",
srcs = ["fast_int_set.go"],
importpath = "github.com/pingcap/tidb/util/intset",
visibility = ["//visibility:public"],
deps = ["@org_golang_x_tools//container/intsets"],
)

go_test(
name = "intset_test",
timeout = "short",
srcs = [
"fast_int_set_bench_test.go",
"fast_int_set_test.go",
],
embed = [":intset"],
flaky = True,
deps = [
"@com_github_stretchr_testify//require",
"@org_golang_x_exp//maps",
"@org_golang_x_exp//slices",
"@org_golang_x_tools//container/intsets",
],
)
Loading

0 comments on commit 6e21798

Please sign in to comment.