kosiew commented on code in PR #24402:
URL: https://github.com/apache/datafusion/pull/24402#discussion_r3812095859
##########
datafusion/physical-expr/src/expressions/case.rs:
##########
@@ -320,6 +334,96 @@ fn is_cheap_and_infallible(expr: &Arc<dyn PhysicalExpr>)
-> bool {
expr.is::<Column>()
}
+/// True when `when` / `then` is the divide-by-zero protection pattern
+/// `CASE WHEN y {>, !=, <} 0 THEN x / y`.
+///
+/// Both divide operands must be cheap (column, literal, or cast of those) so
+/// evaluating them on the full batch cannot introduce errors on skipped rows.
+fn is_divide_by_zero_protection(
+ when_expr: &Arc<dyn PhysicalExpr>,
+ then_expr: &Arc<dyn PhysicalExpr>,
+) -> bool {
+ let Some(checked) = extract_nonzero_checked_operand(when_expr) else {
+ return false;
+ };
+ let Some((numerator, divisor)) = extract_division_operands(then_expr) else
{
+ return false;
+ };
+ unwrap_casts(divisor).eq(unwrap_casts(checked))
+ && is_cheap_for_full_batch(numerator)
+ && is_cheap_for_full_batch(divisor)
+}
+
+/// Operand being tested for a non-zero (or strictly positive / negative)
value.
+///
+/// Matches `y > 0`, `y != 0`, `y < 0` and the swapped forms `0 < y`, `0 != y`,
+/// `0 > y`.
+fn extract_nonzero_checked_operand(
+ expr: &Arc<dyn PhysicalExpr>,
+) -> Option<&Arc<dyn PhysicalExpr>> {
+ let binary = expr.downcast_ref::<BinaryExpr>()?;
+ match binary.op() {
+ Operator::Gt | Operator::NotEq | Operator::Lt
+ if is_literal_zero(binary.right().as_ref()) =>
+ {
+ Some(binary.left())
+ }
+ Operator::Lt | Operator::NotEq | Operator::Gt
+ if is_literal_zero(binary.left().as_ref()) =>
+ {
+ Some(binary.right())
+ }
+ _ => None,
+ }
+}
+
+/// `(numerator, divisor)` from a `/` expression.
+type DivisionOperands<'a> = (&'a Arc<dyn PhysicalExpr>, &'a Arc<dyn
PhysicalExpr>);
+
+fn extract_division_operands(
+ expr: &Arc<dyn PhysicalExpr>,
+) -> Option<DivisionOperands<'_>> {
+ let binary = expr.downcast_ref::<BinaryExpr>()?;
+ if *binary.op() == Operator::Divide {
+ Some((binary.left(), binary.right()))
+ } else {
+ None
+ }
+}
+
+/// Strip `Cast` / `TryCast` layers so `a > 0 THEN x / CAST(a AS …)` still
matches.
+fn unwrap_casts(expr: &Arc<dyn PhysicalExpr>) -> &Arc<dyn PhysicalExpr> {
+ if let Some(cast) = expr.downcast_ref::<CastExpr>() {
+ unwrap_casts(cast.expr())
+ } else if let Some(try_cast) = expr.downcast_ref::<TryCastExpr>() {
+ unwrap_casts(try_cast.expr())
+ } else {
+ expr
+ }
+}
+
+fn is_literal_zero(expr: &dyn PhysicalExpr) -> bool {
+ let Some(lit) = expr.downcast_ref::<Literal>() else {
+ return false;
+ };
+ match ScalarValue::new_zero(&lit.value().data_type()) {
+ Ok(zero) => lit.value() == &zero,
+ Err(_) => false,
+ }
+}
+
+fn is_cheap_for_full_batch(expr: &Arc<dyn PhysicalExpr>) -> bool {
+ if expr.is::<Column>() || expr.is::<Literal>() {
+ true
+ } else if let Some(cast) = expr.downcast_ref::<CastExpr>() {
Review Comment:
I think we need to be a bit more careful with `CastExpr` here.
`is_cheap_for_full_batch` currently treats every cast as safe to evaluate on
rows that CASE would otherwise exclude, but casts can fail.
With this specialization, the divisor is evaluated over the full batch later
on. For example, `CASE WHEN d < 0 THEN CAST(10 AS TINYINT) / CAST(d AS TINYINT)
END` over `(-1), (1000), (0)` should produce `-10, NULL, NULL`. Instead, the
specialized path can error while casting the excluded `1000`, which changes
CASE short-circuiting semantics.
Could we avoid specializing fallible `CastExpr`s here, or only allow casts
that we can prove are infallible? Otherwise this should fall back to the
existing filtered CASE evaluation. It would also be good to add a regression
test with an invalid cast on an excluded row.
##########
datafusion/sqllogictest/test_files/case.slt:
##########
@@ -655,6 +655,38 @@ NULL
NULL
-1
+# Swapped comparison form of the same pattern
Review Comment:
Could we also add an SQL-level test for the other swapped comparison
direction, for example `0 > d` with negative, zero, and positive divisors? The
Rust unit test already covers `<`, and having this here would make sure the
planner-to-physical matching also stays covered for this SQL spelling.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
