perf: improve splitting of the offset vector in `GroupColumn::take_n` implementations, adding a new helper fn by mzabaluev · Pull Request #23730 · apache/datafusion

mzabaluev · 2026-07-20T22:58:55Z

Rationale for this change

The take_n implementation for ByteGroupValueBuilder uses split_vec_min_alloc to split the offset vector, then appends the n-th offset to the prefix vector. In case when the prefix is the newly allocated vector, it was initially reserved for size n, so it's highly likely to reallocate. On the remaining side of the split, the n-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_offsets in datafusion-common, which minimizes allocation in the way of split_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 the split_off-like branch.
To keep results in context, the timed function is GroupValues::emit on 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)

emit_first_small/shift_larger/i32_offsets_grp_1000
                        time:   [4.2893 µs 4.2972 µs 4.3041 µs]
                        change: [−4.9355% −3.8304% −2.6570%] (p = 0.00 < 0.05)
                        Performance has improved.
Found 22 outliers among 200 measurements (11.00%)
  22 (11.00%) low mild
emit_first_small/shift_midsize/i32_offsets_grp_1000
                        time:   [6.1022 µs 6.1104 µs 6.1178 µs]
                        change: [−5.1683% −4.5025% −3.8491%] (p = 0.00 < 0.05)
                        Performance has improved.
emit_first_small/shift_smaller/i32_offsets_grp_1000
                        time:   [7.8567 µs 7.8657 µs 7.8739 µs]
                        change: [−1.4737% −0.9469% −0.4095%] (p = 0.00 < 0.05)
                        Change within noise threshold.
emit_first_small/shift_larger/i64_offsets_grp_1000
                        time:   [4.7588 µs 4.7664 µs 4.7737 µs]
                        change: [−6.7077% −5.5080% −4.3595%] (p = 0.00 < 0.05)
                        Performance has improved.
Found 37 outliers among 200 measurements (18.50%)
  9 (4.50%) low severe
  28 (14.00%) low mild
emit_first_small/shift_midsize/i64_offsets_grp_1000
                        time:   [6.4207 µs 6.4270 µs 6.4328 µs]
                        change: [−6.5177% −5.7360% −5.0481%] (p = 0.00 < 0.05)
                        Performance has improved.
Found 30 outliers among 200 measurements (15.00%)
  30 (15.00%) low mild
emit_first_small/shift_smaller/i64_offsets_grp_1000
                        time:   [7.9585 µs 7.9659 µs 7.9728 µs]
                        change: [−1.8425% −1.3334% −0.8050%] (p = 0.00 < 0.05)
                        Change within noise threshold.
Found 19 outliers among 200 measurements (9.50%)
  19 (9.50%) low mild

emit_first_large/shift_larger/i32_offsets_grp_1000000
                        time:   [7.7722 ms 7.7977 ms 7.8252 ms]
                        change: [−5.4236% −4.9211% −4.4407%] (p = 0.00 < 0.05)
                        Performance has improved.
Found 1 outliers among 50 measurements (2.00%)
  1 (2.00%) high mild
emit_first_large/shift_midsize/i32_offsets_grp_1000000
                        time:   [10.263 ms 10.275 ms 10.289 ms]
                        change: [−5.8561% −5.6391% −5.4269%] (p = 0.00 < 0.05)
                        Performance has improved.
Found 1 outliers among 50 measurements (2.00%)
  1 (2.00%) high mild
emit_first_large/shift_smaller/i32_offsets_grp_1000000
                        time:   [8.0424 ms 8.0776 ms 8.1174 ms]
                        change: [−0.0236% +0.4654% +1.0597%] (p = 0.08 > 0.05)
                        No change in performance detected.
Found 5 outliers among 50 measurements (10.00%)
  4 (8.00%) high mild
  1 (2.00%) high severe
emit_first_large/shift_larger/i64_offsets_grp_1000000
                        time:   [8.5912 ms 8.6398 ms 8.6944 ms]
                        change: [−7.5494% −6.8858% −6.1907%] (p = 0.00 < 0.05)
                        Performance has improved.
Found 2 outliers among 50 measurements (4.00%)
  2 (4.00%) high mild
emit_first_large/shift_midsize/i64_offsets_grp_1000000
                        time:   [11.108 ms 11.159 ms 11.212 ms]
                        change: [−8.0008% −7.4733% −6.9419%] (p = 0.00 < 0.05)
                        Performance has improved.
emit_first_large/shift_smaller/i64_offsets_grp_1000000
                        time:   [8.2358 ms 8.2551 ms 8.2763 ms]
                        change: [−1.4108% −1.0876% −0.7471%] (p = 0.00 < 0.05)
                        Change within noise threshold.
Found 2 outliers among 50 measurements (4.00%)
  2 (4.00%) high mild

I set --turbo-mode=ALL_CORE_MAX in the instance parameters and pinned CPU cores to reduce variation.

Are there any user-facing changes?

utils::take_n_offsets function is added in datafusion-common.

codecov-commenter · 2026-07-20T23:36:58Z

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 80.70%. Comparing base (050046a) to head (16319b8).
⚠️ Report is 2 commits behind head on main.

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.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

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.

mzabaluev · 2026-07-21T00:08:54Z

+/// 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>


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.

The linked issue is over 2 years old and doesn't mention anything about Vec. Did you link the right one?

That issue is in arrow-rs, thanks, fixed.

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

LGTM. Thanks for running benches to refute my theory.

github-actions Bot added common Related to common crate physical-plan Changes to the physical-plan crate labels Jul 20, 2026

mzabaluev mentioned this pull request Jul 20, 2026

feat(physical-plan): GroupColumn support for List / LargeList #23648

Open

mzabaluev-flarion force-pushed the take-n-offsets branch 3 times, most recently from 9bb1d54 to e8a0f93 Compare July 21, 2026 12:02

mzabaluev added 4 commits July 21, 2026 14:15

bench: GroupValues::emit(First(n)) with bytes

2729647

Benchmark the performance of take_n for ByteGroupValueBuilder in the context of the multi-group-by operation where it is used in production.

bench: drop emitted groups in emit_first

56a9d5a

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.

perf(physical-plan): use take_n_offsets utility

a1d95a2

This is faster and could be reused with other GroupColumn implementations.

mzabaluev commented Jul 21, 2026

View reviewed changes

mzabaluev-flarion force-pushed the take-n-offsets branch from e8a0f93 to a1d95a2 Compare July 21, 2026 14:43

docs: fix up argument name

3e6eb87

KonaeAkira suggested changes Jul 21, 2026

View reviewed changes

Comment thread datafusion/common/src/utils/mod.rs Outdated

Comment thread datafusion/common/src/utils/mod.rs

mzabaluev and others added 5 commits July 21, 2026 23:14

refactor: safe way to copy out a Vec

16319b8

The loop does n + 1 iterations instead of n, but who's counting. Co-authored-by: Trương Hoàng Long <longtruong2411@gmail.com>

chore: remove unused import

d552beb

refactor: no need for saturating_mul

0803c90

bench(common): split_vec benchmark

c97e5d1

Created to test the new function take_n_offsets, but benchmarks for the similar split_vec_min_alloc can be added later.

bench(common): tune take_n_offsets sample sizes

440e3a9

KonaeAkira approved these changes Jul 22, 2026

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

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#23730
mzabaluev wants to merge 10 commits into
apache:mainfrom
mzabaluev:take-n-offsets

mzabaluev commented Jul 20, 2026 •

edited

Loading

Uh oh!

codecov-commenter commented Jul 20, 2026 •

edited

Loading

Uh oh!

mzabaluev Jul 21, 2026 •

edited

Loading

Uh oh!

KonaeAkira Jul 21, 2026

Uh oh!

mzabaluev Jul 21, 2026

Uh oh!

Uh oh!

Uh oh!

KonaeAkira left a comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Uh oh!

Conversation

mzabaluev commented Jul 20, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Rationale for this change

What changes are included in this PR?

Are these changes tested?

Are there any user-facing changes?

Uh oh!

codecov-commenter commented Jul 20, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

mzabaluev Jul 21, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

KonaeAkira Jul 21, 2026

Choose a reason for hiding this comment

Uh oh!

mzabaluev Jul 21, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

KonaeAkira left a comment

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

mzabaluev commented Jul 20, 2026 •

edited

Loading

codecov-commenter commented Jul 20, 2026 •

edited

Loading

mzabaluev Jul 21, 2026 •

edited

Loading