shinzoxD commented on code in PR #24402:
URL: https://github.com/apache/datafusion/pull/24402#discussion_r3815165343
##########
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:
Addressed in e91860399: `is_cheap_for_full_batch` no longer treats every
`CastExpr` as safe. A cast is only specialized when it is provably infallible
(safe-mode CAST, TRY_CAST, same-type, or a widening numeric cast via
`CastExpr::check_bigger_cast`). Narrowing casts such as Int32/Int64 -> TinyInt
skip specialization when the source type is known at construction, and column
casts are re-checked against the batch schema at evaluation time so they fall
back to filtered CASE. Added a unit test and a `case.slt` regression for `CASE
WHEN d < 0 THEN CAST(10 AS TINYINT) / CAST(d AS TINYINT) END` over `(-1),
(1000), (0)` => `-10, NULL, NULL` (does not error on the excluded 1000).
--
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]
