From 39f43c9cec1a0dbcea96504b612cef5ab334c317 Mon Sep 17 00:00:00 2001
From: longfang <zhaolongfang1@163.com>
Date: Thu, 26 Dec 2019 01:37:29 -0800
Subject: [PATCH] Add push the limit down to projection.

---
 .../transformation_rules_suite_in.json        |  1 +
 .../transformation_rules_suite_out.json       | 13 +++++
 planner/cascades/transformation_rules.go      | 48 +++++++++++++++++++
 planner/cascades/transformation_rules_test.go |  1 +
 4 files changed, 63 insertions(+)

diff --git a/planner/cascades/testdata/transformation_rules_suite_in.json b/planner/cascades/testdata/transformation_rules_suite_in.json
index f20520b33824b..7442f521dd7ec 100644
--- a/planner/cascades/testdata/transformation_rules_suite_in.json
+++ b/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",
diff --git a/planner/cascades/testdata/transformation_rules_suite_out.json b/planner/cascades/testdata/transformation_rules_suite_out.json
index 495c8a227eace..c3d0e8ea5a093 100644
--- a/planner/cascades/testdata/transformation_rules_suite_out.json
+++ b/planner/cascades/testdata/transformation_rules_suite_out.json
@@ -345,6 +345,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_6 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": [
diff --git a/planner/cascades/transformation_rules.go b/planner/cascades/transformation_rules.go
index d0f7de46863f8..f038ea5ffcc88 100644
--- a/planner/cascades/transformation_rules.go
+++ b/planner/cascades/transformation_rules.go
@@ -64,6 +64,7 @@ var defaultTransformationMap = map[memo.Operand][]Transformation{
 	},
 	memo.OperandLimit: {
 		NewRuleTransformLimitToTopN(),
+		NewRulePushLimitDownProjection(),
 	},
 	memo.OperandProjection: {
 		NewRuleEliminateProjection(),
@@ -665,6 +666,53 @@ 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]
+
+	newLimit := plannercore.LogicalLimit{
+		Offset: limit.Offset,
+		Count:  limit.Count,
+	}.Init(limit.SCtx(), limit.SelectBlockOffset())
+
+	projExpr := memo.NewGroupExpr(proj)
+	limitExpr := memo.NewGroupExpr(newLimit)
+	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
diff --git a/planner/cascades/transformation_rules_test.go b/planner/cascades/transformation_rules_test.go
index f4f6d73ec7cdd..460bc02ed660f 100644
--- a/planner/cascades/transformation_rules_test.go
+++ b/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(),