Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
229 changes: 120 additions & 109 deletions datafusion/optimizer/src/simplify_expressions/expr_simplifier.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1575,7 +1575,9 @@ impl TreeNodeRewriter for Simplifier<'_> {
// ---> (X AND A) OR (Y AND B AND NOT X) OR ... (NOT (X OR Y) AND Q)
//
// Note: the rationale for this rewrite is that the expr can then be further
// simplified using the existing rules for AND/OR
// simplified using the existing rules for AND/OR. Unlike CASE, AND/OR
// do not guarantee branch-local evaluation, so only expose columns
// and literals from conditional branches.
Expr::Case(Case {
expr: None,
when_then_expr,
Expand All @@ -1586,7 +1588,8 @@ impl TreeNodeRewriter for Simplifier<'_> {
// or all thens are literal bools and a small number of them are true
|| (when_then_expr.iter().all(|(_, then)| is_bool_lit(then))
&& when_then_expr.iter().filter(|(_, then)| is_true(then)).count() < 3))
&& info.is_boolean_type(&when_then_expr[0].1)? =>
&& info.is_boolean_type(&when_then_expr[0].1)?
&& can_lower_case_to_boolean(&when_then_expr, else_expr.as_deref()) =>
{
// String disjunction of all the when predicates encountered so far. Not nullable.
let mut filter_expr = lit(false);
Expand Down Expand Up @@ -1643,7 +1646,8 @@ impl TreeNodeRewriter for Simplifier<'_> {
.filter(|(_, then)| is_false(then))
.count()
< 3
&& else_expr.as_deref().is_none_or(is_bool_lit) =>
&& else_expr.as_deref().is_none_or(is_bool_lit)
&& can_lower_case_to_boolean(&when_then_expr, else_expr.as_deref()) =>
{
Transformed::yes(
Expr::Case(Case {
Expand Down Expand Up @@ -2432,6 +2436,20 @@ fn simplify_inlist_set_operation(
}))
}

/// Conservatively checks the inputs whose evaluation can change when lowering
/// CASE to AND/OR. [`Expr`] has no general fallibility analysis: only columns and
/// literals are admitted for all WHEN conditions and branch outputs.
fn can_lower_case_to_boolean(
when_then_expr: &[(Box<Expr>, Box<Expr>)],
else_expr: Option<&Expr>,
) -> bool {
let is_leaf = |expr: &Expr| matches!(expr, Expr::Column(_) | Expr::Literal(..));
when_then_expr
.iter()
.all(|(when, then)| is_leaf(when) && is_leaf(then))
&& else_expr.is_none_or(is_leaf)
}

/// Returns expression testing a boolean `expr` for being exactly `true` (not `false` or NULL).
fn is_exactly_true(expr: Expr, info: &SimplifyContext) -> Result<Expr> {
if !info.nullable(&expr)? {
Expand Down Expand Up @@ -2577,12 +2595,15 @@ mod tests {
// ELSE false
// END
//
// Can be simplified to `i < 5`
// Fold the constant conditions, but preserve CASE because the THEN
// expression is outside the conservative column/literal subset.
let expr = when(col("i").gt(lit(5)).and(lit(false)), col("i").gt(lit(5)))
.when(col("i").lt(lit(5)).and(lit(true)), col("i").lt(lit(5)))
.otherwise(lit(false))
.unwrap();
let expected = col("i").lt(lit(5));
let expected = when(col("i").lt(lit(5)), col("i").lt(lit(5)))
.otherwise(lit(false))
.unwrap();
assert_eq!(expected, simplifier.simplify(expr).unwrap());
}

Expand Down Expand Up @@ -4210,76 +4231,89 @@ mod tests {
Some(Box::new(lit("ready"))),
));

assert_eq!(
simplify(binary_expr(
complex_case.clone(),
// Comparisons still fold into literal outputs, but later WHEN conditions must
// remain conditional: they are not columns or literals.
for (op, value, outputs) in [
(
Operator::Eq,
lit("completed"),
)),
not_distinct_from(col("c1").eq(lit("completed")), lit(true)).and(
distinct_from(col("c1").eq(lit("inboxed")), lit(true))
.and(distinct_from(col("c1").eq(lit("scheduled")), lit(true)))
)
);

assert_eq!(
simplify(binary_expr(
complex_case.clone(),
"completed",
[false, false, true, false, false, false, false],
),
(
Operator::NotEq,
lit("completed"),
)),
distinct_from(col("c1").eq(lit("completed")), lit(true))
.or(not_distinct_from(col("c1").eq(lit("inboxed")), lit(true))
.or(not_distinct_from(col("c1").eq(lit("scheduled")), lit(true))))
);

assert_eq!(
simplify(binary_expr(
complex_case.clone(),
"completed",
[true, true, false, true, true, true, true],
),
(
Operator::Eq,
lit("running"),
)),
not_distinct_from(col("c2"), lit(true)).and(
distinct_from(col("c1").eq(lit("inboxed")), lit(true))
.and(distinct_from(col("c1").eq(lit("scheduled")), lit(true)))
.and(distinct_from(col("c1").eq(lit("completed")), lit(true)))
.and(distinct_from(col("c1").eq(lit("paused")), lit(true)))
)
);

assert_eq!(
simplify(binary_expr(
complex_case.clone(),
"running",
[false, false, false, false, true, false, false],
),
(
Operator::Eq,
lit("ready"),
)),
distinct_from(col("c1").eq(lit("inboxed")), lit(true))
.and(distinct_from(col("c1").eq(lit("scheduled")), lit(true)))
.and(distinct_from(col("c1").eq(lit("completed")), lit(true)))
.and(distinct_from(col("c1").eq(lit("paused")), lit(true)))
.and(distinct_from(col("c2"), lit(true)))
.and(distinct_from(
col("c1").eq(lit("invoked")).and(col("c3").gt(lit(0))),
lit(true)
))
);

assert_eq!(
simplify(binary_expr(
complex_case.clone(),
"ready",
[false, false, false, false, false, false, true],
),
(
Operator::NotEq,
lit("ready"),
"ready",
[true, true, true, true, true, true, false],
),
] {
let Expr::Case(mut expected) = complex_case.clone() else {
unreachable!()
};
for ((_, then), value) in expected.when_then_expr.iter_mut().zip(outputs) {
**then = lit(value);
}
expected.else_expr = Some(Box::new(lit(outputs[6])));
assert_eq!(
simplify(binary_expr(complex_case.clone(), op, lit(value))),
Expr::Case(expected)
);
}
}

#[test]
fn simplify_case_preserves_conditional_expressions() {
let fallible =
Expr::Cast(Cast::new(Box::new(col("c1")), DataType::Int32)).gt(lit(0_i32));
let volatile = Expr::ScalarFunction(ScalarFunction::new_udf(
Arc::new(ScalarUDF::new_from_impl(VolatileUdf::new())),
vec![],
));
for expr in [
Expr::Case(Case::new(
None,
vec![(Box::new(fallible.clone()), Box::new(lit(false)))],
Some(Box::new(lit(false))),
)),
not_distinct_from(col("c1").eq(lit("inboxed")), lit(true))
.or(not_distinct_from(col("c1").eq(lit("scheduled")), lit(true)))
.or(not_distinct_from(col("c1").eq(lit("completed")), lit(true)))
.or(not_distinct_from(col("c1").eq(lit("paused")), lit(true)))
.or(not_distinct_from(col("c2"), lit(true)))
.or(not_distinct_from(
col("c1").eq(lit("invoked")).and(col("c3").gt(lit(0))),
lit(true)
))
);
Expr::Case(Case::new(
None,
vec![(Box::new(col("c2")), Box::new(fallible.clone()))],
Some(Box::new(lit(false))),
)),
Expr::Case(Case::new(
None,
vec![(Box::new(col("c2")), Box::new(lit(true)))],
Some(Box::new(fallible.clone())),
)),
Expr::Case(Case::new(
None,
vec![
(Box::new(col("c2")), Box::new(lit(false))),
(Box::new(fallible), Box::new(lit(true))),
],
Some(Box::new(lit(false))),
)),
Expr::Case(Case::new(
None,
vec![(Box::new(volatile.gt(lit(0_i16))), Box::new(lit(true)))],
Some(Box::new(col("c2"))),
)),
] {
assert_eq!(simplify(expr.clone()), expr);
}
}

#[test]
Expand Down Expand Up @@ -4323,19 +4357,20 @@ mod tests {

// CASE WHEN ISNULL(c2) THEN true ELSE c2
// -->
// ISNULL(c2) OR c2
//
// Need to call simplify 2x due to
// https://github.com/apache/datafusion/issues/1160
// Preserve CASE because the WHEN expression is outside the conservative
// column/literal subset.
let expected = Expr::Case(Case::new(
None,
vec![(Box::new(col("c2").is_null()), Box::new(lit(true)))],
Some(Box::new(col("c2"))),
));
assert_eq!(
simplify(simplify(Expr::Case(Case::new(
None,
vec![(Box::new(col("c2").is_null()), Box::new(lit(true)),)],
Some(Box::new(col("c2"))),
)))),
col("c2")
.is_null()
.or(col("c2").is_not_null().and(col("c2")))
expected
);

// CASE WHEN c1 then true WHEN c2 then false ELSE true
Expand Down Expand Up @@ -4376,29 +4411,21 @@ mod tests {
col("c1_non_null").or(col("c1_non_null").not().and(col("c2_non_null").not()))
);

// CASE WHEN c > 0 THEN true END AS c1
assert_eq!(
simplify(simplify(Expr::Case(Case::new(
// Preserve CASE with a non-leaf WHEN condition, with or without ELSE.
for expr in [
Expr::Case(Case::new(
None,
vec![(Box::new(col("c3").gt(lit(0_i64))), Box::new(lit(true)))],
None,
)))),
not_distinct_from(col("c3").gt(lit(0_i64)), lit(true)).or(distinct_from(
col("c3").gt(lit(0_i64)),
lit(true)
)
.and(lit_bool_null()))
);

// CASE WHEN c > 0 THEN true ELSE false END AS c1
assert_eq!(
simplify(simplify(Expr::Case(Case::new(
)),
Expr::Case(Case::new(
None,
vec![(Box::new(col("c3").gt(lit(0_i64))), Box::new(lit(true)))],
Some(Box::new(lit(false))),
)))),
not_distinct_from(col("c3").gt(lit(0_i64)), lit(true))
);
)),
] {
assert_eq!(simplify(simplify(expr.clone())), expr);
}
}

#[test]
Expand Down Expand Up @@ -4595,22 +4622,6 @@ mod tests {
assert_eq!(simplify(expr.clone()), expr);
}

fn distinct_from(left: impl Into<Expr>, right: impl Into<Expr>) -> Expr {
Expr::BinaryExpr(BinaryExpr {
left: Box::new(left.into()),
op: Operator::IsDistinctFrom,
right: Box::new(right.into()),
})
}

fn not_distinct_from(left: impl Into<Expr>, right: impl Into<Expr>) -> Expr {
Expr::BinaryExpr(BinaryExpr {
left: Box::new(left.into()),
op: Operator::IsNotDistinctFrom,
right: Box::new(right.into()),
})
}

#[test]
fn simplify_expr_bool_or() {
// col || true is always true
Expand Down
63 changes: 63 additions & 0 deletions datafusion/sqllogictest/test_files/case.slt
Original file line number Diff line number Diff line change
Expand Up @@ -906,3 +906,66 @@ drop table floats;
#####
# End of lookup table CASE tests
#####

statement ok
CREATE VIEW guarded_cast AS SELECT * FROM (VALUES ('1'), ('abc'), ('2'), (NULL)) t(s);

# A Boolean CASE must not become an AND that exposes the cast to 'abc'.
# Regression for https://github.com/apache/datafusion/issues/25136.
query T rowsort
SELECT s FROM guarded_cast
WHERE CASE WHEN s ~ '^[0-9]+$' THEN CAST(s AS INT) > 0 ELSE false END;
----
1
2

# Preserve branch evaluation in projections, including both forms of ELSE NULL.
query TBBB rowsort
SELECT s,
CASE WHEN s ~ '^[0-9]+$' THEN CAST(s AS INT) > 0 ELSE true END,
CASE WHEN s ~ '^[0-9]+$' THEN CAST(s AS INT) > 0 ELSE NULL END,
CASE WHEN s ~ '^[0-9]+$' THEN CAST(s AS INT) > 0 END
FROM guarded_cast;
----
1 true true true
2 true true true
NULL true NULL NULL
abc true NULL NULL

# The ELSE expression also runs only on rows that reach it.
query T rowsort
SELECT s FROM guarded_cast
WHERE CASE WHEN s ~ '^[a-z]+$' THEN false ELSE CAST(s AS INT) > 0 END;
----
1
2

# Literal outputs alone are insufficient: later WHEN expressions are conditional.
query T rowsort
SELECT s FROM guarded_cast
WHERE CASE WHEN s ~ '^[a-z]+$' THEN false
WHEN CAST(s AS INT) > 0 THEN true ELSE false END;
----
1
2

# Division is another fallible branch expression, independent of cast handling.
query I rowsort
SELECT n FROM (VALUES (0), (1), (2), (NULL)) t(n)
WHERE CASE WHEN n <> 0 THEN 10 / n > 1 ELSE false END;
----
1
2

# A selected invalid branch must still raise its original error.
query error Cannot cast string 'abc' to value of Int32 type
SELECT CASE WHEN s ~ '^[a-z]+$' THEN CAST(s AS INT) > 0 ELSE false END
FROM guarded_cast;

# The first WHEN must be evaluated even when both branch outputs are identical.
query error Cannot cast string 'abc' to value of Int32 type
SELECT CASE WHEN CAST(s AS INT) > 0 THEN false ELSE false END
FROM (VALUES ('abc')) t(s);

statement ok
DROP VIEW guarded_cast;
Loading