Skip to content

refactor: delegate ANSI integer arithmetic to arrow checked kernels - #5280

Open
kazantsev-maksim wants to merge 74 commits into
apache:mainfrom
kazantsev-maksim:checked_arithmetic
Open

refactor: delegate ANSI integer arithmetic to arrow checked kernels#5280
kazantsev-maksim wants to merge 74 commits into
apache:mainfrom
kazantsev-maksim:checked_arithmetic

Conversation

@kazantsev-maksim

@kazantsev-maksim kazantsev-maksim commented Aug 6, 2026

Copy link
Copy Markdown
Contributor

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?

  1. Added ansi_arithmetic_kernel to dispatch operations directly to arrow::compute::kernels::numeric::{add, sub, mul, div} for integer types when eval_mode is ANSI.
  2. Introduced a strongly typed MathOp enum (Add, Sub, Mul, Div) to replace string literals across arithmetic functions.
  3. Removed unreachable dead code (the redundant else if is_ansi_mode branch) and the obsolete is_div parameter inside checked_binary / try_arithmetic_kernel.
  4. Streamlined checked_arithmetic_internal using match guards (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

image

@kazantsev-maksim kazantsev-maksim changed the title Checked arithmetic refactoring: delegate ANSI integer arithmetic to arrow checked kernels Aug 6, 2026
@kazantsev-maksim kazantsev-maksim changed the title refactoring: delegate ANSI integer arithmetic to arrow checked kernels refactor: delegate ANSI integer arithmetic to arrow checked kernels Aug 6, 2026

@andygrove andygrove left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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_zero is Float64 only, so the integer MathOp::Div arm, which is the newly routed path, has no divide-by-zero coverage. That is exactly where the ArrowError::DivideByZero remap has to be right.
  • A case with a ColumnarValue::Scalar operand, which is the branch the Datum change 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>(

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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>(

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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(

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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>(

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@kazantsev-maksim

Copy link
Copy Markdown
Contributor Author

Thanks @andygrove for very much for the thorough review and great suggestions! I've tried to address all your feedback in the latest update.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

checked_arithmetic: delegate ANSI integer arithmetic to arrow checked kernels

2 participants