Skip to content

Commit

Permalink
feat: add limit push down rule (#16403)
Browse files Browse the repository at this point in the history
* feat: add limit push down rule

* fix

* fix test
  • Loading branch information
xudong963 authored Sep 6, 2024
1 parent dd21d7d commit a9a6630
Show file tree
Hide file tree
Showing 6 changed files with 108 additions and 0 deletions.
1 change: 1 addition & 0 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

2 changes: 2 additions & 0 deletions src/query/sql/src/planner/optimizer/rule/factory.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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())),
Expand Down
2 changes: 2 additions & 0 deletions src/query/sql/src/planner/optimizer/rule/rewrite/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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;
Expand Down
Original file line number Diff line number Diff line change
@@ -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
}
}
3 changes: 3 additions & 0 deletions src/query/sql/src/planner/optimizer/rule/rule.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -83,6 +84,7 @@ pub enum RuleID {
PushDownFilterSort,
PushDownFilterProjectSet,
PushDownFilterWindow,
PushDownLimit,
PushDownLimitUnion,
PushDownLimitOuterJoin,
PushDownLimitEvalScalar,
Expand Down Expand Up @@ -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"),
Expand Down
8 changes: 8 additions & 0 deletions tests/sqllogictests/suites/mode/standalone/explain/limit.test
Original file line number Diff line number Diff line change
Expand Up @@ -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: []

0 comments on commit a9a6630

Please sign in to comment.