nik9000 commented on code in PR #14200:
URL: https://github.com/apache/datafusion/pull/14200#discussion_r1921633344
##########
datafusion/physical-expr/src/expressions/case.rs:
##########
@@ -441,6 +494,66 @@ impl CaseExpr {
Ok(ColumnarValue::Array(zip(&remainder, &else_, &then_value)?))
}
+
+ /// This is a specialization that fuses a greater-than-zero guard with
division so
+ /// so don't need to be lazy. This is a very common guard in TPC-DS as
well as normal
+ /// human expressions.
+ ///
+ /// `CASE WHEN y > 0 THEN x / y ELSE NULL END`
+ ///
+ /// It is executed as though `x / NULL_IF(y <= 0, y)` which is safe
because arrow-rs
+ /// skips evaluating fallible functions if one of their inputs is null.
+ fn div_gt_0(&self, batch: &RecordBatch) -> Result<ColumnarValue> {
+ use arrow::compute::kernels::cmp::lt_eq;
+ use arrow::compute::kernels::numeric::div;
+ use datafusion_physical_expr_common::datum::{
+ apply, apply_cmp, apply_cmp_for_nested,
+ };
+
+ let return_type = self.data_type(&batch.schema())?;
+
+ let when_then = &self.when_then_expr[0];
+ let when = when_then.0.as_any().downcast_ref::<BinaryExpr>().unwrap();
+ let then = when_then.1.as_any().downcast_ref::<BinaryExpr>().unwrap();
+
+ // NOCOMMIT are these ok? This seems like an important assertion.
+ debug_assert_eq!(None, self.else_expr);
+
+ println!("{when:?}");
+ println!("{then:?}");
+
+ let divisor = then.right().evaluate(batch)?;
+ let zero =
ColumnarValue::Scalar(ScalarValue::new_zero(&divisor.data_type())?);
+ println!("{divisor:?}");
+ // NOCOMMIT should I manually launch the lt_eq? That could prevent the
funny to_array below.
+ // NOCOMMIT instead of nullif we could replace `inf` with null?
+ match apply_cmp(&divisor, &zero, lt_eq)? {
+ ColumnarValue::Array(mask) => {
+ let mask = as_boolean_array(&mask)?;
+ // Divisor is already an array if we got here because
comparing it to 0 made an array.
+ let masked = nullif(&divisor.to_array(batch.num_rows())?,
mask)?;
+ apply(
+ &then.left().evaluate(batch)?,
+ &ColumnarValue::Array(masked),
+ div,
+ )
+ }
+ ColumnarValue::Scalar(mask) => {
Review Comment:
Can we only get here if the divisor is a literal? If so, we really don't
need this branch at all - we just won't end up here.
Is there a constant folding phase that I'm dodging in these tests?
--
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]
