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 06beeaa12670c..a82b7f2829969 100644 --- a/planner/cascades/testdata/transformation_rules_suite_out.json +++ b/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": [ diff --git a/planner/cascades/transformation_rules.go b/planner/cascades/transformation_rules.go index 382725ddccf27..0ba53a24d4bab 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(), @@ -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 diff --git a/planner/cascades/transformation_rules_test.go b/planner/cascades/transformation_rules_test.go index 91dd6625209aa..0aaf30ca550a6 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(),