planner/cascades: add transformation rule PushLimitDownProjcetion (#14254)

planner/cascades/testdata/transformation_rules_suite_in.json

@@ -32,6 +32,7 @@
     "name": "TestTopNRules",
     "cases": [
       "select b from t order by a limit 2",
+      "select b from t limit 2",
       "select a+b from t order by a limit 1 offset 2",
       "select c from t order by t.a limit 1",
       "select c from t order by t.a + t.b limit 1",

planner/cascades/testdata/transformation_rules_suite_out.json

@@ -347,6 +347,19 @@
           "    TableScan_6 table:t, pk col:test.t.a"
         ]
       },
+      {
+        "SQL": "select b from t limit 2",
+        "Result": [
+          "Group#0 Schema:[test.t.b]", 
+          "    Projection_2 input:[Group#1], test.t.b", 
+          "Group#1 Schema:[test.t.b]", 
+          "    Limit_3 input:[Group#2], offset:0, count:2", 
+          "Group#2 Schema:[test.t.b]", 
+          "    TiKVSingleGather_5 input:[Group#3], table:t", 
+          "Group#3 Schema:[test.t.b]", 
+          "    TableScan_4 table:t"
+        ]
+      },
       {
         "SQL": "select a+b from t order by a limit 1 offset 2",
         "Result": [

planner/cascades/transformation_rules.go

@@ -64,6 +64,7 @@ var defaultTransformationMap = map[memo.Operand][]Transformation{
 	},
 	memo.OperandLimit: {
 		NewRuleTransformLimitToTopN(),
+		NewRulePushLimitDownProjection(),
 	},
 	memo.OperandProjection: {
 		NewRuleEliminateProjection(),
@@ -666,6 +667,49 @@ func (r *TransformLimitToTopN) OnTransform(old *memo.ExprIter) (newExprs []*memo
 	return []*memo.GroupExpr{topNExpr}, true, false, nil
 }
 
+// PushLimitDownProjection pushes Limit to Projection.
+type PushLimitDownProjection struct {
+	baseRule
+}
+
+// NewRulePushLimitDownProjection creates a new Transformation.
+// The pattern of this rule is `Limit->Projection->X` to `Projection->Limit->X`.
+func NewRulePushLimitDownProjection() Transformation {
+	rule := &PushLimitDownProjection{}
+	rule.pattern = memo.BuildPattern(
+		memo.OperandLimit,
+		memo.EngineTiDBOnly,
+		memo.NewPattern(memo.OperandProjection, memo.EngineTiDBOnly),
+	)
+	return rule
+}
+
+// Match implements Transformation interface.
+func (r *PushLimitDownProjection) Match(expr *memo.ExprIter) bool {
+	proj := expr.Children[0].GetExpr().ExprNode.(*plannercore.LogicalProjection)
+	for _, expr := range proj.Exprs {
+		if expression.HasAssignSetVarFunc(expr) {
+			return false
+		}
+	}
+	return true
+}
+
+// OnTransform implements Transformation interface.
+// This rule tries to pushes the Limit through Projection.
+func (r *PushLimitDownProjection) OnTransform(old *memo.ExprIter) (newExprs []*memo.GroupExpr, eraseOld bool, eraseAll bool, err error) {
+	limit := old.GetExpr().ExprNode.(*plannercore.LogicalLimit)
+	proj := old.Children[0].GetExpr().ExprNode.(*plannercore.LogicalProjection)
+	childGroup := old.Children[0].GetExpr().Children[0]
+
+	projExpr := memo.NewGroupExpr(proj)
+	limitExpr := memo.NewGroupExpr(limit)
+	limitExpr.SetChildren(childGroup)
+	limitGroup := memo.NewGroupWithSchema(limitExpr, childGroup.Prop.Schema)
+	projExpr.SetChildren(limitGroup)
+	return []*memo.GroupExpr{projExpr}, true, false, nil
+}
+
 // PushSelDownJoin pushes Selection through Join.
 type PushSelDownJoin struct {
 	baseRule

planner/cascades/transformation_rules_test.go

@@ -149,6 +149,7 @@ func (s *testTransformationRuleSuite) TestTopNRules(c *C) {
 	s.optimizer.ResetTransformationRules(map[memo.Operand][]Transformation{
 		memo.OperandLimit: {
 			NewRuleTransformLimitToTopN(),
+			NewRulePushLimitDownProjection(),
 		},
 		memo.OperandDataSource: {
 			NewRuleEnumeratePaths(),