Repository navigation
perf: Fix performance regression in unary math UDFs - #26062
Merged
Merged
Conversation
Benchmark `sqrt` and `degrees` on Float64 and Float32 arrays, with and without nulls.
apache#22308 made `sqrt` return an error for negative inputs by switching every unary math function from `unary` to `try_unary` with a per-value validator. `try_unary` can return early on any value, so its loop is not vectorized, and for nullable input it visits valid indices one at a time. That made `sqrt`, which otherwise compiles to vector square-root instructions, several times slower. Add `unary_with_input_check`, which applies the function and checks every value in one branch-free loop, then looks for an error to report only if some value failed the check. `sqrt`'s validator now returns an error message instead of a `Result`, so the check never builds an error. The other unary math functions go back to using `unary`. The error is now reported as an execution error, rather than wrapped in an Arrow compute error.
Contributor
Author
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #26062 +/- ##
========================================
Coverage 82.73% 82.74%
========================================
Files 1147 1147
Lines 449353 449795 +442
Branches 449353 449795 +442
========================================
+ Hits 371783 372184 +401
- Misses 54916 54943 +27
- Partials 22654 22668 +14 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
Contributor
|
run benchmark math_expressions |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing neilc/perf-sqrt (c9a5496) to 8248a57 (merge-base) diff Run configurationrun benchmark math_expressionsResults will be posted here when complete File an issue against this benchmark runner |
alamb
approved these changes
Oct 6, 2026
alamb
left a comment
alamb
left a comment
Contributor
There was a problem hiding this comment.
Looks good to me -- thanks @neilconway -- I reviewed the logic and the comments and kicked off a benchmark
I also took the liberty of pushing a commit to fix the CI error: https://github.com/apache/datafusion/actions/runs/37379821165/job/111998601947?pr=26062
| /// Use this for functions that return an error for some argument values, such | ||
| /// as `sqrt`, which returns an error for negative numbers. `try_unary` can also | ||
| /// return errors, but it can return early on any value, which keeps the | ||
| /// compiler from vectorizing its loop. That makes cheap functions like `sqrt` |
Contributor
There was a problem hiding this comment.
I feel like the mention down here about vectorization is "burying the lead" so to speak -- it might be better if the comment started with a mention of "faster version of try_unary that gives the compiler the best chance to vectorize the check and the operation" or something
Contributor
Author
There was a problem hiding this comment.
Thanks, I revised the comment to change the emphasis and describe the performance tradeoffs more clearly.
| /// slots, in the same loop as `op`; only if some value fails is the array | ||
| /// searched again for an error to report. `input_error` should therefore be a | ||
| /// cheap check, such as a comparison. | ||
| pub(crate) fn unary_with_input_check<T: ArrowPrimitiveType>( |
Contributor
There was a problem hiding this comment.
this seems like it might be a good one to propose porting upstream to arrows rs (or as an example on try_unary 🤔
Contributor
Author
There was a problem hiding this comment.
Yeah, I've been thinking about that. Adding something to Arrow definitely makes sense, although because making the right choice depends on a bunch of factors (e.g., how expensive the op is, whether it is vectorizable, null density, CPU architecture / version of SIMD), we'd want to make sure that users have a clear decision process for which primitive to use.
| .values() | ||
| .iter() | ||
| .map(|&x| { | ||
| any_invalid |= input_error(x).is_some(); |
Contributor
There was a problem hiding this comment.
it is interesting this is mre vectorizable -- I was sot of expecting two loops
| .collect(); | ||
|
|
||
| // The check above also ran on null slots, which can hold any value, so the | ||
| // failure may be spurious. Re-check just the non-null values. |
Contributor
There was a problem hiding this comment.
though in theory this can run and slow down arrays with nulls
Contributor
Author
There was a problem hiding this comment.
Ah, right, this is a very good point! try_unary on NULL-heavy inputs can actually be faster. Measuring on my local machine:
- For cheap functions like
sqrt,degrees,radians,unaryis faster for <= 80% NULLs. - For expensive functions like
exp,sin, andln, the cross-over point is something like 25% NULLs.
In an extreme case of 90% nulls (randomly distributed), exp is about 5 times faster using try_unary than with unary. Whereas with no nulls, unary is about 10% faster.
Intuitively, expensive functions (a) can't be vectorized anyway (b) waste more work on the values in null slots.
I think reverting the blanket switch to try_unary stlll makes sense, because it was probably unintended, and I'd suspect that most people invoking math functions on large data sets won't have NULL-heavy data. But we could certainly try to make this more intelligent, e.g., by applying a heuristic based on NULL density to switch to try_unary, or by having the more expensive functions always use try_unary. Either way I'd be inclined to leave it to a followup PR.
Contributor
Author
There was a problem hiding this comment.
A related behavior is that null slots can hold arbitrary values; if those values happen to fail the error check, we'll need to take the slow path, even if all the valid slots are not erroneous. But we can at least do better here than the initial version of the PR -- in the slow path, we were skipping NULLs on the recheck with a filter(), but it's faster to use a bitmap to mask out null slots. I added that optimization to the PR.
|
🤖 Benchmark completed (GKE) | trigger Instance: Comparing neilc/perf-sqrt (c9a5496) to 8248a57 (merge-base) diff Run configurationrun benchmark math_expressionsCPU Details (lscpu)Details
Resource Usagemath_expressions — base (merge-base)
math_expressions — branch
File an issue against this benchmark runner |
Contributor
Author
|
@alamb FWIW the |
Lead with what the helper is for (vectorizing the input check and `op`) instead of burying it, and say when `try_unary` can still be faster: arrays with many nulls, especially with expensive operations.
If a null slot holds a value that fails the input check, every batch pays for the second pass even though there is no error. That happens easily in practice: arithmetic kernels compute on null slots too, so the null slots of `x - 10` usually hold -10. Build the re-check as a bitmap with `collect_bool` and AND it with the validity bitmap, instead of testing each value's null bit and branching. On an array whose null slots fail the check, this takes the re-check from about 5.5µs to 2.8µs per 8192 rows; arrays that pass the first check are unaffected.
jayzhan211
approved these changes
Oct 8, 2026
jayzhan211
left a comment
jayzhan211
left a comment
Contributor
There was a problem hiding this comment.
Thanks @neilconway , overall LGTM
Contributor
Author
|
@jayzhan211 @alamb Thanks for the reviews! |
Contributor
|
Thank you |
Contributor
|
Nice! |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
6 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Which issue does this PR close?
Rationale for this change
#22308 changed
sqrtto raise an error on negative floating point inputs. In the course of doing that, it changed all of the unary math UDFs to usetry_unaryinstead ofunary. Switching totry_unaryregressed the performance of those UDFs. The effect was the most extreme forsqrt, because #22308 also added a per-value check that inhibited vectorization, but the other UDFs also suffered fromtry_unary's additional overhead.This PR preserves the error-handling change in #22308 but implements it in a different way: we add support for an optional "input check" function to the unary math macro. If supplied, the check function is applied to every input value (including null slot) in a branch-free loop, which doesn't inhibit vectorization. If that loop detects an erroneous input, we do a second branching loop to find the problematic value to report the error. With this scheme, all the math functions can go back to using
unary.Using
unaryis not always a win: for (1) expensive math functions on (2) NULL-heavy data sets, the wasted work from invoking the function on null slots can exceed the overhead of checking the NULL bitmap, and expensive math functions typically prevent vectorization anyway. It would be possible to try to be smarter here about the threshold whentry_unarybecomes faster thanunary, but for now this PR reverts to the pre-#22308 performance.Benchmarks: (x86, AMD EPYC Milan):
What changes are included in this PR?
See above. Also added a benchmark for
sqrtanddegrees(degreesis a trivial math UDF that effectively measures the dispatch overhead). Other, more expensive math UDFs still got slower but their relative slowdown was much less.What is the testing strategy for this PR?
Existing tests pass; added new unit test for new error-check facility.
Are there any user-facing changes?
No.