plan: covert max/min to Limit + Sort operators (pingcap#5105)

plan/aggregation_eliminate.go

@@ -0,0 +1,76 @@
+// Copyright 2017 PingCAP, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+// // Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package plan
+
+import (
+	"github.com/pingcap/tidb/ast"
+	"github.com/pingcap/tidb/context"
+)
+
+// aggEliminater tries to elimiate max/min aggregate function.
+// For SQL like `select max(c) from t;`, we could optimize it to `select c from t order by c desc limit 1;`.
+// For SQL like `select min(c) from t;`, we could optimize it to `select c from t order by c limit 1;`.
+type aggEliminater struct {
+	ctx context.Context
+}
+
+func (a *aggEliminater) optimize(p LogicalPlan, ctx context.Context) (LogicalPlan, error) {
+	a.ctx = ctx
+	return a.eliminateAgg(p), nil
+}
+
+// Try to convert max/min to Limit+Sort operators.
+func (a *aggEliminater) eliminateAgg(p LogicalPlan) LogicalPlan {
+	if agg, ok := p.(*LogicalAggregation); ok {
+		// We only consider case with single max/min function.
+		if len(agg.AggFuncs) != 1 || len(agg.GroupByItems) != 0 {
+			return p
+		}
+		f := agg.AggFuncs[0]
+		if f.GetName() != ast.AggFuncMax && f.GetName() != ast.AggFuncMin {
+			return p
+		}
+
+		// Add Sort and Limit operators.
+		// For max function, the sort order should be desc.
+		desc := f.GetName() == ast.AggFuncMax
+		// Compose Sort operator.
+		sort := LogicalSort{}.init(a.ctx)
+		sort.ByItems = append(sort.ByItems, &ByItems{f.GetArgs()[0], desc})
+		sort.SetSchema(p.Children()[0].Schema().Clone())
+		sort.SetChildren(p.Children()...)
+		// Compose Limit operator.
+		li := LogicalLimit{Count: 1}.init(a.ctx)
+		li.SetSchema(sort.Schema().Clone())
+		li.SetChildren(sort)
+
+		// Add a projection operator here.
+		// During topn_push_down, the sort/limit operator will be converted to a topn operator and the schema of sort/limit will be ignored.
+		// So the schema of the LogicalAggregation will be lost. We add this projection operator to keep the schema unlost.
+		// For SQL like `select * from t where v=(select min(t1.v) from t t1, t t2, t t3 where t1.id=t2.id and t2.id=t3.id and t1.id=t.id);`,
+		// the min(t1.v) will be refered in the outer selection. So we should keep the schema from LogicalAggragation.
+		proj := LogicalProjection{}.init(a.ctx)
+		proj.Exprs = append(proj.Exprs, f.GetArgs()[0])
+		proj.SetSchema(p.Schema().Clone())
+		proj.SetChildren(li)
+		return proj
+	}
+
+	newChildren := make([]Plan, 0, len(p.Children()))
+	for _, child := range p.Children() {
+		newChild := a.eliminateAgg(child.(LogicalPlan))
+		newChildren = append(newChildren, newChild)
+	}
+	p.SetChildren(newChildren...)
+	return p
+}

plan/logical_plan_builder.go

@@ -66,6 +66,10 @@ func (p *LogicalAggregation) collectGroupByColumns() {
 func (b *planBuilder) buildAggregation(p LogicalPlan, aggFuncList []*ast.AggregateFuncExpr, gbyItems []expression.Expression) (LogicalPlan, map[int]int) {
 	b.optFlag = b.optFlag | flagBuildKeyInfo
 	b.optFlag = b.optFlag | flagAggregationOptimize
+	// We may apply aggregation eliminate optimization.
+	// So we add the flagAggEliminate to try to convert max/min to topn and flagPushDownTopN to handle the newly added topn operator.
+	b.optFlag = b.optFlag | flagAggEliminate
+	b.optFlag = b.optFlag | flagPushDownTopN
 
 	agg := LogicalAggregation{AggFuncs: make([]aggregation.Aggregation, 0, len(aggFuncList))}.init(b.ctx)
 	schema := expression.NewSchema(make([]*expression.Column, 0, len(aggFuncList)+p.Schema().Len())...)

plan/logical_plan_test.go

@@ -1622,3 +1622,56 @@ func (s *testPlanSuite) TestNameResolver(c *C) {
 		}
 	}
 }
+
+func (s *testPlanSuite) TestAggEliminater(c *C) {
+	defer func() {
+		testleak.AfterTest(c)()
+	}()
+	tests := []struct {
+		sql  string
+		best string
+	}{
+		// Max to Limit + Sort-Desc.
+		{
+			sql:  "select max(a) from t;",
+			best: "DataScan(t)->TopN([test.t.a true],0,1)->Projection->Projection",
+		},
+		// Min to Limit + Sort.
+		{
+			sql:  "select min(a) from t;",
+			best: "DataScan(t)->TopN([test.t.a],0,1)->Projection->Projection",
+		},
+		// Do nothing to max + firstrow.
+		{
+			sql:  "select max(a), b from t;",
+			best: "DataScan(t)->Aggr(max(test.t.a),firstrow(test.t.b))->Projection",
+		},
+		// Do nothing to max+min.
+		{
+			sql:  "select max(a), min(a) from t;",
+			best: "DataScan(t)->Aggr(max(test.t.a),min(test.t.a))->Projection",
+		},
+		// Do nothing to max with groupby.
+		{
+			sql:  "select max(a) from t group by b;",
+			best: "DataScan(t)->Aggr(max(test.t.a))->Projection",
+		},
+	}
+
+	for _, tt := range tests {
+		comment := Commentf("for %s", tt.sql)
+		stmt, err := s.ParseOneStmt(tt.sql, "", "")
+		c.Assert(err, IsNil, comment)
+		Preprocess(s.ctx, stmt, s.is, false)
+		builder := &planBuilder{
+			ctx:       mockContext(),
+			is:        s.is,
+			colMapper: make(map[*ast.ColumnNameExpr]int),
+		}
+		p := builder.build(stmt).(LogicalPlan)
+		c.Assert(builder.err, IsNil)
+		p, err = logicalOptimize(builder.optFlag, p.(LogicalPlan), builder.ctx)
+		c.Assert(err, IsNil)
+		c.Assert(ToString(p), Equals, tt.best, comment)
+	}
+}

plan/optimizer.go

@@ -36,6 +36,7 @@ const (
 	flagDecorrelate
 	flagPredicatePushDown
 	flagAggregationOptimize
+	flagAggEliminate // Keep AggEliminater before PushDownTopN, for it will add a TopN operator.
 	flagPushDownTopN
 )
 
@@ -46,6 +47,7 @@ var optRuleList = []logicalOptRule{
 	&decorrelateSolver{},
 	&ppdSolver{},
 	&aggregationOptimizer{},
+	&aggEliminater{},
 	&pushDownTopNOptimizer{},
 }