feat(optimizer): push LogicalLimit down LogicalProject, enabled i…

…n batch (risingwavelabs#8971)

Signed-off-by: Clearlove <yifei.c.wei@gmail.com>

src/frontend/planner_test/tests/testdata/limit.yaml

@@ -1,11 +1,15 @@
 # This file is automatically generated. See `src/frontend/planner_test/README.md` for more information.
 - sql: |
     create table t (v int);
-    select * from t limit 4;
+    select v+v from t limit 4;
   logical_plan: |
     LogicalLimit { limit: 4, offset: 0 }
-    └─LogicalProject { exprs: [t.v] }
+    └─LogicalProject { exprs: [(t.v + t.v) as $expr1] }
       └─LogicalScan { table: t, columns: [t.v, t._row_id] }
+  optimized_logical_plan_for_batch: |
+    LogicalProject { exprs: [(t.v + t.v) as $expr1] }
+    └─LogicalLimit { limit: 4, offset: 0 }
+      └─LogicalScan { table: t, columns: [t.v] }
 - sql: |
     create table t (v int);
     select * from t offset 4;
@@ -15,13 +19,18 @@
       └─LogicalScan { table: t, columns: [t.v, t._row_id] }
 - sql: |
     create table t (v int);
-    select * from ( select * from t limit 5 ) limit 4;
+    select * from ( select v+v from t limit 5 ) limit 4;
   logical_plan: |
     LogicalLimit { limit: 4, offset: 0 }
-    └─LogicalProject { exprs: [t.v] }
+    └─LogicalProject { exprs: [$expr1] }
       └─LogicalLimit { limit: 5, offset: 0 }
-        └─LogicalProject { exprs: [t.v] }
+        └─LogicalProject { exprs: [(t.v + t.v) as $expr1] }
           └─LogicalScan { table: t, columns: [t.v, t._row_id] }
+  optimized_logical_plan_for_batch: |
+    LogicalLimit { limit: 4, offset: 0 }
+    └─LogicalProject { exprs: [(t.v + t.v) as $expr1] }
+      └─LogicalLimit { limit: 5, offset: 0 }
+        └─LogicalScan { table: t, columns: [t.v] }
 - sql: |
     create table t (v int);
     select * from t fetch first 4 rows only;

src/frontend/src/optimizer/logical_optimization.rs

@@ -256,6 +256,12 @@ lazy_static! {
         ApplyOrder::TopDown,
     );
 
+    static ref LIMIT_PUSH_DOWN: OptimizationStage = OptimizationStage::new(
+        "Push Down Limit",
+        vec![LimitPushDownRule::create()],
+        ApplyOrder::TopDown,
+    );
+
     static ref PULL_UP_HOP: OptimizationStage = OptimizationStage::new(
         "Pull up hop",
         vec![PullUpHopRule::create()],
@@ -503,6 +509,8 @@ impl LogicalOptimizer {
 
         plan = plan.optimize_by_rules(&TOP_N_AGG_ON_INDEX);
 
+        plan = plan.optimize_by_rules(&LIMIT_PUSH_DOWN);
+
         #[cfg(debug_assertions)]
         InputRefValidator.validate(plan.clone());
 

src/frontend/src/optimizer/rule/limit_push_down_rule.rs

@@ -0,0 +1,40 @@
+//  Copyright 2023 RisingWave Labs
+//
+//  Licensed under the Apache License, Version 2.0 (the "License");
+//  you may not use this file except in compliance with the License.
+//  You may obtain a copy of the License at
+//
+//  http://www.apache.org/licenses/LICENSE-2.0
+//
+//  Unless required by applicable law or agreed to in writing, software
+//  distributed under the License is distributed on an "AS IS" BASIS,
+//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+//  See the License for the specific language governing permissions and
+//  limitations under the License.
+//
+// Copyright (c) 2011-present, Facebook, Inc.  All rights reserved.
+// This source code is licensed under both the GPLv2 (found in the
+// COPYING file in the root directory) and Apache 2.0 License
+// (found in the LICENSE.Apache file in the root directory).
+
+use super::{BoxedRule, Rule};
+use crate::optimizer::plan_node::{LogicalLimit, LogicalProject, PlanTreeNodeUnary};
+use crate::optimizer::PlanRef;
+
+pub struct LimitPushDownRule {}
+
+impl Rule for LimitPushDownRule {
+    fn apply(&self, plan: PlanRef) -> Option<PlanRef> {
+        let limit: &LogicalLimit = plan.as_logical_limit()?;
+        let project: LogicalProject = limit.input().as_logical_project()?.to_owned();
+        let input = project.input();
+        let logical_limit = limit.clone_with_input(input);
+        Some(project.clone_with_input(logical_limit.into()).into())
+    }
+}
+
+impl LimitPushDownRule {
+    pub fn create() -> BoxedRule {
+        Box::new(LimitPushDownRule {})
+    }
+}

src/frontend/src/optimizer/rule/mod.rs

@@ -101,6 +101,8 @@ mod always_false_filter_rule;
 pub use always_false_filter_rule::*;
 mod join_project_transpose_rule;
 pub use join_project_transpose_rule::*;
+mod limit_push_down_rule;
+pub use limit_push_down_rule::*;
 mod pull_up_hop_rule;
 pub use pull_up_hop_rule::*;
 
@@ -148,6 +150,7 @@ macro_rules! for_all_rules {
             , { BushyTreeJoinOrderingRule }
             , { StreamProjectMergeRule }
             , { JoinProjectTransposeRule }
+            , { LimitPushDownRule }
             , { PullUpHopRule }
         }
     };