From 1cf8d63309bde017c0dadb739cb091f2644afa86 Mon Sep 17 00:00:00 2001 From: Ti Chi Robot Date: Fri, 8 Dec 2023 17:17:02 +0800 Subject: [PATCH] planner: fix possible inconsistent output cols among union's children (#48775) (#48895) close pingcap/tidb#48755 --- executor/index_merge_reader_test.go | 49 +++++++++++--------- planner/core/issuetest/BUILD.bazel | 2 +- planner/core/issuetest/planner_issue_test.go | 27 +++++++++++ planner/core/rule_column_pruning.go | 18 +++---- testkit/testkit.go | 11 +++++ 5 files changed, 76 insertions(+), 31 deletions(-) diff --git a/executor/index_merge_reader_test.go b/executor/index_merge_reader_test.go index 8842064e63fa2..10aa1cd12ebbd 100644 --- a/executor/index_merge_reader_test.go +++ b/executor/index_merge_reader_test.go @@ -568,22 +568,26 @@ func TestPessimisticLockOnPartitionForIndexMerge(t *testing.T) { " ├─IndexReader(Build) 3.00 root index:IndexFullScan", " │ └─IndexFullScan 3.00 cop[tikv] table:t2, index:c_datetime(c_datetime) keep order:false", " └─PartitionUnion(Probe) 5545.21 root ", - " ├─IndexMerge 5542.21 root type: union", - " │ ├─IndexRangeScan(Build) 3323.33 cop[tikv] table:t1, partition:p0, index:c1(c1) range:[-inf,10), keep order:false, stats:pseudo", - " │ ├─IndexRangeScan(Build) 3323.33 cop[tikv] table:t1, partition:p0, index:c2(c2) range:[-inf,10), keep order:false, stats:pseudo", - " │ └─TableRowIDScan(Probe) 5542.21 cop[tikv] table:t1, partition:p0 keep order:false, stats:pseudo", - " ├─IndexMerge 1.00 root type: union", - " │ ├─IndexRangeScan(Build) 1.00 cop[tikv] table:t1, partition:p1, index:c1(c1) range:[-inf,10), keep order:false", - " │ ├─IndexRangeScan(Build) 1.00 cop[tikv] table:t1, partition:p1, index:c2(c2) range:[-inf,10), keep order:false", - " │ └─TableRowIDScan(Probe) 1.00 cop[tikv] table:t1, partition:p1 keep order:false", - " ├─IndexMerge 1.00 root type: union", - " │ ├─IndexRangeScan(Build) 1.00 cop[tikv] table:t1, partition:p2, index:c1(c1) range:[-inf,10), keep order:false", - " │ ├─IndexRangeScan(Build) 1.00 cop[tikv] table:t1, partition:p2, index:c2(c2) range:[-inf,10), keep order:false", - " │ └─TableRowIDScan(Probe) 1.00 cop[tikv] table:t1, partition:p2 keep order:false", - " └─IndexMerge 1.00 root type: union", - " ├─IndexRangeScan(Build) 1.00 cop[tikv] table:t1, partition:p3, index:c1(c1) range:[-inf,10), keep order:false", - " ├─IndexRangeScan(Build) 1.00 cop[tikv] table:t1, partition:p3, index:c2(c2) range:[-inf,10), keep order:false", - " └─TableRowIDScan(Probe) 1.00 cop[tikv] table:t1, partition:p3 keep order:false", + " ├─Projection 5542.21 root test.t1.c_datetime, test.t1.c1, test.t1._tidb_rowid, test.t1._tidb_tid", + " │ └─IndexMerge 5542.21 root type: union", + " │ ├─IndexRangeScan(Build) 3323.33 cop[tikv] table:t1, partition:p0, index:c1(c1) range:[-inf,10), keep order:false, stats:pseudo", + " │ ├─IndexRangeScan(Build) 3323.33 cop[tikv] table:t1, partition:p0, index:c2(c2) range:[-inf,10), keep order:false, stats:pseudo", + " │ └─TableRowIDScan(Probe) 5542.21 cop[tikv] table:t1, partition:p0 keep order:false, stats:pseudo", + " ├─Projection 1.00 root test.t1.c_datetime, test.t1.c1, test.t1._tidb_rowid, test.t1._tidb_tid", + " │ └─IndexMerge 1.00 root type: union", + " │ ├─IndexRangeScan(Build) 1.00 cop[tikv] table:t1, partition:p1, index:c1(c1) range:[-inf,10), keep order:false", + " │ ├─IndexRangeScan(Build) 1.00 cop[tikv] table:t1, partition:p1, index:c2(c2) range:[-inf,10), keep order:false", + " │ └─TableRowIDScan(Probe) 1.00 cop[tikv] table:t1, partition:p1 keep order:false", + " ├─Projection 1.00 root test.t1.c_datetime, test.t1.c1, test.t1._tidb_rowid, test.t1._tidb_tid", + " │ └─IndexMerge 1.00 root type: union", + " │ ├─IndexRangeScan(Build) 1.00 cop[tikv] table:t1, partition:p2, index:c1(c1) range:[-inf,10), keep order:false", + " │ ├─IndexRangeScan(Build) 1.00 cop[tikv] table:t1, partition:p2, index:c2(c2) range:[-inf,10), keep order:false", + " │ └─TableRowIDScan(Probe) 1.00 cop[tikv] table:t1, partition:p2 keep order:false", + " └─Projection 1.00 root test.t1.c_datetime, test.t1.c1, test.t1._tidb_rowid, test.t1._tidb_tid", + " └─IndexMerge 1.00 root type: union", + " ├─IndexRangeScan(Build) 1.00 cop[tikv] table:t1, partition:p3, index:c1(c1) range:[-inf,10), keep order:false", + " ├─IndexRangeScan(Build) 1.00 cop[tikv] table:t1, partition:p3, index:c2(c2) range:[-inf,10), keep order:false", + " └─TableRowIDScan(Probe) 1.00 cop[tikv] table:t1, partition:p3 keep order:false", )) tk.MustQuery(`select /*+ use_index_merge(t1) */ c1 from t1 join t2 on t1.c_datetime >= t2.c_datetime @@ -710,20 +714,21 @@ func TestIntersectionWithDifferentConcurrency(t *testing.T) { for _, concurrency := range execCon { tk.MustExec(fmt.Sprintf("set tidb_executor_concurrency = %d", concurrency)) for i := 0; i < 2; i++ { + sql := "select /*+ use_index_merge(t1, primary, c2, c3) */ c1 from t1 where c2 < 1024 and c3 > 1024" if i == 0 { // Dynamic mode. tk.MustExec("set tidb_partition_prune_mode = 'dynamic'") - res := tk.MustQuery("explain select /*+ use_index_merge(t1, primary, c2, c3) */ c1 from t1 where c2 < 1024 and c3 > 1024") - require.Contains(t, res.Rows()[1][0], "IndexMerge") + require.True(t, tk.HasPlan(sql, "IndexMerge")) + require.True(t, tk.HasNoPlan(sql, "PartitionUnion")) } else { tk.MustExec("set tidb_partition_prune_mode = 'static'") - res := tk.MustQuery("explain select /*+ use_index_merge(t1, primary, c2, c3) */ c1 from t1 where c2 < 1024 and c3 > 1024") if tblIdx == 0 { // partition table - require.Contains(t, res.Rows()[1][0], "PartitionUnion") - require.Contains(t, res.Rows()[2][0], "IndexMerge") + require.True(t, tk.HasPlan(sql, "IndexMerge")) + require.True(t, tk.HasPlan(sql, "PartitionUnion")) } else { - require.Contains(t, res.Rows()[1][0], "IndexMerge") + require.True(t, tk.HasPlan(sql, "IndexMerge")) + require.True(t, tk.HasNoPlan(sql, "PartitionUnion")) } } for i := 0; i < queryCnt; i++ { diff --git a/planner/core/issuetest/BUILD.bazel b/planner/core/issuetest/BUILD.bazel index 3674541d90389..a6c9ab8a9e24f 100644 --- a/planner/core/issuetest/BUILD.bazel +++ b/planner/core/issuetest/BUILD.bazel @@ -6,6 +6,6 @@ go_test( srcs = ["planner_issue_test.go"], flaky = True, race = "on", - shard_count = 7, + shard_count = 8, deps = ["//testkit"], ) diff --git a/planner/core/issuetest/planner_issue_test.go b/planner/core/issuetest/planner_issue_test.go index c71a366fc2482..bb4008df94e7a 100644 --- a/planner/core/issuetest/planner_issue_test.go +++ b/planner/core/issuetest/planner_issue_test.go @@ -15,6 +15,8 @@ package issuetest import ( + "sort" + "strconv" "testing" "github.com/pingcap/tidb/testkit" @@ -116,6 +118,31 @@ func TestIssue46083(t *testing.T) { tk.MustExec("INSERT INTO v0 WITH ta2 AS (TABLE v0) TABLE ta2 FOR UPDATE OF ta2;") } +func TestIssue48755(t *testing.T) { + store := testkit.CreateMockStore(t) + tk := testkit.NewTestKit(t, store) + tk.MustExec("use test") + tk.MustExec("set @@tidb_max_chunk_size=32") + tk.MustExec("create table t(a int, b int);") + tk.MustExec("insert into t values(1, 1);") + tk.MustExec("insert into t select a+1, a+1 from t;") + tk.MustExec("insert into t select a+2, a+2 from t;") + tk.MustExec("insert into t select a+4, a+4 from t;") + tk.MustExec("insert into t select a+8, a+8 from t;") + tk.MustExec("insert into t select a+16, a+16 from t;") + tk.MustExec("insert into t select a+32, a+32 from t;") + rs := tk.MustQuery("select a from (select 100 as a, 100 as b union all select * from t) t where b != 0;") + expectedResult := make([]string, 0, 65) + for i := 1; i < 65; i++ { + expectedResult = append(expectedResult, strconv.FormatInt(int64(i), 10)) + } + expectedResult = append(expectedResult, "100") + sort.Slice(expectedResult, func(i, j int) bool { + return expectedResult[i] < expectedResult[j] + }) + rs.Sort().Check(testkit.Rows(expectedResult...)) +} + func TestIssue47881(t *testing.T) { store := testkit.CreateMockStore(t) tk := testkit.NewTestKit(t, store) diff --git a/planner/core/rule_column_pruning.go b/planner/core/rule_column_pruning.go index 0816aea01a943..fb55caa5b8c7e 100644 --- a/planner/core/rule_column_pruning.go +++ b/planner/core/rule_column_pruning.go @@ -247,6 +247,9 @@ func (p *LogicalUnionAll) PruneColumns(parentUsedCols []*expression.Column, opt if !hasBeenUsed { parentUsedCols = make([]*expression.Column, len(p.schema.Columns)) copy(parentUsedCols, p.schema.Columns) + for i := range used { + used[i] = true + } } for _, child := range p.Children() { err := child.PruneColumns(parentUsedCols, opt) @@ -256,16 +259,15 @@ func (p *LogicalUnionAll) PruneColumns(parentUsedCols []*expression.Column, opt } prunedColumns := make([]*expression.Column, 0) + for i := len(used) - 1; i >= 0; i-- { + if !used[i] { + prunedColumns = append(prunedColumns, p.schema.Columns[i]) + p.schema.Columns = append(p.schema.Columns[:i], p.schema.Columns[i+1:]...) + } + } + appendColumnPruneTraceStep(p, prunedColumns, opt) if hasBeenUsed { // keep the schema of LogicalUnionAll same as its children's - used := expression.GetUsedList(p.children[0].Schema().Columns, p.schema) - for i := len(used) - 1; i >= 0; i-- { - if !used[i] { - prunedColumns = append(prunedColumns, p.schema.Columns[i]) - p.schema.Columns = append(p.schema.Columns[:i], p.schema.Columns[i+1:]...) - } - } - appendColumnPruneTraceStep(p, prunedColumns, opt) // It's possible that the child operator adds extra columns to the schema. // Currently, (*LogicalAggregation).PruneColumns() might do this. // But we don't need such columns, so we add an extra Projection to prune this column when this happened. diff --git a/testkit/testkit.go b/testkit/testkit.go index 81778a1cef8c4..492cca63597bb 100644 --- a/testkit/testkit.go +++ b/testkit/testkit.go @@ -236,6 +236,17 @@ func (tk *TestKit) HasPlan(sql string, plan string, args ...interface{}) bool { return false } +// HasNoPlan checks if the result execution plan doesn't contain specific plan. +func (tk *TestKit) HasNoPlan(sql string, plan string, args ...interface{}) bool { + rs := tk.MustQuery("explain "+sql, args...) + for i := range rs.rows { + if strings.Contains(rs.rows[i][0], plan) { + return false + } + } + return true +} + // HasTiFlashPlan checks if the result execution plan contains TiFlash plan. func (tk *TestKit) HasTiFlashPlan(sql string, args ...interface{}) bool { rs := tk.MustQuery("explain "+sql, args...)