Update optimize_children to return Result<Option<LogicalPlan>>

datafusion-examples/examples/rewrite_expr.rs

@@ -79,10 +79,9 @@ impl OptimizerRule for MyRule {
         config: &dyn OptimizerConfig,
     ) -> Result<Option<LogicalPlan>> {
         // recurse down and optimize children first
-        let plan = utils::optimize_children(self, plan, config)?;
-
-        match plan {
-            LogicalPlan::Filter(filter) => {
+        let optimized_plan = utils::optimize_children(self, plan, config)?;
+        match optimized_plan {
+            Some(LogicalPlan::Filter(filter)) => {
                 let mut expr_rewriter = MyExprRewriter {};
                 let predicate = filter.predicate.clone();
                 let predicate = predicate.rewrite(&mut expr_rewriter)?;
@@ -91,7 +90,19 @@ impl OptimizerRule for MyRule {
                     filter.input,
                 )?)))
             }
-            _ => Ok(Some(plan.clone())),
+            Some(optimized_plan) => Ok(Some(optimized_plan)),
+            None => match plan {
+                LogicalPlan::Filter(filter) => {
+                    let mut expr_rewriter = MyExprRewriter {};
+                    let predicate = filter.predicate.clone();
+                    let predicate = predicate.rewrite(&mut expr_rewriter)?;
+                    Ok(Some(LogicalPlan::Filter(Filter::try_new(
+                        predicate,
+                        filter.input.clone(),
+                    )?)))
+                }
+                _ => Ok(None),
+            },
         }
     }
 }

datafusion/core/tests/user_defined_plan.rs

@@ -318,7 +318,7 @@ impl OptimizerRule for TopKOptimizerRule {
 
         // If we didn't find the Limit/Sort combination, recurse as
         // normal and build the result.
-        Ok(Some(optimize_children(self, plan, config)?))
+        optimize_children(self, plan, config)
     }
 
     fn name(&self) -> &str {

datafusion/optimizer/src/common_subexpr_eliminate.rs

@@ -102,7 +102,7 @@ impl OptimizerRule for CommonSubexprEliminate {
         let mut expr_set = ExprSet::new();
 
         let original_schema = plan.schema().clone();
-        let mut optimized_plan = match plan {
+        let optimized_plan = match plan {
             LogicalPlan::Projection(Projection {
                 expr,
                 input,
@@ -115,11 +115,11 @@ impl OptimizerRule for CommonSubexprEliminate {
                 let (mut new_expr, new_input) =
                     self.rewrite_expr(&[expr], &[&arrays], input, &mut expr_set, config)?;
 
-                LogicalPlan::Projection(Projection::try_new_with_schema(
+                Some(LogicalPlan::Projection(Projection::try_new_with_schema(
                     pop_expr(&mut new_expr)?,
                     Arc::new(new_input),
                     schema.clone(),
-                )?)
+                )?))
             }
             LogicalPlan::Filter(filter) => {
                 let input = &filter.input;
@@ -142,7 +142,10 @@ impl OptimizerRule for CommonSubexprEliminate {
                 )?;
 
                 if let Some(predicate) = pop_expr(&mut new_expr)?.pop() {
-                    LogicalPlan::Filter(Filter::try_new(predicate, Arc::new(new_input))?)
+                    Some(LogicalPlan::Filter(Filter::try_new(
+                        predicate,
+                        Arc::new(new_input),
+                    )?))
                 } else {
                     return Err(DataFusionError::Internal(
                         "Failed to pop predicate expr".to_string(),
@@ -165,11 +168,11 @@ impl OptimizerRule for CommonSubexprEliminate {
                     config,
                 )?;
 
-                LogicalPlan::Window(Window {
+                Some(LogicalPlan::Window(Window {
                     input: Arc::new(new_input),
                     window_expr: pop_expr(&mut new_expr)?,
                     schema: schema.clone(),
-                })
+                }))
             }
             LogicalPlan::Aggregate(Aggregate {
                 group_expr,
@@ -194,12 +197,12 @@ impl OptimizerRule for CommonSubexprEliminate {
                 let new_aggr_expr = pop_expr(&mut new_expr)?;
                 let new_group_expr = pop_expr(&mut new_expr)?;
 
-                LogicalPlan::Aggregate(Aggregate::try_new_with_schema(
+                Some(LogicalPlan::Aggregate(Aggregate::try_new_with_schema(
                     Arc::new(new_input),
                     new_group_expr,
                     new_aggr_expr,
                     schema.clone(),
-                )?)
+                )?))
             }
             LogicalPlan::Sort(Sort { expr, input, fetch }) => {
                 let input_schema = Arc::clone(input.schema());
@@ -208,11 +211,11 @@ impl OptimizerRule for CommonSubexprEliminate {
                 let (mut new_expr, new_input) =
                     self.rewrite_expr(&[expr], &[&arrays], input, &mut expr_set, config)?;
 
-                LogicalPlan::Sort(Sort {
+                Some(LogicalPlan::Sort(Sort {
                     expr: pop_expr(&mut new_expr)?,
                     input: Arc::new(new_input),
                     fetch: *fetch,
-                })
+                }))
             }
             LogicalPlan::Join(_)
             | LogicalPlan::CrossJoin(_)
@@ -242,12 +245,16 @@ impl OptimizerRule for CommonSubexprEliminate {
             }
         };
 
-        // add an additional projection if the output schema changed.
-        if optimized_plan.schema() != &original_schema {
-            optimized_plan = build_recover_project_plan(&original_schema, optimized_plan);
+        match optimized_plan {
+            Some(optimized_plan) if optimized_plan.schema() != &original_schema => {
+                // add an additional projection if the output schema changed.
+                Ok(Some(build_recover_project_plan(
+                    &original_schema,
+                    optimized_plan,
+                )))
+            }
+            plan => Ok(plan),
         }
-
-        Ok(Some(optimized_plan))
     }
 
     fn name(&self) -> &str {

datafusion/optimizer/src/eliminate_cross_join.rs

@@ -82,7 +82,7 @@ impl OptimizerRule for EliminateCrossJoin {
                         )?;
                     }
                     _ => {
-                        return Ok(Some(utils::optimize_children(self, plan, config)?));
+                        return utils::optimize_children(self, plan, config);
                     }
                 }
 
@@ -102,7 +102,7 @@ impl OptimizerRule for EliminateCrossJoin {
                     )?;
                 }
 
-                left = utils::optimize_children(self, &left, config)?;
+                left = utils::optimize_children(self, &left, config)?.unwrap_or(left);
 
                 if plan.schema() != left.schema() {
                     left = LogicalPlan::Projection(Projection::new_from_schema(
@@ -128,7 +128,7 @@ impl OptimizerRule for EliminateCrossJoin {
                 }
             }
 
-            _ => Ok(Some(utils::optimize_children(self, plan, config)?)),
+            _ => utils::optimize_children(self, plan, config),
         }
     }
 

datafusion/optimizer/src/push_down_filter.rs

@@ -28,6 +28,7 @@ use datafusion_expr::{
 use std::collections::{HashMap, HashSet};
 use std::iter::once;
 use std::sync::Arc;
+use utils::optimize_children;
 
 /// Push Down Filter optimizer rule pushes filter clauses down the plan
 /// # Introduction
@@ -524,15 +525,14 @@ impl OptimizerRule for PushDownFilter {
             LogicalPlan::Join(join) => {
                 let optimized_plan = push_down_join(plan, join, None)?;
                 return match optimized_plan {
-                    Some(optimized_plan) => Ok(Some(utils::optimize_children(
-                        self,
-                        &optimized_plan,
-                        config,
-                    )?)),
-                    None => Ok(Some(utils::optimize_children(self, plan, config)?)),
+                    Some(optimized_plan) => Ok(Some(
+                        optimize_children(self, &optimized_plan, config)?
+                            .unwrap_or(optimized_plan),
+                    )),
+                    None => optimize_children(self, plan, config),
                 };
             }
-            _ => return Ok(Some(utils::optimize_children(self, plan, config)?)),
+            _ => return optimize_children(self, plan, config),
         };
 
         let child_plan = filter.input.as_ref();
@@ -749,7 +749,9 @@ impl OptimizerRule for PushDownFilter {
             _ => plan.clone(),
         };
 
-        Ok(Some(utils::optimize_children(self, &new_plan, config)?))
+        Ok(Some(
+            optimize_children(self, &new_plan, config)?.unwrap_or(new_plan),
+        ))
     }
 }
 

datafusion/optimizer/src/utils.rs

@@ -37,18 +37,26 @@ use std::sync::Arc;
 /// type. Useful for optimizer rules which want to leave the type
 /// of plan unchanged but still apply to the children.
 /// This also handles the case when the `plan` is a [`LogicalPlan::Explain`].
+///
+/// Returning `Ok(None)` indicates that the plan can't be optimized by the `optimizer`.
 pub fn optimize_children(
     optimizer: &impl OptimizerRule,
     plan: &LogicalPlan,
     config: &dyn OptimizerConfig,
-) -> Result<LogicalPlan> {
+) -> Result<Option<LogicalPlan>> {
     let new_exprs = plan.expressions();
     let mut new_inputs = Vec::with_capacity(plan.inputs().len());
+    let mut plan_is_changed = false;
     for input in plan.inputs() {
         let new_input = optimizer.try_optimize(input, config)?;
+        plan_is_changed = plan_is_changed || new_input.is_some();
         new_inputs.push(new_input.unwrap_or_else(|| input.clone()))
     }
-    from_plan(plan, &new_exprs, &new_inputs)
+    if plan_is_changed {
+        Ok(Some(from_plan(plan, &new_exprs, &new_inputs)?))
+    } else {
+        Ok(None)
+    }
 }
 
 /// Splits a conjunctive [`Expr`] such as `A AND B AND C` => `[A, B, C]`