expression,planner: reuse the PropConstSolver to improve performance

pkg/expression/constant_propagation.go

@@ -15,6 +15,9 @@
 package expression
 
 import (
+	"slices"
+	"sync"
+
 	"github.com/pingcap/tidb/pkg/expression/exprctx"
 	"github.com/pingcap/tidb/pkg/parser/ast"
 	"github.com/pingcap/tidb/pkg/parser/mysql"
@@ -23,14 +26,14 @@ import (
 	"github.com/pingcap/tidb/pkg/util/chunk"
 	"github.com/pingcap/tidb/pkg/util/collate"
 	"github.com/pingcap/tidb/pkg/util/disjointset"
+	"github.com/pingcap/tidb/pkg/util/intest"
 	"github.com/pingcap/tidb/pkg/util/logutil"
 	"go.uber.org/zap"
 )
 
 // MaxPropagateColsCnt means the max number of columns that can participate propagation.
 var MaxPropagateColsCnt = 100
 
-// nolint:structcheck
 type basePropConstSolver struct {
 	colMapper map[int64]int             // colMapper maps column to its index
 	eqList    []*Constant               // if eqList[i] != nil, it means col_i = eqList[i]
@@ -39,15 +42,34 @@ type basePropConstSolver struct {
 	ctx       exprctx.ExprContext
 }
 
+func newBasePropConstSolver() basePropConstSolver {
+	return basePropConstSolver{
+		colMapper: make(map[int64]int, 4),
+		eqList:    make([]*Constant, 0, 4),
+		columns:   make([]*Column, 0, 4),
+		unionSet:  disjointset.NewIntSet(4),
+	}
+}
+
+func (s *basePropConstSolver) Clear() {
+	clear(s.colMapper)
+	s.eqList = s.eqList[:0]
+	s.columns = s.columns[:0]
+	s.unionSet.Clear()
+	s.ctx = nil
+}
+
 func (s *basePropConstSolver) getColID(col *Column) int {
 	return s.colMapper[col.UniqueID]
 }
 
-func (s *basePropConstSolver) insertCol(col *Column) {
-	_, ok := s.colMapper[col.UniqueID]
-	if !ok {
-		s.colMapper[col.UniqueID] = len(s.colMapper)
-		s.columns = append(s.columns, col)
+func (s *basePropConstSolver) insertCols(cols ...*Column) {
+	for _, col := range cols {
+		_, ok := s.colMapper[col.UniqueID]
+		if !ok {
+			s.colMapper[col.UniqueID] = len(s.colMapper)
+			s.columns = append(s.columns, col)
+		}
 	}
 }
 
@@ -208,26 +230,60 @@ func tryToReplaceCond(ctx BuildContext, src *Column, tgt *Column, cond Expressio
 	return false, false, cond
 }
 
+var propConstSolverPool = sync.Pool{
+	New: func() any {
+		solver := &propConstSolver{
+			basePropConstSolver: newBasePropConstSolver(),
+			conditions:          make([]Expression, 0, 4),
+		}
+		return solver
+	},
+}
+
 type propConstSolver struct {
 	basePropConstSolver
 	conditions []Expression
 }
 
+// newPropConstSolver returns a PropagateConstantSolver.
+func newPropConstSolver() PropagateConstantSolver {
+	solver := propConstSolverPool.Get().(*propConstSolver)
+	return solver
+}
+
+// PropagateConstant propagate constant values of deterministic predicates in a condition.
+func (s *propConstSolver) PropagateConstant(ctx exprctx.ExprContext, conditions []Expression) []Expression {
+	s.ctx = ctx
+	return s.solve(conditions)
+}
+
+func (s *propConstSolver) Clear() {
+	s.basePropConstSolver.Clear()
+	s.conditions = s.conditions[:0]
+	propConstSolverPool.Put(s)
+}
+
 // propagateConstantEQ propagates expressions like 'column = constant' by substituting the constant for column, the
 // procedure repeats multiple times. An example runs as following:
 // a = d & b * 2 = c & c = d + 2 & b = 1 & a = 4, we pick eq cond b = 1 and a = 4
 // d = 4 & 2 = c & c = d + 2 & b = 1 & a = 4, we propagate b = 1 and a = 4 and pick eq cond c = 2 and d = 4
 // d = 4 & 2 = c & false & b = 1 & a = 4, we propagate c = 2 and d = 4, and do constant folding: c = d + 2 will be folded as false.
 func (s *propConstSolver) propagateConstantEQ() {
-	s.eqList = make([]*Constant, len(s.columns))
+	intest.Assert(len(s.eqList) == 0)
+	s.eqList = slices.Grow(s.eqList, len(s.columns))
+	for range s.columns {
+		s.eqList = append(s.eqList, nil)
+	}
 	visited := make([]bool, len(s.conditions))
+	cols := make([]*Column, 0, 4)
+	cons := make([]Expression, 0, 4)
 	for range MaxPropagateColsCnt {
 		mapper := s.pickNewEQConds(visited)
 		if len(mapper) == 0 {
 			return
 		}
-		cols := make([]*Column, 0, len(mapper))
-		cons := make([]Expression, 0, len(mapper))
+		cols = slices.Grow(cols, len(mapper))
+		cons = slices.Grow(cons, len(mapper))
 		for id, con := range mapper {
 			cols = append(cols, s.columns[id])
 			cons = append(cons, con)
@@ -237,6 +293,8 @@ func (s *propConstSolver) propagateConstantEQ() {
 				s.conditions[i] = ColumnSubstitute(s.ctx, cond, NewSchema(cols...), cons)
 			}
 		}
+		cols = cols[:0]
+		cons = cons[:0]
 	}
 }
 
@@ -260,7 +318,11 @@ func (s *propConstSolver) propagateConstantEQ() {
 // We maintain a unionSet representing the equivalent for every two columns.
 func (s *propConstSolver) propagateColumnEQ() {
 	visited := make([]bool, len(s.conditions))
-	s.unionSet = disjointset.NewIntSet(len(s.columns))
+	if s.unionSet == nil {
+		s.unionSet = disjointset.NewIntSet(len(s.columns))
+	} else {
+		s.unionSet.GrowNewIntSet(len(s.columns))
+	}
 	for i := range s.conditions {
 		if fun, ok := s.conditions[i].(*ScalarFunction); ok && fun.FuncName.L == ast.EQ {
 			lCol, lOk := fun.GetArgs()[0].(*Column)
@@ -306,10 +368,11 @@ func (s *propConstSolver) setConds2ConstFalse() {
 	if MaybeOverOptimized4PlanCache(s.ctx, s.conditions...) {
 		s.ctx.SetSkipPlanCache("some parameters may be overwritten when constant propagation")
 	}
-	s.conditions = []Expression{&Constant{
+	s.conditions = s.conditions[:0]
+	s.conditions = append(s.conditions, &Constant{
 		Value:   types.NewDatum(false),
 		RetType: types.NewFieldType(mysql.TypeTiny),
-	}}
+	})
 }
 
 // pickNewEQConds tries to pick new equal conds and puts them to retMapper.
@@ -357,13 +420,10 @@ func (s *propConstSolver) pickNewEQConds(visited []bool) (retMapper map[int]*Con
 }
 
 func (s *propConstSolver) solve(conditions []Expression) []Expression {
-	cols := make([]*Column, 0, len(conditions))
+	s.conditions = slices.Grow(s.conditions, len(conditions))
 	for _, cond := range conditions {
 		s.conditions = append(s.conditions, SplitCNFItems(cond)...)
-		cols = append(cols, ExtractColumns(cond)...)
-	}
-	for _, col := range cols {
-		s.insertCol(col)
+		s.insertCols(ExtractColumns(cond)...)
 	}
 	if len(s.columns) > MaxPropagateColsCnt {
 		logutil.BgLogger().Warn("too many columns in a single CNF",
@@ -374,15 +434,22 @@ func (s *propConstSolver) solve(conditions []Expression) []Expression {
 	}
 	s.propagateConstantEQ()
 	s.propagateColumnEQ()
-	s.conditions = propagateConstantDNF(s.ctx, s.conditions)
+	s.conditions = propagateConstantDNF(s.ctx, s.conditions...)
 	s.conditions = RemoveDupExprs(s.conditions)
-	return s.conditions
+	return slices.Clone(s.conditions)
 }
 
 // PropagateConstant propagate constant values of deterministic predicates in a condition.
 // This is a constant propagation logic for expression list such as ['a=1', 'a=b']
-func PropagateConstant(ctx exprctx.ExprContext, conditions []Expression) []Expression {
-	return newPropConstSolver().PropagateConstant(exprctx.WithConstantPropagateCheck(ctx), conditions)
+func PropagateConstant(ctx exprctx.ExprContext, conditions ...Expression) []Expression {
+	if len(conditions) == 0 {
+		return conditions
+	}
+	solver := newPropConstSolver()
+	defer func() {
+		solver.Clear()
+	}()
+	return solver.PropagateConstant(exprctx.WithConstantPropagateCheck(ctx), conditions)
 }
 
 type propOuterJoinConstSolver struct {
@@ -613,7 +680,11 @@ func (s *propOuterJoinConstSolver) propagateColumnEQ() {
 		return
 	}
 	visited := make([]bool, 2*len(s.joinConds)+len(s.filterConds))
-	s.unionSet = disjointset.NewIntSet(len(s.columns))
+	if s.unionSet == nil {
+		s.unionSet = disjointset.NewIntSet(len(s.columns))
+	} else {
+		s.unionSet.GrowNewIntSet(len(s.columns))
+	}
 	var outerCol, innerCol *Column
 	// Only consider column equal condition in joinConds.
 	// If we have column equal in filter condition, the outer join should have been simplified already.
@@ -659,17 +730,13 @@ func (s *propOuterJoinConstSolver) propagateColumnEQ() {
 }
 
 func (s *propOuterJoinConstSolver) solve(joinConds, filterConds []Expression) ([]Expression, []Expression) {
-	cols := make([]*Column, 0, len(joinConds)+len(filterConds))
 	for _, cond := range joinConds {
 		s.joinConds = append(s.joinConds, SplitCNFItems(cond)...)
-		cols = append(cols, ExtractColumns(cond)...)
+		s.insertCols(ExtractColumns(cond)...)
 	}
 	for _, cond := range filterConds {
 		s.filterConds = append(s.filterConds, SplitCNFItems(cond)...)
-		cols = append(cols, ExtractColumns(cond)...)
-	}
-	for _, col := range cols {
-		s.insertCol(col)
+		s.insertCols(ExtractColumns(cond)...)
 	}
 	if len(s.columns) > MaxPropagateColsCnt {
 		logutil.BgLogger().Warn("too many columns",
@@ -680,18 +747,18 @@ func (s *propOuterJoinConstSolver) solve(joinConds, filterConds []Expression) ([
 	}
 	s.propagateConstantEQ()
 	s.propagateColumnEQ()
-	s.joinConds = propagateConstantDNF(s.ctx, s.joinConds)
-	s.filterConds = propagateConstantDNF(s.ctx, s.filterConds)
+	s.joinConds = propagateConstantDNF(s.ctx, s.joinConds...)
+	s.filterConds = propagateConstantDNF(s.ctx, s.filterConds...)
 	return s.joinConds, s.filterConds
 }
 
 // propagateConstantDNF find DNF item from CNF, and propagate constant inside DNF.
-func propagateConstantDNF(ctx exprctx.ExprContext, conds []Expression) []Expression {
+func propagateConstantDNF(ctx exprctx.ExprContext, conds ...Expression) []Expression {
 	for i, cond := range conds {
 		if dnf, ok := cond.(*ScalarFunction); ok && dnf.FuncName.L == ast.LogicOr {
 			dnfItems := SplitDNFItems(cond)
 			for j, item := range dnfItems {
-				dnfItems[j] = ComposeCNFCondition(ctx, PropagateConstant(ctx, []Expression{item})...)
+				dnfItems[j] = ComposeCNFCondition(ctx, PropagateConstant(ctx, item)...)
 			}
 			conds[i] = ComposeDNFCondition(ctx, dnfItems...)
 		}
@@ -720,17 +787,5 @@ func PropConstOverOuterJoin(ctx exprctx.ExprContext, joinConds, filterConds []Ex
 // PropagateConstantSolver is a constant propagate solver.
 type PropagateConstantSolver interface {
 	PropagateConstant(ctx exprctx.ExprContext, conditions []Expression) []Expression
-}
-
-// newPropConstSolver returns a PropagateConstantSolver.
-func newPropConstSolver() PropagateConstantSolver {
-	solver := &propConstSolver{}
-	solver.colMapper = make(map[int64]int)
-	return solver
-}
-
-// PropagateConstant propagate constant values of deterministic predicates in a condition.
-func (s *propConstSolver) PropagateConstant(ctx exprctx.ExprContext, conditions []Expression) []Expression {
-	s.ctx = ctx
-	return s.solve(conditions)
+	Clear()
 }

pkg/expression/scalar_function.go

@@ -729,6 +729,7 @@ func (sf *ScalarFunction) Equals(other any) bool {
 // ReHashCode is used after we change the argument in place.
 func ReHashCode(sf *ScalarFunction) {
 	sf.hashcode = sf.hashcode[:0]
+	sf.hashcode = slices.Grow(sf.hashcode, 1+len(sf.FuncName.L)+len(sf.GetArgs())*8+1)
 	sf.hashcode = append(sf.hashcode, scalarFunctionFlag)
 	sf.hashcode = codec.EncodeCompactBytes(sf.hashcode, hack.Slice(sf.FuncName.L))
 	for _, arg := range sf.GetArgs() {

pkg/expression/util.go

@@ -1399,16 +1399,15 @@ func IsImmutableFunc(expr Expression) bool {
 // are mutable or have side effects, we cannot remove it even if it has duplicates;
 // if the plan is going to be cached, we cannot remove expressions containing `?` neither.
 func RemoveDupExprs(exprs []Expression) []Expression {
-	res := make([]Expression, 0, len(exprs))
 	exists := make(map[string]struct{}, len(exprs))
-	for _, expr := range exprs {
+	return slices.DeleteFunc(exprs, func(expr Expression) bool {
 		key := string(expr.HashCode())
 		if _, ok := exists[key]; !ok || IsMutableEffectsExpr(expr) {
-			res = append(res, expr)
 			exists[key] = struct{}{}
+			return false
 		}
-	}
-	return res
+		return true
+	})
 }
 
 // GetUint64FromConstant gets a uint64 from constant expression.

pkg/planner/cascades/old/transformation_rules.go

@@ -870,7 +870,7 @@ func (*pushDownJoin) predicatePushDown(
 		tempCond = append(tempCond, join.OtherConditions...)
 		tempCond = append(tempCond, predicates...)
 		tempCond = expression.ExtractFiltersFromDNFs(sctx.GetExprCtx(), tempCond)
-		tempCond = expression.PropagateConstant(sctx.GetExprCtx(), tempCond)
+		tempCond = expression.PropagateConstant(sctx.GetExprCtx(), tempCond...)
 		// Return table dual when filter is constant false or null.
 		dual := logicalop.Conds2TableDual(join, tempCond)
 		if dual != nil {

pkg/planner/core/constraint/exprs.go

@@ -21,6 +21,9 @@ import (
 
 // DeleteTrueExprs deletes the surely true expressions
 func DeleteTrueExprs(buildCtx expression.BuildContext, stmtCtx *stmtctx.StatementContext, conds []expression.Expression) []expression.Expression {
+	if len(conds) == 0 {
+		return conds
+	}
 	newConds := make([]expression.Expression, 0, len(conds))
 	for _, cond := range conds {
 		con, ok := cond.(*expression.Constant)

pkg/planner/core/operator/logicalop/expression_util.go

@@ -21,6 +21,9 @@ import (
 
 // Conds2TableDual builds a LogicalTableDual if cond is constant false or null.
 func Conds2TableDual(p base.LogicalPlan, conds []expression.Expression) base.LogicalPlan {
+	if len(conds) == 0 {
+		return nil
+	}
 	exprCtx := p.SCtx().GetExprCtx()
 	for _, cond := range conds {
 		if expression.IsConstNull(cond) {

pkg/planner/core/operator/logicalop/logical_join.go

@@ -255,7 +255,7 @@ func (p *LogicalJoin) PredicatePushDown(predicates []expression.Expression, opt
 		tempCond = append(tempCond, p.OtherConditions...)
 		tempCond = append(tempCond, predicates...)
 		tempCond = expression.ExtractFiltersFromDNFs(p.SCtx().GetExprCtx(), tempCond)
-		tempCond = expression.PropagateConstant(p.SCtx().GetExprCtx(), tempCond)
+		tempCond = expression.PropagateConstant(p.SCtx().GetExprCtx(), tempCond...)
 		// Return table dual when filter is constant false or null.
 		dual := Conds2TableDual(p, tempCond)
 		if dual != nil {
@@ -270,7 +270,7 @@ func (p *LogicalJoin) PredicatePushDown(predicates []expression.Expression, opt
 		leftCond = leftPushCond
 		rightCond = rightPushCond
 	case AntiSemiJoin:
-		predicates = expression.PropagateConstant(p.SCtx().GetExprCtx(), predicates)
+		predicates = expression.PropagateConstant(p.SCtx().GetExprCtx(), predicates...)
 		// Return table dual when filter is constant false or null.
 		dual := Conds2TableDual(p, predicates)
 		if dual != nil {

pkg/planner/core/operator/logicalop/logical_plans_misc.go

@@ -89,7 +89,7 @@ func addSelection(p base.LogicalPlan, child base.LogicalPlan, conditions []expre
 		return
 	}
 	exprCtx := p.SCtx().GetExprCtx()
-	conditions = expression.PropagateConstant(exprCtx, conditions)
+	conditions = expression.PropagateConstant(exprCtx, conditions...)
 	// Return table dual when filter is constant false or null.
 	dual := Conds2TableDual(child, conditions)
 	if dual != nil {

pkg/planner/core/operator/logicalop/logical_selection.go

@@ -108,7 +108,7 @@ func (p *LogicalSelection) PredicatePushDown(predicates []expression.Expression,
 	retConditions, child = p.Children()[0].PredicatePushDown(append(canBePushDown, predicates...), opt)
 	retConditions = append(retConditions, canNotBePushDown...)
 	if len(retConditions) > 0 {
-		p.Conditions = expression.PropagateConstant(exprCtx, retConditions)
+		p.Conditions = expression.PropagateConstant(exprCtx, retConditions...)
 		// Return table dual when filter is constant false or null.
 		dual := Conds2TableDual(p, p.Conditions)
 		if dual != nil {

pkg/planner/core/rule_predicate_simplification.go

@@ -179,7 +179,7 @@ func applyPredicateSimplification(sctx base.PlanContext, predicates []expression
 	removeRedundantORBranch(sctx, simplifiedPredicate)
 	pruneEmptyORBranches(sctx, simplifiedPredicate)
 	exprCtx := sctx.GetExprCtx()
-	simplifiedPredicate = expression.PropagateConstant(exprCtx, simplifiedPredicate)
+	simplifiedPredicate = expression.PropagateConstant(exprCtx, simplifiedPredicate...)
 	simplifiedPredicate = constraint.DeleteTrueExprs(exprCtx, sctx.GetSessionVars().StmtCtx, simplifiedPredicate)
 	return simplifiedPredicate
 }
@@ -267,7 +267,7 @@ func unsatisfiable(ctx base.PlanContext, p1, p2 expression.Expression) bool {
 	}
 	newPredList := make([]expression.Expression, 0, 1)
 	newPredList = append(newPredList, newPred)
-	newPredList = expression.PropagateConstant(ctx.GetExprCtx(), newPredList)
+	newPredList = expression.PropagateConstant(ctx.GetExprCtx(), newPredList...)
 	return unsatisfiableExpression(ctx, newPredList[0])
 }