Skip to content
166 changes: 158 additions & 8 deletions datafusion/expr/src/expr.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2988,6 +2988,36 @@ fn rewrite_placeholder_from_subquery(
}

struct SchemaDisplay<'a>(&'a Expr);

fn write_schema_display_binary_child(
f: &mut Formatter<'_>,
expr: &Expr,
precedence: u8,
is_right: bool,
) -> fmt::Result {
match expr {
Expr::BinaryExpr(child) => {
let child_precedence = child.op.precedence();
if child_precedence == 0
|| child_precedence < precedence
|| (is_right && child_precedence == precedence)
{
write!(f, "({})", SchemaDisplay(expr))
} else {
write!(f, "{}", SchemaDisplay(expr))
}
}
_ => write!(f, "{}", SchemaDisplay(expr)),
}
}

fn write_schema_display_unary_child(f: &mut Formatter<'_>, expr: &Expr) -> fmt::Result {
match expr {
Expr::BinaryExpr(_) => write!(f, "({})", SchemaDisplay(expr)),
_ => write!(f, "{}", SchemaDisplay(expr)),
}
}

impl Display for SchemaDisplay<'_> {
fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
match self.0 {
Expand Down Expand Up @@ -3042,7 +3072,10 @@ impl Display for SchemaDisplay<'_> {
}
}
Expr::BinaryExpr(BinaryExpr { left, op, right }) => {
write!(f, "{} {op} {}", SchemaDisplay(left), SchemaDisplay(right),)
let precedence = op.precedence();
write_schema_display_binary_child(f, left, precedence, false)?;
write!(f, " {op} ")?;
write_schema_display_binary_child(f, right, precedence, true)
}
Expr::Case(Case {
expr,
Expand Down Expand Up @@ -3152,8 +3185,15 @@ impl Display for SchemaDisplay<'_> {

Ok(())
}
Expr::Negative(expr) => write!(f, "(- {})", SchemaDisplay(expr)),
Expr::Not(expr) => write!(f, "NOT {}", SchemaDisplay(expr)),
Expr::Negative(expr) => {
write!(f, "(- ")?;
write_schema_display_unary_child(f, expr)?;
write!(f, ")")
}
Expr::Not(expr) => {
write!(f, "NOT ")?;
write_schema_display_unary_child(f, expr)
}
Expr::Unnest(Unnest { expr, outer }) => {
let name = if *outer { "UNNEST_OUTER" } else { "UNNEST" };
write!(f, "{name}({})", SchemaDisplay(expr))
Expand Down Expand Up @@ -3294,6 +3334,35 @@ impl Display for SchemaDisplay<'_> {
/// A helper struct for displaying an `Expr` as an SQL-like string.
struct SqlDisplay<'a>(&'a Expr);

fn write_sql_display_binary_child(
f: &mut Formatter<'_>,
expr: &Expr,
precedence: u8,
is_right: bool,
) -> fmt::Result {
match expr {
Expr::BinaryExpr(child) => {
let child_precedence = child.op.precedence();
if child_precedence == 0
|| child_precedence < precedence
|| (is_right && child_precedence == precedence)
{
write!(f, "({})", SqlDisplay(expr))
} else {
write!(f, "{}", SqlDisplay(expr))
}
}
_ => write!(f, "{}", SqlDisplay(expr)),
}
}

fn write_sql_display_unary_child(f: &mut Formatter<'_>, expr: &Expr) -> fmt::Result {
match expr {
Expr::BinaryExpr(_) => write!(f, "({})", SqlDisplay(expr)),
_ => write!(f, "{}", SqlDisplay(expr)),
}
}

impl Display for SqlDisplay<'_> {
fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
match self.0 {
Expand Down Expand Up @@ -3324,7 +3393,10 @@ impl Display for SqlDisplay<'_> {
}
}
Expr::BinaryExpr(BinaryExpr { left, op, right }) => {
write!(f, "{} {op} {}", SqlDisplay(left), SqlDisplay(right),)
let precedence = op.precedence();
write_sql_display_binary_child(f, left, precedence, false)?;
write!(f, " {op} ")?;
write_sql_display_binary_child(f, right, precedence, true)
}
Expr::Case(Case {
expr,
Expand Down Expand Up @@ -3428,8 +3500,15 @@ impl Display for SqlDisplay<'_> {

Ok(())
}
Expr::Negative(expr) => write!(f, "(- {})", SqlDisplay(expr)),
Expr::Not(expr) => write!(f, "NOT {}", SqlDisplay(expr)),
Expr::Negative(expr) => {
write!(f, "(- ")?;
write_sql_display_unary_child(f, expr)?;
write!(f, ")")
}
Expr::Not(expr) => {
write!(f, "NOT ")?;
write_sql_display_unary_child(f, expr)
}
Expr::Unnest(Unnest { expr, outer }) => {
let name = if *outer { "UNNEST_OUTER" } else { "UNNEST" };
write!(f, "{name}({})", SqlDisplay(expr))
Expand Down Expand Up @@ -3557,6 +3636,13 @@ pub fn schema_name_from_sorts(sorts: &[Sort]) -> Result<String, fmt::Error> {
pub const OUTER_REFERENCE_COLUMN_PREFIX: &str = "outer_ref";
pub const UNNEST_COLUMN_PREFIX: &str = "UNNEST";

fn write_expr_display_unary_child(f: &mut Formatter<'_>, expr: &Expr) -> fmt::Result {
match expr {
Expr::BinaryExpr(_) => write!(f, "({expr})"),
_ => write!(f, "{expr}"),
}
}

/// Format expressions for display as part of a logical plan. In many cases, this will produce
/// similar output to `Expr.name()` except that column names will be prefixed with '#'.
impl Display for Expr {
Expand Down Expand Up @@ -3597,8 +3683,15 @@ impl Display for Expr {
format_type_and_metadata(field.data_type(), Some(field.metadata()));
write!(f, "TRY_CAST({expr} AS {formatted})")
}
Expr::Not(expr) => write!(f, "NOT {expr}"),
Expr::Negative(expr) => write!(f, "(- {expr})"),
Expr::Not(expr) => {
write!(f, "NOT ")?;
write_expr_display_unary_child(f, expr)
}
Expr::Negative(expr) => {
write!(f, "(- ")?;
write_expr_display_unary_child(f, expr)?;
write!(f, ")")
}
Expr::IsNull(expr) => write!(f, "{expr} IS NULL"),
Expr::IsNotNull(expr) => write!(f, "{expr} IS NOT NULL"),
Expr::IsTrue(expr) => write!(f, "{expr} IS TRUE"),
Expand Down Expand Up @@ -4229,6 +4322,63 @@ mod test {
assert_eq!("NULL", null_expr.human_display().to_string());
}

#[test]
fn format_nested_binary_exprs_with_parentheses() {
let one_plus_two = binary_expr(lit(1i64), Operator::Plus, lit(2i64));
let expr = binary_expr(one_plus_two.clone(), Operator::Multiply, lit(3i64));

assert_eq!(
"(Int64(1) + Int64(2)) * Int64(3)",
expr.schema_name().to_string()
);
assert_eq!("(1 + 2) * 3", expr.human_display().to_string());

let nested_subtraction = binary_expr(
lit(1i64),
Operator::Minus,
binary_expr(lit(2i64), Operator::Minus, lit(3i64)),
);
assert_eq!(
"Int64(1) - (Int64(2) - Int64(3))",
nested_subtraction.schema_name().to_string()
);
assert_eq!(
"1 - (2 - 3)",
nested_subtraction.human_display().to_string()
);

let nested_multiplication = binary_expr(
lit(6i64),
Operator::Divide,
binary_expr(lit(2i64), Operator::Multiply, lit(3i64)),
);
assert_eq!(
"Int64(6) / (Int64(2) * Int64(3))",
nested_multiplication.schema_name().to_string()
);
assert_eq!(
"6 / (2 * 3)",
nested_multiplication.human_display().to_string()
);

let negative_expr = Expr::Negative(Box::new(one_plus_two.clone()));
assert_eq!(
"(- (Int64(1) + Int64(2)))",
negative_expr.schema_name().to_string()
);
assert_eq!("(- (1 + 2))", negative_expr.human_display().to_string());
assert_eq!("(- (Int64(1) + Int64(2)))", negative_expr.to_string());

let not_expr =
Expr::Not(Box::new(binary_expr(lit(1i64), Operator::Eq, lit(2i64))));
assert_eq!(
"NOT (Int64(1) = Int64(2))",
not_expr.schema_name().to_string()
);
assert_eq!("NOT (1 = 2)", not_expr.human_display().to_string());
assert_eq!("NOT (Int64(1) = Int64(2))", not_expr.to_string());
}

#[test]
fn test_partial_ord() {
// Test validates that partial ord is defined for Expr, not
Expand Down
45 changes: 41 additions & 4 deletions datafusion/optimizer/src/common_subexpr_eliminate.rs
Original file line number Diff line number Diff line change
Expand Up @@ -85,9 +85,19 @@ impl CommonSubexprEliminate {
schema,
..
} = projection;
let name_preserver = NamePreserver::new_for_projection();
let saved_names = expr
.iter()
.map(|expr| name_preserver.save(expr))
.collect::<Vec<_>>();
let input = Arc::unwrap_or_clone(input);
self.try_unary_plan(expr, input, config)?
.map_data(|(new_expr, new_input)| {
let new_expr = new_expr
.into_iter()
.zip(saved_names)
.map(|(expr, saved_name)| saved_name.restore(expr))
.collect();
Projection::try_new_with_schema(new_expr, Arc::new(new_input), schema)
.map(LogicalPlan::Projection)
})
Expand Down Expand Up @@ -852,6 +862,7 @@ mod test {

use super::*;
use crate::assert_optimized_plan_eq_snapshot;
use crate::optimize_projections::OptimizeProjections;
use crate::optimizer::OptimizerContext;
use crate::test::udfs::leaf_udf_expr;
use crate::test::*;
Expand Down Expand Up @@ -912,7 +923,7 @@ mod test {
assert_optimized_plan_equal!(
plan,
@ r"
Aggregate: groupBy=[[]], aggr=[[sum(__common_expr_1 AS test.a * Int32(1) - test.b), sum(__common_expr_1 AS test.a * Int32(1) - test.b * (Int32(1) + test.c))]]
Aggregate: groupBy=[[]], aggr=[[sum(__common_expr_1 AS test.a * (Int32(1) - test.b)), sum(__common_expr_1 AS test.a * (Int32(1) - test.b) * (Int32(1) + test.c))]]
Projection: test.a * (Int32(1) - test.b) AS __common_expr_1, test.a, test.b, test.c
TableScan: test
"
Expand All @@ -933,13 +944,39 @@ mod test {
assert_optimized_plan_equal!(
plan,
@ r"
Projection: __common_expr_1 - test.c AS alias1 * __common_expr_1 AS test.a + test.b, __common_expr_1 AS test.a + test.b
Projection: __common_expr_1 - test.c AS alias1 * __common_expr_1 AS test.a + test.b AS alias1 * (test.a + test.b), __common_expr_1 AS test.a + test.b
Projection: test.a + test.b AS __common_expr_1, test.a, test.b, test.c
TableScan: test
"
)
}

#[test]
fn projection_name_preserved_for_nested_common_expression() -> Result<()> {
let table_scan = test_table_scan()?;
let common_expr = col("a") * col("b");
let plan = LogicalPlanBuilder::from(table_scan)
.project(vec![
common_expr.clone() + Expr::Negative(Box::new(common_expr)),
])?
.build()?;

let rules: Vec<Arc<dyn OptimizerRule + Send + Sync>> = vec![
Arc::new(CommonSubexprEliminate::new()),
Arc::new(OptimizeProjections::new()),
];
assert_optimized_plan_eq_snapshot!(
OptimizerContext::new(),
rules,
plan,
@ r"
Projection: __common_expr_1 AS test.a * test.b + (- __common_expr_1 AS test.a * test.b) AS test.a * test.b + (- (test.a * test.b))
Projection: test.a * test.b AS __common_expr_1
TableScan: test projection=[a, b]
"
)
}

#[test]
fn aggregate() -> Result<()> {
let table_scan = test_table_scan()?;
Expand Down Expand Up @@ -1770,8 +1807,8 @@ mod test {
assert_optimized_plan_equal!(
plan,
@ r"
Projection: __common_expr_1 AS NOT test.a = test.b, __common_expr_1 AS NOT test.b = test.a
Projection: NOT test.a = test.b AS __common_expr_1, test.a, test.b, test.c
Projection: __common_expr_1 AS NOT (test.a = test.b), __common_expr_1 AS NOT (test.b = test.a)
Projection: NOT (test.a = test.b) AS __common_expr_1, test.a, test.b, test.c
TableScan: test
"
)?;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -647,7 +647,7 @@ mod tests {
.build()?;

let actual = get_optimized_plan_formatted(plan, &time);
let expected = "Projection: NOT test.a AS Boolean(true) OR Boolean(false) != test.a\
let expected = "Projection: NOT test.a AS (Boolean(true) OR Boolean(false)) != test.a\
\n TableScan: test";

assert_eq!(expected, actual);
Expand Down
2 changes: 1 addition & 1 deletion datafusion/sql/tests/sql_integration.rs
Original file line number Diff line number Diff line change
Expand Up @@ -3963,7 +3963,7 @@ fn negative_sum_intervals_in_projection() {
assert_snapshot!(
plan,
@r#"
Projection: (- IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 2, nanoseconds: 0 }") + IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 5, nanoseconds: 0 }") + (- IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 4, nanoseconds: 0 }") + IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 7, nanoseconds: 0 }")))
Projection: (- (IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 2, nanoseconds: 0 }") + IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 5, nanoseconds: 0 }") + (- (IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 4, nanoseconds: 0 }") + IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 7, nanoseconds: 0 }")))))
EmptyRelation: rows=1
"#
);
Expand Down
Loading
Loading