andygrove commented on code in PR #5280:
URL: https://github.com/apache/datafusion-comet/pull/5280#discussion_r3728431647
##########
native/spark-expr/src/math_funcs/checked_arithmetic.rs:
##########
@@ -29,40 +30,66 @@ use datafusion::common::DataFusionError;
use datafusion::physical_plan::ColumnarValue;
use std::sync::Arc;
-pub fn try_arithmetic_kernel<T>(
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+enum MathOp {
+ Add,
+ Sub,
+ Mul,
+ Div,
+}
+
+fn try_arithmetic_kernel<T>(
left: &PrimitiveArray<T>,
right: &PrimitiveArray<T>,
- op: &str,
is_ansi_mode: bool,
+ op: MathOp,
) -> Result<ArrayRef, DataFusionError>
where
T: ArrowPrimitiveType,
{
match op {
- "checked_add" => checked_binary(left, right, is_ansi_mode, false, |l,
r| l.add_checked(r)),
- "checked_sub" => checked_binary(left, right, is_ansi_mode, false, |l,
r| l.sub_checked(r)),
- "checked_mul" => checked_binary(left, right, is_ansi_mode, false, |l,
r| l.mul_checked(r)),
- "checked_div" => checked_binary(left, right, is_ansi_mode, true, |l,
r| l.div_checked(r)),
- _ => Err(DataFusionError::Internal(format!(
- "Unsupported operation: {:?}",
- op
- ))),
+ MathOp::Add => checked_binary(left, right, is_ansi_mode, |l, r|
l.add_checked(r)),
+ MathOp::Sub => checked_binary(left, right, is_ansi_mode, |l, r|
l.sub_checked(r)),
+ MathOp::Mul => checked_binary(left, right, is_ansi_mode, |l, r|
l.mul_checked(r)),
+ MathOp::Div => checked_binary(left, right, is_ansi_mode, |l, r|
l.div_checked(r)),
}
}
+fn ansi_arithmetic_kernel<T>(
Review Comment:
`ArrayRef` already implements `Datum` via `impl<T: Array> Datum for T` in
`arrow-array/src/scalar.rs`, so this function does not need a type parameter at
all. If you change the signature to take `&dyn Datum`, the four integer match
arms below collapse into a single one:
```rust
DataType::Int8 | DataType::Int16 | DataType::Int32 | DataType::Int64
if is_ansi_mode =>
{
ansi_arithmetic_kernel(&left_arr, &right_arr, op)
}
```
I tried this locally and it compiles clean with all six tests passing. It
drops about 30 lines, which feels worth it given that the goal here is removing
duplication. It also sets up the `Datum` change from the issue nicely, since
`&dyn Datum` is what you would pass a `Scalar` through.
##########
native/spark-expr/src/math_funcs/checked_arithmetic.rs:
##########
@@ -189,50 +206,68 @@ fn checked_arithmetic_internal(
(ColumnarValue::Scalar(l), ColumnarValue::Scalar(r)) =>
(l.to_array()?, r.to_array()?),
};
- // Rust only supports checked_arithmetic on numeric types
let result_array = match data_type {
- DataType::Int8 => try_arithmetic_kernel::<Int8Type>(
+ DataType::Int8 if is_ansi_mode => ansi_arithmetic_kernel(
left_arr.as_primitive::<Int8Type>(),
right_arr.as_primitive::<Int8Type>(),
op,
+ ),
+ DataType::Int8 => try_arithmetic_kernel::<Int8Type>(
+ left_arr.as_primitive::<Int8Type>(),
+ right_arr.as_primitive::<Int8Type>(),
is_ansi_mode,
+ op,
),
- DataType::Int16 => try_arithmetic_kernel::<Int16Type>(
+ DataType::Int16 if is_ansi_mode => ansi_arithmetic_kernel(
left_arr.as_primitive::<Int16Type>(),
right_arr.as_primitive::<Int16Type>(),
op,
+ ),
+ DataType::Int16 => try_arithmetic_kernel::<Int16Type>(
+ left_arr.as_primitive::<Int16Type>(),
+ right_arr.as_primitive::<Int16Type>(),
is_ansi_mode,
+ op,
),
- DataType::Int32 => try_arithmetic_kernel::<Int32Type>(
+ DataType::Int32 if is_ansi_mode => ansi_arithmetic_kernel(
left_arr.as_primitive::<Int32Type>(),
right_arr.as_primitive::<Int32Type>(),
op,
+ ),
+ DataType::Int32 => try_arithmetic_kernel::<Int32Type>(
+ left_arr.as_primitive::<Int32Type>(),
+ right_arr.as_primitive::<Int32Type>(),
is_ansi_mode,
+ op,
),
- DataType::Int64 => try_arithmetic_kernel::<Int64Type>(
+ DataType::Int64 if is_ansi_mode => ansi_arithmetic_kernel(
left_arr.as_primitive::<Int64Type>(),
right_arr.as_primitive::<Int64Type>(),
op,
+ ),
+ DataType::Int64 => try_arithmetic_kernel::<Int64Type>(
+ left_arr.as_primitive::<Int64Type>(),
+ right_arr.as_primitive::<Int64Type>(),
is_ansi_mode,
+ op,
),
- // Spark always casts division operands to floats
- DataType::Float16 if (op == "checked_div") =>
try_arithmetic_kernel::<Float16Type>(
+ DataType::Float16 if op == MathOp::Div =>
try_arithmetic_kernel::<Float16Type>(
Review Comment:
Could we keep the `// Spark always casts division operands to floats`
comment? It is the thing that explains why these three arms are `Div`-only.
It would also help to spell out why the float arms deliberately do *not* go
through `ansi_arithmetic_kernel`, because it is quite subtle. Arrow's
`float_op` handles `Op::Div` with the infallible `op!` macro using
`div_wrapping`, so `numeric::div` on floats returns inf instead of raising,
whereas `div_checked` on floats does error on a zero divisor. Sending floats
through the numeric kernel would silently break ANSI DIVIDE_BY_ZERO for double
division. Right now the only thing guarding against that is one assertion in
`test_checked_div_by_zero`, and I think the next person tidying this file would
reasonably try to finish the refactor and hit it.
##########
native/spark-expr/src/math_funcs/checked_arithmetic.rs:
##########
@@ -189,50 +206,68 @@ fn checked_arithmetic_internal(
(ColumnarValue::Scalar(l), ColumnarValue::Scalar(r)) =>
(l.to_array()?, r.to_array()?),
};
- // Rust only supports checked_arithmetic on numeric types
let result_array = match data_type {
- DataType::Int8 => try_arithmetic_kernel::<Int8Type>(
+ DataType::Int8 if is_ansi_mode => ansi_arithmetic_kernel(
Review Comment:
These four arms are near-identical and only differ in the primitive type
used for the downcast. See my note on `ansi_arithmetic_kernel` above. Taking
`&dyn Datum` lets this become one arm and removes the downcasts entirely.
##########
native/spark-expr/src/math_funcs/checked_arithmetic.rs:
##########
@@ -29,40 +30,66 @@ use datafusion::common::DataFusionError;
use datafusion::physical_plan::ColumnarValue;
use std::sync::Arc;
-pub fn try_arithmetic_kernel<T>(
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+enum MathOp {
+ Add,
+ Sub,
+ Mul,
+ Div,
+}
+
+fn try_arithmetic_kernel<T>(
left: &PrimitiveArray<T>,
right: &PrimitiveArray<T>,
- op: &str,
is_ansi_mode: bool,
+ op: MathOp,
) -> Result<ArrayRef, DataFusionError>
where
T: ArrowPrimitiveType,
{
match op {
- "checked_add" => checked_binary(left, right, is_ansi_mode, false, |l,
r| l.add_checked(r)),
- "checked_sub" => checked_binary(left, right, is_ansi_mode, false, |l,
r| l.sub_checked(r)),
- "checked_mul" => checked_binary(left, right, is_ansi_mode, false, |l,
r| l.mul_checked(r)),
- "checked_div" => checked_binary(left, right, is_ansi_mode, true, |l,
r| l.div_checked(r)),
- _ => Err(DataFusionError::Internal(format!(
- "Unsupported operation: {:?}",
- op
- ))),
+ MathOp::Add => checked_binary(left, right, is_ansi_mode, |l, r|
l.add_checked(r)),
+ MathOp::Sub => checked_binary(left, right, is_ansi_mode, |l, r|
l.sub_checked(r)),
+ MathOp::Mul => checked_binary(left, right, is_ansi_mode, |l, r|
l.mul_checked(r)),
+ MathOp::Div => checked_binary(left, right, is_ansi_mode, |l, r|
l.div_checked(r)),
}
}
+fn ansi_arithmetic_kernel<T>(
+ left: &PrimitiveArray<T>,
+ right: &PrimitiveArray<T>,
+ op: MathOp,
+) -> Result<ArrayRef, DataFusionError>
+where
+ T: ArrowPrimitiveType,
+{
+ let result_array = match op {
+ MathOp::Add => numeric::add(left, right),
+ MathOp::Sub => numeric::sub(left, right),
+ MathOp::Mul => numeric::mul(left, right),
+ MathOp::Div => numeric::div(left, right),
+ };
+
+ result_array.map_err(|e| match e {
+ ArrowError::DivideByZero => divide_by_zero_error().into(),
+ _ => DataFusionError::from(SparkError::ArithmeticOverflow {
+ from_type: String::from("integer"),
+ }),
+ })
+}
+
fn checked_binary<T, F>(
Review Comment:
After this change, is the `is_ansi_mode` branch in here only reachable from
the three float division arms? If so, we now have two ANSI implementations in
the file, and the name does not really hint that this one is
float-division-only. Would it be clearer to pull that branch out into a small
dedicated helper, or at least note the float-only reachability in a comment?
--
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]
