Skip to content

Commit

Permalink
Avoid copies in TypeCoercion
Browse files Browse the repository at this point in the history
  • Loading branch information
alamb committed Apr 23, 2024
1 parent a98adc7 commit f2e2c3b
Show file tree
Hide file tree
Showing 2 changed files with 25 additions and 10 deletions.
1 change: 1 addition & 0 deletions datafusion/expr/src/logical_plan/plan.rs
Original file line number Diff line number Diff line change
Expand Up @@ -815,6 +815,7 @@ impl LogicalPlan {
}
}
}

/// Replaces placeholder param values (like `$1`, `$2`) in [`LogicalPlan`]
/// with the specified `param_values`.
///
Expand Down
34 changes: 24 additions & 10 deletions datafusion/optimizer/src/analyzer/type_coercion.rs
Original file line number Diff line number Diff line change
Expand Up @@ -106,16 +106,30 @@ fn analyze_internal(
};

let preserver = NamePreserver::new(&plan);
Ok(plan
.map_expressions(|expr| {
// ensure aggregate names don't change:
// https://github.com/apache/datafusion/issues/3555
let original_name = preserver.save(&expr)?;
expr.rewrite(&mut expr_rewrite)?
.map_data(|expr| original_name.restore(expr))
})?
// propagate the the transformation information from children
.update_transformed(children_transformed))
let transformed_plan = plan.map_expressions(|expr| {
// ensure aggregate names don't change:
// https://github.com/apache/datafusion/issues/3555
let original_name = preserver.save(&expr)?;
expr.rewrite(&mut expr_rewrite)?
.map_data(|expr| original_name.restore(expr))
})?;

// if any of the expressions were rewritten, we need to recreate the plan to
// recalculate the schema. At the moment this requires a copy
if transformed_plan.transformed || children_transformed {
// TODO avoid this copy
let plan = transformed_plan.data;
let new_inputs = plan
.inputs()
.into_iter()
.map(|input| input.clone())
.collect::<Vec<_>>();

plan.with_new_exprs(plan.expressions(), new_inputs)
.map(Transformed::yes)
} else {
Ok(transformed_plan)
}
}

pub(crate) struct TypeCoercionRewriter {
Expand Down

0 comments on commit f2e2c3b

Please sign in to comment.