Repository navigation
perf(coalesce): directly copy sparse UInt32 selections - #11174
Merged
Merged
Conversation
Rich-T-kid
commented
Sep 23, 2026
Rich-T-kid
left a comment
Rich-T-kid
left a comment
Contributor
Author
There was a problem hiding this comment.
left a few thoughts
Contributor
Author
|
at the risk of scope creep I think its worth introducing a couple things in this PR
follow up PR will implement
we will keep the default path for non-specialized arrays. |
Rich-T-kid
marked this pull request as ready for review
September 23, 2026 17:50
Rich-T-kid
force-pushed
the
rich-T-kid/coalescer-interleave
branch
from
September 24, 2026 17:28
b8cda8d to
8bb4889
Compare
Contributor
Author
|
PR description and commits are cleaned up. ready for review 🚀 |
Rich-T-kid
changed the title
Contributor
Author
|
@alamb @Dandandan bumping this as its very related to apache/datafusion#7001 (comment) |
Contributor
Author
|
I think im going to revamp this PR. I dont like the current API thats being used & I dont like the scope creep. Ill explain more in the refrenced issue. |
Rich-T-kid
marked this pull request as draft
October 4, 2026 19:14
Rich-T-kid
force-pushed
the
rich-T-kid/coalescer-interleave
branch
2 times, most recently
from
October 7, 2026 03:06
e497cc2 to
25a2c8e
Compare
Rich-T-kid
commented
Oct 7, 2026
Rich-T-kid
left a comment
Rich-T-kid
left a comment
Contributor
Author
There was a problem hiding this comment.
Should add benchmark + test for the new methods + trait
Rich-T-kid
commented
Oct 7, 2026
Rich-T-kid
force-pushed
the
rich-T-kid/coalescer-interleave
branch
from
October 12, 2026 04:30
25a2c8e to
03facf1
Compare
Rich-T-kid
changed the title
github-actions
Bot
added
the
arrow-kernels
Rich-T-kid
commented
Oct 13, 2026
Rich-T-kid
force-pushed
the
rich-T-kid/coalescer-interleave
branch
from
October 14, 2026 15:27
c57c595 to
ab144b3
Compare
Contributor
Author
|
run benchmark coalesce_kernels |
|
🤖 Benchmark starting (GKE) | trigger Target: coalesce_kernels Sharding factor: 1 (1 workers). Comparing rich-T-kid/coalescer-interleave (ab144b3) to 0706d31 (merge-base) diff Run configurationrun benchmark coalesce_kernels
env:
BENCH_FILTER: take
shards: 1
Results will be posted when all workers finish. File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Target: coalesce_kernels Comparing rich-T-kid/coalescer-interleave (ab144b3) to 0706d31 (merge-base) diff Run configurationrun benchmark coalesce_kernels
env:
BENCH_FILTER: take
shards: 1
DetailsPer-runner informationcoalesce_kernels — shard 1/1Node: gk3-benchmark-cluster-nap-1wzwczws-b8f8a7e3-njkg Instance: c4a-highmem-16 (12 vCPU / 65 GiB) uname: BENCH_COMMAND: cargo bench --features=arrow,async,test_common,experimental,object_store --bench coalesce_kernelsCPU Details (lscpu)Resource Usagebase (merge-base)
branch
File an issue against this benchmark runner |
Rich-T-kid
force-pushed
the
rich-T-kid/coalescer-interleave
branch
from
October 14, 2026 17:41
ab144b3 to
a766c82
Compare
Scatter sparse non-null UInt32 selections directly into specialized primitive and byte-view builders, avoiding the intermediate taken RecordBatch. Validate indices before builder mutation and keep dense selections on the existing take path. Compact sparse ByteView selections into bounded data blocks, while reusing source buffers for dense selections and repeated chunks. This keeps offsets valid when selected non-inline data exceeds a single Arrow buffer. Closes apache#8957
Keep the existing primitive cutoff, but use the materialized take path above 1/16 selectivity when a batch contains Utf8View or BinaryView columns.
Hoist UInt32 OOB validation to the public entry so the error is uniform across copy paths, drop the double validation, and keep the InProgressArray trait on UInt32Array with the per-chunk alloc removed.
Rich-T-kid
force-pushed
the
rich-T-kid/coalescer-interleave
branch
from
November 2, 2026 01:31
1d35f33 to
cb107b9
Compare
Contributor
Author
|
run benchmark coalesce_kernels env: |
|
🤖 Benchmark starting (GKE) | trigger Target: coalesce_kernels Sharding factor: 1 (1 workers). Comparing rich-T-kid/coalescer-interleave (a8a97b1) to b65bab9 (merge-base) diff Run configurationrun benchmark coalesce_kernels
env:
BENCH_FILTER: '^take:'
shards: 1
Results will be posted when all workers finish. File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Target: coalesce_kernels Comparing rich-T-kid/coalescer-interleave (a8a97b1) to b65bab9 (merge-base) diff Run configurationrun benchmark coalesce_kernels
env:
BENCH_FILTER: '^take:'
shards: 1
DetailsPer-runner informationcoalesce_kernels — shard 1/1Node: gk3-benchmark-cluster-nap-w5rqtefr-378755d1-ld4w Instance: c4a-highmem-16 (12 vCPU / 65 GiB) uname: BENCH_COMMAND: cargo bench --features=arrow,async,test_common,experimental,object_store --bench coalesce_kernelsCPU Details (lscpu)Resource Usagebase (merge-base)
branch
File an issue against this benchmark runner |
Keep threshold boundaries and out-of-bounds atomicity coverage in table-driven tests while removing duplicate setup.
Contributor
Author
|
run benchmark coalesce_kernels env: |
|
🤖 Benchmark starting (GKE) | trigger Target: coalesce_kernels Sharding factor: 1 (1 workers). Comparing rich-T-kid/coalescer-interleave (9da5e89) to b65bab9 (merge-base) diff Run configurationrun benchmark coalesce_kernels
env:
BENCH_FILTER: '^take: (primitive|.*(utf8view|binaryview))'
shards: 1
Results will be posted when all workers finish. File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Target: coalesce_kernels Comparing rich-T-kid/coalescer-interleave (9da5e89) to b65bab9 (merge-base) diff Run configurationrun benchmark coalesce_kernels
env:
BENCH_FILTER: '^take: (primitive|.*(utf8view|binaryview))'
shards: 1
DetailsPer-runner informationcoalesce_kernels — shard 1/1Node: gk3-benchmark-cluster-nap-1pkbfl73-8fc9a82d-84xd Instance: c4a-highmem-16 (12 vCPU / 65 GiB) uname: BENCH_COMMAND: cargo bench --features=arrow,async,test_common,experimental,object_store --bench coalesce_kernelsCPU Details (lscpu)Resource Usagebase (merge-base)
branch
File an issue against this benchmark runner |
alamb
approved these changes
Nov 12, 2026
alamb
left a comment
alamb
left a comment
Contributor
There was a problem hiding this comment.
Thank you @Rich-T-kid and @mbutrovich -- this is a really neat PR and I thank you for sticking with it. The performance report is pretty incredible
I had some small minor comments but nothing I think is required
Contributor
Author
|
thanks @alamb @mbutrovich |
Contributor
|
Let the race begin |
Surajsuthar
pushed a commit
to Surajsuthar/arrow-rs
that referenced
this pull request
Nov 14, 2026
# Which issue does this PR close? Closes apache#12210. # Rationale for this change The sparse ByteView indexed-copy path added in apache#11174 checks the destination buffer boundary for every selected non-inline value, even though it already computes the total selected byte length before copying. # What changes are included in this PR? Before this change, one loop handled both cases. For every selected non-inline value it computed the destination end, checked whether the value crossed the buffer limit, potentially finalized the buffer and allocated another one, and then copied the value. The caller already knows the total selected byte length. This change uses that total once before the loop: - If all selected bytes fit, the common path copies them into one buffer without checking the split boundary for every value. - If they do not fit, it calls a cold helper containing the existing split-buffer loop. The diff looks larger because it adds the specialized common loop and moves the old general loop into the helper; the split-buffer behavior is unchanged. In generated code, the common loop goes from two destination checks to one: the split-buffer threshold branch is removed, while `Vec::extend_from_slice`'s capacity guard remains. <details> <summary>Assembly comparison and reproduction</summary> Verified at base `b8d7b68bfc` and head `e619f21624` with `rustc 1.99.1` on `aarch64-apple-darwin`. Base, before the copy: ```asm add x8, x27, x10 lsr x8, x8, apache#31 cbz x8, ... // continue when dst_len + value_len fits ``` Both revisions retain the `Vec` capacity guard around `memcpy`: ```asm sub x8, x8, x21 cmp x27, x8 b.hi ... // grow Vec when remaining capacity is too small ``` At the PR head, the first sequence is absent from the common loop. To check either revision: ```bash RUSTFLAGS='-C target-cpu=native' cargo rustc \ -p arrow --bench coalesce_kernels --release \ --features test_utils -- --emit=asm BIN=$(rg --files target/release/deps | \ rg '/coalesce_kernels-[^.]+$' | head -1) objdump -d "$BIN" | rustfilt | less ``` Search the disassembly for `append_indexed_views_and_copy_strings`. Run the commands at `b8d7b68bfc` and `e619f21624` to compare the two loops. Install `rustfilt` with `cargo install rustfilt` if needed. </details> The existing overflow checks and maximum buffer size are preserved. # Are these changes tested? `test_copy_rows_by_indices_compacts_into_single_buffer` sets `max_buffer_size` equal to the selected byte length, covering the exact-fit fast path and validating the output layout. The existing split-buffer test covers `selected_bytes > max_buffer_size`, so both sides of the new dispatch are exercised. The final reruns compare the PR head `e619f21624` with its parent `b8d7b68bfc`. | Benchmark | Result | Runs | | --- | --- | --- | | `coalesce_kernels` (ARM64) | 1% selections are 7-13% faster; 0.1% selections are 2-5% faster | [results](apache#12212 (comment)) | | `coalesce_kernels` (AMD64) | Seven cases are 2-18% faster; BinaryView with no nulls at 1% is 3-4% slower in two runs | [full run](apache#12212 (comment)), [rerun](apache#12212 (comment)) | | Focused Utf8View rerun (ARM64) | All four cases are 3-10% faster | [results](apache#12212 (comment)) | | Focused Utf8View rerun (AMD64) | All four cases are 7-15% faster | [results](apache#12212 (comment)) | # Are there any significant user-facing changes? No.
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?
Closes #8957.
Rationale for this change
BatchCoalescer::push_batch_with_indicesnormally materializes a selectedRecordBatchwithtake_record_batchbefore copying it into the coalescer. Sparse selections can copy supported columns directly and avoid that intermediate batch.What changes are included in this PR?
Non-null
UInt32indices are validated once before any builder is mutated. Sparse selections then copy primitive and byte-view values directly; other columns use their existing per-column take path. Byte-view columns use a stricter density cutoff and reuse or compact source buffers as appropriate.Dense selections, nullable indices, and other index types retain the materialized path. Direct copies still honor target batch boundaries and the configured large-batch bypass.
Are these changes tested?
The tests compare direct copies with
takeacross primitive, byte-view, and mixed schemas. They cover nulls, reordered and repeated indices, buffer reuse and remapping, density cutoffs, batch splitting, empty selections, schema mismatches, large-batch bypass, and out-of-bounds rejection before partial mutation.Hosted benchmark results show the sparse-selection gains. That run predates the final safety-contract cleanup, which has not been benchmarked separately.
Are there any significant user-facing changes?
No public API changes.