perf: optimize nanvl in datafusion-functions#23458
Conversation
comphead
left a comment
comphead
left a comment
There was a problem hiding this comment.
Nice improvement overall — the null-free fast path clearly removes a lot of Option and null bookkeeping overhead, and the benchmark wins in the PR description are compelling.
One additional optimization opportunity: the current fallback path handles all null-containing cases together, so it still pays for x.is_valid(i) and, in the NaN case, y.is_valid(i) checks inside a single generic loop. I think this can be tightened further by splitting the fallback into separate cases for:
(None, Some(_))(Some(_), None)(Some(_), Some(_))
That should reduce branching and avoid unnecessary validity checks for partially-null inputs while preserving the same semantics.
Concretely, I’d expect something along these lines in nanvl_impl:
match (x.nulls(), y.nulls()) {
(None, None) => { /* current fast path */ }
(None, Some(_)) => { /* only y validity matters when x_value.is_nan() */ }
(Some(_), None) => { /* only x validity matters; y validity checks disappear */ }
(Some(_), Some(_)) => { /* current generic null-aware path */ }
}Also, a couple of smaller follow-ups:
- the PR description says “raw f64 value slices”, but the implementation is now generic across
Float16/32/64 - it would be useful to add benchmarks for partially-null inputs too, since the current numbers mostly validate the no-null fast path
Adds benchmarks for the three null-configuration match arms (only-x-null, only-y-null, both-null) and unit tests covering the single-null-input cases. These benchmarks were used to evaluate splitting the null-aware path into per-configuration arms; measurements showed that split regresses performance by 8-16%, so only the coverage additions are kept.
I tried out the main suggestion here and saw worse performance. Curious to know if you can reproduce this? |
I addressed these points. Thanks. |
I ran the benchmarks again and see a ~11-15% regression with the proposed approach. Claude tells me that this is due to a "whole-function codegen penalty". |
Thanks for trying that, its just what caught my eye, anyway current PR already brings up a benefit |
| /// splits out a null-free fast path that iterates the raw value slices, | ||
| /// skipping per-element validity checks and `Option` handling. The null-aware | ||
| /// path builds its null buffer lazily via [`NullBufferBuilder`] exactly as the | ||
| /// `FromIterator` implementation does, so the result is byte-for-byte the same. |
There was a problem hiding this comment.
probably can remove reference to previous version (like how it mentions result is byte-for-byte the same)
| let mut values = Vec::with_capacity(len); | ||
| for i in 0..len { | ||
| // `y` is only consulted when `x` is a (non-null) NaN, matching | ||
| // the original short-circuiting match. |
There was a problem hiding this comment.
same here, no need to mention original short-circuiting match considering the code will be gone
4 participants
Which issue does this PR close?
N/A
Rationale for this change
Improve performance of existing function.
What changes are included in this PR?
Added a null-free fast path to the nanvl array kernel that iterates the raw value slices (generic across
Float16/Float32/Float64) instead of per-elementOptioniteration plus collect, eliminating null-bookkeeping overhead on the common no-null input.Are these changes tested?
Existing tests + one new unit test
Benchmark (criterion):
Wins:
Within noise:
Partially-null inputs (base =
main, cand = this PR). Even on inputscontaining nulls the rewritten kernel is faster than
main:Also added benchmarks and unit tests for partially-null inputs (only-x-null, only-y-null, both-null). Splitting the null-aware path into per-configuration match arms was evaluated with these benchmarks but regressed performance by 8-16% (the larger function body penalizes even the unchanged both-null arm), so only the coverage additions were kept.
Are there any user-facing changes?
No