diff --git a/datafusion/expr/src/expr_rewriter/mod.rs b/datafusion/expr/src/expr_rewriter/mod.rs index f0f02f398642..05ee3e112d62 100644 --- a/datafusion/expr/src/expr_rewriter/mod.rs +++ b/datafusion/expr/src/expr_rewriter/mod.rs @@ -318,21 +318,20 @@ impl NamePreserver { Self { use_alias: true } } - pub fn save(&self, expr: &Expr) -> Result { - let original_name = if self.use_alias { + pub fn save(&self, expr: &Expr) -> SavedName { + if self.use_alias { let (relation, name) = expr.qualified_name(); SavedName::Saved { relation, name } } else { SavedName::None - }; - Ok(original_name) + } } } impl SavedName { /// Ensures the qualified name of the rewritten expression is preserved - pub fn restore(self, expr: Expr) -> Result { - let expr = match self { + pub fn restore(self, expr: Expr) -> Expr { + match self { SavedName::Saved { relation, name } => { let (new_relation, new_name) = expr.qualified_name(); if new_relation != relation || new_name != name { @@ -342,8 +341,7 @@ impl SavedName { } } SavedName::None => expr, - }; - Ok(expr) + } } } @@ -543,9 +541,9 @@ mod test { let mut rewriter = TestRewriter { rewrite_to: rewrite_to.clone(), }; - let saved_name = NamePreserver { use_alias: true }.save(&expr_from).unwrap(); + let saved_name = NamePreserver { use_alias: true }.save(&expr_from); let new_expr = expr_from.clone().rewrite(&mut rewriter).unwrap().data; - let new_expr = saved_name.restore(new_expr).unwrap(); + let new_expr = saved_name.restore(new_expr); let original_name = expr_from.qualified_name(); let new_name = new_expr.qualified_name(); diff --git a/datafusion/expr/src/logical_plan/plan.rs b/datafusion/expr/src/logical_plan/plan.rs index 18a624dd9cb2..20adeb7cc808 100644 --- a/datafusion/expr/src/logical_plan/plan.rs +++ b/datafusion/expr/src/logical_plan/plan.rs @@ -1441,7 +1441,7 @@ impl LogicalPlan { let schema = Arc::clone(plan.schema()); let name_preserver = NamePreserver::new(&plan); plan.map_expressions(|e| { - let original_name = name_preserver.save(&e)?; + let original_name = name_preserver.save(&e); let transformed_expr = e.infer_placeholder_types(&schema)?.transform_up(|e| { if let Expr::Placeholder(Placeholder { id, .. }) = e { @@ -1452,7 +1452,7 @@ impl LogicalPlan { } })?; // Preserve name to avoid breaking column references to this expression - transformed_expr.map_data(|expr| original_name.restore(expr)) + Ok(transformed_expr.update_data(|expr| original_name.restore(expr))) }) }) .map(|res| res.data) diff --git a/datafusion/optimizer/src/analyzer/count_wildcard_rule.rs b/datafusion/optimizer/src/analyzer/count_wildcard_rule.rs index 0036f6df43f6..32ca790b0094 100644 --- a/datafusion/optimizer/src/analyzer/count_wildcard_rule.rs +++ b/datafusion/optimizer/src/analyzer/count_wildcard_rule.rs @@ -76,7 +76,7 @@ fn is_count_star_window_aggregate(window_function: &WindowFunction) -> bool { fn analyze_internal(plan: LogicalPlan) -> Result> { let name_preserver = NamePreserver::new(&plan); plan.map_expressions(|expr| { - let original_name = name_preserver.save(&expr)?; + let original_name = name_preserver.save(&expr); let transformed_expr = expr.transform_up(|expr| match expr { Expr::WindowFunction(mut window_function) if is_count_star_window_aggregate(&window_function) => @@ -94,7 +94,7 @@ fn analyze_internal(plan: LogicalPlan) -> Result> { } _ => Ok(Transformed::no(expr)), })?; - transformed_expr.map_data(|data| original_name.restore(data)) + Ok(transformed_expr.update_data(|data| original_name.restore(data))) }) } diff --git a/datafusion/optimizer/src/analyzer/function_rewrite.rs b/datafusion/optimizer/src/analyzer/function_rewrite.rs index 94f5657b899b..ec3626b2c899 100644 --- a/datafusion/optimizer/src/analyzer/function_rewrite.rs +++ b/datafusion/optimizer/src/analyzer/function_rewrite.rs @@ -61,7 +61,7 @@ impl ApplyFunctionRewrites { let name_preserver = NamePreserver::new(&plan); plan.map_expressions(|expr| { - let original_name = name_preserver.save(&expr)?; + let original_name = name_preserver.save(&expr); // recursively transform the expression, applying the rewrites at each step let transformed_expr = expr.transform_up(|expr| { @@ -74,7 +74,7 @@ impl ApplyFunctionRewrites { Ok(result) })?; - transformed_expr.map_data(|expr| original_name.restore(expr)) + Ok(transformed_expr.update_data(|expr| original_name.restore(expr))) }) } } diff --git a/datafusion/optimizer/src/analyzer/type_coercion.rs b/datafusion/optimizer/src/analyzer/type_coercion.rs index 284650c3d64e..a0b5c9552c83 100644 --- a/datafusion/optimizer/src/analyzer/type_coercion.rs +++ b/datafusion/optimizer/src/analyzer/type_coercion.rs @@ -119,9 +119,9 @@ fn analyze_internal( let name_preserver = NamePreserver::new(&plan); // apply coercion rewrite all expressions in the plan individually plan.map_expressions(|expr| { - let original_name = name_preserver.save(&expr)?; - expr.rewrite(&mut expr_rewrite)? - .map_data(|expr| original_name.restore(expr)) + let original_name = name_preserver.save(&expr); + expr.rewrite(&mut expr_rewrite) + .map(|transformed| transformed.update_data(|e| original_name.restore(e))) })? // some plans need extra coercion after their expressions are coerced .map_data(|plan| expr_rewrite.coerce_plan(plan))? diff --git a/datafusion/optimizer/src/common_subexpr_eliminate.rs b/datafusion/optimizer/src/common_subexpr_eliminate.rs index 583c6cf50de3..14b23a4524fb 100644 --- a/datafusion/optimizer/src/common_subexpr_eliminate.rs +++ b/datafusion/optimizer/src/common_subexpr_eliminate.rs @@ -414,9 +414,9 @@ impl CommonSubexprEliminate { exprs .iter() .map(|expr| name_preserver.save(expr)) - .collect::>>() + .collect::>() }) - .collect::>>()?; + .collect::>(); new_window_expr_list.into_iter().zip(saved_names).try_rfold( new_input, |plan, (new_window_expr, saved_names)| { @@ -426,7 +426,7 @@ impl CommonSubexprEliminate { .map(|(new_window_expr, saved_name)| { saved_name.restore(new_window_expr) }) - .collect::>>()?; + .collect::>(); Window::try_new(new_window_expr, Arc::new(plan)) .map(LogicalPlan::Window) }, @@ -604,14 +604,14 @@ impl CommonSubexprEliminate { let saved_names = aggr_expr .iter() .map(|expr| name_perserver.save(expr)) - .collect::>>()?; + .collect::>(); let new_aggr_expr = rewritten_aggr_expr .into_iter() - .zip(saved_names.into_iter()) + .zip(saved_names) .map(|(new_expr, saved_name)| { saved_name.restore(new_expr) }) - .collect::>>()?; + .collect::>(); // Since `group_expr` may have changed, schema may also. // Use `try_new()` method. diff --git a/datafusion/optimizer/src/optimize_projections/mod.rs b/datafusion/optimizer/src/optimize_projections/mod.rs index 96772d3f2864..0623be504b9b 100644 --- a/datafusion/optimizer/src/optimize_projections/mod.rs +++ b/datafusion/optimizer/src/optimize_projections/mod.rs @@ -497,7 +497,7 @@ fn merge_consecutive_projections(proj: Projection) -> Result Result>>()?; + .collect::>(); Projection::try_new(new_exprs, prev_projection.input).map(Transformed::yes) } else { // not rewritten, so put the projection back together diff --git a/datafusion/optimizer/src/simplify_expressions/simplify_exprs.rs b/datafusion/optimizer/src/simplify_expressions/simplify_exprs.rs index 877989c02d74..c4dba4c8df73 100644 --- a/datafusion/optimizer/src/simplify_expressions/simplify_exprs.rs +++ b/datafusion/optimizer/src/simplify_expressions/simplify_exprs.rs @@ -123,10 +123,10 @@ impl SimplifyExpressions { // Preserve expression names to avoid changing the schema of the plan. let name_preserver = NamePreserver::new(&plan); plan.map_expressions(|e| { - let original_name = name_preserver.save(&e)?; + let original_name = name_preserver.save(&e); let new_e = simplifier .simplify(e) - .and_then(|expr| original_name.restore(expr))?; + .map(|expr| original_name.restore(expr))?; // TODO it would be nice to have a way to know if the expression was simplified // or not. For now conservatively return Transformed::yes Ok(Transformed::yes(new_e)) diff --git a/datafusion/optimizer/src/unwrap_cast_in_comparison.rs b/datafusion/optimizer/src/unwrap_cast_in_comparison.rs index 6043f0d7c8b5..318badc08823 100644 --- a/datafusion/optimizer/src/unwrap_cast_in_comparison.rs +++ b/datafusion/optimizer/src/unwrap_cast_in_comparison.rs @@ -117,9 +117,9 @@ impl OptimizerRule for UnwrapCastInComparison { let name_preserver = NamePreserver::new(&plan); plan.map_expressions(|expr| { - let original_name = name_preserver.save(&expr)?; - expr.rewrite(&mut expr_rewriter)? - .map_data(|expr| original_name.restore(expr)) + let original_name = name_preserver.save(&expr); + expr.rewrite(&mut expr_rewriter) + .map(|transformed| transformed.update_data(|e| original_name.restore(e))) }) } }