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
@@ -47,11 +56,7 @@ func (s *Schema) String() string {
 	}
 	ukStrs := make([]string, 0, len(s.Keys))
 	for _, key := range s.Keys {
-		ukColStrs := make([]string, 0, len(key))
-		for _, col := range key {
-			ukColStrs = append(ukColStrs, col.String())
-		}
-		ukStrs = append(ukStrs, "["+strings.Join(ukColStrs, ",")+"]")
+		ukStrs = append(ukStrs, key.String())
 	}
 	return "Column: [" + strings.Join(colStrs, ",") + "] Unique key: [" + strings.Join(ukStrs, ",") + "]"
 }

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,27 @@ 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])
+			appendAggregationEliminateTraceStep(agg, uniqueKey, opt)
 			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,28 +92,49 @@ 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(), "", "")
+					appendDistinctEliminateTraceStep(agg, uniqueKey, af, opt)
 				}
 			}
 		}
 	}
 }
 
+func appendAggregationEliminateTraceStep(agg *LogicalAggregation, uniqueKey expression.KeyInfo, opt *logicalOptimizeOp) {
+	if opt.tracer == nil {
+		return
+	}
+	opt.appendStepToCurrent(agg.ID(), agg.TP(),
+		fmt.Sprintf("%s is unique key", uniqueKey.String()),
+		"aggregation eliminated to projection")
+}
+
+func appendDistinctEliminateTraceStep(agg *LogicalAggregation, uniqueKey expression.KeyInfo, af *aggregation.AggFuncDesc,
+	opt *logicalOptimizeOp) {
+	if opt.tracer == nil {
+		return
+	}
+	opt.appendStepToCurrent(agg.ID(), agg.TP(),
+		fmt.Sprintf("%s is unique key", uniqueKey.String()),
+		fmt.Sprintf("agg[%s]'s distinct eliminated", af.Name))
+}
+
 // ConvertAggToProj convert aggregation to projection.
 func ConvertAggToProj(agg *LogicalAggregation, schema *expression.Schema) (bool, *LogicalProjection) {
 	proj := LogicalProjection{
@@ -196,7 +223,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