Repository navigation
perf: Optimize MutableBuffer::extend for exact-size iterators - #11656
Open
neilconway wants to merge 7 commits into
Open
neilconway wants to merge 7 commits into
neilconway wants to merge 7 commits into
Conversation
…ata` benchmarks, with new inputs each iteration `MutableBuffer::extend` had one benchmark, which copies `u32` slices. Add an arrow-buffer benchmark that appends 32Ki items, one `extend` call per chunk, with chunks of 16 or 1024 items on average. It covers these iterator shapes: - copied slices of `u8`, and mapped slices of `u8`, `i32` and `u128`. - an `i32` range. - `repeat_n`. - a `Chain` of two slices. - a `Filter`, whose size hint, unlike the others', is not exact. `mutable_array` only covered `Utf8`. Add `Utf8View` and dictionary cases, with the existing long random slices and with slices of 1 to 8 rows. Every case now starts with room for all the rows it copies, so its time goes to the copies rather than to growing the buffers from 5 rows. `zip_kernels` and `mutable_array` ran every iteration on the same input. The CPU's branch predictor then learns that input's exact sequence of branch outcomes, such as which string views are stored inline or how long each copied run is. A query sees each batch once, so it never gets that benefit, and code that branches on the data looks faster than it would be in a query. On an AMD EPYC Milan server, zipping two arrays of 8192 string views of 3 to 19 bytes with an all-true mask (one copy, with a branch per view on whether it is inline) took about 3 branch misses per call when every call used the same arrays. Cycling through 32 or more pairs of arrays raised that to about 3,400, and the call took 3.4x as long. The predictor memorized up to 32Ki to 64Ki views, depending on the build, and almost none at 128Ki. With lists of 1024 slices of 1 to 8 rows, it memorized up to 8 lists; from 32 lists on, it missed about once per slice. Every data input now changes each iteration, cycling through: - `zip_kernels`: 64 arrays for each array input and 64 masks for each random mask case (512Ki rows). - `mutable_array`: 64 arrays and 64 lists of slices. - `mutable_buffer_extend`: 16 splits of the items into chunks, and 16 arrays of 32Ki items for each case that reads input data (512Ki items). Scalars and the constant masks stay the same on every iteration, as they would across the batches of a query. `mutable_buffer_extend` chunks have random lengths, 1 to 31 or 1 to 2047 items. With fixed lengths, the predictor would learn when each call's loop ends, and a loop that handles several items per pass would never have items left over.
`MutableBuffer::extend` checks the remaining capacity before writing
each item, then pushes any items left once the capacity runs out. When
the iterator's size hint gives an exact length, reserve that many items
and zip the iterator with `0..len` instead. For slice iterators and
other iterators that support random access, `zip` then walks both sides
by index, and the compiler can often vectorize the loop. Iterators
without an exact size hint, such as `Filter`, keep the existing loop.
The new loop sets the buffer's length from the loop index after each
item. Adding the item size to the length instead made the compiler keep
a running total in vector registers: on x86 the vectorized loop then
handled half as many `i32` items per pass, and 1024-item calls on
mapped `i32` slices were 38% to 45% slower than before.
`MutableArrayData` copies string views and dictionary keys through this
method, and so do `BufferBuilder::append_n` and collecting into a
`MutableBuffer`. Its string view copy returned inline views (12 bytes
or fewer) unchanged and added the buffer offset to the others. On an
Apple M4 Max, the new loop turned that choice into a branch per view,
which the CPU often guesses wrong when short and long strings are
mixed, and zipping string view arrays took up to 3.3x as long as
before. The copy now adds 0 to inline views instead, so every view
takes the same path. The existing loop already had that branch on x86,
so view copies get faster there too.
Benchmarks, each version run twice in alternating order. Changes are
against the benchmark commit. The two runs of a version were within 2%
of each other for over 90% of the cases; larger gaps that matter for a
result are noted.
- AMD EPYC Milan:
- `zip_kernels` string views: -16% to -69%. With an all-true or
all-false mask, a zip took 9 µs instead of 28 µs, and its branch
misses per 8192 views fell from 3,464 to 13. The two 50% masks
with arrays on both sides, which `zip` copies with `interleave`,
are -2%. `Utf8` is within 1.5%.
- `mutable_array`: `str_view` -18% with short slices and +12% with
long slices, whose strings of 0 to 399 bytes are almost all longer
than 12 bytes, so the CPU always guessed the branch right. `dict`
-10% with long slices and -14% with short ones. `Utf8` unchanged:
with nulls, one run without this change took 29 ms, while the
other three, and four more runs of each version made separately,
took 35.5 to 37 ms.
- `mutable_buffer_extend`: calls of 16 items -9% to -31%, except
`Filter` +2%. Calls of 1024 items: mapped slices -5% (`i32`, with
runs 5% apart) to -21%, copied `u8` -6%, the others within 3%.
- Apple M4 Max:
- `zip_kernels` string views: -3% to -22%. The two 50% masks with
arrays on both sides are within 1%. `Utf8` is within 2.1%.
- `mutable_array`: `str_view` -17% with short slices and -6% with long
slices. `dict` -14% with long slices and -31% with short ones.
`Utf8` unchanged.
- `mutable_buffer_extend`: calls of 16 items -22% to -57%, except
copied `u8` (-5%, with runs 14% apart) and `Filter` (+11%, 0.4 to
1.4 ns per call; slower in every run with this change, including
earlier runs of it). Calls of 1024 items: mapped `u8` -31%, mapped
`u128` -19%, `Chain` -9% (with runs 6% apart), the others within 5%
except `Filter`, whose runs were up to 24% apart.
An iterator whose size hint reports an exact length but that yields
more items now has the extra items dropped; before, they were appended.
Such an iterator breaks the `Iterator::size_hint` contract. An iterator
that yields fewer items than it reports still has all of them appended.
A lower bound whose size in bytes overflows `usize` is now ignored.
Before, the multiplication panicked in debug builds and wrapped in
release builds.
github-actions
Bot
added
the
arrow-core
mbutrovich
reviewed
Oct 17, 2026
mbutrovich
left a comment
mbutrovich
left a comment
Contributor
There was a problem hiding this comment.
Thanks @neilconway. Cycling through inputs in the benchmarks makes the branch-prediction results much easier to trust.
| @@ -226,13 +226,12 @@ fn build_extend_view(array: &ArrayData, buffer_offset: u32) -> Extend<'_> { | |||
| mutable | |||
| .buffer1 | |||
| .extend(views[start..start + len].iter().map(|v| { | |||
Contributor
There was a problem hiding this comment.
Do you know where the extra time goes in mutable str_view 1024 on x86? The description shows 318.2 us to 356.8 us (+12.1%) there, while the other x86 mutable_array cases improve. On an M5 Max I see the same benchmark improve by 3.6% (168.2 us to 162.1 us). This is the long-slice case with strings of 0 to 399 bytes. Few of those strings are inline, so the old branch was well predicted. If the cost comes from the branchless select, it would help to know before trading it for the gains in the short-slice cases.
Contributor
Author
There was a problem hiding this comment.
Yeah, this is a good call out; the branchless view copy is indeed the culprit here. I had Claude explore a bunch of alternatives and it seems that rather than unpacking and repacking the view, we can use this trick instead:
--- a/arrow-data/src/transform/mod.rs
+++ b/arrow-data/src/transform/mod.rs
@@ -20,7 +20,7 @@
//! Provides utilities for creating, manipulating, and converting Arrow arrays
//! made of primitive types, strings, and nested types.
-use super::{ArrayData, ArrayDataBuilder, ByteView, data::new_buffers};
+use super::{ArrayData, ArrayDataBuilder, data::new_buffers};
use crate::bit_mask::set_bits;
use arrow_buffer::buffer::{BooleanBuffer, NullBuffer};
use arrow_buffer::{ArrowNativeType, Buffer, IntervalMonthDayNano, MutableBuffer, bit_util, i256};
@@ -229,9 +229,10 @@ fn build_extend_view(array: &ArrayData, buffer_offset: u32) -> Extend<'_> {
// A view of 12 bytes or fewer stores the string inline and
// has no buffer index, so adding 0 leaves it unchanged.
let offset = if *v as u32 > 12 { buffer_offset } else { 0 };
- let mut view = ByteView::from(*v);
- view.buffer_index += offset;
- u128::from(view)
+ // The buffer index is bits 64 to 95 of a view. The new index
+ // is less than the number of data buffers, so the sum never
+ // carries into the bits above it.
+ v + ((offset as u128) << 64)
That results in turning the x86 regression into a nice win (~20%) on both x86 and ARM.
Contributor
There was a problem hiding this comment.
the branchless view copy is indeed the culprit here. I had Claude explore a bunch of alternatives and it seems that rather than unpacking and repacking the view, we can use this trick instead
The codegen agrees with that. I built arrow-data with the release profile for x86_64-unknown-linux-gnu (rustc 1.99.0). At e2eecf1, the repack sends the high half of each view through an SSE register to clear the old buffer index:
movq xmm1, qword ptr [r8 + 8] # high 64 bits of the view
...
add r11d, dword ptr [r8 + 8] # new buffer index
pand xmm1, xmm0 # clear the old buffer index
movq rbx, xmm1
or rbx, r11 # merge in the new oneAt 2edf0bb, the offset is added straight into the high word from memory. There is no SSE round trip, and LLVM unrolls the loop by two:
cmp r15d, 13
mov r12d, 0
cmovae r12d, r11d # buffer_offset, or 0 for an inline view
add r12, qword ptr [r8 + r14 - 8] # high 64 bits + offsetOn aarch64 the and / add w / orr repack becomes a single add x13, x15, x13. That takes the loop from 11 to 9 instructions per view.
Contributor
Author
There was a problem hiding this comment.
Thanks for checking this!
`MutableArrayData` adds a buffer offset to the buffer index of each string view it copies, and adds 0 to views of inline strings. It unpacked each view into a `ByteView`, added the offset to the `buffer_index` field, and packed the fields back into a `u128`. Now it shifts the offset to the position of the buffer index, bits 64 to 95, and adds it to the whole view. The new index is less than the number of data buffers, so the sum never carries into the bits above it. After the previous commit, x86 copies of long runs of views that all hold buffer indices were up to 18% slower than before it. A null-free column of strings longer than 12 bytes has only such views. The shifted add makes those copies 12% to 23% faster than before the previous commit. Benchmarks on an AMD EPYC Milan server, each version run twice in alternating order. Changes are against the benchmark commit, then against the previous commit: - `mutable_array` `str_view` with long slices, whose strings of 0 to 399 bytes are almost all longer than 12 bytes: -23% (the previous commit: +10%), -30% against the previous commit. With short slices: -12%, +3% against the previous commit. - `zip_kernels` string views, two arrays: -36% to -77%, and -27% to +1% against the previous commit. The exception is 10% true rows, where one of the two runs of this commit was 11% slower than the other; three more runs of each version put this commit 1% faster than the previous one. - `merge` of null-free string view arrays, with a benchmark that is not part of this change: every case is 10% to 81% faster. The seven cases that the previous commit made 7.5% to 18% slower, copies of 13- to 40-byte strings in long runs, are 12% to 16% faster. Against the previous commit: -27% to +2.6%. - `take` on lists of 0 to 8 string views, also with a benchmark that is not part of this change: -38%, +2.4% against the previous commit.
When an iterator's size hint reports an exact length, `MutableBuffer::extend` appends at most that many items, so it drops the extra items of an iterator whose size hint is wrong. The `Extend` and `FromIterator` impls of `MutableBuffer` and `BufferBuilder` all append this way, and their docs now say so. The string view copy compares a view's length with `MAX_INLINE_VIEW_LEN` instead of a literal 12, as `ByteView` does.
Contributor
Author
|
Thanks for the review, @mbutrovich ! |
Contributor
Sure thing, I'll take another pass today. |
mbutrovich
approved these changes
Oct 19, 2026
mbutrovich
left a comment
mbutrovich
left a comment
Contributor
There was a problem hiding this comment.
This LGTM, thanks @neilconway!
Rich-T-kid
reviewed
Oct 20, 2026
Rich-T-kid
left a comment
Rich-T-kid
left a comment
Contributor
There was a problem hiding this comment.
@neilconway a large portion of this PR seems to be benchmarks. Could you split this into a separate PR?
Contributor
Author
Sure, done in #11742 @Rich-T-kid |
The comment on the shifted add said that "the new index is less than the number of data buffers", without saying which index or why that rules out a carry. It now says that adding `buffer_offset` gives the string's buffer index in the output, which fits in 32 bits: `MutableArrayData::with_capacities` panics if the total number of data buffers in its input arrays doesn't fit in 32 bits. The comment on the length check now speaks of a view of a string of 12 bytes or fewer, since every view is 16 bytes.
Jefffrey
pushed a commit
to mbutrovich/arrow-rs
that referenced
this pull request
Oct 23, 2026
…e#11742) # Which issue does this PR close? - N/A. Split out of apache#11656, which uses these benchmarks. # Rationale for this change 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. These improvements are intended to evaluate apache#11656 more accurately. # What changes are included in this PR? - Add an arrow-buffer `mutable_buffer_extend` benchmark that appends 32Ki items, one `extend` call per chunk, with chunks of random length (16 or 1024 items on average). It covers copied and mapped slices, an `i32` range, `repeat_n`, a `Chain` of two slices, and a `Filter`, whose size hint is not exact. - Add `Utf8View` and dictionary cases to `mutable_array`, with the existing long random slices and with slices of 1 to 8 rows. Every case now starts with room for all the rows it copies, so its time goes to the copies rather than to growing the buffers. - Cycle `zip_kernels`, `mutable_array` and `mutable_buffer_extend` through 16 to 64 inputs instead of reusing one. Scalars and the constant masks stay the same, as they would across the batches of a query. # Are these changes tested? Benchmark-only change. # Are there any significant user-facing changes? No. # AI usage Claude Code (Opus 5.5), reviewed with Codex (Astra 6.1). Co-authored-by: Ed Seidl <etseidl@users.noreply.github.com>
Contributor
|
run benchmark mutable_buffer_extend |
|
🤖 Benchmark starting (GKE) | trigger Target: mutable_buffer_extend Sharding factor: 1 (1 workers). Comparing neilc/mutable-buffer-extend-exact-size (be88f77) to df4dc3e (merge-base) diff Run configurationrun benchmark mutable_buffer_extend
shards: 1
Results will be posted when all workers finish. File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Target: mutable_buffer_extend Comparing neilc/mutable-buffer-extend-exact-size (be88f77) to df4dc3e (merge-base) diff Run configurationrun benchmark mutable_buffer_extend
shards: 1
DetailsPer-runner informationmutable_buffer_extend — shard 1/1Node: gk3-benchmark-cluster-nap-m22gfx50-6285b8fc-jhdt Instance: c4a-highmem-16 (12 vCPU / 65 GiB) uname: BENCH_COMMAND: cargo bench --features=arrow,async,test_common,experimental,object_store --bench mutable_buffer_extendCPU Details (lscpu)Resource Usagebase (merge-base)
branch
File an issue against this benchmark runner |
Rich-T-kid
self-requested a review
October 23, 2026 15:26
Contributor
Author
|
Interesting bot results, with a significant regression on one of the |
…extend-exact-size
This branch has not been deployed
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
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
MutableBuffer::extendchecks the remaining capacity before writing each item from the iterator to the buffer. Checking buffer capacity for every item has non-trivial overhead and can also inhibit vectorization. When given an exact-size iterator (size hint returnslower == upper), we can skip the capacity check, which also improves the compiler's ability to vectorize the loop.The most important in-tree caller of
MutableBuffer::extendisMutableArrayData::try_extend, when copying string views and dictionary keys. This optimization initially regressed the string view case by about 3x on ARM, because the newMutableBuffer::extendlogic 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. We also add the buffer offset to each view with a shift instead of unpacking and repacking the view, which avoids an x86 slowdown when copying long runs of strings longer than 12 bytes.In turn, an important in-tree caller of
MutableArrayData::try_extendiszip. The two optimizations above improvezipfor string views by 15-77% on x86 and 8-21% on ARM, depending on the data distribution. This should in turn helpCASE,COALESCE, andNVL2in DataFusion, for certain shapes and data types.Along the way, we extend the benchmarks for
MutableBuffer::extendto cover more iterator shapes, and extend theMutableArrayDatabenchmarks to cover more data types. We also revise how data is generated for thezip_kernelsandmutable_arraybenchmarks: 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, each version run twice in alternating order. "Runs X% apart" marks cases whose runs differed by more than 2%.
What changes are included in this PR?
MutableBuffer::extendfor exact size iterators; add tests for buggy iterator hintsMutableArrayData::try_extendwhen copying string viewsMutableBufferandMutableArrayDatazipandMutableArrayDatabenchmarks to make branch prediction behavior more realisticAre these changes tested?
Yes; existing tests pass, new tests added.
Are there any significant user-facing changes?
When
MutableBuffer::extendis passed an iterator whose size hint reports an exact length (the lower bound equals the upper bound) but that actually yields more items than the hint reports, the extra items are silently dropped. Previously, those items were appended. I think this is acceptable: the iterator's size hint is buggy, and including the extra items is not clearly better than omitting them.AI usage
Claude Code (Opus 5.5), reviewed with Codex (Astra 6.1).