[NoahKusaba]

Which issue does this PR close?

Part of #2578. First of three PRs; no behaviour change.

Rationale for this change

The writer works on BufferedBatches' internals in three places: encode and
spill each rebuild the same gather loop, and the input loop tracks the
buffer's memory by hand. Moving both into BufferedBatches keeps the gather
and the memory accounting in one place, behind a small interface (push_batch,
drain, memory_size, is_empty).

#2578's follow-up then adds its second buffer layout inside BufferedBatches,
without touching the writer's encode, spill or memory accounting.

What changes are included in this PR?

• drain(f) empties the buffer, passing each output batch (at most
  batch_size rows) and its partition to f. Encode and spill call it.
• memory_size() returns the buffered batches' size plus the index lists'
  capacity, kept as a running total. The writer passes it to try_resize.
• new takes batch_size.

memory_size() equals the number the writer computed before, so every spill
decision, and every output batch, is unchanged.

Are there any user-facing changes?

No behaviour change. BufferedBatches was re-exported from sort_shuffle
but has no users outside the writer, so it is now crate-private. Its API
changes:

• new(num_partitions, schema) → new(num_partitions, schema, batch_size)
• take(), BufferedTake, indices_for(), indices_allocated_size() and
  num_buffered_rows() are replaced by drain() and memory_size().

How was this tested?

• Buffer unit tests rewritten for the new API.
• Existing writer tests, including the spill tests, pass unchanged.
• cargo test -p ballista-core -p ballista-executor -p ballista, clippy and
  fmt are clean.

[andygrove]

Thanks @NoahKusaba, this is a nice cleanup. Having encode and spill share one gather loop reads much better, and BufferedBatches is the right home for the memory accounting. I checked that memory_size() matches the old number. extend grows the index lists exactly like the old reserve and push loop, and try_resize ends up at the same reservation the try_grow calls added up to, so spill decisions don't change. The encoded bytes are identical too.

I can't approve it as is though, because it raises memory use in encode_buffered_partitions. Every output partition's IPC writer now stays open until the drain finishes, and with zstd each one holds its own compression context. At 200 output partitions that's about 140 MB more per input partition, and the memory pool never sees any of it. A refactor that's meant to be a no-op shouldn't change the memory profile, so this needs to go back to one open writer at a time. Details inline, plus a couple of smaller things.

Since this is on the shuffle write path, could you also share a before/after against main once that's fixed? The sort_shuffle bench is a quick option and covers both encode and spill at 200 partitions:

git checkout main
cargo bench -p ballista-benchmarks --bench sort_shuffle -- --save-baseline main
git checkout refactor/sort-shuffle-buffer-drain
cargo bench -p ballista-benchmarks --bench sort_shuffle -- --baseline main

If you have TPC-H SF10 handy, a full run from benchmarks/ would be great too:

cargo run --release --bin tpch -- benchmark ballista \
  --host localhost --port 50050 \
  --iterations 3 --path <tpch parquet dir> --format parquet

No need for an EXPLAIN diff since the plan doesn't change.

[NoahKusaba]

Peak RSS growth during the write, zstd, one 8192-row batch per partition:

partitions	main	before the fix	after the fix
200	36 MB	139 MB	36 MB
1000	140 MB	655 MB	140 MB

I also wrote the same 37-partition shuffle (40 batches with a string column)
on main and on this branch, with lz4 and zstd, with and without spilling.
The data files, index files and reported stats are byte-identical, and the
spill counts match.

sort_shuffle bench against main (laptop, on AC, criterion's 10 samples):

benchmark	main	this branch	change
no spill, 10 batches, 200 partitions	71.0 ms	65.1 ms	−8.8%
no spill, 50 batches, 200 partitions	159.4 ms	140.4 ms	−10.6%
spill, 8 MB limit, 200 partitions	1.23 s	1.00 s	−18.6%
spill, 2 MB limit, 200 partitions	2.27 s	1.77 s	−22.3%

No regression. The branch does save a little work: push_batch extends each
index list in one call instead of pushing row by row, and the memory total is
kept as it grows instead of summing every partition's index capacity twice
per input batch. That's well under a millisecond per run here, though, so I
don't think it's a significant speedup. The main baseline ran first, and
two runs of this branch differed by 5–6%, so I expect most of the gap is
machine noise.

I skipped the TPC-H run since the plan and the output bytes don't change, but
I'm happy to run it if you'd like.

[andygrove]

Disclaimer: this review was generated by an LLM (Claude Code) at my request, against e01542e0. File/line references were checked against the code, but please treat the conclusions as suggestions and push back where it's wrong.

Thanks @NoahKusaba, everything from the last round is addressed. Encode is back to one open StreamWriter at a time, and your RSS numbers match main. Empty partitions are skipped, the batch refs are built once per drain, and the upgrade note and the reallocation test are in. Thanks for running the benchmarks too. One optional suggestion inline.

[NoahKusaba]

Note on the "Retrigger CI" commit (d273b17): the failure on 47fc2f9 was a flaky chaos test, not this change.

test linux crates failed in ballista-chaos --test ha, on baseline_matches_local_datafusion::case_2_aqe_on, with cluster must start: timed out waiting for 2 executors to register. The other 15 ha tests passed, and so did the rest of the workflow.

The executor logs show both executors got the same Flight port (127.0.0.1:45565). Executor 0 then failed with AddrInUse and shut down, so only one executor was left to register:

executor-0: Built-in arrow flight server listening on: 127.0.0.1:45565
executor-0: Could not start built-in arrow flight server.
executor-0: executor services stopped with reason Err(TonicError(... AddrInUse ...))
executor-1: Built-in arrow flight server listening on: 127.0.0.1:45565

This is a race in the test harness. spawn_executor calls free_ports() once per executor (chaos-testing/src/cluster.rs). That function binds to :0 to find a free port, then releases it before the child process binds. Within one call the ports are distinct, but nothing stops a later call from getting the same number. Executor 0 starts its Flight server only after gRPC setup and registration, so executor 1's free_ports() got 45565 back from the kernel before executor 0 had bound it.

This PR doesn't touch the chaos harness or executor startup, and the same suite passed on the rerun.