neilconway opened a new pull request, #11656:
URL: https://github.com/apache/arrow-rs/pull/11656
# Which issue does this PR close?
- N/A
# Rationale for this change
`MutableBuffer::extend` checks the remaining capacity before writing each
item from the iterator to the buffer. The per-element capacity check has
non-trivial overhead and can also inhibit vectorization. When given a
fixed-size iterator (size hint returns `lower == upper`), we can skip the
per-element capacity check. This improves the compiler's ability to vectorize
the loop.
The most important in-tree caller of `MutableBuffer::extend` is
`MutableArrayData::try_extend`, when copying string views and dictionary keys.
This optimization initially regressed the string view case by about 3x on ARM,
because the new `extend` logic resulted in the compiler inserting a branch on
the length of the string that is poorly predicted on data containing a mix of
inline and out-of-line strings. x86 didn't regress because it always had a
poorly-predicted branch. We address the ARM regression by adjusting the
view-copying logic to avoid the branch, which speeds up both ARM and x86.
In turn, an important in-tree caller of `MutableArrayData::try_extend` is
`zip`. The two optimizations above improve `zip` for string views by 6-69% on
x86 and 3-22% on ARM, depending on the data distribution. This should in turn
help `CASE`, `COALESCE`, and `NVL2` in DataFusion, for certain shapes and data
types.
Along the way, we extend the benchmarks for `MutableBuffer::extend` to cover
more iterator shapes, and extend the `MutableArrayData` benchmarks to cover
more data types. We also revise how data is generated for the `zip_kernels` and
`mutable_array` benchmarks: the previous approach repeatedly passed the same
input data to the functions on every iteration. This allows the branch
predictor to do much better than it would on real-world data, which initially
hid the ARM regression described above.
Benchmarks:
- AMD EPYC Milan (x86)
- zip_kernels, string views, 8192 rows of 3 to 19 bytes
- array vs scalar
- all_true: 28.30 µs → 8.67 µs (−69.4%)
- 99pct_true: 33.84 µs → 12.69 µs (−62.5%)
- 90pct_true: 63.86 µs → 39.99 µs (−37.4%)
- 50pct_true: 137.2 µs → 107.5 µs (−21.7%)
- true_then_false: 62.55 µs → 45.20 µs (−27.7%)
- 10pct_true: 115.3 µs → 95.30 µs (−17.3%)
- 1pct_true: 99.62 µs → 82.61 µs (−17.1%)
- all_false: 96.48 µs → 80.40 µs (−16.7%)
- 50pct_nulls: 128.9 µs → 104.6 µs (−18.9%)
- scalar vs array
- all_true: 103.4 µs → 85.00 µs (−17.8%) (runs 5% apart)
- 99pct_true: 104.5 µs → 87.84 µs (−15.9%)
- 90pct_true: 121.7 µs → 99.56 µs (−18.2%)
- 50pct_true: 148.5 µs → 114.0 µs (−23.2%)
- true_then_false: 67.21 µs → 49.05 µs (−27.0%)
- 10pct_true: 72.49 µs → 47.27 µs (−34.8%)
- 1pct_true: 39.15 µs → 17.89 µs (−54.3%) (runs 3% apart)
- all_false: 27.84 µs → 8.88 µs (−68.1%) (runs 2.6% apart)
- 50pct_nulls: 110.5 µs → 80.31 µs (−27.4%)
- array vs array
- all_true: 28.72 µs → 9.07 µs (−68.4%)
- 99pct_true: 35.69 µs → 14.11 µs (−60.5%)
- 90pct_true: 72.65 µs → 44.72 µs (−38.4%)
- 50pct_true: 83.63 µs → 82.14 µs (−1.8%)
- true_then_false: 29.10 µs → 9.20 µs (−68.4%)
- true_then_false: 29.10 µs → 9.20 µs (−68.4%)
- 10pct_true: 73.22 µs → 45.08 µs (−38.4%)
- 1pct_true: 35.52 µs → 13.83 µs (−61.1%) (runs 2.9% apart)
- all_false: 28.65 µs → 9.02 µs (−68.5%)
- 50pct_nulls: 83.94 µs → 82.45 µs (−1.8%)
- zip_kernels, Utf8: all 27 cases between −1.5% and +0.5%
- mutable_array
- mutable str 1024: 93.19 ms → 92.80 ms (−0.4%)
- mutable str nulls 1024: 32.56 ms → 35.87 ms (+10.2%) (runs 19% apart).
Noise, not the change: one base run took 29 ms, and the other three runs took
36 ms. In a separate test of four runs each, both versions took 35.5 to 37 ms.
- mutable str_view 1024: 318.2 µs → 356.8 µs (+12.1%)
- mutable str_view short slices 1024: 19.08 µs → 15.70 µs (−17.7%)
- mutable dict 1024: 56.62 µs → 51.12 µs (−9.7%)
- mutable dict short slices 1024: 19.01 µs → 16.37 µs (−13.9%)
- mutable_buffer_extend
- chunks of 16 items on average
- copied u8: 32.88 µs → 24.08 µs (−26.8%)
- map u8: 36.69 µs → 30.00 µs (−18.2%)
- map i32: 35.05 µs → 28.43 µs (−18.9%)
- map u128: 39.44 µs → 35.73 µs (−9.4%) (runs 2.6% apart)
- range i32: 34.38 µs → 24.44 µs (−28.9%) (runs 3% apart)
- repeat_n u32: 33.60 µs → 23.29 µs (−30.7%)
- chain i32: 57.15 µs → 49.34 µs (−13.7%)
- filter i32: 32.68 µs → 33.42 µs (+2.3%)
- chunks of 1024 items on average
- copied u8: 0.89 µs → 0.83 µs (−6.2%)
- map u8: 1.34 µs → 1.06 µs (−20.5%)
- map i32: 3.74 µs → 3.53 µs (−5.4%) (runs 5% apart)
- map u128: 24.73 µs → 19.51 µs (−21.1%)
- range i32: 2.57 µs → 2.53 µs (−1.5%)
- repeat_n u32: 2.58 µs → 2.50 µs (−3.0%)
- chain i32: 27.48 µs → 27.61 µs (+0.5%)
- filter i32: 18.52 µs → 18.48 µs (−0.2%)
- Apple M4 Max (ARM)
- zip_kernels, string views, 8192 rows of 3 to 19 bytes
- array vs scalar
- all_true: 4.25 µs → 4.07 µs (−4.2%)
- 99pct_true: 7.22 µs → 6.62 µs (−8.3%)
- 90pct_true: 23.63 µs → 20.99 µs (−11.2%)
- 50pct_true: 61.19 µs → 54.88 µs (−10.3%)
- true_then_false: 25.16 µs → 19.82 µs (−21.2%)
- 10pct_true: 53.25 µs → 43.30 µs (−18.7%)
- 1pct_true: 49.16 µs → 39.02 µs (−20.6%)
- all_false: 45.80 µs → 35.51 µs (−22.5%)
- 50pct_nulls: 57.88 µs → 49.16 µs (−15.1%)
- scalar vs array
- all_true: 49.08 µs → 39.41 µs (−19.7%)
- 99pct_true: 52.81 µs → 43.12 µs (−18.4%)
- 90pct_true: 58.19 µs → 48.51 µs (−16.6%)
- 50pct_true: 70.12 µs → 63.38 µs (−9.6%)
- true_then_false: 28.22 µs → 22.94 µs (−18.7%) (runs 2.6% apart)
- 10pct_true: 28.37 µs → 26.12 µs (−7.9%)
- 1pct_true: 10.65 µs → 10.22 µs (−4.0%)
- all_false: 4.27 µs → 4.11 µs (−3.7%)
- 50pct_nulls: 48.32 µs → 43.11 µs (−10.8%)
- array vs array
- all_true: 4.37 µs → 4.23 µs (−3.2%)
- 99pct_true: 7.58 µs → 7.00 µs (−7.6%)
- 90pct_true: 25.27 µs → 22.70 µs (−10.2%)
- 50pct_true: 49.18 µs → 49.31 µs (+0.3%)
- true_then_false: 4.42 µs → 4.25 µs (−4.0%)
- 10pct_true: 25.48 µs → 23.31 µs (−8.5%)
- 1pct_true: 7.57 µs → 7.05 µs (−6.9%)
- all_false: 4.34 µs → 4.19 µs (−3.5%)
- 50pct_nulls: 49.25 µs → 48.97 µs (−0.6%)
- zip_kernels, Utf8: all 27 cases between −2.1% and +1.9%
- mutable_array
- mutable str 1024: 2.40 ms → 2.39 ms (−0.4%)
- mutable str nulls 1024: 1.26 ms → 1.27 ms (+0.1%)
- mutable str_view 1024: 179.4 µs → 168.6 µs (−6.0%)
- mutable str_view short slices 1024: 10.32 µs → 8.53 µs (−17.3%)
- mutable dict 1024: 36.90 µs → 31.65 µs (−14.2%)
- mutable dict short slices 1024: 11.66 µs → 7.99 µs (−31.5%)
- mutable_buffer_extend
- chunks of 16 items on average
- copied u8: 12.21 µs → 11.54 µs (−5.5%) (runs 14% apart)
- map u8: 19.97 µs → 8.52 µs (−57.4%) (runs 5% apart)
- map i32: 19.78 µs → 12.36 µs (−37.5%) (runs 5% apart)
- map u128: 21.01 µs → 13.07 µs (−37.8%) (runs 2.4% apart)
- range i32: 17.22 µs → 11.10 µs (−35.6%)
- repeat_n u32: 17.71 µs → 8.90 µs (−49.7%)
- chain i32: 30.84 µs → 23.98 µs (−22.2%)
- filter i32: 17.79 µs → 19.83 µs (+11.5%) (runs 8% apart)
- chunks of 1024 items on average
- copied u8: 0.51 µs → 0.51 µs (−1.5%)
- map u8: 0.77 µs → 0.53 µs (−31.2%) (runs 9% apart)
- map i32: 1.88 µs → 1.87 µs (−0.3%)
- map u128: 12.00 µs → 9.74 µs (−18.8%)
- range i32: 1.77 µs → 1.77 µs (+0.2%)
- repeat_n u32: 1.77 µs → 1.70 µs (−4.3%) (runs 2.4% apart)
- chain i32: 14.55 µs → 13.29 µs (−8.7%) (runs 6% apart)
- filter i32: 9.36 µs → 10.29 µs (+9.9%) (runs 24% apart)
# What changes are included in this PR?
* Optimize `MutableBuffer::extend` for exact size iterators; add tests for
poorly behaving iterator hints
* Optimize `MutableArrayData::try_extend` when copying string views
* Extend benchmarks for `MutableBuffer` and `MutableArrayData`
* Revise data generation for `zip` and `MutableArrayData` to make branch
prediction behavior more realistic
# Are these changes tested?
Yes; existing tests pass, new tests added.
# Are there any significant user-facing changes?
No.
# AI usage
Claude Code (Opus 5.5), reviewed with Codex (Astra 6.1).
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
