diff --git a/planner/core/issuetest/planner_issue_test.go b/planner/core/issuetest/planner_issue_test.go index dd6f4e6534934..1362f871c3f45 100644 --- a/planner/core/issuetest/planner_issue_test.go +++ b/planner/core/issuetest/planner_issue_test.go @@ -15,6 +15,8 @@ package issuetest import ( + "slices" + "strconv" "testing" "github.com/pingcap/tidb/testkit" @@ -115,3 +117,26 @@ func TestIssue46083(t *testing.T) { tk.MustExec("CREATE TEMPORARY TABLE v0(v1 int)") 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") + slices.Sort(expectedResult) + rs.Sort().Check(testkit.Rows(expectedResult...)) +} 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.