From 4785dc4a622f030dc08b159005522f5f6eb93199 Mon Sep 17 00:00:00 2001
From: Zhou Kunqin <25057648+time-and-fate@users.noreply.github.com>
Date: Wed, 1 Dec 2021 20:35:53 +0800
Subject: [PATCH] planner: fix inconsistent schema between UnionAll and child
 operator (#30231)

---
 planner/core/integration_test.go    | 22 ++++++++++++++++++++++
 planner/core/rule_column_pruning.go | 17 +++++++++++++++++
 2 files changed, 39 insertions(+)

diff --git a/planner/core/integration_test.go b/planner/core/integration_test.go
index 90cfb8932aabb..e58579c5bcb58 100644
--- a/planner/core/integration_test.go
+++ b/planner/core/integration_test.go
@@ -4713,6 +4713,12 @@ func (s *testIntegrationSerialSuite) TestPushDownGroupConcatToTiFlash(c *C) {
 
 func (s *testIntegrationSuite) TestIssue27797(c *C) {
 	tk := testkit.NewTestKit(c, s.store)
+	origin := tk.MustQuery("SELECT @@session.tidb_partition_prune_mode")
+	originStr := origin.Rows()[0][0].(string)
+	defer func() {
+		tk.MustExec("set @@session.tidb_partition_prune_mode = '" + originStr + "'")
+	}()
+	tk.MustExec("set @@session.tidb_partition_prune_mode = 'static'")
 	tk.MustExec("use test")
 	tk.MustExec("drop table if exists t27797")
 	tk.MustExec("create table t27797(a int, b int, c int, d int) " +
@@ -4948,6 +4954,22 @@ func (s *testIntegrationSuite) TestIssue30094(c *C) {
 	))
 }
 
+func (s *testIntegrationSuite) TestIssue29705(c *C) {
+	tk := testkit.NewTestKit(c, s.store)
+	origin := tk.MustQuery("SELECT @@session.tidb_partition_prune_mode")
+	originStr := origin.Rows()[0][0].(string)
+	defer func() {
+		tk.MustExec("set @@session.tidb_partition_prune_mode = '" + originStr + "'")
+	}()
+	tk.MustExec("set @@session.tidb_partition_prune_mode = 'static'")
+	tk.MustExec("use test")
+	tk.MustExec("drop table if exists t;")
+	tk.MustExec("create table t(id int) partition by hash(id) partitions 4;")
+	tk.MustExec("insert into t values(1);")
+	result := tk.MustQuery("SELECT COUNT(1) FROM ( SELECT COUNT(1) FROM t b GROUP BY id) a;")
+	result.Check(testkit.Rows("1"))
+}
+
 func (s *testIntegrationSerialSuite) TestIssue30271(c *C) {
 	defer collate.SetNewCollationEnabledForTest(false)
 	collate.SetNewCollationEnabledForTest(true)
diff --git a/planner/core/rule_column_pruning.go b/planner/core/rule_column_pruning.go
index 77c727f4fb5ce..1d144c72807f9 100644
--- a/planner/core/rule_column_pruning.go
+++ b/planner/core/rule_column_pruning.go
@@ -241,6 +241,23 @@ func (p *LogicalUnionAll) PruneColumns(parentUsedCols []*expression.Column) erro
 				p.schema.Columns = append(p.schema.Columns[:i], p.schema.Columns[i+1:]...)
 			}
 		}
+		// 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.
+		for i, child := range p.Children() {
+			if p.schema.Len() < child.Schema().Len() {
+				schema := p.schema.Clone()
+				exprs := make([]expression.Expression, len(p.schema.Columns))
+				for j, col := range schema.Columns {
+					exprs[j] = col
+				}
+				proj := LogicalProjection{Exprs: exprs, AvoidColumnEvaluator: true}.Init(p.ctx, p.blockOffset)
+				proj.SetSchema(schema)
+
+				proj.SetChildren(child)
+				p.children[i] = proj
+			}
+		}
 	}
 	return nil
 }