Skip to content

Commit

Permalink
planner: let optimizer trace support gcSubstitute (#40815)
Browse files Browse the repository at this point in the history 
close #40662
  • Loading branch information
@Yisaer
Yisaer authored Jan 30, 2023
1 parent 28d0325 commit 4ed4562
Showing 1 changed file with 35 additions and 20 deletions.
55 changes: 35 additions & 20 deletions planner/core/rule_generate_column_substitute.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,12 +15,11 @@
package core

import (
"bytes"
"context"

"github.com/pingcap/tidb/expression"
"github.com/pingcap/tidb/parser/ast"
"github.com/pingcap/tidb/sessionctx"
"github.com/pingcap/tidb/sessionctx/stmtctx"
"github.com/pingcap/tidb/types"
)

Expand All @@ -37,13 +36,13 @@ 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, _ *logicalOptimizeOp) (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 {
return lp, nil
}
return gc.substitute(ctx, lp, exprToColumn), nil
return gc.substitute(ctx, lp, exprToColumn, opt), nil
}

// collectGenerateColumn collect the generate column and save them to a map from their expressions to themselves.
Expand Down Expand Up @@ -74,18 +73,33 @@ func collectGenerateColumn(lp LogicalPlan, exprToColumn ExprColumnMap) {
}
}

func tryToSubstituteExpr(expr *expression.Expression, sctx sessionctx.Context, candidateExpr expression.Expression, tp types.EvalType, schema *expression.Schema, col *expression.Column) {
if (*expr).Equal(sctx, candidateExpr) && candidateExpr.GetType().EvalType() == tp &&
func tryToSubstituteExpr(expr *expression.Expression, lp LogicalPlan, candidateExpr expression.Expression, tp types.EvalType, schema *expression.Schema, col *expression.Column, opt *logicalOptimizeOp) {
if (*expr).Equal(lp.SCtx(), candidateExpr) && candidateExpr.GetType().EvalType() == tp &&
schema.ColumnIndex(col) != -1 {
*expr = col
appendSubstituteColumnStep(lp, candidateExpr, col, opt)
}
}

func substituteExpression(cond expression.Expression, sctx *stmtctx.StatementContext, sessionCtx sessionctx.Context, exprToColumn ExprColumnMap, schema *expression.Schema) {
func appendSubstituteColumnStep(lp LogicalPlan, candidateExpr expression.Expression, col *expression.Column, opt *logicalOptimizeOp) {
reason := func() string { return "" }
action := func() string {
buffer := bytes.NewBufferString("expression:")
buffer.WriteString(candidateExpr.String())
buffer.WriteString(" substituted by")
buffer.WriteString(" column:")
buffer.WriteString(col.String())
return buffer.String()
}
opt.appendStepToCurrent(lp.ID(), lp.TP(), reason, action)
}

func substituteExpression(cond expression.Expression, lp LogicalPlan, exprToColumn ExprColumnMap, schema *expression.Schema, opt *logicalOptimizeOp) {
sf, ok := cond.(*expression.ScalarFunction)
if !ok {
return
}
sctx := lp.SCtx().GetSessionVars().StmtCtx
defer func() {
// If the argument is not changed, hash code doesn't need to recount again.
// But we always do it to keep the code simple and stupid.
Expand All @@ -96,10 +110,10 @@ func substituteExpression(cond expression.Expression, sctx *stmtctx.StatementCon
switch sf.FuncName.L {
case ast.EQ, ast.LT, ast.LE, ast.GT, ast.GE:
for candidateExpr, column := range exprToColumn {
tryToSubstituteExpr(&sf.GetArgs()[1], sessionCtx, candidateExpr, sf.GetArgs()[0].GetType().EvalType(), schema, column)
tryToSubstituteExpr(&sf.GetArgs()[1], lp, candidateExpr, sf.GetArgs()[0].GetType().EvalType(), schema, column, opt)
}
for candidateExpr, column := range exprToColumn {
tryToSubstituteExpr(&sf.GetArgs()[0], sessionCtx, candidateExpr, sf.GetArgs()[1].GetType().EvalType(), schema, column)
tryToSubstituteExpr(&sf.GetArgs()[0], lp, candidateExpr, sf.GetArgs()[1].GetType().EvalType(), schema, column, opt)
}
case ast.In:
expr = &sf.GetArgs()[0]
Expand All @@ -115,43 +129,42 @@ func substituteExpression(cond expression.Expression, sctx *stmtctx.StatementCon
}
if canSubstitute {
for candidateExpr, column := range exprToColumn {
tryToSubstituteExpr(expr, sessionCtx, candidateExpr, tp, schema, column)
tryToSubstituteExpr(expr, lp, candidateExpr, tp, schema, column, opt)
}
}
case ast.Like:
expr = &sf.GetArgs()[0]
tp = sf.GetArgs()[1].GetType().EvalType()
for candidateExpr, column := range exprToColumn {
tryToSubstituteExpr(expr, sessionCtx, candidateExpr, tp, schema, column)
tryToSubstituteExpr(expr, lp, candidateExpr, tp, schema, column, opt)
}
case ast.LogicOr, ast.LogicAnd:
substituteExpression(sf.GetArgs()[0], sctx, sessionCtx, exprToColumn, schema)
substituteExpression(sf.GetArgs()[1], sctx, sessionCtx, exprToColumn, schema)
substituteExpression(sf.GetArgs()[0], lp, exprToColumn, schema, opt)
substituteExpression(sf.GetArgs()[1], lp, exprToColumn, schema, opt)
case ast.UnaryNot:
substituteExpression(sf.GetArgs()[0], sctx, sessionCtx, exprToColumn, schema)
substituteExpression(sf.GetArgs()[0], lp, exprToColumn, schema, opt)
}
}

func (gc *gcSubstituter) substitute(ctx context.Context, lp LogicalPlan, exprToColumn ExprColumnMap) LogicalPlan {
sctx := lp.SCtx().GetSessionVars().StmtCtx
func (gc *gcSubstituter) substitute(ctx context.Context, lp LogicalPlan, exprToColumn ExprColumnMap, opt *logicalOptimizeOp) LogicalPlan {
var tp types.EvalType
switch x := lp.(type) {
case *LogicalSelection:
for _, cond := range x.Conditions {
substituteExpression(cond, sctx, lp.SCtx(), exprToColumn, x.Schema())
substituteExpression(cond, lp, exprToColumn, x.Schema(), opt)
}
case *LogicalProjection:
for i := range x.Exprs {
tp = x.Exprs[i].GetType().EvalType()
for candidateExpr, column := range exprToColumn {
tryToSubstituteExpr(&x.Exprs[i], lp.SCtx(), candidateExpr, tp, x.children[0].Schema(), column)
tryToSubstituteExpr(&x.Exprs[i], lp, candidateExpr, tp, x.children[0].Schema(), column, opt)
}
}
case *LogicalSort:
for i := range x.ByItems {
tp = x.ByItems[i].Expr.GetType().EvalType()
for candidateExpr, column := range exprToColumn {
tryToSubstituteExpr(&x.ByItems[i].Expr, lp.SCtx(), candidateExpr, tp, x.Schema(), column)
tryToSubstituteExpr(&x.ByItems[i].Expr, lp, candidateExpr, tp, x.Schema(), column, opt)
}
}
case *LogicalAggregation:
Expand All @@ -162,6 +175,7 @@ func (gc *gcSubstituter) substitute(ctx context.Context, lp LogicalPlan, exprToC
if aggFunc.Args[i].Equal(lp.SCtx(), candidateExpr) && candidateExpr.GetType().EvalType() == tp &&
x.Schema().ColumnIndex(column) != -1 {
aggFunc.Args[i] = column
appendSubstituteColumnStep(lp, candidateExpr, column, opt)
}
}
}
Expand All @@ -172,12 +186,13 @@ func (gc *gcSubstituter) substitute(ctx context.Context, lp LogicalPlan, exprToC
if x.GroupByItems[i].Equal(lp.SCtx(), candidateExpr) && candidateExpr.GetType().EvalType() == tp &&
x.Schema().ColumnIndex(column) != -1 {
x.GroupByItems[i] = column
appendSubstituteColumnStep(lp, candidateExpr, column, opt)
}
}
}
}
for _, child := range lp.Children() {
gc.substitute(ctx, child, exprToColumn)
gc.substitute(ctx, child, exprToColumn, opt)
}
return lp
}
Expand Down

0 comments on commit 4ed4562

Please sign in to comment.