Skip to content

Commit

Permalink
Perf: avoid unnecessary allocations when transforming Expr
Browse files Browse the repository at this point in the history
(cherry picked from commit c7347f8)
  • Loading branch information
sadboy authored and peter-toth committed Dec 20, 2023
1 parent 8fa80e7 commit cc15a71
Showing 1 changed file with 14 additions and 8 deletions.
22 changes: 14 additions & 8 deletions datafusion/expr/src/tree_node/expr.rs
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,8 @@

//! Tree node implementation for logical expr
use std::mem;

use crate::expr::{
AggregateFunction, AggregateFunctionDefinition, Alias, Between, BinaryExpr, Case,
Cast, GetIndexedField, GroupingSet, InList, InSubquery, Like, Placeholder,
Expand Down Expand Up @@ -429,15 +431,14 @@ impl TreeNode for Expr {
}
}

fn transform_boxed<F>(boxed_expr: Box<Expr>, transform: &mut F) -> Result<Box<Expr>>
fn transform_boxed<F>(mut boxed_expr: Box<Expr>, transform: &mut F) -> Result<Box<Expr>>
where
F: FnMut(Expr) -> Result<Expr>,
{
// TODO:
// It might be possible to avoid an allocation (the Box::new) below by reusing the box.
let expr: Expr = *boxed_expr;
let rewritten_expr = transform(expr)?;
Ok(Box::new(rewritten_expr))
// We reuse the existing Box to avoid an allocation:
let t = mem::replace(&mut *boxed_expr, Expr::Wildcard { qualifier: None });
let _ = mem::replace(&mut *boxed_expr, transform(t)?);
Ok(boxed_expr)
}

fn transform_option_box<F>(
Expand Down Expand Up @@ -468,9 +469,14 @@ where
}

/// &mut transform a `Vec` of `Expr`s
fn transform_vec<F>(v: Vec<Expr>, transform: &mut F) -> Result<Vec<Expr>>
fn transform_vec<F>(mut v: Vec<Expr>, transform: &mut F) -> Result<Vec<Expr>>
where
F: FnMut(Expr) -> Result<Expr>,
{
v.into_iter().map(transform).collect()
// Perform an in-place mutation of the Vec to avoid allocation:
for expr in v.iter_mut() {
let t = mem::replace(expr, Expr::Wildcard { qualifier: None });
let _ = mem::replace(expr, transform(t)?);
}
Ok(v)
}

0 comments on commit cc15a71

Please sign in to comment.