Skip to content

Commit 37d0c7a

Browse files
committed
Reduce cloning in LogicalPlanBuilder
- Migrate function arguments from `LogicalPlan` to `impl Into<Arc<LogicalPlan>>` - Update usages (mostly in tests)
1 parent 102a162 commit 37d0c7a

10 files changed

Lines changed: 194 additions & 194 deletions

File tree

‎datafusion/core/src/physical_planner.rs‎

Lines changed: 3 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1377,17 +1377,15 @@ impl DefaultPhysicalPlanner {
13771377
let (left, left_col_keys, left_projected) =
13781378
wrap_projection_for_join_if_necessary(
13791379
&left_keys,
1380-
original_left.as_ref().clone(),
1380+
Arc::clone(original_left),
13811381
)?;
13821382
let (right, right_col_keys, right_projected) =
13831383
wrap_projection_for_join_if_necessary(
13841384
&right_keys,
1385-
original_right.as_ref().clone(),
1385+
Arc::clone(original_right),
13861386
)?;
1387-
let column_on = (left_col_keys, right_col_keys);
13881387

1389-
let left = Arc::new(left);
1390-
let right = Arc::new(right);
1388+
let column_on = (left_col_keys, right_col_keys);
13911389
let (new_join, requalified) = Join::try_new_with_project_input(
13921390
node,
13931391
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
@@ -259,26 +259,30 @@ pub fn strip_outer_reference(expr: Expr) -> Expr {
259259
/// Returns plan with expressions coerced to types compatible with
260260
/// schema types
261261
pub fn coerce_plan_expr_for_schema(
262-
plan: LogicalPlan,
262+
plan: impl Into<Arc<LogicalPlan>>,
263263
schema: &DFSchema,
264-
) -> Result<LogicalPlan> {
265-
match plan {
264+
) -> Result<Arc<LogicalPlan>> {
265+
let plan = plan.into();
266+
if matches!(plan.as_ref(), LogicalPlan::Projection(_)) {
266267
// special case Projection to avoid adding multiple projections
267-
LogicalPlan::Projection(Projection { expr, input, .. }) => {
268-
let new_exprs = coerce_exprs_for_schema(expr, input.schema(), schema)?;
269-
let projection = Projection::try_new(new_exprs, input)?;
270-
Ok(LogicalPlan::Projection(projection))
271-
}
272-
_ => {
273-
let exprs: Vec<Expr> = plan.schema().iter().map(Expr::from).collect();
274-
let new_exprs = coerce_exprs_for_schema(exprs, plan.schema(), schema)?;
275-
let add_project = new_exprs.iter().any(|expr| expr.try_as_col().is_none());
276-
if add_project {
277-
let projection = Projection::try_new(new_exprs, Arc::new(plan))?;
278-
Ok(LogicalPlan::Projection(projection))
279-
} else {
280-
Ok(plan)
281-
}
268+
let LogicalPlan::Projection(Projection { expr, input, .. }) =
269+
Arc::unwrap_or_clone(plan)
270+
else {
271+
unreachable!()
272+
};
273+
274+
let new_exprs = coerce_exprs_for_schema(expr, input.schema(), schema)?;
275+
let projection = Projection::try_new(new_exprs, input)?;
276+
Ok(Arc::new(LogicalPlan::Projection(projection)))
277+
} else {
278+
let exprs: Vec<Expr> = plan.schema().iter().map(Expr::from).collect();
279+
let new_exprs = coerce_exprs_for_schema(exprs, plan.schema(), schema)?;
280+
let add_project = new_exprs.iter().any(|expr| expr.try_as_col().is_none());
281+
if add_project {
282+
let projection = Projection::try_new(new_exprs, plan)?;
283+
Ok(Arc::new(LogicalPlan::Projection(projection)))
284+
} else {
285+
Ok(plan)
282286
}
283287
}
284288
}
@@ -482,7 +486,7 @@ mod test {
482486
fn normalize_cols() {
483487
let expr = col("a") + col("b") + col("c");
484488

485-
// Schemas with some matching and some non matching cols
489+
// Schemas with some matching and some non-matching cols
486490
let schema_a = make_schema_with_empty_metadata(
487491
vec![Some("tableA".into()), Some("tableA".into())],
488492
vec!["a", "aa"],

0 commit comments

Comments
 (0)