planner/core: separate aggPrune from aggPushDown

plan/logical_plan_builder.go

@@ -75,6 +75,7 @@ func (b *planBuilder) buildAggregation(p LogicalPlan, aggFuncList []*ast.Aggrega
 	b.optFlag = b.optFlag | flagPushDownTopN
 	// when we eliminate the max and min we may add `is not null` filter.
 	b.optFlag = b.optFlag | flagPredicatePushDown
+	b.optFlag = b.optFlag | flagEliminateAgg
 
 	plan4Agg := LogicalAggregation{AggFuncs: make([]*aggregation.AggFuncDesc, 0, len(aggFuncList))}.init(b.ctx)
 	schema4Agg := expression.NewSchema(make([]*expression.Column, 0, len(aggFuncList)+p.Schema().Len())...)

plan/logical_plan_test.go

@@ -875,6 +875,10 @@ func (s *testPlanSuite) TestEagerAggregation(c *C) {
 			sql:  "select max(a.c) from t a join t b on a.a=b.a and a.b=b.b group by a.b",
 			best: "Join{DataScan(a)->DataScan(b)}(a.a,b.a)(a.b,b.b)->Aggr(max(a.c))->Projection",
 		},
+		{
+			sql:  "select t1.a, count(t2.b) from t t1, t t2 where t1.a = t2.a group by t1.a",
+			best: "Join{DataScan(t1)->DataScan(t2)}(t1.a,t2.a)->Projection->Projection",
+		},
 	}
 	s.ctx.GetSessionVars().AllowAggPushDown = true
 	for _, tt := range tests {
@@ -1315,10 +1319,6 @@ func (s *testPlanSuite) TestAggPrune(c *C) {
 			sql:  "select sum(b) from t group by c, d, e",
 			best: "DataScan(t)->Aggr(sum(test.t.b))->Projection",
 		},
-		{
-			sql:  "select t1.a, count(t2.b) from t t1, t t2 where t1.a = t2.a group by t1.a",
-			best: "Join{DataScan(t1)->DataScan(t2)}(t1.a,t2.a)->Projection->Projection",
-		},
 		{
 			sql:  "select tt.a, sum(tt.b) from (select a, b from t) tt group by tt.a",
 			best: "DataScan(t)->Projection->Projection->Projection",
@@ -1328,7 +1328,6 @@ func (s *testPlanSuite) TestAggPrune(c *C) {
 			best: "DataScan(t)->Projection->Projection->Projection->Projection",
 		},
 	}
-	s.ctx.GetSessionVars().AllowAggPushDown = true
 	for _, tt := range tests {
 		comment := Commentf("for %s", tt.sql)
 		stmt, err := s.ParseOneStmt(tt.sql, "", "")
@@ -1337,11 +1336,10 @@ func (s *testPlanSuite) TestAggPrune(c *C) {
 		p, err := BuildLogicalPlan(s.ctx, stmt, s.is)
 		c.Assert(err, IsNil)
 
-		p, err = logicalOptimize(flagPredicatePushDown|flagPrunColumns|flagBuildKeyInfo|flagAggregationOptimize, p.(LogicalPlan))
+		p, err = logicalOptimize(flagPredicatePushDown|flagPrunColumns|flagBuildKeyInfo|flagEliminateAgg, p.(LogicalPlan))
 		c.Assert(err, IsNil)
 		c.Assert(ToString(p), Equals, tt.best, comment)
 	}
-	s.ctx.GetSessionVars().AllowAggPushDown = false
 }
 
 func (s *testPlanSuite) TestVisitInfo(c *C) {

plan/optimizer.go

@@ -32,6 +32,7 @@ const (
 	flagPrunColumns uint64 = 1 << iota
 	flagEliminateProjection
 	flagBuildKeyInfo
+	flagEliminateAgg
 	flagDecorrelate
 	flagMaxMinEliminate
 	flagPredicatePushDown
@@ -44,11 +45,12 @@ var optRuleList = []logicalOptRule{
 	&columnPruner{},
 	&projectionEliminater{},
 	&buildKeySolver{},
+	&aggregationRecursiveEliminater{},
 	&decorrelateSolver{},
 	&maxMinEliminator{},
 	&ppdSolver{},
 	&partitionProcessor{},
-	&aggregationOptimizer{},
+	&aggregationPushDownSolver{},
 	&pushDownTopNOptimizer{},
 }
 

plan/rule_aggregation_elimination.go

@@ -0,0 +1,129 @@
+// Copyright 2018 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/expression"
+	"github.com/pingcap/tidb/expression/aggregation"
+	"github.com/pingcap/tidb/mysql"
+	"github.com/pingcap/tidb/sessionctx"
+	"github.com/pingcap/tidb/types"
+)
+
+type aggregationRecursiveEliminater struct {
+	aggregationEliminateChecker
+}
+
+type aggregationEliminateChecker struct {
+}
+
+// tryToEliminateAggregation will eliminate aggregation grouped by unique key.
+// e.g. select min(b) from t group by a. If a is a unique key, then this sql is equal to `select b from t group by a`.
+// For count(expr), sum(expr), avg(expr), count(distinct expr, [expr...]) we may need to rewrite the expr. Details are shown below.
+// If we can eliminate agg successful, we return a projection. Else we return a nil pointer.
+func (a *aggregationEliminateChecker) tryToEliminateAggregation(agg *LogicalAggregation) *LogicalProjection {
+	schemaByGroupby := expression.NewSchema(agg.groupByCols...)
+	coveredByUniqueKey := false
+	for _, key := range agg.children[0].Schema().Keys {
+		if schemaByGroupby.ColumnsIndices(key) != nil {
+			coveredByUniqueKey = true
+			break
+		}
+	}
+	if coveredByUniqueKey {
+		// GroupByCols has unique key, so this aggregation can be removed.
+		proj := a.convertAggToProj(agg)
+		proj.SetChildren(agg.children[0])
+		return proj
+	}
+	return nil
+}
+
+func (a *aggregationEliminateChecker) convertAggToProj(agg *LogicalAggregation) *LogicalProjection {
+	proj := LogicalProjection{
+		Exprs: make([]expression.Expression, 0, len(agg.AggFuncs)),
+	}.init(agg.ctx)
+	for _, fun := range agg.AggFuncs {
+		expr := a.rewriteExpr(agg.ctx, fun)
+		proj.Exprs = append(proj.Exprs, expr)
+	}
+	proj.SetSchema(agg.schema.Clone())
+	return proj
+}
+
+// rewriteExpr will rewrite the aggregate function to expression doesn't contain aggregate function.
+func (a *aggregationEliminateChecker) rewriteExpr(ctx sessionctx.Context, aggFunc *aggregation.AggFuncDesc) expression.Expression {
+	switch aggFunc.Name {
+	case ast.AggFuncCount:
+		if aggFunc.Mode == aggregation.FinalMode {
+			return a.rewriteSumOrAvg(ctx, aggFunc.Args)
+		}
+		return a.rewriteCount(ctx, aggFunc.Args)
+	case ast.AggFuncSum, ast.AggFuncAvg:
+		return a.rewriteSumOrAvg(ctx, aggFunc.Args)
+	default:
+		// Default we do nothing about expr.
+		return aggFunc.Args[0]
+	}
+}
+
+func (a *aggregationEliminateChecker) rewriteCount(ctx sessionctx.Context, exprs []expression.Expression) expression.Expression {
+	// If is count(expr), we will change it to if(isnull(expr), 0, 1).
+	// If is count(distinct x, y, z) we will change it to if(isnull(x) or isnull(y) or isnull(z), 0, 1).
+	isNullExprs := make([]expression.Expression, 0, len(exprs))
+	for _, expr := range exprs {
+		isNullExpr := expression.NewFunctionInternal(ctx, ast.IsNull, types.NewFieldType(mysql.TypeTiny), expr)
+		isNullExprs = append(isNullExprs, isNullExpr)
+	}
+	innerExpr := expression.ComposeDNFCondition(ctx, isNullExprs...)
+	newExpr := expression.NewFunctionInternal(ctx, ast.If, types.NewFieldType(mysql.TypeLonglong), innerExpr, expression.Zero, expression.One)
+	return newExpr
+}
+
+// See https://dev.mysql.com/doc/refman/5.7/en/group-by-functions.html
+// The SUM() and AVG() functions return a DECIMAL value for exact-value arguments (integer or DECIMAL),
+// and a DOUBLE value for approximate-value arguments (FLOAT or DOUBLE).
+func (a *aggregationEliminateChecker) rewriteSumOrAvg(ctx sessionctx.Context, exprs []expression.Expression) expression.Expression {
+	// FIXME: Consider the case that avg is final mode.
+	expr := exprs[0]
+	switch expr.GetType().Tp {
+	// Integer type should be cast to decimal.
+	case mysql.TypeTiny, mysql.TypeShort, mysql.TypeInt24, mysql.TypeLong, mysql.TypeLonglong:
+		return expression.BuildCastFunction(ctx, expr, types.NewFieldType(mysql.TypeNewDecimal))
+		// Double and Decimal doesn't need to be cast.
+	case mysql.TypeDouble, mysql.TypeNewDecimal:
+		return expr
+		// Float should be cast to double. And other non-numeric type should be cast to double too.
+	default:
+		return expression.BuildCastFunction(ctx, expr, types.NewFieldType(mysql.TypeDouble))
+	}
+}
+
+func (a *aggregationRecursiveEliminater) optimize(p LogicalPlan) (LogicalPlan, error) {
+	newChildren := make([]LogicalPlan, 0, len(p.Children()))
+	for _, child := range p.Children() {
+		newChild, _ := a.optimize(child)
+		newChildren = append(newChildren, newChild)
+	}
+	p.SetChildren(newChildren...)
+	agg, ok := p.(*LogicalAggregation)
+	if !ok {
+		return p, nil
+	}
+	if proj := a.tryToEliminateAggregation(agg); proj != nil {
+		return proj, nil
+	}
+	return p, nil
+}