apache · tustvold · Dec 15, 2022 · Dec 13, 2022
diff --git a/datafusion-examples/examples/rewrite_expr.rs b/datafusion-examples/examples/rewrite_expr.rs
@@ -73,11 +73,11 @@ impl OptimizerRule for MyRule {
         "my_rule"
     }
 
-    fn optimize(
+    fn try_optimize(
         &self,
         plan: &LogicalPlan,
         _config: &mut OptimizerConfig,
-    ) -> Result<LogicalPlan> {
+    ) -> Result<Option<LogicalPlan>> {
         // recurse down and optimize children first
         let plan = utils::optimize_children(self, plan, _config)?;
 
@@ -86,12 +86,12 @@ impl OptimizerRule for MyRule {
                 let mut expr_rewriter = MyExprRewriter {};
                 let predicate = filter.predicate().clone();
                 let predicate = predicate.rewrite(&mut expr_rewriter)?;
-                Ok(LogicalPlan::Filter(Filter::try_new(
+                Ok(Some(LogicalPlan::Filter(Filter::try_new(
                     predicate,
                     filter.input().clone(),
-                )?))
+                )?)))
             }
-            _ => Ok(plan.clone()),
+            _ => Ok(Some(plan.clone())),
         }
     }
 }

diff --git a/datafusion/core/tests/user_defined_plan.rs b/datafusion/core/tests/user_defined_plan.rs
@@ -282,11 +282,11 @@ impl QueryPlanner for TopKQueryPlanner {
 struct TopKOptimizerRule {}
 impl OptimizerRule for TopKOptimizerRule {
     // Example rewrite pass to insert a user defined LogicalPlanNode
-    fn optimize(
+    fn try_optimize(
         &self,
         plan: &LogicalPlan,
         optimizer_config: &mut OptimizerConfig,
-    ) -> Result<LogicalPlan> {
+    ) -> Result<Option<LogicalPlan>> {
         // Note: this code simply looks for the pattern of a Limit followed by a
         // Sort and replaces it by a TopK node. It does not handle many
         // edge cases (e.g multiple sort columns, sort ASC / DESC), etc.
@@ -304,20 +304,22 @@ impl OptimizerRule for TopKOptimizerRule {
             {
                 if expr.len() == 1 {
                     // we found a sort with a single sort expr, replace with a a TopK
-                    return Ok(LogicalPlan::Extension(Extension {
+                    return Ok(Some(LogicalPlan::Extension(Extension {
                         node: Arc::new(TopKPlanNode {
                             k: *fetch,
-                            input: self.optimize(input.as_ref(), optimizer_config)?,
+                            input: self
+                                .try_optimize(input.as_ref(), optimizer_config)?
+                                .unwrap_or_else(|| input.as_ref().clone()),
                             expr: expr[0].clone(),
                         }),
-                    }));
+                    })));
                 }
             }
         }
 
         // If we didn't find the Limit/Sort combination, recurse as
         // normal and build the result.
-        optimize_children(self, plan, optimizer_config)
+        Ok(Some(optimize_children(self, plan, optimizer_config)?))
     }
 
     fn name(&self) -> &str {

diff --git a/datafusion/optimizer/src/common_subexpr_eliminate.rs b/datafusion/optimizer/src/common_subexpr_eliminate.rs
@@ -91,16 +91,6 @@ impl CommonSubexprEliminate {
 }
 
 impl OptimizerRule for CommonSubexprEliminate {
-    fn optimize(
-        &self,
-        plan: &LogicalPlan,
-        optimizer_config: &mut OptimizerConfig,
-    ) -> Result<LogicalPlan> {
-        Ok(self
-            .try_optimize(plan, optimizer_config)?
-            .unwrap_or_else(|| plan.clone()))
-    }
-
     fn try_optimize(
         &self,
         plan: &LogicalPlan,
@@ -595,7 +585,8 @@ mod test {
     fn assert_optimized_plan_eq(expected: &str, plan: &LogicalPlan) {
         let optimizer = CommonSubexprEliminate {};
         let optimized_plan = optimizer
-            .optimize(plan, &mut OptimizerConfig::new())
+            .try_optimize(plan, &mut OptimizerConfig::new())
+            .unwrap()
             .expect("failed to optimize plan");
         let formatted_plan = format!("{:?}", optimized_plan);
         assert_eq!(expected, formatted_plan);
@@ -839,7 +830,10 @@ mod test {
             .build()
             .unwrap();
         let rule = CommonSubexprEliminate {};
-        let optimized_plan = rule.optimize(&plan, &mut OptimizerConfig::new()).unwrap();
+        let optimized_plan = rule
+            .try_optimize(&plan, &mut OptimizerConfig::new())
+            .unwrap()
+            .unwrap();
 
         let schema = optimized_plan.schema();
         let fields_with_datatypes: Vec<_> = schema

diff --git a/datafusion/optimizer/src/decorrelate_where_exists.rs b/datafusion/optimizer/src/decorrelate_where_exists.rs
@@ -74,16 +74,6 @@ impl DecorrelateWhereExists {
 }
 
 impl OptimizerRule for DecorrelateWhereExists {
-    fn optimize(
-        &self,
-        plan: &LogicalPlan,
-        optimizer_config: &mut OptimizerConfig,
-    ) -> Result<LogicalPlan> {
-        Ok(self
-            .try_optimize(plan, optimizer_config)?
-            .unwrap_or_else(|| plan.clone()))
-    }
-
     fn try_optimize(
         &self,
         plan: &LogicalPlan,

diff --git a/datafusion/optimizer/src/decorrelate_where_in.rs b/datafusion/optimizer/src/decorrelate_where_in.rs
@@ -78,16 +78,6 @@ impl DecorrelateWhereIn {
 }
 
 impl OptimizerRule for DecorrelateWhereIn {
-    fn optimize(
-        &self,
-        plan: &LogicalPlan,
-        optimizer_config: &mut OptimizerConfig,
-    ) -> datafusion_common::Result<LogicalPlan> {
-        Ok(self
-            .try_optimize(plan, optimizer_config)?
-            .unwrap_or_else(|| plan.clone()))
-    }
-
     fn try_optimize(
         &self,
         plan: &LogicalPlan,

diff --git a/datafusion/optimizer/src/eliminate_cross_join.rs b/datafusion/optimizer/src/eliminate_cross_join.rs
@@ -51,20 +51,10 @@ impl EliminateCrossJoin {
 /// This fix helps to improve the performance of TPCH Q19. issue#78
 ///
 impl OptimizerRule for EliminateCrossJoin {
-    fn optimize(
-        &self,
-        plan: &LogicalPlan,
-        _optimizer_config: &mut OptimizerConfig,
-    ) -> Result<LogicalPlan> {
-        Ok(self
-            .try_optimize(plan, _optimizer_config)?
-            .unwrap_or_else(|| plan.clone()))
-    }
-
     fn try_optimize(
         &self,
         plan: &LogicalPlan,
-        _optimizer_config: &mut OptimizerConfig,
+        optimizer_config: &mut OptimizerConfig,
     ) -> Result<Option<LogicalPlan>> {
         match plan {
             LogicalPlan::Filter(filter) => {
@@ -91,7 +81,7 @@ impl OptimizerRule for EliminateCrossJoin {
                         return Ok(Some(utils::optimize_children(
                             self,
                             plan,
-                            _optimizer_config,
+                            optimizer_config,
                         )?));
                     }
                 }
@@ -112,7 +102,7 @@ impl OptimizerRule for EliminateCrossJoin {
                     )?;
                 }
 
-                left = utils::optimize_children(self, &left, _optimizer_config)?;
+                left = utils::optimize_children(self, &left, optimizer_config)?;
 
                 if plan.schema() != left.schema() {
                     left = LogicalPlan::Projection(Projection::new_from_schema(
@@ -141,7 +131,7 @@ impl OptimizerRule for EliminateCrossJoin {
             _ => Ok(Some(utils::optimize_children(
                 self,
                 plan,
-                _optimizer_config,
+                optimizer_config,
             )?)),
         }
     }
@@ -399,7 +389,8 @@ mod tests {
     fn assert_optimized_plan_eq(plan: &LogicalPlan, expected: Vec<&str>) {
         let rule = EliminateCrossJoin::new();
         let optimized_plan = rule
-            .optimize(plan, &mut OptimizerConfig::new())
+            .try_optimize(plan, &mut OptimizerConfig::new())
+            .unwrap()
             .expect("failed to optimize plan");
         let formatted = optimized_plan.display_indent_schema().to_string();
         let actual: Vec<&str> = formatted.trim().lines().collect();

diff --git a/datafusion/optimizer/src/eliminate_filter.rs b/datafusion/optimizer/src/eliminate_filter.rs
@@ -37,20 +37,10 @@ impl EliminateFilter {
 }
 
 impl OptimizerRule for EliminateFilter {
-    fn optimize(
-        &self,
-        plan: &LogicalPlan,
-        _optimizer_config: &mut OptimizerConfig,
-    ) -> Result<LogicalPlan> {
-        Ok(self
-            .try_optimize(plan, _optimizer_config)?
-            .unwrap_or_else(|| plan.clone()))
-    }
-
     fn try_optimize(
         &self,
         plan: &LogicalPlan,
-        _optimizer_config: &mut OptimizerConfig,
+        optimizer_config: &mut OptimizerConfig,
     ) -> Result<Option<LogicalPlan>> {
         let predicate_and_input = match plan {
             LogicalPlan::Filter(filter) => match filter.predicate() {
@@ -63,7 +53,7 @@ impl OptimizerRule for EliminateFilter {
         };
 
         match predicate_and_input {
-            Some((true, input)) => self.try_optimize(input, _optimizer_config),
+            Some((true, input)) => self.try_optimize(input, optimizer_config),
             Some((false, input)) => Ok(Some(LogicalPlan::EmptyRelation(EmptyRelation {
                 produce_one_row: false,
                 schema: input.schema().clone(),
@@ -73,7 +63,7 @@ impl OptimizerRule for EliminateFilter {
                 Ok(Some(utils::optimize_children(
                     self,
                     plan,
-                    _optimizer_config,
+                    optimizer_config,
                 )?))
             }
         }
@@ -93,7 +83,8 @@ mod tests {
     fn assert_optimized_plan_eq(plan: &LogicalPlan, expected: &str) {
         let rule = EliminateFilter::new();
         let optimized_plan = rule
-            .optimize(plan, &mut OptimizerConfig::new())
+            .try_optimize(plan, &mut OptimizerConfig::new())
+            .unwrap()
             .expect("failed to optimize plan");
         let formatted_plan = format!("{:?}", optimized_plan);
         assert_eq!(formatted_plan, expected);

diff --git a/datafusion/optimizer/src/eliminate_limit.rs b/datafusion/optimizer/src/eliminate_limit.rs
@@ -37,16 +37,6 @@ impl EliminateLimit {
 }
 
 impl OptimizerRule for EliminateLimit {
-    fn optimize(
-        &self,
-        plan: &LogicalPlan,
-        optimizer_config: &mut OptimizerConfig,
-    ) -> Result<LogicalPlan> {
-        Ok(self
-            .try_optimize(plan, optimizer_config)?
-            .unwrap_or_else(|| plan.clone()))
-    }
-
     fn try_optimize(
         &self,
         plan: &LogicalPlan,
@@ -100,7 +90,8 @@ mod tests {
 
     fn assert_optimized_plan_eq(plan: &LogicalPlan, expected: &str) -> Result<()> {
         let optimized_plan = EliminateLimit::new()
-            .optimize(plan, &mut OptimizerConfig::new())
+            .try_optimize(plan, &mut OptimizerConfig::new())
+            .unwrap()
             .expect("failed to optimize plan");
         let formatted_plan = format!("{:?}", optimized_plan);
         assert_eq!(formatted_plan, expected);
@@ -113,10 +104,12 @@ mod tests {
         expected: &str,
     ) -> Result<()> {
         let optimized_plan = PushDownLimit::new()
-            .optimize(plan, &mut OptimizerConfig::new())
+            .try_optimize(plan, &mut OptimizerConfig::new())
+            .unwrap()
             .expect("failed to optimize plan");
         let optimized_plan = EliminateLimit::new()
-            .optimize(&optimized_plan, &mut OptimizerConfig::new())
+            .try_optimize(&optimized_plan, &mut OptimizerConfig::new())
+            .unwrap()
             .expect("failed to optimize plan");
         let formatted_plan = format!("{:?}", optimized_plan);
         assert_eq!(formatted_plan, expected);

diff --git a/datafusion/optimizer/src/eliminate_outer_join.rs b/datafusion/optimizer/src/eliminate_outer_join.rs
@@ -62,16 +62,6 @@ impl EliminateOuterJoin {
 
 /// Attempt to eliminate outer joins.
 impl OptimizerRule for EliminateOuterJoin {
-    fn optimize(
-        &self,
-        plan: &LogicalPlan,
-        optimizer_config: &mut OptimizerConfig,
-    ) -> Result<LogicalPlan> {
-        Ok(self
-            .try_optimize(plan, optimizer_config)?
-            .unwrap_or_else(|| plan.clone()))
-    }
-
     fn try_optimize(
         &self,
         plan: &LogicalPlan,
@@ -330,7 +320,8 @@ mod tests {
     fn assert_optimized_plan_eq(plan: &LogicalPlan, expected: &str) -> Result<()> {
         let rule = EliminateOuterJoin::new();
         let optimized_plan = rule
-            .optimize(plan, &mut OptimizerConfig::new())
+            .try_optimize(plan, &mut OptimizerConfig::new())
+            .unwrap()
             .expect("failed to optimize plan");
         let formatted_plan = format!("{:?}", optimized_plan);
         assert_eq!(formatted_plan, expected);

diff --git a/datafusion/optimizer/src/filter_null_join_keys.rs b/datafusion/optimizer/src/filter_null_join_keys.rs
@@ -35,16 +35,6 @@ use std::sync::Arc;
 pub struct FilterNullJoinKeys {}
 
 impl OptimizerRule for FilterNullJoinKeys {
-    fn optimize(
-        &self,
-        plan: &LogicalPlan,
-        optimizer_config: &mut OptimizerConfig,
-    ) -> datafusion_common::Result<LogicalPlan> {
-        Ok(self
-            .try_optimize(plan, optimizer_config)?
-            .unwrap_or_else(|| plan.clone()))
-    }
-
     fn try_optimize(
         &self,
         plan: &LogicalPlan,
@@ -165,7 +155,8 @@ mod tests {
 
     fn optimize_plan(plan: &LogicalPlan) -> LogicalPlan {
         let rule = FilterNullJoinKeys::default();
-        rule.optimize(plan, &mut OptimizerConfig::new())
+        rule.try_optimize(plan, &mut OptimizerConfig::new())
+            .unwrap()
             .expect("failed to optimize plan")
     }
 

diff --git a/datafusion/optimizer/src/inline_table_scan.rs b/datafusion/optimizer/src/inline_table_scan.rs
@@ -35,20 +35,10 @@ impl InlineTableScan {
 }
 
 impl OptimizerRule for InlineTableScan {
-    fn optimize(
-        &self,
-        plan: &LogicalPlan,
-        _optimizer_config: &mut OptimizerConfig,
-    ) -> Result<LogicalPlan> {
-        Ok(self
-            .try_optimize(plan, _optimizer_config)?
-            .unwrap_or_else(|| plan.clone()))
-    }
-
     fn try_optimize(
         &self,
         plan: &LogicalPlan,
-        _optimizer_config: &mut OptimizerConfig,
+        optimizer_config: &mut OptimizerConfig,
     ) -> Result<Option<LogicalPlan>> {
         match plan {
             // Match only on scans without filter / projection / fetch
@@ -63,7 +53,7 @@ impl OptimizerRule for InlineTableScan {
                 if let Some(sub_plan) = source.get_logical_plan() {
                     // Recursively apply optimization
                     let plan =
-                        utils::optimize_children(self, sub_plan, _optimizer_config)?;
+                        utils::optimize_children(self, sub_plan, optimizer_config)?;
                     let plan = LogicalPlanBuilder::from(plan)
                         .project(vec![Expr::Wildcard])?
                         .alias(table_name)?;
@@ -80,7 +70,7 @@ impl OptimizerRule for InlineTableScan {
                 Ok(Some(utils::optimize_children(
                     self,
                     plan,
-                    _optimizer_config,
+                    optimizer_config,
                 )?))
             }
         }
@@ -168,7 +158,8 @@ mod tests {
         let plan = scan.filter(col("x.a").eq(lit(1))).unwrap().build().unwrap();
 
         let optimized_plan = rule
-            .optimize(&plan, &mut OptimizerConfig::new())
+            .try_optimize(&plan, &mut OptimizerConfig::new())
+            .unwrap()
             .expect("failed to optimize plan");
         let formatted_plan = format!("{:?}", optimized_plan);
         let expected = "Filter: x.a = Int32(1)\