planner/cascades: introduce TransformationRuleBatch & add rule Inject…

…ProjectionBelowTopN (pingcap#14504)

planner/cascades/optimize.go

@@ -30,22 +30,22 @@ var DefaultOptimizer = NewOptimizer()
 
 // Optimizer is the struct for cascades optimizer.
 type Optimizer struct {
-	transformationRuleMap map[memo.Operand][]Transformation
-	implementationRuleMap map[memo.Operand][]ImplementationRule
+	transformationRuleBatches []TransformationRuleBatch
+	implementationRuleMap     map[memo.Operand][]ImplementationRule
 }
 
 // NewOptimizer returns a cascades optimizer with default transformation
 // rules and implementation rules.
 func NewOptimizer() *Optimizer {
 	return &Optimizer{
-		transformationRuleMap: defaultTransformationMap,
-		implementationRuleMap: defaultImplementationMap,
+		transformationRuleBatches: DefaultRuleBatches,
+		implementationRuleMap:     defaultImplementationMap,
 	}
 }
 
-// ResetTransformationRules resets the transformationRuleMap of the optimizer, and returns the optimizer.
-func (opt *Optimizer) ResetTransformationRules(rules map[memo.Operand][]Transformation) *Optimizer {
-	opt.transformationRuleMap = rules
+// ResetTransformationRules resets the transformationRuleBatches of the optimizer, and returns the optimizer.
+func (opt *Optimizer) ResetTransformationRules(ruleBatches ...TransformationRuleBatch) *Optimizer {
+	opt.transformationRuleBatches = ruleBatches
 	return opt
 }
 
@@ -55,12 +55,6 @@ func (opt *Optimizer) ResetImplementationRules(rules map[memo.Operand][]Implemen
 	return opt
 }
 
-// GetTransformationRules gets the all the candidate Transformation rules of the optimizer
-// based on the logical plan node.
-func (opt *Optimizer) GetTransformationRules(node plannercore.LogicalPlan) []Transformation {
-	return opt.transformationRuleMap[memo.GetOperand(node)]
-}
-
 // GetImplementationRules gets all the candidate implementation rules of the optimizer
 // for the logical plan node.
 func (opt *Optimizer) GetImplementationRules(node plannercore.LogicalPlan) []ImplementationRule {
@@ -129,38 +123,40 @@ func (opt *Optimizer) onPhasePreprocessing(sctx sessionctx.Context, plan planner
 }
 
 func (opt *Optimizer) onPhaseExploration(sctx sessionctx.Context, g *memo.Group) error {
-	for !g.Explored {
-		err := opt.exploreGroup(g)
-		if err != nil {
-			return err
+	for round, ruleBatch := range opt.transformationRuleBatches {
+		for !g.Explored(round) {
+			err := opt.exploreGroup(g, round, ruleBatch)
+			if err != nil {
+				return err
+			}
 		}
 	}
 	return nil
 }
 
-func (opt *Optimizer) exploreGroup(g *memo.Group) error {
-	if g.Explored {
+func (opt *Optimizer) exploreGroup(g *memo.Group, round int, ruleBatch TransformationRuleBatch) error {
+	if g.Explored(round) {
 		return nil
 	}
-	g.Explored = true
+	g.SetExplored(round)
 
 	for elem := g.Equivalents.Front(); elem != nil; elem = elem.Next() {
 		curExpr := elem.Value.(*memo.GroupExpr)
-		if curExpr.Explored {
+		if curExpr.Explored(round) {
 			continue
 		}
-		curExpr.Explored = true
+		curExpr.SetExplored(round)
 
 		// Explore child groups firstly.
 		for _, childGroup := range curExpr.Children {
-			for !childGroup.Explored {
-				if err := opt.exploreGroup(childGroup); err != nil {
+			for !childGroup.Explored(round) {
+				if err := opt.exploreGroup(childGroup, round, ruleBatch); err != nil {
 					return err
 				}
 			}
 		}
 
-		eraseCur, err := opt.findMoreEquiv(g, elem)
+		eraseCur, err := opt.findMoreEquiv(g, elem, round, ruleBatch)
 		if err != nil {
 			return err
 		}
@@ -172,11 +168,12 @@ func (opt *Optimizer) exploreGroup(g *memo.Group) error {
 }
 
 // findMoreEquiv finds and applies the matched transformation rules.
-func (opt *Optimizer) findMoreEquiv(g *memo.Group, elem *list.Element) (eraseCur bool, err error) {
+func (opt *Optimizer) findMoreEquiv(g *memo.Group, elem *list.Element, round int, ruleBatch TransformationRuleBatch) (eraseCur bool, err error) {
 	expr := elem.Value.(*memo.GroupExpr)
-	for _, rule := range opt.GetTransformationRules(expr.ExprNode) {
+	operand := memo.GetOperand(expr.ExprNode)
+	for _, rule := range ruleBatch[operand] {
 		pattern := rule.GetPattern()
-		if !pattern.Operand.Match(memo.GetOperand(expr.ExprNode)) {
+		if !pattern.Operand.Match(operand) {
 			continue
 		}
 		// Create a binding of the current Group expression and the pattern of
@@ -198,7 +195,7 @@ func (opt *Optimizer) findMoreEquiv(g *memo.Group, elem *list.Element) (eraseCur
 					g.Insert(e)
 				}
 				// If we delete all of the other GroupExprs, we can break the search.
-				g.Explored = true
+				g.SetExplored(round)
 				return false, nil
 			}
 
@@ -210,7 +207,7 @@ func (opt *Optimizer) findMoreEquiv(g *memo.Group, elem *list.Element) (eraseCur
 				// If the new Group expression is successfully inserted into the
 				// current Group, mark the Group as unexplored to enable the exploration
 				// on the new Group expressions.
-				g.Explored = false
+				g.SetUnexplored(round)
 			}
 		}
 	}

planner/cascades/optimize_test.go

@@ -106,7 +106,7 @@ func (s *testCascadesSuite) TestPreparePossibleProperties(c *C) {
 		},
 	})
 	defer func() {
-		s.optimizer.ResetTransformationRules(defaultTransformationMap)
+		s.optimizer.ResetTransformationRules(DefaultRuleBatches...)
 	}()
 	stmt, err := s.ParseOneStmt("select f, sum(a) from t group by f", "", "")
 	c.Assert(err, IsNil)
@@ -187,7 +187,7 @@ func (s *testCascadesSuite) TestAppliedRuleSet(c *C) {
 		},
 	})
 	defer func() {
-		s.optimizer.ResetTransformationRules(defaultTransformationMap)
+		s.optimizer.ResetTransformationRules(DefaultRuleBatches...)
 	}()
 	stmt, err := s.ParseOneStmt("select 1", "", "")
 	c.Assert(err, IsNil)

planner/cascades/stringer_test.go

@@ -61,7 +61,7 @@ func (s *testStringerSuite) TestGroupStringer(c *C) {
 		},
 	})
 	defer func() {
-		s.optimizer.ResetTransformationRules(defaultTransformationMap)
+		s.optimizer.ResetTransformationRules(DefaultRuleBatches...)
 	}()
 	var input []string
 	var output []struct {

planner/cascades/testdata/integration_suite_in.json

@@ -53,7 +53,8 @@
       "select b from t order by b limit 3",
       "select a from t order by a limit 1 offset 2",
       "select * from ((select a as aa from t t1) union all (select b as aa from t t2)) as t3 order by aa",
-      "select a, b, lag(a,1) over (order by b) from t order by b"
+      "select a, b, lag(a,1) over (order by b) from t order by b",
+      "select * from (select a+1 as c, a+b as d from t) as t1 order by c+d limit 10"
     ]
   },
   {