perf: vectorize float/decimal to narrow-integer casts#4941
Conversation
The four macros for float-to-int and decimal-to-int narrowing casts (cast_float_to_int16_down, cast_float_to_int32_up, cast_decimal_to_int16_down, cast_decimal_to_int32_up) built the output with a per-element iterator-collect over Option/Result. Replace that with Arrow's unary (legacy) and try_unary (ANSI) kernels, which map the values buffer in one pass and carry the null buffer over, following the same pattern used by cast_int_to_int_macro. The decimal macros also hoist the constant scale divisor out of the per-element loop. Overflow, NaN, saturation, and wrap-through-Int semantics are preserved unchanged; only the iteration mechanism changes. Add Rust unit tests for the float-to-Byte and decimal-to-Int/Byte legacy wrap paths and the ANSI overflow error paths (previously covered only by Scala tests), plus a benchmark. Non-overflow casts are 49-91% faster with no regression. Part of apache#4936.
mbutrovich
left a comment
There was a problem hiding this comment.
First pass, thanks @andygrove!
| assert_eq!(decimal_array.value(1), -10000); // -100 * 10^2 | ||
| assert!(decimal_array.is_null(2)); | ||
| } | ||
|
|
There was a problem hiding this comment.
(test_cast_float64_to_int8_legacy_wraps) covers 300.7 -> 300 -> 44, which exercises the Byte narrowing wrap but keeps the intermediate inside i32 range. The distinctive Spark behavior these macros exist to replicate is the double narrowing: float overflows i32 first, then that already-wrapped Int narrows to Byte/Short. The legacy float path is (value as i32) as $rust_dest_type, and value as i32 saturates on overflow in Rust (3e9_f64 as i32 == i32::MAX), so a value like 3e9 produces i32::MAX then narrows to -1 for i8. That saturate-then-narrow is exactly the fragile path. Suggested change: add a case with Some(3e9) (or Some(f64::INFINITY)) to test_cast_float64_to_int8_legacy_wraps and assert the wrapped byte, so the legacy overflow-then-narrow behavior is pinned and cannot silently drift if the as i32 step is ever refactored.
| }) | ||
| .collect::<Result<$dest_array_type, _>>()?, | ||
| })?, | ||
| _ => cast_array.unary::<_, $dest_arrow_type>(|value| (value as i32) as $rust_dest_type), |
There was a problem hiding this comment.
native/spark-expr/src/conversion_funcs/numeric.rs:334 (float legacy arm cast_array.unary::<_, $dest_arrow_type>(|value| (value as i32) as $rust_dest_type)), :373 (float32_up legacy), :428 and :473 (decimal legacy arms with value / divisor).
PrimitiveArray::unary applies the op to every slot including nulls (arrow-array/src/array/primitive_array.rs:880-903: "Applies the function for all values, including those on null slots ... requires that the operation must be infallible (not error/panic) for any value"). The old iterator-collect only ran the body for Some values, so this PR widens the set of inputs the closure sees to include whatever garbage sits in null slots.
I verified this is safe here. Float value as i32 and as $rust_dest_type are saturating as casts that never panic on any bit pattern including NaN. Decimal value / divisor cannot panic because divisor = 10^scale is always positive (the only i128 division panic is i128::MIN / -1). So there is no regression. The risk is a future edit adding a fallible op (a checked_*, an index, an unwrap) into one of these closures without realizing it now runs on null slots. The existing comments explain the Spark cast-through-Int semantics but not this invariant. Suggested change: add one line to the legacy-arm comment in each macro, e.g. // unary runs op on null slots too; the as-cast / positive-divisor division here is infallible for any bit pattern.
Which issue does this PR close?
Part of #4936.
Rationale for this change
The float-to-int and decimal-to-int narrowing casts built their output with a per-element
iter().map(...).collect::<Result<...>>()overOption/Result, the same slow pattern thatspark_cast_int_to_intmoved off of (that change was up to 100x faster).What changes are included in this PR?
Rewrites the four macros (
cast_float_to_int16_down,cast_float_to_int32_up,cast_decimal_to_int16_down,cast_decimal_to_int32_up) to use Arrow'sunary(legacy/try) andtry_unary(ANSI) kernels, which map the values buffer in one pass and carry the null buffer over, following the existingcast_int_to_int_macro. Each macro gains a destinationArrowPrimitiveTypeparameter, threaded through the 12 call sites. The decimal macros also hoist the constant10^scaledivisor out of the per-element loop.Overflow, NaN, saturation, and the cast-through-Int wrap semantics are preserved exactly; only the iteration mechanism changes.
How are these changes tested?
Added Rust unit tests (these paths previously had only Scala coverage): float64-to-Byte legacy wrap, float64-to-Int ANSI ok + overflow error, decimal-to-Int legacy, decimal-to-Byte legacy wrap, and decimal-to-Int ANSI overflow error. Output is bit-identical to
main.Benchmark (criterion), baseline
mainvs this branch, 8192-row columns: