perf: improve splitting of the offset vector in GroupColumn::take_n implementations, adding a new helper fn#23730
perf: improve splitting of the offset vector in GroupColumn::take_n implementations, adding a new helper fn#23730mzabaluev wants to merge 12 commits into
GroupColumn::take_n implementations, adding a new helper fn#23730Conversation
github-actions
Bot
added
common
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #23730 +/- ##
==========================================
- Coverage 80.71% 80.70% -0.01%
==========================================
Files 1089 1089
Lines 368750 368765 +15
Branches 368750 368765 +15
==========================================
- Hits 297631 297630 -1
- Misses 53375 53386 +11
- Partials 17744 17749 +5 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
9bb1d54 to
e8a0f93
Compare
A refinement of split_vec_min_alloc for splitting offset buffers. To improve performance of the current use in `BytesGroupValues::take_n`, this function adjusts the remaining offsets when they are copied to the new location, rather than two passes of copying then subtracting in place. It also accounts for n + 1 length of the resulting prefix vector, which current split_vec_min_alloc falls on the wrong side of.
Benchmark the performance of take_n for ByteGroupValueBuilder in the context of the multi-group-by operation where it is used in production.
Drop the returned groups from emit(First(n)) in the benchmarked function, to eliminate accumulating allocation behavior that's not characteristic of the use sites.
This is faster and could be reused with other GroupColumn implementations.
| /// Allocates for whichever side is smaller, so the new allocation is | ||
| /// `min(n + 1, vec.len() - n)`. This matters when the split emits a prefix | ||
| /// under memory pressure, where `n` can be close to `vec.len()`. | ||
| pub fn take_n_offsets<O>(offsets: &mut Vec<O>, n: usize) -> Vec<O> |
There was a problem hiding this comment.
I thought of adding this as a method to OffsetBufferBuilder to make the return value correct-by-construction rather than just assuming the offsets can be subtracted the cut-off value from. But that would go against the general push towards Vec in arrow-rs#10245.
There was a problem hiding this comment.
The linked issue is over 2 years old and doesn't mention anything about Vec. Did you link the right one?
There was a problem hiding this comment.
That issue is in arrow-rs, thanks, fixed.
e8a0f93 to
a1d95a2
Compare
The loop does n + 1 iterations instead of n, but who's counting. Co-authored-by: Trương Hoàng Long <longtruong2411@gmail.com>
Created to test the new function take_n_offsets, but benchmarks for the similar split_vec_min_alloc can be added later.
KonaeAkira
left a comment
KonaeAkira
left a comment
There was a problem hiding this comment.
LGTM. Thanks for running benches to refute my theory.
|
As seen in the |
| O: OffsetSizeTrait, | ||
| { | ||
| let cut_offset = offsets[n]; | ||
| if n < offsets.len() - n { |
There was a problem hiding this comment.
| if n < offsets.len() - n { | |
| if n <= offsets.len() / 2 { |
Why not this?
There was a problem hiding this comment.
This also pushes one item above the promise to allocate for the smaller part.
In the take_n_offsets drain-and-shift branch. The rewritten, safe version vectorizes just as well and performs the same in the benchmarks.
| let new_len = offsets.len() - n; | ||
| // Shift the retained tail down to the front and rebase it to start at 0 | ||
| // in a single fused pass. | ||
| for i in 0..new_len { | ||
| offsets[i] = offsets[n + i] - cut_offset; | ||
| } |
There was a problem hiding this comment.
I have rewritten the loop of contention now.
However, there is a near 100% performance cliff at the threshold point between this and the else branch, for large array lengths that don't fit into cache (for small arrays, there is a smaller bump in the other way). So it's tempting to make the split-off threshold at a lower fraction of the length, at the cost of making larger allocations. I suspect that the real fix, however, would be in implementing something proposed in #19906 and #19649.
4 participants
Rationale for this change
The
take_nimplementation forByteGroupValueBuilderusessplit_vec_min_allocto split the offset vector, then appends then-th offset to the prefix vector. In case when the prefix is the newly allocated vector, it was initially reserved for sizen, so it's highly likely to reallocate. On the remaining side of the split, then-th offset is subtracted from the remaining offsets in place after the values have been copied to the new location. This may be slower than a single read-subtract-write pass.The implementation in #23648 will benefit from the same optimization.
What changes are included in this PR?
Add a utility function named
take_n_offsetsin datafusion-common, which minimizes allocation in the way ofsplit_vec_min_alloc, but reserves the right capacity at the outset and shifts the remaining offsets in a single pass with the copying.Are these changes tested?
Added a micro-benchmark that shows improvement on the
drain-like branch of the split, and no change on thesplit_off-like branch.To keep results in context, the timed function is
GroupValues::emiton a setup with two grouped columns of byte array types; one of them a string column, to try and catch any regressions in future specialization or refactoring.Benchmark results (on a GCE instance type C4, Intel Xeon Emerald Rapids host CPU)
I set
--turbo-mode=ALL_CORE_MAXin the instance parameters and pinned CPU cores to reduce variation.Are there any user-facing changes?
utils::take_n_offsetsfunction is added in datafusion-common.