Skip to content

Commit 2aaeb76

Browse files
committed
Reduce cloning in LogicalPlanBuilder
- Migrate function arguments from `LogicalPlan` to `impl Into<Arc<LogicalPlan>>` - Update usages (mostly in tests)
1 parent 50d74a7 commit 2aaeb76

10 files changed

Lines changed: 199 additions & 199 deletions

File tree

‎datafusion/core/src/physical_planner.rs‎

Lines changed: 3 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1364,17 +1364,15 @@ impl DefaultPhysicalPlanner {
13641364
let (left, left_col_keys, left_projected) =
13651365
wrap_projection_for_join_if_necessary(
13661366
&left_keys,
1367-
original_left.as_ref().clone(),
1367+
Arc::clone(original_left),
13681368
)?;
13691369
let (right, right_col_keys, right_projected) =
13701370
wrap_projection_for_join_if_necessary(
13711371
&right_keys,
1372-
original_right.as_ref().clone(),
1372+
Arc::clone(original_right),
13731373
)?;
1374-
let column_on = (left_col_keys, right_col_keys);
13751374

1376-
let left = Arc::new(left);
1377-
let right = Arc::new(right);
1375+
let column_on = (left_col_keys, right_col_keys);
13781376
let (new_join, requalified) = Join::try_new_with_project_input(
13791377
node,
13801378
Arc::clone(&left),

‎datafusion/expr/src/expr_rewriter/mod.rs‎

Lines changed: 25 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -41,15 +41,15 @@ pub use order_by::rewrite_sort_cols_by_aggs;
4141

4242
/// Trait for rewriting [`Expr`]s into function calls.
4343
///
44-
/// This trait is used with `FunctionRegistry::register_function_rewrite` to
44+
/// This trait is used with `FunctionRegistry::register_function_rewrite`
4545
/// to evaluating `Expr`s using functions that may not be built in to DataFusion
4646
///
4747
/// For example, concatenating arrays `a || b` is represented as
4848
/// `Operator::ArrowAt`, but can be implemented by calling a function
4949
/// `array_concat` from the `functions-nested` crate.
5050
// This is not used in datafusion internally, but it is still helpful for downstream project so don't remove it.
5151
pub trait FunctionRewrite: Debug {
52-
/// Return a human readable name for this rewrite
52+
/// Return a human-readable name for this rewrite
5353
fn name(&self) -> &str;
5454

5555
/// Potentially rewrite `expr` to some other expression
@@ -219,26 +219,30 @@ pub fn strip_outer_reference(expr: Expr) -> Expr {
219219
/// Returns plan with expressions coerced to types compatible with
220220
/// schema types
221221
pub fn coerce_plan_expr_for_schema(
222-
plan: LogicalPlan,
222+
plan: impl Into<Arc<LogicalPlan>>,
223223
schema: &DFSchema,
224-
) -> Result<LogicalPlan> {
225-
match plan {
224+
) -> Result<Arc<LogicalPlan>> {
225+
let plan = plan.into();
226+
if matches!(plan.as_ref(), LogicalPlan::Projection(_)) {
226227
// special case Projection to avoid adding multiple projections
227-
LogicalPlan::Projection(Projection { expr, input, .. }) => {
228-
let new_exprs = coerce_exprs_for_schema(expr, input.schema(), schema)?;
229-
let projection = Projection::try_new(new_exprs, input)?;
230-
Ok(LogicalPlan::Projection(projection))
231-
}
232-
_ => {
233-
let exprs: Vec<Expr> = plan.schema().iter().map(Expr::from).collect();
234-
let new_exprs = coerce_exprs_for_schema(exprs, plan.schema(), schema)?;
235-
let add_project = new_exprs.iter().any(|expr| expr.try_as_col().is_none());
236-
if add_project {
237-
let projection = Projection::try_new(new_exprs, Arc::new(plan))?;
238-
Ok(LogicalPlan::Projection(projection))
239-
} else {
240-
Ok(plan)
241-
}
228+
let LogicalPlan::Projection(Projection { expr, input, .. }) =
229+
Arc::unwrap_or_clone(plan)
230+
else {
231+
unreachable!()
232+
};
233+
234+
let new_exprs = coerce_exprs_for_schema(expr, input.schema(), schema)?;
235+
let projection = Projection::try_new(new_exprs, input)?;
236+
Ok(Arc::new(LogicalPlan::Projection(projection)))
237+
} else {
238+
let exprs: Vec<Expr> = plan.schema().iter().map(Expr::from).collect();
239+
let new_exprs = coerce_exprs_for_schema(exprs, plan.schema(), schema)?;
240+
let add_project = new_exprs.iter().any(|expr| expr.try_as_col().is_none());
241+
if add_project {
242+
let projection = Projection::try_new(new_exprs, plan)?;
243+
Ok(Arc::new(LogicalPlan::Projection(projection)))
244+
} else {
245+
Ok(plan)
242246
}
243247
}
244248
}
@@ -442,7 +446,7 @@ mod test {
442446
fn normalize_cols() {
443447
let expr = col("a") + col("b") + col("c");
444448

445-
// Schemas with some matching and some non matching cols
449+
// Schemas with some matching and some non-matching cols
446450
let schema_a = make_schema_with_empty_metadata(
447451
vec![Some("tableA".into()), Some("tableA".into())],
448452
vec!["a", "aa"],

0 commit comments

Comments
 (0)