diff --git a/datafusion/optimizer/src/eliminate_filter.rs b/datafusion/optimizer/src/eliminate_filter.rs index 2bf5cfa303904..368e94c92c8d9 100644 --- a/datafusion/optimizer/src/eliminate_filter.rs +++ b/datafusion/optimizer/src/eliminate_filter.rs @@ -17,13 +17,12 @@ //! [`EliminateFilter`] replaces `where false` or `where null` with an empty relation. -use crate::optimizer::ApplyOrder; -use datafusion_common::{Result, ScalarValue}; -use datafusion_expr::{ - logical_plan::{EmptyRelation, LogicalPlan}, - Expr, Filter, -}; +use datafusion_common::tree_node::Transformed; +use datafusion_common::{internal_err, Result, ScalarValue}; +use datafusion_expr::logical_plan::tree_node::unwrap_arc; +use datafusion_expr::{EmptyRelation, Expr, Filter, LogicalPlan}; +use crate::optimizer::ApplyOrder; use crate::{OptimizerConfig, OptimizerRule}; /// Optimization rule that eliminate the scalar value (true/false/null) filter @@ -44,31 +43,10 @@ impl EliminateFilter { impl OptimizerRule for EliminateFilter { fn try_optimize( &self, - plan: &LogicalPlan, + _plan: &LogicalPlan, _config: &dyn OptimizerConfig, ) -> Result> { - match plan { - LogicalPlan::Filter(Filter { - predicate: Expr::Literal(ScalarValue::Boolean(v)), - input, - .. - }) => { - match *v { - // input also can be filter, apply again - Some(true) => Ok(Some( - self.try_optimize(input, _config)? - .unwrap_or_else(|| input.as_ref().clone()), - )), - Some(false) | None => { - Ok(Some(LogicalPlan::EmptyRelation(EmptyRelation { - produce_one_row: false, - schema: input.schema().clone(), - }))) - } - } - } - _ => Ok(None), - } + internal_err!("Should have called EliminateFilter::rewrite") } fn name(&self) -> &str { @@ -78,17 +56,47 @@ impl OptimizerRule for EliminateFilter { fn apply_order(&self) -> Option { Some(ApplyOrder::TopDown) } + + fn supports_rewrite(&self) -> bool { + true + } + + fn rewrite( + &self, + plan: LogicalPlan, + _config: &dyn OptimizerConfig, + ) -> Result> { + match plan { + LogicalPlan::Filter(Filter { + predicate: Expr::Literal(ScalarValue::Boolean(v)), + input, + .. + }) => match v { + Some(true) => Ok(Transformed::yes( + self.rewrite(unwrap_arc(input), _config)?.data, + )), + Some(false) | None => Ok(Transformed::yes(LogicalPlan::EmptyRelation( + EmptyRelation { + produce_one_row: false, + schema: unwrap_arc(input).schema().clone(), + }, + ))), + }, + _ => Ok(Transformed::no(plan)), + } + } } #[cfg(test)] mod tests { - use crate::eliminate_filter::EliminateFilter; + use std::sync::Arc; + use datafusion_common::{Result, ScalarValue}; use datafusion_expr::{ col, lit, logical_plan::builder::LogicalPlanBuilder, sum, Expr, LogicalPlan, }; - use std::sync::Arc; + use crate::eliminate_filter::EliminateFilter; use crate::test::*; fn assert_optimized_plan_equal(plan: LogicalPlan, expected: &str) -> Result<()> {