This is an automated email from the ASF dual-hosted git repository.
alamb pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new a5ce56831a implement rewrite for EliminateLimit (#10253)
a5ce56831a is described below
commit a5ce56831a9ec61e634d1c285c1b28d8c3891503
Author: Kevin Mingtarja <[email protected]>
AuthorDate: Sat Apr 27 01:08:26 2024 +0700
implement rewrite for EliminateLimit (#10253)
---
datafusion/optimizer/src/eliminate_limit.rs | 64 +++++++++++++++++------------
1 file changed, 38 insertions(+), 26 deletions(-)
diff --git a/datafusion/optimizer/src/eliminate_limit.rs
b/datafusion/optimizer/src/eliminate_limit.rs
index 39231d784e..1b0907d973 100644
--- a/datafusion/optimizer/src/eliminate_limit.rs
+++ b/datafusion/optimizer/src/eliminate_limit.rs
@@ -18,8 +18,9 @@
//! [`EliminateLimit`] eliminates `LIMIT` when possible
use crate::optimizer::ApplyOrder;
use crate::{OptimizerConfig, OptimizerRule};
-use datafusion_common::Result;
-use datafusion_expr::logical_plan::{EmptyRelation, LogicalPlan};
+use datafusion_common::tree_node::Transformed;
+use datafusion_common::{internal_err, Result};
+use datafusion_expr::logical_plan::{tree_node::unwrap_arc, EmptyRelation,
LogicalPlan};
/// Optimizer rule to replace `LIMIT 0` or `LIMIT` whose ancestor LIMIT's skip
is
/// greater than or equal to current's fetch
@@ -41,32 +42,10 @@ impl EliminateLimit {
impl OptimizerRule for EliminateLimit {
fn try_optimize(
&self,
- plan: &LogicalPlan,
+ _plan: &LogicalPlan,
_config: &dyn OptimizerConfig,
) -> Result<Option<LogicalPlan>> {
- if let LogicalPlan::Limit(limit) = plan {
- match limit.fetch {
- Some(fetch) => {
- if fetch == 0 {
- return
Ok(Some(LogicalPlan::EmptyRelation(EmptyRelation {
- produce_one_row: false,
- schema: limit.input.schema().clone(),
- })));
- }
- }
- None => {
- if limit.skip == 0 {
- let input = limit.input.as_ref();
- // input also can be Limit, so we should apply again.
- return Ok(Some(
- self.try_optimize(input, _config)?
- .unwrap_or_else(|| input.clone()),
- ));
- }
- }
- }
- }
- Ok(None)
+ internal_err!("Should have called EliminateLimit::rewrite")
}
fn name(&self) -> &str {
@@ -76,6 +55,39 @@ impl OptimizerRule for EliminateLimit {
fn apply_order(&self) -> Option<ApplyOrder> {
Some(ApplyOrder::BottomUp)
}
+
+ fn supports_rewrite(&self) -> bool {
+ true
+ }
+
+ fn rewrite(
+ &self,
+ plan: LogicalPlan,
+ _config: &dyn OptimizerConfig,
+ ) -> Result<
+ datafusion_common::tree_node::Transformed<LogicalPlan>,
+ datafusion_common::DataFusionError,
+ > {
+ match plan {
+ LogicalPlan::Limit(limit) => {
+ if let Some(fetch) = limit.fetch {
+ if fetch == 0 {
+ return Ok(Transformed::yes(LogicalPlan::EmptyRelation(
+ EmptyRelation {
+ produce_one_row: false,
+ schema: limit.input.schema().clone(),
+ },
+ )));
+ }
+ } else if limit.skip == 0 {
+ // input also can be Limit, so we should apply again.
+ return Ok(self.rewrite(unwrap_arc(limit.input),
_config).unwrap());
+ }
+ Ok(Transformed::no(LogicalPlan::Limit(limit)))
+ }
+ _ => Ok(Transformed::no(plan)),
+ }
+ }
}
#[cfg(test)]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
