perf(core): write one sort-shuffle file per task instead of one per input

[Dandandan]

Rebased onto main now that #2315 has merged; this PR is a single commit.
The change only matters when a task owns several input partitions, which is
what #2315 made the default.

Summary

A task that owns several input partitions wrote one data.arrow plus index
per input partition. Every downstream reader fetching partition k
therefore opened one file per input partition, and a stage left M files
behind, where M is the stage's input partition count.

The writer now emits one file per task.

This matters more since ballista.scheduler.max_partitions_per_task defaults
to 0 (#2315): a task now routinely owns several input partitions, so M
files per stage became M/P.

Design

Each input partition still buckets and spills concurrently, and still encodes
its own buckets to IPC bytes on its own task, so the interleave, framing
and compression stay parallel across inputs. The coordinator — which already
awaited every input before responding — then concatenates the finished buffers
into one file and writes one index.

The file is laid out partition-major: the schema header, then output
partition 0's bytes from every input in turn, then partition 1's, and so on.

data.arrow:       [ bucket 0 ][ bucket 1 ][ bucket 2 ]
                    \_ in0,in1 _/
data.arrow.index: { 0 -> off0, 1 -> off1, 2 -> off2 }

Keeping each output partition contiguous is what lets the index stay one
offset per partition, which is why the reader is unchanged:
create_shuffle_path already resolves a sort-shuffle summary to
{stage_id}/{file_id}/data.arrow, and MultiStreamPartitionStream already
crosses concatenated IPC streams inside a byte range.

Spill directories move from {stage}/{file_id}/spill to
{stage}/{task_id}/spill-{input}, so a task owns exactly one directory under
the stage and cleanup no longer strands an empty directory per input partition.

[avantgardnerio]

Benchmarked this at SF1000 on EKS. Short version: it's a solid win, and I think the mechanism is more interesting than the headline number.

Setup

TPC-H SF1000 (ZSTD parquet, 32 MiB row groups) on S3, 32 executor pods x 8 vcores = 256 vcores across 4x r6i.24xlarge, 48 GiB memory pool per pod, target_partitions=256, max_partitions_per_task=0 so every task owns 8 input partitions, lz4 shuffle compression, AQE on, 1 iteration.

Baseline is main at this PR's merge-base, so the only difference between the two builds is the two files in this PR. I crossed that with ballista.shuffle.sort_based.memory_limit_per_task_bytes at 0 (what benchmarking.md uses) and at 268435456 (the shipping default), because it turns out that setting interacts with this change.

Seven queries: the four heaviest sort-shuffle writers, plus Q1 and Q6 as controls. Q6 plans ShuffleWriterExec rather than SortShuffleWriterExec, so nothing in this PR can affect it, which makes it a useful noise gauge.

Wall clock, seconds

Config	Q6	Q1	Q3	Q10	Q18	Q21	Q9	Total
main, spill off	7.64	12.60	23.50	27.73	46.58	67.25	81.69	267.00
main, spill off (repeat)	7.90	13.07	23.72	29.63	48.76	69.08	80.00	272.16
main, spill 256 MiB	7.83	12.78	26.90	30.95	49.15	69.83	61.00	258.43
this PR, spill off	7.31	12.39	21.08	23.52	46.93	63.97	64.37	239.58
this PR, spill 256 MiB	8.45	13.21	20.91	23.89	48.19	60.26	62.83	237.73

At the shipping default the PR is 8.0% faster overall, with Q3 -22.3%, Q10 -22.8%, Q21 -13.7%.

Spill volumes come out byte-identical between the two builds at each setting (Q9: 416.9 GB over 1790 events either way), which is a nice confirmation that this changes file layout only, not bucketing or spill decisions.

Hypothesis: the win is parallel encoding, not throughput

What the gather + IPC framing + lz4 work has to cross is the end-of-task serial path. There look to be three ways it can avoid that:

1. Spilling does it incrementally, while the input stream is still arriving.
2. This PR does it on the 8 per-input tasks, in parallel.
3. Nothing else does.

That predicts the PR should help most exactly where spilling wasn't already helping, and the numbers line up:

• Q18 is the one query that spills even at memory_limit_per_task_bytes=0 (35.5 GB, via memory-pool pressure rather than the per-task budget). It's also the one query the PR doesn't move: -1.6%, inside noise.
• Q9 is where spilling helps main most on its own (81.7s -> 61.0s). It's the other query the PR doesn't improve on (+3.0%, also inside noise).
• Q3, Q10 and Q21 never spill at 0, and they're where the PR wins double digits.

The part I find most useful is what this does to the spill budget as a tuning knob. On main it's a real trade: turning it on gains 25% on Q9 but costs 15% on Q3 and 12% on Q10. With this PR the setting stops mattering much (totals 239.58 vs 237.73, under 1%). Operators currently have to pick a side; after this they mostly don't. That reads to me as a better argument for the change than the 8%.

Two caveats on reading the above

One iteration per cell. Q6 ranged 7.31-8.45s across the five runs, so the noise floor is around 8%. The double-digit per-query deltas clear that comfortably; the 8.0% total does not, and neither Q9's +3.0% nor Q18's -1.6% should be read as anything. Happy to rerun with 3 iterations if that's worth having.

Also worth flagging because I nearly quoted it: the write_time counter shows a 96% drop here and that figure is meaningless. On main the timer wraps finalize_output, which includes the interleave, IPC framing and compression. Here it wraps write_task_consolidated, which is copies only, since the encode moved into encode_buffered_partitions and isn't timed. On top of that main accumulates the counter once per input partition (8 per task) versus once per task here, so the aggregates aren't comparable either. Timing encode_buffered_partitions under the same counter would keep before/after comparisons meaningful for whoever benchmarks the shuffle writer next.

[avantgardnerio]

Nice work!

I originally tried to encode to a single file when implementing MPTs, but it ran much slower because I had encoding on the serial path. This PR keeps the encoding parallel, regardless of spill, and that's a big win (see perf numbers below).

I gave it a thorough review, and the only nit is that it would be worth revisiting the metrics at some point, but that doesn't have to be in this PR.

[Dandandan]

Nice work!

I originally tried to encode to a single file when implementing MPTs, but it ran much slower because I had encoding on the serial path. This PR keeps the encoding parallel, regardless of spill, and that's a big win (see perf numbers below).

I gave it a thorough review, and the only nit is that it would be worth revisiting the metrics at some point, but that doesn't have to be in this PR.

Thanks for the review! 🚀