feat: add limit push down rule (#16403)

* feat: add limit push down rule

* fix

* fix test

Author: xudong963

Cargo.lock

@@ -34,6 +34,7 @@ use crate::optimizer::rule::rewrite::RulePushDownFilterProjectSet;
 use crate::optimizer::rule::rewrite::RulePushDownFilterScan;
 use crate::optimizer::rule::rewrite::RulePushDownFilterSort;
 use crate::optimizer::rule::rewrite::RulePushDownFilterUnion;
+use crate::optimizer::rule::rewrite::RulePushDownLimit;
 use crate::optimizer::rule::rewrite::RulePushDownLimitOuterJoin;
 use crate::optimizer::rule::rewrite::RulePushDownLimitScan;
 use crate::optimizer::rule::rewrite::RulePushDownLimitSort;
@@ -63,6 +64,7 @@ impl RuleFactory {
             RuleID::PushDownFilterScan => Ok(Box::new(RulePushDownFilterScan::new(metadata))),
             RuleID::PushDownFilterSort => Ok(Box::new(RulePushDownFilterSort::new())),
             RuleID::PushDownFilterProjectSet => Ok(Box::new(RulePushDownFilterProjectSet::new())),
+            RuleID::PushDownLimit => Ok(Box::new(RulePushDownLimit::new(metadata))),
             RuleID::PushDownLimitUnion => Ok(Box::new(RulePushDownLimitUnion::new())),
             RuleID::PushDownLimitScan => Ok(Box::new(RulePushDownLimitScan::new())),
             RuleID::PushDownSortScan => Ok(Box::new(RulePushDownSortScan::new())),

src/query/sql/src/planner/optimizer/rule/factory.rs

@@ -30,6 +30,7 @@ mod rule_push_down_filter_scan;
 mod rule_push_down_filter_sort;
 mod rule_push_down_filter_union;
 mod rule_push_down_filter_window;
+mod rule_push_down_limit;
 mod rule_push_down_limit_aggregate;
 mod rule_push_down_limit_expression;
 mod rule_push_down_limit_join;
@@ -60,6 +61,7 @@ pub use rule_push_down_filter_scan::RulePushDownFilterScan;
 pub use rule_push_down_filter_sort::RulePushDownFilterSort;
 pub use rule_push_down_filter_union::RulePushDownFilterUnion;
 pub use rule_push_down_filter_window::RulePushDownFilterWindow;
+pub use rule_push_down_limit::RulePushDownLimit;
 pub use rule_push_down_limit_aggregate::RulePushDownLimitAggregate;
 pub use rule_push_down_limit_expression::RulePushDownLimitEvalScalar;
 pub use rule_push_down_limit_join::RulePushDownLimitOuterJoin;

src/query/sql/src/planner/optimizer/rule/rewrite/mod.rs

@@ -0,0 +1,92 @@
+// Copyright 2021 Datafuse 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.
+
+use std::sync::Arc;
+
+use databend_common_exception::Result;
+use databend_common_expression::Column;
+use databend_common_expression::DataField;
+use databend_common_expression::DataSchemaRefExt;
+
+use crate::optimizer::extract::Matcher;
+use crate::optimizer::rule::Rule;
+use crate::optimizer::rule::TransformResult;
+use crate::optimizer::RelExpr;
+use crate::optimizer::RuleID;
+use crate::optimizer::SExpr;
+use crate::plans::ConstantTableScan;
+use crate::plans::Limit;
+use crate::plans::Operator;
+use crate::plans::RelOp;
+use crate::plans::RelOperator;
+use crate::MetadataRef;
+
+pub struct RulePushDownLimit {
+    id: RuleID,
+    metadata: MetadataRef,
+    matchers: Vec<Matcher>,
+}
+
+impl RulePushDownLimit {
+    pub fn new(metadata: MetadataRef) -> Self {
+        Self {
+            id: RuleID::PushDownLimit,
+            metadata,
+            matchers: vec![Matcher::MatchOp {
+                op_type: RelOp::Limit,
+                children: vec![Matcher::Leaf],
+            }],
+        }
+    }
+}
+
+impl Rule for RulePushDownLimit {
+    fn id(&self) -> RuleID {
+        self.id
+    }
+
+    fn apply(&self, s_expr: &SExpr, state: &mut TransformResult) -> Result<()> {
+        let limit: Limit = s_expr.plan().clone().try_into()?;
+        if let Some(limit_val) = limit.limit
+            && limit_val == 0
+        {
+            let output_columns = limit
+                .derive_relational_prop(&RelExpr::with_s_expr(s_expr))?
+                .output_columns
+                .clone();
+            let metadata = self.metadata.read();
+            let mut fields = Vec::with_capacity(output_columns.len());
+            for col in output_columns.iter() {
+                fields.push(DataField::new(
+                    &col.to_string(),
+                    metadata.column(*col).data_type(),
+                ));
+            }
+            let empty_scan = ConstantTableScan {
+                values: vec![Column::Null { len: 0 }; output_columns.len()],
+                num_rows: 0,
+                schema: DataSchemaRefExt::create(fields),
+                columns: output_columns,
+            };
+            let result = SExpr::create_leaf(Arc::new(RelOperator::ConstantTableScan(empty_scan)));
+            state.add_result(result);
+        }
+
+        Ok(())
+    }
+
+    fn matchers(&self) -> &[Matcher] {
+        &self.matchers
+    }
+}

src/query/sql/src/planner/optimizer/rule/rewrite/rule_push_down_limit.rs

@@ -34,6 +34,7 @@ pub static DEFAULT_REWRITE_RULES: LazyLock<Vec<RuleID>> = LazyLock::new(|| {
         RuleID::PushDownFilterUnion,
         RuleID::PushDownFilterAggregate,
         RuleID::PushDownFilterWindow,
+        RuleID::PushDownLimit,
         RuleID::PushDownLimitUnion,
         RuleID::PushDownLimitEvalScalar,
         RuleID::PushDownLimitSort,
@@ -83,6 +84,7 @@ pub enum RuleID {
     PushDownFilterSort,
     PushDownFilterProjectSet,
     PushDownFilterWindow,
+    PushDownLimit,
     PushDownLimitUnion,
     PushDownLimitOuterJoin,
     PushDownLimitEvalScalar,
@@ -118,6 +120,7 @@ impl Display for RuleID {
             RuleID::PushDownFilterScan => write!(f, "PushDownFilterScan"),
             RuleID::PushDownFilterSort => write!(f, "PushDownFilterSort"),
             RuleID::PushDownFilterProjectSet => write!(f, "PushDownFilterProjectSet"),
+            RuleID::PushDownLimit => write!(f, "PushDownLimit"),
             RuleID::PushDownLimitUnion => write!(f, "PushDownLimitUnion"),
             RuleID::PushDownLimitOuterJoin => write!(f, "PushDownLimitOuterJoin"),
             RuleID::PushDownLimitEvalScalar => write!(f, "PushDownLimitEvalScalar"),

src/query/sql/src/planner/optimizer/rule/rule.rs

@@ -220,3 +220,11 @@ Limit
                     ├── partitions scanned: 1
                     ├── push downs: [filters: [], limit: NONE]
                     └── estimated rows: 2.00
+
+
+query T
+explain select c1 from (select count(t1.number) as c1 from numbers(1) as t1 group by number) as t3 left join (select count(t.number) as c from numbers(2) as t group by number) as t4 on t3.c1=t4.c order by t3.c1 limit 0
+----
+ConstantTableScan
+├── output columns: [count(t1.number) (#2)]
+└── column 0: []