diff --git a/datafusion/optimizer/src/eliminate_limit.rs b/datafusion/optimizer/src/eliminate_limit.rs index 39231d784e00..1b0907d9736d 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> { - 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 { Some(ApplyOrder::BottomUp) } + + fn supports_rewrite(&self) -> bool { + true + } + + fn rewrite( + &self, + plan: LogicalPlan, + _config: &dyn OptimizerConfig, + ) -> Result< + datafusion_common::tree_node::Transformed, + 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)]