planner: add optimizer trace framework for logicalOptimize (pingcap#2…

…9559)

executor/executor.go

@@ -1682,6 +1682,8 @@ func ResetContextOfStmt(ctx sessionctx.Context, s ast.StmtNode) (err error) {
 	sc.CTEStorageMap = map[int]*CTEStorages{}
 	sc.IsStaleness = false
 	sc.LockTableIDs = make(map[int64]struct{})
+	sc.EnableOptimizeTrace = false
+	sc.LogicalOptimizeTrace = nil
 
 	sc.InitMemTracker(memory.LabelForSQLText, vars.MemQuotaQuery)
 	sc.InitDiskTracker(memory.LabelForSQLText, -1)

planner/core/logical_plan_test.go

@@ -2071,3 +2071,58 @@ 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().StmtCtx.EnableOptimizeTrace = 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/optimizer.go

@@ -35,6 +35,7 @@ import (
 	"github.com/pingcap/tidb/types"
 	utilhint "github.com/pingcap/tidb/util/hint"
 	"github.com/pingcap/tidb/util/set"
+	"github.com/pingcap/tidb/util/tracing"
 	"go.uber.org/atomic"
 )
 
@@ -83,9 +84,44 @@ var optRuleList = []logicalOptRule{
 	&columnPruner{}, // column pruning again at last, note it will mess up the results of buildKeySolver
 }
 
+type logicalOptimizeOp struct {
+	// tracer is goring to track optimize steps during rule optimizing
+	tracer *tracing.LogicalOptimizeTracer
+}
+
+func defaultLogicalOptimizeOption() *logicalOptimizeOp {
+	return &logicalOptimizeOp{}
+}
+
+func (op *logicalOptimizeOp) withEnableOptimizeTracer(tracer *tracing.LogicalOptimizeTracer) *logicalOptimizeOp {
+	op.tracer = tracer
+	return op
+}
+
+func (op *logicalOptimizeOp) appendBeforeRuleOptimize(name string, before LogicalPlan) {
+	if op.tracer == nil {
+		return
+	}
+	op.tracer.AppendRuleTracerBeforeRuleOptimize(name, before.buildLogicalPlanTrace())
+}
+
+func (op *logicalOptimizeOp) appendStepToCurrent(id int, tp, reason, action string) {
+	if op.tracer == nil {
+		return
+	}
+	op.tracer.AppendRuleTracerStepToCurrent(id, tp, reason, action)
+}
+
+func (op *logicalOptimizeOp) trackAfterRuleOptimize(after LogicalPlan) {
+	if op.tracer == nil {
+		return
+	}
+	op.tracer.TrackLogicalPlanAfterRuleOptimize(after.buildLogicalPlanTrace())
+}
+
 // logicalOptRule means a logical optimizing rule, which contains decorrelate, ppd, column pruning, etc.
 type logicalOptRule interface {
-	optimize(context.Context, LogicalPlan) (LogicalPlan, error)
+	optimize(context.Context, LogicalPlan, *logicalOptimizeOp) (LogicalPlan, error)
 	name() string
 }
 
@@ -335,6 +371,15 @@ func enableParallelApply(sctx sessionctx.Context, plan PhysicalPlan) PhysicalPla
 }
 
 func logicalOptimize(ctx context.Context, flag uint64, logic LogicalPlan) (LogicalPlan, error) {
+	opt := defaultLogicalOptimizeOption()
+	stmtCtx := logic.SCtx().GetSessionVars().StmtCtx
+	if stmtCtx.EnableOptimizeTrace {
+		tracer := &tracing.LogicalOptimizeTracer{Steps: make([]*tracing.LogicalRuleOptimizeTracer, 0)}
+		opt = opt.withEnableOptimizeTracer(tracer)
+		defer func() {
+			stmtCtx.LogicalOptimizeTrace = tracer
+		}()
+	}
 	var err error
 	for i, rule := range optRuleList {
 		// The order of flags is same as the order of optRule in the list.
@@ -343,10 +388,12 @@ func logicalOptimize(ctx context.Context, flag uint64, logic LogicalPlan) (Logic
 		if flag&(1<<uint(i)) == 0 || isLogicalRuleDisabled(rule) {
 			continue
 		}
-		logic, err = rule.optimize(ctx, logic)
+		opt.appendBeforeRuleOptimize(rule.name(), logic)
+		logic, err = rule.optimize(ctx, logic, opt)
 		if err != nil {
 			return nil, err
 		}
+		opt.trackAfterRuleOptimize(logic)
 	}
 	return logic, err
 }

planner/core/plan.go

@@ -28,6 +28,7 @@ import (
 	"github.com/pingcap/tidb/sessionctx"
 	"github.com/pingcap/tidb/types"
 	"github.com/pingcap/tidb/util/stringutil"
+	"github.com/pingcap/tidb/util/tracing"
 	"github.com/pingcap/tipb/go-tipb"
 )
 
@@ -305,6 +306,9 @@ type LogicalPlan interface {
 
 	// canPushToCop check if we might push this plan to a specific store.
 	canPushToCop(store kv.StoreType) bool
+
+	// buildLogicalPlanTrace clone necessary information from LogicalPlan
+	buildLogicalPlanTrace() *tracing.LogicalPlanTrace
 }
 
 // PhysicalPlan is a tree of the physical operators.
@@ -377,6 +381,15 @@ func (p *baseLogicalPlan) ExplainInfo() string {
 	return ""
 }
 
+// buildLogicalPlanTrace implements LogicalPlan
+func (p *baseLogicalPlan) buildLogicalPlanTrace() *tracing.LogicalPlanTrace {
+	planTrace := &tracing.LogicalPlanTrace{ID: p.ID(), TP: p.TP()}
+	for _, child := range p.Children() {
+		planTrace.Children = append(planTrace.Children, child.buildLogicalPlanTrace())
+	}
+	return planTrace
+}
+
 type basePhysicalPlan struct {
 	basePlan
 

planner/core/rule_aggregation_elimination.go

@@ -70,7 +70,7 @@ func (a *aggregationEliminateChecker) tryToEliminateAggregation(agg *LogicalAggr
 	return nil
 }
 
-func (a *aggregationEliminateChecker) tryToEliminateDistinct(agg *LogicalAggregation) {
+func (a *aggregationEliminateChecker) tryToEliminateDistinct(agg *LogicalAggregation, opt *logicalOptimizeOp) {
 	for _, af := range agg.AggFuncs {
 		if af.HasDistinct {
 			cols := make([]*expression.Column, 0, len(af.Args))
@@ -100,6 +100,8 @@ func (a *aggregationEliminateChecker) tryToEliminateDistinct(agg *LogicalAggrega
 				}
 				if distinctByUniqueKey {
 					af.HasDistinct = false
+					// TODO: fulfill in future pr
+					opt.appendStepToCurrent(agg.ID(), agg.TP(), "", "")
 				}
 			}
 		}
@@ -179,10 +181,10 @@ func wrapCastFunction(ctx sessionctx.Context, arg expression.Expression, targetT
 	return expression.BuildCastFunction(ctx, arg, targetTp)
 }
 
-func (a *aggregationEliminator) optimize(ctx context.Context, p LogicalPlan) (LogicalPlan, error) {
+func (a *aggregationEliminator) optimize(ctx context.Context, p LogicalPlan, opt *logicalOptimizeOp) (LogicalPlan, error) {
 	newChildren := make([]LogicalPlan, 0, len(p.Children()))
 	for _, child := range p.Children() {
-		newChild, err := a.optimize(ctx, child)
+		newChild, err := a.optimize(ctx, child, opt)
 		if err != nil {
 			return nil, err
 		}
@@ -193,7 +195,7 @@ func (a *aggregationEliminator) optimize(ctx context.Context, p LogicalPlan) (Lo
 	if !ok {
 		return p, nil
 	}
-	a.tryToEliminateDistinct(agg)
+	a.tryToEliminateDistinct(agg, opt)
 	if proj := a.tryToEliminateAggregation(agg); proj != nil {
 		return proj, nil
 	}

planner/core/rule_aggregation_push_down.go

@@ -367,7 +367,7 @@ func (a *aggregationPushDownSolver) pushAggCrossUnion(agg *LogicalAggregation, u
 	return newAgg, nil
 }
 
-func (a *aggregationPushDownSolver) optimize(ctx context.Context, p LogicalPlan) (LogicalPlan, error) {
+func (a *aggregationPushDownSolver) optimize(ctx context.Context, p LogicalPlan, opt *logicalOptimizeOp) (LogicalPlan, error) {
 	return a.aggPushDown(p)
 }
 

planner/core/rule_build_key_info.go

@@ -25,9 +25,9 @@ import (
 
 type buildKeySolver struct{}
 
-func (s *buildKeySolver) optimize(ctx context.Context, lp LogicalPlan) (LogicalPlan, error) {
-	buildKeyInfo(lp)
-	return lp, nil
+func (s *buildKeySolver) optimize(ctx context.Context, p LogicalPlan, opt *logicalOptimizeOp) (LogicalPlan, error) {
+	buildKeyInfo(p)
+	return p, nil
 }
 
 // buildKeyInfo recursively calls LogicalPlan's BuildKeyInfo method.

planner/core/rule_column_pruning.go

@@ -29,7 +29,7 @@ import (
 type columnPruner struct {
 }
 
-func (s *columnPruner) optimize(ctx context.Context, lp LogicalPlan) (LogicalPlan, error) {
+func (s *columnPruner) optimize(ctx context.Context, lp LogicalPlan, opt *logicalOptimizeOp) (LogicalPlan, error) {
 	err := lp.PruneColumns(lp.Schema().Columns)
 	return lp, err
 }

planner/core/rule_decorrelate.go

@@ -119,7 +119,7 @@ func (s *decorrelateSolver) aggDefaultValueMap(agg *LogicalAggregation) map[int]
 }
 
 // optimize implements logicalOptRule interface.
-func (s *decorrelateSolver) optimize(ctx context.Context, p LogicalPlan) (LogicalPlan, error) {
+func (s *decorrelateSolver) optimize(ctx context.Context, p LogicalPlan, opt *logicalOptimizeOp) (LogicalPlan, error) {
 	if apply, ok := p.(*LogicalApply); ok {
 		outerPlan := apply.children[0]
 		innerPlan := apply.children[1]
@@ -139,12 +139,12 @@ func (s *decorrelateSolver) optimize(ctx context.Context, p LogicalPlan) (Logica
 			apply.AttachOnConds(newConds)
 			innerPlan = sel.children[0]
 			apply.SetChildren(outerPlan, innerPlan)
-			return s.optimize(ctx, p)
+			return s.optimize(ctx, p, opt)
 		} else if m, ok := innerPlan.(*LogicalMaxOneRow); ok {
 			if m.children[0].MaxOneRow() {
 				innerPlan = m.children[0]
 				apply.SetChildren(outerPlan, innerPlan)
-				return s.optimize(ctx, p)
+				return s.optimize(ctx, p, opt)
 			}
 		} else if proj, ok := innerPlan.(*LogicalProjection); ok {
 			for i, expr := range proj.Exprs {
@@ -157,14 +157,14 @@ func (s *decorrelateSolver) optimize(ctx context.Context, p LogicalPlan) (Logica
 				proj.SetSchema(apply.Schema())
 				proj.Exprs = append(expression.Column2Exprs(outerPlan.Schema().Clone().Columns), proj.Exprs...)
 				apply.SetSchema(expression.MergeSchema(outerPlan.Schema(), innerPlan.Schema()))
-				np, err := s.optimize(ctx, p)
+				np, err := s.optimize(ctx, p, opt)
 				if err != nil {
 					return nil, err
 				}
 				proj.SetChildren(np)
 				return proj, nil
 			}
-			return s.optimize(ctx, p)
+			return s.optimize(ctx, p, opt)
 		} else if agg, ok := innerPlan.(*LogicalAggregation); ok {
 			if apply.canPullUpAgg() && agg.canPullUp() {
 				innerPlan = agg.children[0]
@@ -214,7 +214,7 @@ func (s *decorrelateSolver) optimize(ctx context.Context, p LogicalPlan) (Logica
 					newAggFuncs = append(newAggFuncs, desc)
 				}
 				agg.AggFuncs = newAggFuncs
-				np, err := s.optimize(ctx, p)
+				np, err := s.optimize(ctx, p, opt)
 				if err != nil {
 					return nil, err
 				}
@@ -283,7 +283,7 @@ func (s *decorrelateSolver) optimize(ctx context.Context, p LogicalPlan) (Logica
 							proj.SetChildren(apply)
 							p = proj
 						}
-						return s.optimize(ctx, p)
+						return s.optimize(ctx, p, opt)
 					}
 					sel.Conditions = originalExpr
 					apply.CorCols = extractCorColumnsBySchema4LogicalPlan(apply.children[1], apply.children[0].Schema())
@@ -294,12 +294,12 @@ func (s *decorrelateSolver) optimize(ctx context.Context, p LogicalPlan) (Logica
 			// the top level Sort has no effect on the subquery's result.
 			innerPlan = sort.children[0]
 			apply.SetChildren(outerPlan, innerPlan)
-			return s.optimize(ctx, p)
+			return s.optimize(ctx, p, opt)
 		}
 	}
 	newChildren := make([]LogicalPlan, 0, len(p.Children()))
 	for _, child := range p.Children() {
-		np, err := s.optimize(ctx, child)
+		np, err := s.optimize(ctx, child, opt)
 		if err != nil {
 			return nil, err
 		}

planner/core/rule_eliminate_projection.go

@@ -145,7 +145,7 @@ type projectionEliminator struct {
 }
 
 // optimize implements the logicalOptRule interface.
-func (pe *projectionEliminator) optimize(ctx context.Context, lp LogicalPlan) (LogicalPlan, error) {
+func (pe *projectionEliminator) optimize(ctx context.Context, lp LogicalPlan, opt *logicalOptimizeOp) (LogicalPlan, error) {
 	root := pe.eliminate(lp, make(map[string]*expression.Column), false)
 	return root, nil
 }

planner/core/rule_generate_column_substitute.go

@@ -37,7 +37,7 @@ type ExprColumnMap map[expression.Expression]*expression.Column
 // For example: select a+1 from t order by a+1, with a virtual generate column c as (a+1) and
 // an index on c. We need to replace a+1 with c so that we can use the index on c.
 // See also https://dev.mysql.com/doc/refman/8.0/en/generated-column-index-optimizations.html
-func (gc *gcSubstituter) optimize(ctx context.Context, lp LogicalPlan) (LogicalPlan, error) {
+func (gc *gcSubstituter) optimize(ctx context.Context, lp LogicalPlan, opt *logicalOptimizeOp) (LogicalPlan, error) {
 	exprToColumn := make(ExprColumnMap)
 	collectGenerateColumn(lp, exprToColumn)
 	if len(exprToColumn) == 0 {

planner/core/rule_join_elimination.go

@@ -225,8 +225,9 @@ func (o *outerJoinEliminator) doOptimize(p LogicalPlan, aggCols []*expression.Co
 	return p, nil
 }
 
-func (o *outerJoinEliminator) optimize(ctx context.Context, p LogicalPlan) (LogicalPlan, error) {
-	return o.doOptimize(p, nil, nil)
+func (o *outerJoinEliminator) optimize(ctx context.Context, p LogicalPlan, opt *logicalOptimizeOp) (LogicalPlan, error) {
+	p, err := o.doOptimize(p, nil, nil)
+	return p, err
 }
 
 func (*outerJoinEliminator) name() string {

planner/core/rule_join_reorder.go

@@ -55,7 +55,7 @@ type jrNode struct {
 	cumCost float64
 }
 
-func (s *joinReOrderSolver) optimize(ctx context.Context, p LogicalPlan) (LogicalPlan, error) {
+func (s *joinReOrderSolver) optimize(ctx context.Context, p LogicalPlan, opt *logicalOptimizeOp) (LogicalPlan, error) {
 	return s.optimizeRecursive(p.SCtx(), p)
 }
 

planner/core/rule_max_min_eliminate.go

@@ -34,7 +34,7 @@ import (
 type maxMinEliminator struct {
 }
 
-func (a *maxMinEliminator) optimize(ctx context.Context, p LogicalPlan) (LogicalPlan, error) {
+func (a *maxMinEliminator) optimize(ctx context.Context, p LogicalPlan, opt *logicalOptimizeOp) (LogicalPlan, error) {
 	return a.eliminateMaxMin(p), nil
 }
 

planner/core/rule_partition_processor.go

@@ -57,8 +57,9 @@ const FullRange = -1
 // partitionProcessor is here because it's easier to prune partition after predicate push down.
 type partitionProcessor struct{}
 
-func (s *partitionProcessor) optimize(ctx context.Context, lp LogicalPlan) (LogicalPlan, error) {
-	return s.rewriteDataSource(lp)
+func (s *partitionProcessor) optimize(ctx context.Context, lp LogicalPlan, opt *logicalOptimizeOp) (LogicalPlan, error) {
+	p, err := s.rewriteDataSource(lp)
+	return p, err
 }
 
 func (s *partitionProcessor) rewriteDataSource(lp LogicalPlan) (LogicalPlan, error) {

planner/core/rule_predicate_push_down.go

@@ -27,7 +27,7 @@ import (
 
 type ppdSolver struct{}
 
-func (s *ppdSolver) optimize(ctx context.Context, lp LogicalPlan) (LogicalPlan, error) {
+func (s *ppdSolver) optimize(ctx context.Context, lp LogicalPlan, opt *logicalOptimizeOp) (LogicalPlan, error) {
 	_, p := lp.PredicatePushDown(nil)
 	return p, nil
 }