planner: implement aggregation eliminate optimize trace

expression/schema.go

@@ -30,6 +30,15 @@ func (ki KeyInfo) Clone() KeyInfo {
 	return result
 }
 
+// String implements fmt.Stringer interface.
+func (ki KeyInfo) String() string {
+	ukColStrs := make([]string, 0, len(ki))
+	for _, col := range ki {
+		ukColStrs = append(ukColStrs, col.String())
+	}
+	return "[" + strings.Join(ukColStrs, ",") + "]"
+}
+
 // Schema stands for the row schema and unique key information get from input.
 type Schema struct {
 	Columns []*Column

planner/core/explain.go

@@ -34,6 +34,7 @@ import (
 	"github.com/pingcap/tidb/types"
 	"github.com/pingcap/tidb/util/plancodec"
 	"github.com/pingcap/tidb/util/stringutil"
+	"github.com/pingcap/tidb/util/tracing"
 	"github.com/pingcap/tipb/go-tipb"
 )
 
@@ -890,11 +891,23 @@ func (p *LogicalAggregation) ExplainInfo() string {
 	return buffer.String()
 }
 
+func (p *LogicalAggregation) buildLogicalPlanTrace() *tracing.LogicalPlanTrace {
+	lpt := p.baseLogicalPlan.buildLogicalPlanTrace()
+	lpt.ExplainInfo = p.ExplainInfo()
+	return lpt
+}
+
 // ExplainInfo implements Plan interface.
 func (p *LogicalProjection) ExplainInfo() string {
 	return expression.ExplainExpressionList(p.Exprs, p.schema)
 }
 
+func (p *LogicalProjection) buildLogicalPlanTrace() *tracing.LogicalPlanTrace {
+	lpt := p.baseLogicalPlan.buildLogicalPlanTrace()
+	lpt.ExplainInfo = p.ExplainInfo()
+	return lpt
+}
+
 // ExplainInfo implements Plan interface.
 func (p *LogicalSelection) ExplainInfo() string {
 	return string(expression.SortedExplainExpressionList(p.Conditions))
@@ -930,6 +943,12 @@ func (ds *DataSource) ExplainInfo() string {
 	return buffer.String()
 }
 
+func (ds *DataSource) buildLogicalPlanTrace() *tracing.LogicalPlanTrace {
+	lpt := ds.baseLogicalPlan.buildLogicalPlanTrace()
+	lpt.ExplainInfo = ds.ExplainInfo()
+	return lpt
+}
+
 // ExplainInfo implements Plan interface.
 func (p *PhysicalExchangeSender) ExplainInfo() string {
 	buffer := bytes.NewBufferString("ExchangeType: ")

planner/core/logical_plan_test.go

@@ -2071,58 +2071,3 @@ func (s *testPlanSuite) TestWindowLogicalPlanAmbiguous(c *C) {
 		}
 	}
 }
-
-func (s *testPlanSuite) TestLogicalOptimizeWithTraceEnabled(c *C) {
-	sql := "select * from t where a in (1,2)"
-	defer testleak.AfterTest(c)()
-	tt := []struct {
-		flags []uint64
-		steps int
-	}{
-		{
-			flags: []uint64{
-				flagEliminateAgg,
-				flagPushDownAgg},
-			steps: 2,
-		},
-		{
-			flags: []uint64{
-				flagEliminateAgg,
-				flagPushDownAgg,
-				flagPrunColumns,
-				flagBuildKeyInfo,
-			},
-			steps: 4,
-		},
-		{
-			flags: []uint64{},
-			steps: 0,
-		},
-	}
-
-	for i, tc := range tt {
-		comment := Commentf("case:%v sql:%s", i, sql)
-		stmt, err := s.ParseOneStmt(sql, "", "")
-		c.Assert(err, IsNil, comment)
-		err = Preprocess(s.ctx, stmt, WithPreprocessorReturn(&PreprocessorReturn{InfoSchema: s.is}))
-		c.Assert(err, IsNil, comment)
-		sctx := MockContext()
-		sctx.GetSessionVars().EnableStmtOptimizeTrace = true
-		builder, _ := NewPlanBuilder().Init(sctx, s.is, &hint.BlockHintProcessor{})
-		domain.GetDomain(sctx).MockInfoCacheAndLoadInfoSchema(s.is)
-		ctx := context.TODO()
-		p, err := builder.Build(ctx, stmt)
-		c.Assert(err, IsNil)
-		flag := uint64(0)
-		for _, f := range tc.flags {
-			flag = flag | f
-		}
-		p, err = logicalOptimize(ctx, flag, p.(LogicalPlan))
-		c.Assert(err, IsNil)
-		_, ok := p.(*LogicalProjection)
-		c.Assert(ok, IsTrue)
-		otrace := sctx.GetSessionVars().StmtCtx.LogicalOptimizeTrace
-		c.Assert(otrace, NotNil)
-		c.Assert(len(otrace.Steps), Equals, tc.steps)
-	}
-}

planner/core/logical_plan_trace_test.go

@@ -0,0 +1,147 @@
+// Copyright 2021 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,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package core
+
+import (
+	"context"
+
+	. "github.com/pingcap/check"
+	"github.com/pingcap/tidb/domain"
+	"github.com/pingcap/tidb/util/hint"
+	"github.com/pingcap/tidb/util/testleak"
+)
+
+func (s *testPlanSuite) TestLogicalOptimizeWithTraceEnabled(c *C) {
+	sql := "select * from t where a in (1,2)"
+	defer testleak.AfterTest(c)()
+	tt := []struct {
+		flags []uint64
+		steps int
+	}{
+		{
+			flags: []uint64{
+				flagEliminateAgg,
+				flagPushDownAgg},
+			steps: 2,
+		},
+		{
+			flags: []uint64{
+				flagEliminateAgg,
+				flagPushDownAgg,
+				flagPrunColumns,
+				flagBuildKeyInfo,
+			},
+			steps: 4,
+		},
+		{
+			flags: []uint64{},
+			steps: 0,
+		},
+	}
+
+	for i, tc := range tt {
+		comment := Commentf("case:%v sql:%s", i, sql)
+		stmt, err := s.ParseOneStmt(sql, "", "")
+		c.Assert(err, IsNil, comment)
+		err = Preprocess(s.ctx, stmt, WithPreprocessorReturn(&PreprocessorReturn{InfoSchema: s.is}))
+		c.Assert(err, IsNil, comment)
+		sctx := MockContext()
+		sctx.GetSessionVars().EnableStmtOptimizeTrace = true
+		builder, _ := NewPlanBuilder().Init(sctx, s.is, &hint.BlockHintProcessor{})
+		domain.GetDomain(sctx).MockInfoCacheAndLoadInfoSchema(s.is)
+		ctx := context.TODO()
+		p, err := builder.Build(ctx, stmt)
+		c.Assert(err, IsNil)
+		flag := uint64(0)
+		for _, f := range tc.flags {
+			flag = flag | f
+		}
+		p, err = logicalOptimize(ctx, flag, p.(LogicalPlan))
+		c.Assert(err, IsNil)
+		_, ok := p.(*LogicalProjection)
+		c.Assert(ok, IsTrue)
+		otrace := sctx.GetSessionVars().StmtCtx.LogicalOptimizeTrace
+		c.Assert(otrace, NotNil)
+		c.Assert(len(otrace.Steps), Equals, tc.steps)
+	}
+}
+
+func (s *testPlanSuite) TestSingleRuleTraceStep(c *C) {
+	defer testleak.AfterTest(c)()
+	tt := []struct {
+		sql             string
+		flags           []uint64
+		assertRuleName  string
+		assertRuleSteps []assertTraceStep
+	}{
+		{
+			sql:            "select min(distinct a) from t group by a",
+			flags:          []uint64{flagBuildKeyInfo, flagEliminateAgg},
+			assertRuleName: "aggregation_eliminate",
+			assertRuleSteps: []assertTraceStep{
+				{
+					assertReason: ".*unique.*",
+					assertAction: ".*distinct eliminated.*",
+				},
+				{
+					assertReason: ".*unique.*",
+					assertAction: ".*aggregation eliminated.*",
+				},
+			},
+		},
+	}
+
+	for i, tc := range tt {
+		sql := tc.sql
+		comment := Commentf("case:%v sql:%s", i, sql)
+		stmt, err := s.ParseOneStmt(sql, "", "")
+		c.Assert(err, IsNil, comment)
+		err = Preprocess(s.ctx, stmt, WithPreprocessorReturn(&PreprocessorReturn{InfoSchema: s.is}))
+		c.Assert(err, IsNil, comment)
+		sctx := MockContext()
+		sctx.GetSessionVars().EnableStmtOptimizeTrace = true
+		builder, _ := NewPlanBuilder().Init(sctx, s.is, &hint.BlockHintProcessor{})
+		domain.GetDomain(sctx).MockInfoCacheAndLoadInfoSchema(s.is)
+		ctx := context.TODO()
+		p, err := builder.Build(ctx, stmt)
+		c.Assert(err, IsNil)
+		flag := uint64(0)
+		for _, f := range tc.flags {
+			flag = flag | f
+		}
+		p, err = logicalOptimize(ctx, flag, p.(LogicalPlan))
+		c.Assert(err, IsNil)
+		_, ok := p.(*LogicalProjection)
+		c.Assert(ok, IsTrue)
+		otrace := sctx.GetSessionVars().StmtCtx.LogicalOptimizeTrace
+		c.Assert(otrace, NotNil)
+		assert := false
+		for _, step := range otrace.Steps {
+			if step.RuleName == tc.assertRuleName {
+				assert = true
+				for i, ruleStep := range step.Steps {
+					c.Assert(ruleStep.Action, Matches, tc.assertRuleSteps[i].assertAction)
+					c.Assert(ruleStep.Reason, Matches, tc.assertRuleSteps[i].assertReason)
+				}
+			}
+		}
+		c.Assert(assert, IsTrue)
+	}
+}
+
+type assertTraceStep struct {
+	assertReason string
+	assertAction string
+}

planner/core/rule_aggregation_elimination.go

@@ -16,6 +16,7 @@ package core
 
 import (
 	"context"
+	"fmt"
 	"math"
 
 	"github.com/pingcap/tidb/expression"
@@ -37,7 +38,7 @@ type aggregationEliminateChecker struct {
 // 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 {
+func (a *aggregationEliminateChecker) tryToEliminateAggregation(agg *LogicalAggregation, opt *logicalOptimizeOp) *LogicalProjection {
 	for _, af := range agg.AggFuncs {
 		// TODO(issue #9968): Actually, we can rewrite GROUP_CONCAT when all the
 		// arguments it accepts are promised to be NOT-NULL.
@@ -54,22 +55,29 @@ func (a *aggregationEliminateChecker) tryToEliminateAggregation(agg *LogicalAggr
 	}
 	schemaByGroupby := expression.NewSchema(agg.GetGroupByCols()...)
 	coveredByUniqueKey := false
+	var uniqueKey expression.KeyInfo
 	for _, key := range agg.children[0].Schema().Keys {
 		if schemaByGroupby.ColumnsIndices(key) != nil {
 			coveredByUniqueKey = true
+			uniqueKey = key
 			break
 		}
 	}
 	if coveredByUniqueKey {
 		// GroupByCols has unique key, so this aggregation can be removed.
 		if ok, proj := ConvertAggToProj(agg, agg.schema); ok {
 			proj.SetChildren(agg.children[0])
+			opt.appendStepToCurrent(agg.ID(), agg.TP(),
+				fmt.Sprintf("%s is unique key", uniqueKey.String()),
+				"aggregation eliminated to projection")
 			return proj
 		}
 	}
 	return nil
 }
 
+// tryToEliminateDistinct will eliminate distinct in the aggregation function if the aggregation args
+// have unique key column. see detail example in https://github.com/pingcap/tidb/issues/23436
 func (a *aggregationEliminateChecker) tryToEliminateDistinct(agg *LogicalAggregation, opt *logicalOptimizeOp) {
 	for _, af := range agg.AggFuncs {
 		if af.HasDistinct {
@@ -86,22 +94,26 @@ func (a *aggregationEliminateChecker) tryToEliminateDistinct(agg *LogicalAggrega
 			if canEliminate {
 				distinctByUniqueKey := false
 				schemaByDistinct := expression.NewSchema(cols...)
+				var uniqueKey expression.KeyInfo
 				for _, key := range agg.children[0].Schema().Keys {
 					if schemaByDistinct.ColumnsIndices(key) != nil {
 						distinctByUniqueKey = true
+						uniqueKey = key
 						break
 					}
 				}
 				for _, key := range agg.children[0].Schema().UniqueKeys {
 					if schemaByDistinct.ColumnsIndices(key) != nil {
 						distinctByUniqueKey = true
+						uniqueKey = key
 						break
 					}
 				}
 				if distinctByUniqueKey {
 					af.HasDistinct = false
-					// TODO: fulfill in future pr
-					opt.appendStepToCurrent(agg.ID(), agg.TP(), "", "")
+					opt.appendStepToCurrent(agg.ID(), agg.TP(),
+						fmt.Sprintf("%s is unique key", uniqueKey.String()),
+						fmt.Sprintf("agg[%s]'s distinct eliminated", af.Name))
 				}
 			}
 		}
@@ -196,7 +208,7 @@ func (a *aggregationEliminator) optimize(ctx context.Context, p LogicalPlan, opt
 		return p, nil
 	}
 	a.tryToEliminateDistinct(agg, opt)
-	if proj := a.tryToEliminateAggregation(agg); proj != nil {
+	if proj := a.tryToEliminateAggregation(agg, opt); proj != nil {
 		return proj, nil
 	}
 	return p, nil

planner/core/rule_aggregation_push_down.go

@@ -368,7 +368,7 @@ func (a *aggregationPushDownSolver) pushAggCrossUnion(agg *LogicalAggregation, u
 }
 
 func (a *aggregationPushDownSolver) optimize(ctx context.Context, p LogicalPlan, opt *logicalOptimizeOp) (LogicalPlan, error) {
-	return a.aggPushDown(p)
+	return a.aggPushDown(p, opt)
 }
 
 func (a *aggregationPushDownSolver) tryAggPushDownForUnion(union *LogicalUnionAll, agg *LogicalAggregation) error {
@@ -395,9 +395,9 @@ func (a *aggregationPushDownSolver) tryAggPushDownForUnion(union *LogicalUnionAl
 }
 
 // aggPushDown tries to push down aggregate functions to join paths.
-func (a *aggregationPushDownSolver) aggPushDown(p LogicalPlan) (_ LogicalPlan, err error) {
+func (a *aggregationPushDownSolver) aggPushDown(p LogicalPlan, opt *logicalOptimizeOp) (_ LogicalPlan, err error) {
 	if agg, ok := p.(*LogicalAggregation); ok {
-		proj := a.tryToEliminateAggregation(agg)
+		proj := a.tryToEliminateAggregation(agg, opt)
 		if proj != nil {
 			p = proj
 		} else {
@@ -428,7 +428,7 @@ func (a *aggregationPushDownSolver) aggPushDown(p LogicalPlan) (_ LogicalPlan, e
 					join.SetChildren(lChild, rChild)
 					join.SetSchema(expression.MergeSchema(lChild.Schema(), rChild.Schema()))
 					buildKeyInfo(join)
-					proj := a.tryToEliminateAggregation(agg)
+					proj := a.tryToEliminateAggregation(agg, opt)
 					if proj != nil {
 						p = proj
 					}
@@ -487,7 +487,7 @@ func (a *aggregationPushDownSolver) aggPushDown(p LogicalPlan) (_ LogicalPlan, e
 	}
 	newChildren := make([]LogicalPlan, 0, len(p.Children()))
 	for _, child := range p.Children() {
-		newChild, err := a.aggPushDown(child)
+		newChild, err := a.aggPushDown(child, opt)
 		if err != nil {
 			return nil, err
 		}