Skip to content

perf(core): reuse shuffle partition-index buffers across batches - #2359

Merged
Dandandan merged 1 commit into
apache:mainfrom
NoahKusaba:perf/reuse-shuffle-partition-index-buffers
Aug 24, 2026
Merged

perf(core): reuse shuffle partition-index buffers across batches#2359
Dandandan merged 1 commit into
apache:mainfrom
NoahKusaba:perf/reuse-shuffle-partition-index-buffers

Conversation

@NoahKusaba

@NoahKusaba NoahKusaba commented Aug 24, 2026

Copy link
Copy Markdown
Contributor

Which issue does this PR close?

Closes #.

Rationale for this change

compute_partition_indices allocated a fresh Vec<Vec<u32>> per input batch
num_partitions vectors grown from zero, then dropped once push_batch copied
the indices out. The churn scales with the output partition count.

What changes are included in this PR?

Copying the hash_buffer implementation that already exists in compute_partition_indices,
it now writes into a caller-owned buffer whose inner Vecs are cleared rather
than dropped, so capacity survives into the next batch. Partition assignment is
unchanged.

AI Benchmark:
100 batches x 8192 rows, hash on l_orderkey, release, best of 5 — timing
hash + index build + push_batch:

partitions before after speedup
16 74.8 µs/batch 61.5 µs/batch 1.22x
64 92.2 µs/batch 59.4 µs/batch 1.55x
200 125.8 µs/batch 62.0 µs/batch 2.03x
400 162.4 µs/batch 71.0 µs/batch 2.29x
800 207.4 µs/batch 76.1 µs/batch 2.73x

Are there any user-facing changes?

No

rows.clear();
}
for (row, &h) in hash_buffer.iter().enumerate() {
out[(h % num_partitions as u64) as usize].push(row as u32);

@Dandandan Dandandan Aug 24, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is also a slow path (modulo is extremely slow on CPUs), in DataFusion I went so far to replace it with strength reduce apache/datafusion#21900

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Interesting, I'll try to incorporate that too!

@NoahKusaba NoahKusaba Aug 24, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actually I'll make a follow up PR today / tomorrow for that change, to make it easier for maintainers to review/approve. Thanks for the suggestion and an implementation to reference :)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Amazing!

@Dandandan
Dandandan merged commit 50d7322 into apache:main Aug 24, 2026
28 checks passed
@NoahKusaba
NoahKusaba deleted the perf/reuse-shuffle-partition-index-buffers branch August 24, 2026 15:05
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants