From be18447c1ca98d4029293a323ecad5a8c98328f1 Mon Sep 17 00:00:00 2001 From: hanmc Date: Fri, 8 Nov 2019 16:13:25 +0800 Subject: [PATCH 1/2] add transformation rule TransformLimitToTopN --- .../transformation_rules_suite_in.json | 7 ++++ .../transformation_rules_suite_out.json | 35 ++++++++++++++++ planner/cascades/transformation_rules.go | 40 ++++++++++++++++++ planner/cascades/transformation_rules_test.go | 42 +++++++++++++++++++ 4 files changed, 124 insertions(+) diff --git a/planner/cascades/testdata/transformation_rules_suite_in.json b/planner/cascades/testdata/transformation_rules_suite_in.json index e4057e5b2693e..4d591e8b847d3 100644 --- a/planner/cascades/testdata/transformation_rules_suite_in.json +++ b/planner/cascades/testdata/transformation_rules_suite_in.json @@ -10,5 +10,12 @@ "select a, max(b) from t group by a having a > 1", "select a, avg(b) from t group by a having a > 1 and max(b) > 10" ] + }, + { + "name": "TestTopNRules", + "cases": [ + "select b from t order by a limit 2", + "select a+b from t order by a limit 1 offset 2" + ] } ] diff --git a/planner/cascades/testdata/transformation_rules_suite_out.json b/planner/cascades/testdata/transformation_rules_suite_out.json index a07e9f5eb3d1e..e7edaca1cfa27 100644 --- a/planner/cascades/testdata/transformation_rules_suite_out.json +++ b/planner/cascades/testdata/transformation_rules_suite_out.json @@ -110,5 +110,40 @@ ] } ] + }, + { + "Name": "TestTopNRules", + "Cases": [ + { + "SQL": "select b from t order by a limit 2", + "Result": [ + "Group#0 Schema:[Column#15]", + " Projection_5 input:[Group#1], Column#13", + "Group#1 Schema:[Column#13,Column#14], UniqueKey:[Column#14]", + " TopN_8 input:[Group#2], Column#14:asc, offset:0, count:2", + "Group#2 Schema:[Column#13,Column#14], UniqueKey:[Column#14]", + " Projection_2 input:[Group#3], Column#2, Column#1", + "Group#3 Schema:[Column#1,Column#2], UniqueKey:[Column#1]", + " TableGather_7 input:[Group#4]", + "Group#4 Schema:[Column#1,Column#2], UniqueKey:[Column#1]", + " TableScan_6 table:t, pk col:Column#1" + ] + }, + { + "SQL": "select a+b from t order by a limit 1 offset 2", + "Result": [ + "Group#0 Schema:[Column#15]", + " Projection_5 input:[Group#1], Column#13", + "Group#1 Schema:[Column#13,Column#14], UniqueKey:[Column#14]", + " TopN_8 input:[Group#2], Column#14:asc, offset:2, count:1", + "Group#2 Schema:[Column#13,Column#14], UniqueKey:[Column#14]", + " Projection_2 input:[Group#3], plus(Column#1, Column#2), Column#1", + "Group#3 Schema:[Column#1,Column#2], UniqueKey:[Column#1]", + " TableGather_7 input:[Group#4]", + "Group#4 Schema:[Column#1,Column#2], UniqueKey:[Column#1]", + " TableScan_6 table:t, pk col:Column#1" + ] + } + ] } ] diff --git a/planner/cascades/transformation_rules.go b/planner/cascades/transformation_rules.go index 2d4675c858a7b..65c47f2c7673c 100644 --- a/planner/cascades/transformation_rules.go +++ b/planner/cascades/transformation_rules.go @@ -52,6 +52,7 @@ const ( rulePushSelDownProjection rulePushSelDownAggregation ruleEnumeratePaths + ruleTransformLimitToTopN ) var transformationRuleList = []Transformation{ @@ -61,6 +62,7 @@ var transformationRuleList = []Transformation{ &PushSelDownProjection{}, &PushSelDownAggregation{}, &EnumeratePaths{}, + &TransformLimitToTopN{}, } var defaultTransformationMap = map[memo.Operand][]TransformationID{ @@ -74,6 +76,9 @@ var defaultTransformationMap = map[memo.Operand][]TransformationID{ memo.OperandDataSource: { ruleEnumeratePaths, }, + memo.OperandLimit: { + ruleTransformLimitToTopN, + }, } var patternMap []*memo.Pattern @@ -411,3 +416,38 @@ func (r *PushSelDownAggregation) OnTransform(old *memo.ExprIter) (newExprs []*me remainedGroupExpr.SetChildren(aggGroup) return []*memo.GroupExpr{remainedGroupExpr}, true, false, nil } + +// TransformLimitToTopN transforms Limit+Sort to TopN. +type TransformLimitToTopN struct { +} + +// GetPattern implements Transformation interface. +// The pattern of this rule is `Limit -> Sort`. +func (r *TransformLimitToTopN) GetPattern() *memo.Pattern { + return memo.BuildPattern( + memo.OperandLimit, + memo.EngineTiDBOnly, + memo.NewPattern(memo.OperandSort, memo.EngineTiDBOnly), + ) +} + +// Match implements Transformation interface. +func (r *TransformLimitToTopN) Match(expr *memo.ExprIter) bool { + return true +} + +// OnTransform implements Transformation interface. +// This rule will transform `Limit -> Sort -> x` to `TopN -> x`. +func (r *TransformLimitToTopN) OnTransform(old *memo.ExprIter) (newExprs []*memo.GroupExpr, eraseOld bool, eraseAll bool, err error) { + limit := old.GetExpr().ExprNode.(*plannercore.LogicalLimit) + sort := old.Children[0].GetExpr().ExprNode.(*plannercore.LogicalSort) + childGroup := old.Children[0].GetExpr().Children[0] + topN := plannercore.LogicalTopN{ + ByItems: sort.ByItems, + Offset: limit.Offset, + Count: limit.Count, + }.Init(limit.SCtx(), limit.SelectBlockOffset()) + topNExpr := memo.NewGroupExpr(topN) + topNExpr.SetChildren(childGroup) + return []*memo.GroupExpr{topNExpr}, true, false, nil +} diff --git a/planner/cascades/transformation_rules_test.go b/planner/cascades/transformation_rules_test.go index ec12a86a2c455..db0b262e01dca 100644 --- a/planner/cascades/transformation_rules_test.go +++ b/planner/cascades/transformation_rules_test.go @@ -50,6 +50,30 @@ func (s *testTransformationRuleSuite) TearDownSuite(c *C) { c.Assert(s.testData.GenerateOutputIfNeeded(), IsNil) } +func testGroupToString(input []string, output []struct { + SQL string + Result []string +}, s *testTransformationRuleSuite, c *C) { + for i, sql := range input { + stmt, err := s.ParseOneStmt(sql, "", "") + c.Assert(err, IsNil) + p, _, err := plannercore.BuildLogicalPlan(context.Background(), s.sctx, stmt, s.is) + c.Assert(err, IsNil) + logic, ok := p.(plannercore.LogicalPlan) + c.Assert(ok, IsTrue) + logic, err = s.optimizer.onPhasePreprocessing(s.sctx, logic) + c.Assert(err, IsNil) + group := convert2Group(logic) + err = s.optimizer.onPhaseExploration(s.sctx, group) + c.Assert(err, IsNil) + s.testData.OnRecord(func() { + output[i].SQL = sql + output[i].Result = ToString(group) + }) + c.Assert(ToString(group), DeepEquals, output[i].Result) + } +} + func (s *testTransformationRuleSuite) TestPredicatePushDown(c *C) { s.optimizer.ResetTransformationRules(map[memo.Operand][]TransformationID{ memo.OperandSelection: { @@ -91,3 +115,21 @@ func (s *testTransformationRuleSuite) TestPredicatePushDown(c *C) { c.Assert(ToString(group), DeepEquals, output[i].Result) } } + +func (s *testTransformationRuleSuite) TestTopNRules(c *C) { + s.optimizer.ResetTransformationRules(map[memo.Operand][]TransformationID{ + memo.OperandLimit: { + ruleTransformLimitToTopN, + }, + memo.OperandDataSource: { + ruleEnumeratePaths, + }, + }) + var input []string + var output []struct { + SQL string + Result []string + } + s.testData.GetTestCases(c, &input, &output) + testGroupToString(input, output, s, c) +} From 394819267f49ccffd2f6a9a004679abb4337e5c0 Mon Sep 17 00:00:00 2001 From: hanmc Date: Fri, 8 Nov 2019 16:17:17 +0800 Subject: [PATCH 2/2] simplify the test --- planner/cascades/transformation_rules_test.go | 19 +------------------ 1 file changed, 1 insertion(+), 18 deletions(-) diff --git a/planner/cascades/transformation_rules_test.go b/planner/cascades/transformation_rules_test.go index db0b262e01dca..2db958ae153dd 100644 --- a/planner/cascades/transformation_rules_test.go +++ b/planner/cascades/transformation_rules_test.go @@ -96,24 +96,7 @@ func (s *testTransformationRuleSuite) TestPredicatePushDown(c *C) { Result []string } s.testData.GetTestCases(c, &input, &output) - for i, sql := range input { - stmt, err := s.ParseOneStmt(sql, "", "") - c.Assert(err, IsNil) - p, _, err := plannercore.BuildLogicalPlan(context.Background(), s.sctx, stmt, s.is) - c.Assert(err, IsNil) - logic, ok := p.(plannercore.LogicalPlan) - c.Assert(ok, IsTrue) - logic, err = s.optimizer.onPhasePreprocessing(s.sctx, logic) - c.Assert(err, IsNil) - group := convert2Group(logic) - err = s.optimizer.onPhaseExploration(s.sctx, group) - c.Assert(err, IsNil) - s.testData.OnRecord(func() { - output[i].SQL = sql - output[i].Result = ToString(group) - }) - c.Assert(ToString(group), DeepEquals, output[i].Result) - } + testGroupToString(input, output, s, c) } func (s *testTransformationRuleSuite) TestTopNRules(c *C) {