planner: move flagPrunColumnsAgain to the last (#14791)

cmd/explaintest/r/explain_easy.result

@@ -505,8 +505,8 @@ PRIMARY KEY (`id`)
 EXPLAIN SELECT COUNT(1) FROM (SELECT COALESCE(b.region_name, '不详') region_name, SUM(a.registration_num) registration_num FROM (SELECT stat_date, show_date, region_id, 0 registration_num FROM test01 WHERE period = 1 AND stat_date >= 20191202 AND stat_date <= 20191202 UNION ALL SELECT stat_date, show_date, region_id, registration_num registration_num FROM test01 WHERE period = 1 AND stat_date >= 20191202 AND stat_date <= 20191202) a LEFT JOIN test02 b ON a.region_id = b.id WHERE registration_num > 0 AND a.stat_date >= '20191202' AND a.stat_date <= '20191202' GROUP BY a.stat_date , a.show_date , COALESCE(b.region_name, '不详') ) JLS;
 id	count	task	operator info
 StreamAgg_22	1.00	root	funcs:count(1)->Column#22
-└─HashAgg_25	1.00	root	group by:Column#30, Column#31, Column#32, funcs:firstrow(1)->Column#29
-  └─Projection_48	0.02	root	Column#14, Column#15, coalesce(test.test02.region_name, 不详)->Column#32
+└─HashAgg_25	1.00	root	group by:Column#32, Column#33, Column#34, funcs:firstrow(1)->Column#31
+  └─Projection_48	0.02	root	Column#14, Column#15, coalesce(test.test02.region_name, 不详)->Column#34
     └─IndexMergeJoin_32	0.02	root	left outer join, inner:TableReader_30, outer key:Column#16, inner key:test.test02.id
       ├─Union_37	0.02	root	
       │ ├─Projection_38	0.01	root	test.test01.stat_date, test.test01.show_date, test.test01.region_id
@@ -676,15 +676,15 @@ id	count	task	operator info
 Sort_13	2.00	root	Column#3:asc
 └─HashAgg_17	2.00	root	group by:Column#3, funcs:firstrow(Column#6)->Column#3
   └─Union_18	2.00	root	
-    ├─HashAgg_19	1.00	root	group by:0, funcs:firstrow(0)->Column#6, funcs:firstrow(0)->Column#3
+    ├─HashAgg_19	1.00	root	group by:1, funcs:firstrow(0)->Column#6, funcs:firstrow(0)->Column#3
     │ └─TableDual_22	1.00	root	rows:1
     └─HashAgg_25	1.00	root	group by:1, funcs:firstrow(1)->Column#6, funcs:firstrow(1)->Column#3
       └─TableDual_28	1.00	root	rows:1
 explain SELECT 0 AS a FROM dual UNION (SELECT 1 AS a FROM dual ORDER BY a);
 id	count	task	operator info
 HashAgg_15	2.00	root	group by:Column#3, funcs:firstrow(Column#6)->Column#3
 └─Union_16	2.00	root	
-  ├─HashAgg_17	1.00	root	group by:0, funcs:firstrow(0)->Column#6, funcs:firstrow(0)->Column#3
+  ├─HashAgg_17	1.00	root	group by:1, funcs:firstrow(0)->Column#6, funcs:firstrow(0)->Column#3
   │ └─TableDual_20	1.00	root	rows:1
   └─StreamAgg_27	1.00	root	group by:Column#1, funcs:firstrow(Column#1)->Column#6, funcs:firstrow(Column#1)->Column#3
     └─Projection_32	1.00	root	1->Column#1

cmd/explaintest/r/topn_push_down.result

@@ -207,7 +207,7 @@ Apply_17	9990.00	root	semi join, inner:Selection_21, equal:[eq(test.t1.a, test.t
 │ └─Selection_19	9990.00	cop[tikv]	not(isnull(test.t1.a))
 │   └─TableFullScan_18	10000.00	cop[tikv]	table:t1, keep order:false, stats:pseudo
 └─Selection_21	0.80	root	not(isnull(test.t2.a))
-  └─Projection_22	1.00	root	test.t2.a, test.t1.b
+  └─Projection_22	1.00	root	test.t2.a
     └─Limit_23	1.00	root	offset:0, count:1
       └─TableReader_29	1.00	root	data:Limit_28
         └─Limit_28	1.00	cop[tikv]	offset:0, count:1

planner/core/logical_plan_test.go

@@ -81,7 +81,7 @@ func (s *testPlanSuite) TestPredicatePushDown(c *C) {
 		c.Assert(err, IsNil, comment)
 		p, _, err := BuildLogicalPlan(ctx, s.ctx, stmt, s.is)
 		c.Assert(err, IsNil)
-		p, err = logicalOptimize(context.TODO(), flagPredicatePushDown|flagDecorrelate|flagPrunColumns, p.(LogicalPlan))
+		p, err = logicalOptimize(context.TODO(), flagPredicatePushDown|flagDecorrelate|flagPrunColumns|flagPrunColumnsAgain, p.(LogicalPlan))
 		c.Assert(err, IsNil)
 		s.testData.OnRecord(func() {
 			output[ith] = ToString(p)
@@ -108,7 +108,7 @@ func (s *testPlanSuite) TestJoinPredicatePushDown(c *C) {
 		c.Assert(err, IsNil, comment)
 		p, _, err := BuildLogicalPlan(ctx, s.ctx, stmt, s.is)
 		c.Assert(err, IsNil, comment)
-		p, err = logicalOptimize(context.TODO(), flagPredicatePushDown|flagDecorrelate|flagPrunColumns, p.(LogicalPlan))
+		p, err = logicalOptimize(context.TODO(), flagPredicatePushDown|flagDecorrelate|flagPrunColumns|flagPrunColumnsAgain, p.(LogicalPlan))
 		c.Assert(err, IsNil, comment)
 		proj, ok := p.(*LogicalProjection)
 		c.Assert(ok, IsTrue, comment)
@@ -147,7 +147,7 @@ func (s *testPlanSuite) TestOuterWherePredicatePushDown(c *C) {
 		c.Assert(err, IsNil, comment)
 		p, _, err := BuildLogicalPlan(ctx, s.ctx, stmt, s.is)
 		c.Assert(err, IsNil, comment)
-		p, err = logicalOptimize(context.TODO(), flagPredicatePushDown|flagDecorrelate|flagPrunColumns, p.(LogicalPlan))
+		p, err = logicalOptimize(context.TODO(), flagPredicatePushDown|flagDecorrelate|flagPrunColumns|flagPrunColumnsAgain, p.(LogicalPlan))
 		c.Assert(err, IsNil, comment)
 		proj, ok := p.(*LogicalProjection)
 		c.Assert(ok, IsTrue, comment)
@@ -192,7 +192,7 @@ func (s *testPlanSuite) TestSimplifyOuterJoin(c *C) {
 		c.Assert(err, IsNil, comment)
 		p, _, err := BuildLogicalPlan(ctx, s.ctx, stmt, s.is)
 		c.Assert(err, IsNil, comment)
-		p, err = logicalOptimize(context.TODO(), flagPredicatePushDown|flagPrunColumns, p.(LogicalPlan))
+		p, err = logicalOptimize(context.TODO(), flagPredicatePushDown|flagPrunColumns|flagPrunColumnsAgain, p.(LogicalPlan))
 		c.Assert(err, IsNil, comment)
 		planString := ToString(p)
 		s.testData.OnRecord(func() {
@@ -232,7 +232,7 @@ func (s *testPlanSuite) TestAntiSemiJoinConstFalse(c *C) {
 		c.Assert(err, IsNil, comment)
 		p, _, err := BuildLogicalPlan(ctx, s.ctx, stmt, s.is)
 		c.Assert(err, IsNil, comment)
-		p, err = logicalOptimize(context.TODO(), flagDecorrelate|flagPredicatePushDown|flagPrunColumns, p.(LogicalPlan))
+		p, err = logicalOptimize(context.TODO(), flagDecorrelate|flagPredicatePushDown|flagPrunColumns|flagPrunColumnsAgain, p.(LogicalPlan))
 		c.Assert(err, IsNil, comment)
 		c.Assert(ToString(p), Equals, ca.best, comment)
 		join, _ := p.(LogicalPlan).Children()[0].(*LogicalJoin)
@@ -259,7 +259,7 @@ func (s *testPlanSuite) TestDeriveNotNullConds(c *C) {
 		c.Assert(err, IsNil, comment)
 		p, _, err := BuildLogicalPlan(ctx, s.ctx, stmt, s.is)
 		c.Assert(err, IsNil, comment)
-		p, err = logicalOptimize(context.TODO(), flagPredicatePushDown|flagPrunColumns|flagDecorrelate, p.(LogicalPlan))
+		p, err = logicalOptimize(context.TODO(), flagPredicatePushDown|flagPrunColumns|flagPrunColumnsAgain|flagDecorrelate, p.(LogicalPlan))
 		c.Assert(err, IsNil, comment)
 		s.testData.OnRecord(func() {
 			output[i].Plan = ToString(p)
@@ -419,7 +419,7 @@ func (s *testPlanSuite) TestTablePartition(c *C) {
 		})
 		p, _, err := BuildLogicalPlan(ctx, s.ctx, stmt, isChoices[ca.IsIdx])
 		c.Assert(err, IsNil)
-		p, err = logicalOptimize(context.TODO(), flagDecorrelate|flagPrunColumns|flagPredicatePushDown|flagPartitionProcessor, p.(LogicalPlan))
+		p, err = logicalOptimize(context.TODO(), flagDecorrelate|flagPrunColumns|flagPrunColumnsAgain|flagPredicatePushDown|flagPartitionProcessor, p.(LogicalPlan))
 		c.Assert(err, IsNil)
 		planString := ToString(p)
 		s.testData.OnRecord(func() {
@@ -444,7 +444,7 @@ func (s *testPlanSuite) TestSubquery(c *C) {
 		p, _, err := BuildLogicalPlan(ctx, s.ctx, stmt, s.is)
 		c.Assert(err, IsNil)
 		if lp, ok := p.(LogicalPlan); ok {
-			p, err = logicalOptimize(context.TODO(), flagBuildKeyInfo|flagDecorrelate|flagPrunColumns, lp)
+			p, err = logicalOptimize(context.TODO(), flagBuildKeyInfo|flagDecorrelate|flagPrunColumns|flagPrunColumnsAgain, lp)
 			c.Assert(err, IsNil)
 		}
 		s.testData.OnRecord(func() {
@@ -469,7 +469,7 @@ func (s *testPlanSuite) TestPlanBuilder(c *C) {
 		p, _, err := BuildLogicalPlan(ctx, s.ctx, stmt, s.is)
 		c.Assert(err, IsNil)
 		if lp, ok := p.(LogicalPlan); ok {
-			p, err = logicalOptimize(context.TODO(), flagPrunColumns, lp)
+			p, err = logicalOptimize(context.TODO(), flagPrunColumns|flagPrunColumnsAgain, lp)
 			c.Assert(err, IsNil)
 		}
 		s.testData.OnRecord(func() {
@@ -516,7 +516,7 @@ func (s *testPlanSuite) TestEagerAggregation(c *C) {
 
 		p, _, err := BuildLogicalPlan(ctx, s.ctx, stmt, s.is)
 		c.Assert(err, IsNil)
-		p, err = logicalOptimize(context.TODO(), flagBuildKeyInfo|flagPredicatePushDown|flagPrunColumns|flagPushDownAgg, p.(LogicalPlan))
+		p, err = logicalOptimize(context.TODO(), flagBuildKeyInfo|flagPredicatePushDown|flagPrunColumns|flagPrunColumnsAgain|flagPushDownAgg, p.(LogicalPlan))
 		c.Assert(err, IsNil)
 		s.testData.OnRecord(func() {
 			output[ith] = ToString(p)
@@ -542,7 +542,7 @@ func (s *testPlanSuite) TestColumnPruning(c *C) {
 
 		p, _, err := BuildLogicalPlan(ctx, s.ctx, stmt, s.is)
 		c.Assert(err, IsNil)
-		lp, err := logicalOptimize(ctx, flagPredicatePushDown|flagPrunColumns, p.(LogicalPlan))
+		lp, err := logicalOptimize(ctx, flagPredicatePushDown|flagPrunColumns|flagPrunColumnsAgain, p.(LogicalPlan))
 		c.Assert(err, IsNil)
 		s.testData.OnRecord(func() {
 			output[i] = make(map[int][]string)
@@ -571,7 +571,7 @@ func (s *testPlanSuite) TestProjectionEliminator(c *C) {
 
 		p, _, err := BuildLogicalPlan(ctx, s.ctx, stmt, s.is)
 		c.Assert(err, IsNil)
-		p, err = logicalOptimize(context.TODO(), flagBuildKeyInfo|flagPrunColumns|flagEliminateProjection, p.(LogicalPlan))
+		p, err = logicalOptimize(context.TODO(), flagBuildKeyInfo|flagPrunColumns|flagPrunColumnsAgain|flagEliminateProjection, p.(LogicalPlan))
 		c.Assert(err, IsNil)
 		c.Assert(ToString(p), Equals, tt.best, Commentf("for %s %d", tt.sql, ith))
 	}
@@ -853,7 +853,7 @@ func (s *testPlanSuite) TestAggPrune(c *C) {
 		p, _, err := BuildLogicalPlan(ctx, s.ctx, stmt, s.is)
 		c.Assert(err, IsNil)
 
-		p, err = logicalOptimize(context.TODO(), flagPredicatePushDown|flagPrunColumns|flagBuildKeyInfo|flagEliminateAgg|flagEliminateProjection, p.(LogicalPlan))
+		p, err = logicalOptimize(context.TODO(), flagPredicatePushDown|flagPrunColumns|flagPrunColumnsAgain|flagBuildKeyInfo|flagEliminateAgg|flagEliminateProjection, p.(LogicalPlan))
 		c.Assert(err, IsNil)
 		planString := ToString(p)
 		s.testData.OnRecord(func() {

planner/core/optimizer.go

@@ -47,12 +47,12 @@ const (
 	flagEliminateProjection
 	flagMaxMinEliminate
 	flagPredicatePushDown
-	flagPrunColumnsAgain
 	flagEliminateOuterJoin
 	flagPartitionProcessor
 	flagPushDownAgg
 	flagPushDownTopN
 	flagJoinReOrder
+	flagPrunColumnsAgain
 )
 
 var optRuleList = []logicalOptRule{
@@ -64,12 +64,12 @@ var optRuleList = []logicalOptRule{
 	&projectionEliminator{},
 	&maxMinEliminator{},
 	&ppdSolver{},
-	&columnPruner{}, // column pruning again after predicate push down
 	&outerJoinEliminator{},
 	&partitionProcessor{},
 	&aggregationPushDownSolver{},
 	&pushDownTopNOptimizer{},
 	&joinReOrderSolver{},
+	&columnPruner{}, // column pruning again at last, note it will mess up the results of buildKeySolver
 }
 
 // logicalOptRule means a logical optimizing rule, which contains decorrelate, ppd, column pruning, etc.

planner/core/planbuilder.go

@@ -389,7 +389,8 @@ func NewPlanBuilder(sctx sessionctx.Context, is infoschema.InfoSchema, processor
 func (b *PlanBuilder) Build(ctx context.Context, node ast.Node) (Plan, error) {
 	b.optFlag = flagPrunColumns
 	defer func() {
-		if b.optFlag&flagPredicatePushDown > 0 {
+		// if there is something after flagPrunColumns, do flagPrunColumnsAgain
+		if b.optFlag&flagPrunColumns > 0 && b.optFlag-flagPrunColumns > flagPrunColumns {
 			b.optFlag |= flagPrunColumnsAgain
 		}
 	}()

planner/core/rule_column_pruning.go

@@ -151,10 +151,30 @@ func (ls *LogicalSort) PruneColumns(parentUsedCols []*expression.Column) error {
 	return child.PruneColumns(parentUsedCols)
 }
 
+// PruneColumns implements LogicalPlan interface.
+// If any expression can view as a constant in execution stage, such as correlated column, constant,
+// we do prune them. Note that we can't prune the expressions contain non-deterministic functions, such as rand().
+func (lt *LogicalTopN) PruneColumns(parentUsedCols []*expression.Column) error {
+	child := lt.children[0]
+	for i := len(lt.ByItems) - 1; i >= 0; i-- {
+		cols := expression.ExtractColumns(lt.ByItems[i].Expr)
+		if len(cols) == 0 {
+			if expression.IsMutableEffectsExpr(lt.ByItems[i].Expr) {
+				continue
+			}
+			lt.ByItems = append(lt.ByItems[:i], lt.ByItems[i+1:]...)
+		} else if lt.ByItems[i].Expr.GetType().Tp == mysql.TypeNull {
+			lt.ByItems = append(lt.ByItems[:i], lt.ByItems[i+1:]...)
+		} else {
+			parentUsedCols = append(parentUsedCols, cols...)
+		}
+	}
+	return child.PruneColumns(parentUsedCols)
+}
+
 // PruneColumns implements LogicalPlan interface.
 func (p *LogicalUnionAll) PruneColumns(parentUsedCols []*expression.Column) error {
 	used := expression.GetUsedList(parentUsedCols, p.schema)
-
 	hasBeenUsed := false
 	for i := range used {
 		hasBeenUsed = hasBeenUsed || used[i]
@@ -165,20 +185,23 @@ func (p *LogicalUnionAll) PruneColumns(parentUsedCols []*expression.Column) erro
 	if !hasBeenUsed {
 		parentUsedCols = make([]*expression.Column, len(p.schema.Columns))
 		copy(parentUsedCols, p.schema.Columns)
-	} else {
-		// Issue 10341: p.schema.Columns might contain table name (AsName), but p.Children()0].Schema().Columns does not.
-		for i := len(used) - 1; i >= 0; i-- {
-			if !used[i] {
-				p.schema.Columns = append(p.schema.Columns[:i], p.schema.Columns[i+1:]...)
-			}
-		}
 	}
 	for _, child := range p.Children() {
 		err := child.PruneColumns(parentUsedCols)
 		if err != nil {
 			return err
 		}
 	}
+
+	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] {
+				p.schema.Columns = append(p.schema.Columns[:i], p.schema.Columns[i+1:]...)
+			}
+		}
+	}
 	return nil
 }
 

planner/core/rule_partition_processor.go

@@ -195,7 +195,7 @@ func (s *partitionProcessor) pruneHashPartition(ds *DataSource, pi *model.Partit
 	}
 	unionAll := LogicalUnionAll{}.Init(ds.SCtx(), ds.blockOffset)
 	unionAll.SetChildren(children...)
-	unionAll.SetSchema(ds.schema)
+	unionAll.SetSchema(ds.schema.Clone())
 	return unionAll, nil
 }
 
@@ -305,7 +305,7 @@ func (s *partitionProcessor) prune(ds *DataSource) (LogicalPlan, error) {
 	}
 	unionAll := LogicalUnionAll{}.Init(ds.SCtx(), ds.blockOffset)
 	unionAll.SetChildren(children...)
-	unionAll.SetSchema(ds.schema)
+	unionAll.SetSchema(ds.schema.Clone())
 	return unionAll, nil
 }
 

planner/core/testdata/plan_suite_unexported_out.json

@@ -96,7 +96,7 @@
     "Cases": [
       "Join{DataScan(t)->DataScan(s)}(test.t.a,test.t.a)->Projection",
       "Join{DataScan(t)->Aggr(count(test.t.c),firstrow(test.t.a))->DataScan(s)}(test.t.a,test.t.a)->Projection->Projection",
-      "Join{DataScan(t)->Aggr(count(test.t.c),firstrow(test.t.a))->DataScan(s)}(test.t.a,test.t.a)->Aggr(firstrow(Column#13),firstrow(test.t.a),count(test.t.b))->Projection->Projection",
+      "Join{DataScan(t)->Aggr(count(test.t.c),firstrow(test.t.a))->DataScan(s)}(test.t.a,test.t.a)->Aggr(firstrow(Column#13),count(test.t.b))->Projection->Projection",
       "Apply{DataScan(t)->DataScan(s)->Sel([eq(test.t.a, test.t.a)])->Aggr(count(test.t.b))}->Projection",
       "Join{DataScan(t)->DataScan(s)->Aggr(count(test.t.b),firstrow(test.t.a))}(test.t.a,test.t.a)->Projection->Projection->Projection",
       "Join{Join{DataScan(t1)->DataScan(t2)}->DataScan(s)->Aggr(count(test.t.b),firstrow(test.t.a))}(test.t.a,test.t.a)->Projection->Projection->Projection",