refactor: delegate ANSI integer arithmetic to arrow checked kernels - #5280
refactor: delegate ANSI integer arithmetic to arrow checked kernels#5280kazantsev-maksim wants to merge 74 commits into
Conversation
This reverts commit 768b3e9.
andygrove
left a comment
There was a problem hiding this comment.
Thanks for picking this up. The MathOp enum, removing the unreachable else if is_ansi_mode branch, and dropping the is_div parameter are all nice cleanups.
I checked the equivalence of the integer ANSI path locally and it holds up. In arrow-arith, numeric::add on two same-typed integer arrays goes arithmetic_op -> integer_op::<T> -> try_op! -> arity::try_binary(add_checked), which is exactly what the old code did. try_binary uses try_for_each_valid_idx, so null slots are skipped in both versions, and the DivideByZero remap lines up. I applied the diff and all six unit tests in the file pass.
A few things I would like to work through before this goes in.
The Datum half of #5092 is not implemented
The issue asked for two changes, and this PR does the first. The issue also asked to pass Datum operands so that checked_arithmetic_internal can stop materializing scalars with to_array_of_size. Right now scalar operands are still expanded to full-length arrays and then downcast, so arrow takes the array-array try_binary branch and never reaches the try_unary scalar fast path. That is the part of the issue with an actual runtime win, so as written the refactor is roughly performance-neutral.
Would you be up for handling the scalar case here too? If you would rather keep this PR focused, could you open a follow-up issue and link it from the description? Otherwise it will get lost when this PR closes #5092.
Test coverage for the newly routed paths
This moves all four integer widths onto a new kernel, and the existing tests only exercise Int32 and Float64 with array-array inputs. It might be worth extending them a little:
- An ANSI integer divide-by-zero case.
test_checked_div_by_zerois Float64 only, so the integerMathOp::Divarm, which is the newly routed path, has no divide-by-zero coverage. That is exactly where theArrowError::DivideByZeroremap has to be right. - A case with a
ColumnarValue::Scalaroperand, which is the branch theDatumchange above would touch. - An ANSI Int64 overflow case, since only Int32 is covered today.
The benchmark in native/spark-expr/benches/checked_arithmetic.rs is array-array only as well, so it would not be able to show a scalar-path improvement if the Datum change lands.
CI
No CI has run on this yet. Head sha 68eb53df has zero check runs, so the workflows will need triggering before we can evaluate it.
| } | ||
| } | ||
|
|
||
| fn ansi_arithmetic_kernel<T>( |
There was a problem hiding this 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:
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.
| }) | ||
| } | ||
|
|
||
| fn checked_binary<T, F>( |
There was a problem hiding this 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?
There was a problem hiding this comment.
Great catch! You're completely right. To address this and make the separation crystal clear, I've pulled that float-specific ANSI path out into a dedicated helper function called ansi_float_div.
| // 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( |
There was a problem hiding this 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.
| ), | ||
| // 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>( |
There was a problem hiding this 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.
|
Thanks @andygrove for very much for the thorough review and great suggestions! I've tried to address all your feedback in the latest update. |
Which issue does this PR close?
Closes #5092
Rationale for this change
This PR refactors the module to delegate integer arithmetic in ANSI mode directly to the native arrow-rs numeric kernels, eliminating code duplication and improving maintainability while preserving the expected Spark error behavior (DIVIDE_BY_ZERO and ARITHMETIC_OVERFLOW).
What changes are included in this PR?
arrow::compute::kernels::numeric::{add, sub, mul, div}for integer types when eval_mode is ANSI.MathOp enum (Add, Sub, Mul, Div)to replace string literals across arithmetic functions.checked_binary/try_arithmetic_kernel.(DataType::Int... if is_ansi_mode)to cleanly separate ANSI and Try code paths without nested conditionals.How are these changes tested?
Tested with existing unit tests