diff --git a/datafusion-examples/examples/rewrite_expr.rs b/datafusion-examples/examples/rewrite_expr.rs index 7a2ee6d93cd1c..216a6932c8d16 100644 --- 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 { + ) -> Result> { // 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 index 1b18e4ebe404d..4b5a7e66f18a3 100644 --- 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 { + ) -> Result> { // 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 index b541aa0ee76ca..6885a0bfbb68d 100644 --- 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 { - 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 index 6b5bd184f5042..54e346a08a685 100644 --- 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 { - 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 index 4614b0699cc7a..dac5ff03b29cd 100644 --- 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 { - 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 index b431742f0acf1..a00362d70e38a 100644 --- 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 { - 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> { 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 index 3e1c91a700d3a..046d7b11f01d9 100644 --- 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 { - 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> { 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 index b60b16e83bc5a..45844e120426d 100644 --- 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 { - 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 index 4d1fecad70e89..a78d61dace929 100644 --- 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 { - 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 index f1a4f13937cd9..8563c13a875a8 100644 --- 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 { - 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 index e0238a7513023..d8283bf7148d3 100644 --- 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 { - 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> { 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)\ diff --git a/datafusion/optimizer/src/optimizer.rs b/datafusion/optimizer/src/optimizer.rs index dd4783ceb3b1c..4e9eadc47e5b6 100644 --- a/datafusion/optimizer/src/optimizer.rs +++ b/datafusion/optimizer/src/optimizer.rs @@ -55,17 +55,7 @@ pub trait OptimizerRule { &self, plan: &LogicalPlan, optimizer_config: &mut OptimizerConfig, - ) -> Result> { - self.optimize(plan, optimizer_config).map(Some) - } - - /// Rewrite `plan` to an optimized form. This method will eventually be deprecated and - /// replace by `try_optimize`. - fn optimize( - &self, - plan: &LogicalPlan, - optimizer_config: &mut OptimizerConfig, - ) -> Result; + ) -> Result>; /// A human readable name for this optimizer rule fn name(&self) -> &str; @@ -407,11 +397,11 @@ mod tests { struct BadRule {} impl OptimizerRule for BadRule { - fn optimize( + fn try_optimize( &self, _plan: &LogicalPlan, _optimizer_config: &mut OptimizerConfig, - ) -> datafusion_common::Result { + ) -> datafusion_common::Result> { Err(DataFusionError::Plan("rule failed".to_string())) } @@ -424,13 +414,13 @@ mod tests { struct GetTableScanRule {} impl OptimizerRule for GetTableScanRule { - fn optimize( + fn try_optimize( &self, _plan: &LogicalPlan, _optimizer_config: &mut OptimizerConfig, - ) -> datafusion_common::Result { + ) -> datafusion_common::Result> { let table_scan = test_table_scan()?; - LogicalPlanBuilder::from(table_scan).build() + Ok(Some(LogicalPlanBuilder::from(table_scan).build()?)) } fn name(&self) -> &str { diff --git a/datafusion/optimizer/src/propagate_empty_relation.rs b/datafusion/optimizer/src/propagate_empty_relation.rs index cfdb354d4cb97..9b5869366d5ff 100644 --- a/datafusion/optimizer/src/propagate_empty_relation.rs +++ b/datafusion/optimizer/src/propagate_empty_relation.rs @@ -34,16 +34,6 @@ impl PropagateEmptyRelation { } impl OptimizerRule for PropagateEmptyRelation { - fn optimize( - &self, - plan: &LogicalPlan, - optimizer_config: &mut OptimizerConfig, - ) -> Result { - Ok(self - .try_optimize(plan, optimizer_config)? - .unwrap_or_else(|| plan.clone())) - } - fn try_optimize( &self, plan: &LogicalPlan, @@ -227,7 +217,8 @@ mod tests { fn assert_optimized_plan_eq(plan: &LogicalPlan, expected: &str) { let rule = PropagateEmptyRelation::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); @@ -236,10 +227,12 @@ mod tests { fn assert_together_optimized_plan_eq(plan: &LogicalPlan, expected: &str) { let optimize_one = EliminateFilter::new() - .optimize(plan, &mut OptimizerConfig::new()) + .try_optimize(plan, &mut OptimizerConfig::new()) + .unwrap() .expect("failed to optimize plan"); let optimize_two = PropagateEmptyRelation::new() - .optimize(&optimize_one, &mut OptimizerConfig::new()) + .try_optimize(&optimize_one, &mut OptimizerConfig::new()) + .unwrap() .expect("failed to optimize plan"); let formatted_plan = format!("{:?}", optimize_two); assert_eq!(formatted_plan, expected); diff --git a/datafusion/optimizer/src/push_down_filter.rs b/datafusion/optimizer/src/push_down_filter.rs index 40ed00ebda04c..fc2a2d84ad361 100644 --- a/datafusion/optimizer/src/push_down_filter.rs +++ b/datafusion/optimizer/src/push_down_filter.rs @@ -503,16 +503,6 @@ impl OptimizerRule for PushDownFilter { "push_down_filter" } - fn optimize( - &self, - plan: &LogicalPlan, - optimizer_config: &mut OptimizerConfig, - ) -> Result { - Ok(self - .try_optimize(plan, optimizer_config)? - .unwrap_or_else(|| plan.clone())) - } - fn try_optimize( &self, plan: &LogicalPlan, @@ -808,7 +798,8 @@ mod tests { fn assert_optimized_plan_eq(plan: &LogicalPlan, expected: &str) -> Result<()> { let optimized_plan = PushDownFilter::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!(plan.schema(), optimized_plan.schema()); @@ -1955,8 +1946,9 @@ mod tests { table_scan_with_pushdown_provider(TableProviderFilterPushDown::Inexact)?; let optimised_plan = PushDownFilter::new() - .optimize(&plan, &mut OptimizerConfig::new()) - .expect("failed to optimize plan"); + .try_optimize(&plan, &mut OptimizerConfig::new()) + .expect("failed to optimize plan") + .unwrap(); let expected = "\ Filter: a = Int64(1)\ @@ -2306,7 +2298,8 @@ mod tests { // Originally global state which can help to avoid duplicate Filters been generated and pushed down. // Now the global state is removed. Need to double confirm that avoid duplicate Filters. let optimized_plan = PushDownFilter::new() - .optimize(&plan, &mut OptimizerConfig::new()) + .try_optimize(&plan, &mut OptimizerConfig::new()) + .unwrap() .expect("failed to optimize plan"); assert_optimized_plan_eq(&optimized_plan, expected) } diff --git a/datafusion/optimizer/src/push_down_limit.rs b/datafusion/optimizer/src/push_down_limit.rs index a404762b69ad8..22e837a6c439f 100644 --- a/datafusion/optimizer/src/push_down_limit.rs +++ b/datafusion/optimizer/src/push_down_limit.rs @@ -75,16 +75,6 @@ fn push_down_join( /// Push down Limit. impl OptimizerRule for PushDownLimit { - fn optimize( - &self, - plan: &LogicalPlan, - optimizer_config: &mut OptimizerConfig, - ) -> Result { - Ok(self - .try_optimize(plan, optimizer_config)? - .unwrap_or_else(|| plan.clone())) - } - fn try_optimize( &self, plan: &LogicalPlan, @@ -281,7 +271,8 @@ mod test { fn assert_optimized_plan_eq(plan: &LogicalPlan, 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 formatted_plan = format!("{:?}", optimized_plan); diff --git a/datafusion/optimizer/src/push_down_projection.rs b/datafusion/optimizer/src/push_down_projection.rs index 0238148c3ec15..642aa188e1671 100644 --- a/datafusion/optimizer/src/push_down_projection.rs +++ b/datafusion/optimizer/src/push_down_projection.rs @@ -46,16 +46,6 @@ use std::{ pub struct PushDownProjection {} impl OptimizerRule for PushDownProjection { - fn optimize( - &self, - plan: &LogicalPlan, - optimizer_config: &mut OptimizerConfig, - ) -> Result { - Ok(self - .try_optimize(plan, optimizer_config)? - .unwrap_or_else(|| plan.clone())) - } - fn try_optimize( &self, plan: &LogicalPlan, @@ -1029,6 +1019,8 @@ mod tests { fn optimize(plan: &LogicalPlan) -> Result { let rule = PushDownProjection::new(); - rule.optimize(plan, &mut OptimizerConfig::new()) + Ok(rule + .try_optimize(plan, &mut OptimizerConfig::new())? + .unwrap()) } } diff --git a/datafusion/optimizer/src/rewrite_disjunctive_predicate.rs b/datafusion/optimizer/src/rewrite_disjunctive_predicate.rs index 3f6cc763bb30b..079046273f0fd 100644 --- a/datafusion/optimizer/src/rewrite_disjunctive_predicate.rs +++ b/datafusion/optimizer/src/rewrite_disjunctive_predicate.rs @@ -124,20 +124,10 @@ impl RewriteDisjunctivePredicate { } impl OptimizerRule for RewriteDisjunctivePredicate { - fn optimize( - &self, - plan: &LogicalPlan, - _optimizer_config: &mut OptimizerConfig, - ) -> Result { - 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> { match plan { LogicalPlan::Filter(filter) => { @@ -146,13 +136,15 @@ impl OptimizerRule for RewriteDisjunctivePredicate { let rewritten_expr = normalize_predicate(rewritten_predicate); Ok(Some(LogicalPlan::Filter(Filter::try_new( rewritten_expr, - Arc::new(Self::optimize(self, filter.input(), _optimizer_config)?), + self.try_optimize(filter.input(), optimizer_config)? + .map(Arc::new) + .unwrap_or_else(|| filter.input().clone()), )?))) } _ => Ok(Some(utils::optimize_children( self, plan, - _optimizer_config, + optimizer_config, )?)), } } diff --git a/datafusion/optimizer/src/scalar_subquery_to_join.rs b/datafusion/optimizer/src/scalar_subquery_to_join.rs index 00e4d89c8df4b..3ee46c8d6c47f 100644 --- a/datafusion/optimizer/src/scalar_subquery_to_join.rs +++ b/datafusion/optimizer/src/scalar_subquery_to_join.rs @@ -68,9 +68,10 @@ impl ScalarSubqueryToJoin { Ok(subquery) => subquery, _ => return Ok(()), }; - let subquery = - self.optimize(&subquery.subquery, optimizer_config)?; - let subquery = Arc::new(subquery); + let subquery = self + .try_optimize(&subquery.subquery, optimizer_config)? + .map(Arc::new) + .unwrap_or_else(|| subquery.subquery.clone()); let subquery = Subquery { subquery }; let res = SubqueryInfo::new(subquery, expr, *op, lhs); subqueries.push(res); @@ -89,16 +90,6 @@ impl ScalarSubqueryToJoin { } impl OptimizerRule for ScalarSubqueryToJoin { - fn optimize( - &self, - plan: &LogicalPlan, - optimizer_config: &mut OptimizerConfig, - ) -> Result { - Ok(self - .try_optimize(plan, optimizer_config)? - .unwrap_or_else(|| plan.clone())) - } - fn try_optimize( &self, plan: &LogicalPlan, @@ -107,7 +98,9 @@ impl OptimizerRule for ScalarSubqueryToJoin { match plan { LogicalPlan::Filter(filter) => { // Apply optimizer rule to current input - let optimized_input = self.optimize(filter.input(), optimizer_config)?; + let optimized_input = self + .try_optimize(filter.input(), optimizer_config)? + .unwrap_or_else(|| filter.input().as_ref().clone()); let (subqueries, other_exprs) = self.extract_subquery_exprs(filter.predicate(), optimizer_config)?; diff --git a/datafusion/optimizer/src/simplify_expressions/simplify_exprs.rs b/datafusion/optimizer/src/simplify_expressions/simplify_exprs.rs index d2c57f3a79c46..f56776b2720f9 100644 --- a/datafusion/optimizer/src/simplify_expressions/simplify_exprs.rs +++ b/datafusion/optimizer/src/simplify_expressions/simplify_exprs.rs @@ -43,15 +43,15 @@ impl OptimizerRule for SimplifyExpressions { "simplify_expressions" } - fn optimize( + fn try_optimize( &self, plan: &LogicalPlan, optimizer_config: &mut OptimizerConfig, - ) -> Result { + ) -> Result> { let mut execution_props = ExecutionProps::new(); execution_props.query_execution_start_time = optimizer_config.query_execution_start_time(); - Self::optimize_internal(plan, &execution_props) + Ok(Some(Self::optimize_internal(plan, &execution_props)?)) } } @@ -172,7 +172,8 @@ mod tests { fn assert_optimized_plan_eq(plan: &LogicalPlan, expected: &str) -> Result<()> { let rule = SimplifyExpressions::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); @@ -384,7 +385,7 @@ mod tests { let rule = SimplifyExpressions::new(); let err = rule - .optimize(plan, &mut config) + .try_optimize(plan, &mut config) .expect_err("expected optimization to fail"); err.to_string() @@ -399,7 +400,8 @@ mod tests { let rule = SimplifyExpressions::new(); let optimized_plan = rule - .optimize(plan, &mut config) + .try_optimize(plan, &mut config) + .unwrap() .expect("failed to optimize plan"); format!("{:?}", optimized_plan) } diff --git a/datafusion/optimizer/src/single_distinct_to_groupby.rs b/datafusion/optimizer/src/single_distinct_to_groupby.rs index 1e13223952670..e4878bd16441a 100644 --- a/datafusion/optimizer/src/single_distinct_to_groupby.rs +++ b/datafusion/optimizer/src/single_distinct_to_groupby.rs @@ -83,20 +83,10 @@ fn contains_grouping_set(expr: &[Expr]) -> bool { } impl OptimizerRule for SingleDistinctToGroupBy { - fn optimize( - &self, - plan: &LogicalPlan, - optimizer_config: &mut OptimizerConfig, - ) -> Result { - 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> { match plan { LogicalPlan::Aggregate(Aggregate { @@ -166,7 +156,7 @@ impl OptimizerRule for SingleDistinctToGroupBy { Vec::new(), )?); let inner_agg = - utils::optimize_children(self, &grouped_aggr, _optimizer_config)?; + utils::optimize_children(self, &grouped_aggr, optimizer_config)?; let outer_aggr_schema = Arc::new(DFSchema::new_with_metadata( outer_group_exprs @@ -213,14 +203,14 @@ impl OptimizerRule for SingleDistinctToGroupBy { Ok(Some(utils::optimize_children( self, plan, - _optimizer_config, + optimizer_config, )?)) } } _ => Ok(Some(utils::optimize_children( self, plan, - _optimizer_config, + optimizer_config, )?)), } } @@ -242,7 +232,8 @@ mod tests { fn assert_optimized_plan_eq(plan: &LogicalPlan, expected: &str) { let rule = SingleDistinctToGroupBy::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.display_indent_schema()); diff --git a/datafusion/optimizer/src/subquery_filter_to_join.rs b/datafusion/optimizer/src/subquery_filter_to_join.rs index ccf6931d3e98f..36c18ea749634 100644 --- a/datafusion/optimizer/src/subquery_filter_to_join.rs +++ b/datafusion/optimizer/src/subquery_filter_to_join.rs @@ -49,16 +49,6 @@ impl SubqueryFilterToJoin { } impl OptimizerRule for SubqueryFilterToJoin { - fn optimize( - &self, - plan: &LogicalPlan, - optimizer_config: &mut OptimizerConfig, - ) -> Result { - Ok(self - .try_optimize(plan, optimizer_config)? - .unwrap_or_else(|| plan.clone())) - } - fn try_optimize( &self, plan: &LogicalPlan, @@ -67,7 +57,9 @@ impl OptimizerRule for SubqueryFilterToJoin { match plan { LogicalPlan::Filter(filter) => { // Apply optimizer rule to current input - let optimized_input = self.optimize(filter.input(), optimizer_config)?; + let optimized_input = self + .try_optimize(filter.input(), optimizer_config)? + .unwrap_or_else(|| filter.input().as_ref().clone()); // Splitting filter expression into components by AND let filters = utils::split_conjunction(filter.predicate()); @@ -104,10 +96,10 @@ impl OptimizerRule for SubqueryFilterToJoin { subquery, negated, } => { - let right_input = self.optimize( + let right_input = self.try_optimize( &subquery.subquery, optimizer_config - )?; + )?.unwrap_or_else(||subquery.subquery.as_ref().clone()); let right_schema = right_input.schema(); if right_schema.fields().len() != 1 { return Err(DataFusionError::Plan( @@ -220,7 +212,8 @@ mod tests { fn assert_optimized_plan_eq(plan: &LogicalPlan, expected: &str) { let rule = SubqueryFilterToJoin::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.display_indent_schema()); assert_eq!(formatted_plan, expected); diff --git a/datafusion/optimizer/src/test/mod.rs b/datafusion/optimizer/src/test/mod.rs index f5f93517e4921..1f51b38f68106 100644 --- a/datafusion/optimizer/src/test/mod.rs +++ b/datafusion/optimizer/src/test/mod.rs @@ -107,7 +107,8 @@ pub fn assert_optimized_plan_eq( expected: &str, ) { 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.display_indent_schema()); assert_eq!(formatted_plan, expected); @@ -118,9 +119,9 @@ pub fn assert_optimizer_err( plan: &LogicalPlan, expected: &str, ) { - let res = rule.optimize(plan, &mut OptimizerConfig::new()); + let res = rule.try_optimize(plan, &mut OptimizerConfig::new()); match res { - Ok(plan) => assert_eq!(format!("{}", plan.display_indent()), "An error"), + Ok(plan) => assert_eq!(format!("{}", plan.unwrap().display_indent()), "An error"), Err(ref e) => { let actual = format!("{}", e); if expected.is_empty() || !actual.contains(expected) { @@ -131,7 +132,10 @@ pub fn assert_optimizer_err( } pub fn assert_optimization_skipped(rule: &dyn OptimizerRule, plan: &LogicalPlan) { - let new_plan = rule.optimize(plan, &mut OptimizerConfig::new()).unwrap(); + let new_plan = rule + .try_optimize(plan, &mut OptimizerConfig::new()) + .unwrap() + .unwrap(); assert_eq!( format!("{}", plan.display_indent()), format!("{}", new_plan.display_indent()) diff --git a/datafusion/optimizer/src/type_coercion.rs b/datafusion/optimizer/src/type_coercion.rs index 27e06782afc56..a7b68cc18d44d 100644 --- a/datafusion/optimizer/src/type_coercion.rs +++ b/datafusion/optimizer/src/type_coercion.rs @@ -55,16 +55,6 @@ impl OptimizerRule for TypeCoercion { "type_coercion" } - fn optimize( - &self, - plan: &LogicalPlan, - optimizer_config: &mut OptimizerConfig, - ) -> Result { - Ok(self - .try_optimize(plan, optimizer_config)? - .unwrap_or_else(|| plan.clone())) - } - fn try_optimize( &self, plan: &LogicalPlan, @@ -617,7 +607,7 @@ mod test { fn assert_optimized_plan_eq(plan: &LogicalPlan, expected: &str) -> Result<()> { let rule = TypeCoercion::new(); let mut config = OptimizerConfig::default(); - let plan = rule.optimize(plan, &mut config)?; + let plan = rule.try_optimize(plan, &mut config)?.unwrap(); assert_eq!(expected, &format!("{:?}", plan)); Ok(()) } diff --git a/datafusion/optimizer/src/unwrap_cast_in_comparison.rs b/datafusion/optimizer/src/unwrap_cast_in_comparison.rs index bfcb0f85aba2a..687e15b965c9f 100644 --- a/datafusion/optimizer/src/unwrap_cast_in_comparison.rs +++ b/datafusion/optimizer/src/unwrap_cast_in_comparison.rs @@ -79,16 +79,6 @@ impl UnwrapCastInComparison { } impl OptimizerRule for UnwrapCastInComparison { - fn optimize( - &self, - plan: &LogicalPlan, - _optimizer_config: &mut OptimizerConfig, - ) -> Result { - Ok(self - .try_optimize(plan, _optimizer_config)? - .unwrap_or_else(|| plan.clone())) - } - fn try_optimize( &self, plan: &LogicalPlan,