[NoahKusaba]

Which issue does this PR close?

No issue. Found while measuring the sort-shuffle writer's memory accounting
for #2774.

Rationale for this change

At the end of input, the sort-shuffle writer encodes its buffered rows into one
IPC buffer per output partition. The buffered batches stay alive until the last
partition is gathered, so the writer holds the batches and the encoded output
at the same time. Only the batches were reserved in the memory pool, so the
pool could grant the encoded output's memory to other consumers.

Measured with a counting allocator, 200 output partitions and 2M rows with a
string column: the writer held 287 MiB at its peak but had reserved 205 MiB,
82 MiB short. With this change it reserves 263 MiB at its peak. Most of the
remaining gap is allocator size-class rounding, which no reservation counts.

What changes are included in this PR?

encode_buffered_partitions reserves each partition's encoded bytes as soon as
that partition is encoded:

• If the pool grants them, the partition stays in memory.
• If the pool refuses them, the encoded stream is appended to the partition's
  spill file instead, with a new SpillManager::spill_stream. Since perf: default the shuffle compression codec to zstd #2805, the
  in-memory encode and spill_batches both write through write_ipc_stream,
  so this is byte for byte what spilling the rows would have written.
• Afterwards the reservation holds only the encoded bytes, until the
  consolidated write consumes them.

So only the partition just encoded is ever held beyond what the pool grants,
and only until the pool's answer. Every partition is handled the same way, so a
small partition after a refused one can still stay in memory.

A spill during encoding counts as one memory-pool spill event, with its own
debug log line, and is timed in spill_time. spill_batches and the new
spill_stream share opening the spill file and updating the counts.

Alternatives considered

• Reserve an estimate before encoding. The only estimate available before
  encoding is the buffered size. With compression that is several times the
  encoded bytes, so the writer would spill whole buffers that would have fit.
• Keep a refused partition in memory and spill only the partitions after
  it. Simpler, but it leaves one partition's encoded bytes over the pool's
  limit until the shuffle file is written.

Are there any user-facing changes?

No API or configuration changes. spill_stream is pub(super). Two behaviour
changes:

• Under memory pressure, the end of input can now spill. Before, it never
  did. spill_count and disk I/O can therefore go up for jobs that run short
  of pool memory.
• After encoding, the reservation drops from the buffered size to the encoded
  size, which is usually smaller with compression. Other inputs of the same
  task can use that memory sooner.

Without memory pressure the write path is unchanged apart from one pool call per
output partition.

How was this tested?

• encoding_spills_the_partitions_the_pool_refuses: encodes 4 partitions of
  1000, 2000, 100 and 900 rows with room for partitions 0 and 2 only. It checks
  which partitions stay in memory and which spill, each partition's encoded
  bytes, and the bytes left reserved. So the small partition 2 stays in memory
  after partition 1 is refused. An unlimited pool keeps every partition in
  memory.
• spills_when_the_pool_cannot_hold_the_encoded_output: runs the writer
  through the existing run_round_trip helper with a 193 KiB pool. That pool
  holds the 192 KiB of buffered rows but not the encoded output as well, so
  only encoding spills. The helper checks that the writer spills, that every
  row comes back, and that the pool ends empty.
• spill_stream_writes_what_spill_batches_would: two spills each through
  spill_batches and through spill_stream produce the same spill file bytes
  and the same counts.
• Removing the pool check, removing the final resize, spilling every partition
  after the first refusal, or skipping spill_stream's write or its counting
  each makes at least one of these tests fail.
• cargo fmt, cargo clippy --all-targets --workspace --exclude ballista-cli --all-features -- -D warnings, cargo doc --document-private-items with
  -D warnings, and cargo test -p ballista-core --lib (490 passed).

TPC-H

SF10 parquet, one local scheduler and one executor (12 vcores, --memory-pool-size 4GB), --partitions 12 --iterations 3, release builds of main at 04455a1 and this PR. Average of 3 iterations, in seconds:

Query	main	PR	change
1	2.458	2.340	-4.8%
2	0.814	0.649	-20.3%
3	2.898	2.482	-14.4%
4	1.155	0.927	-19.7%
5	1.703	1.720	+1.0%
6	0.744	0.762	+2.4%
7	4.389	4.209	-4.1%
8	1.571	1.586	+1.0%
9	4.118	4.215	+2.4%
10	2.414	2.482	+2.8%
11	0.374	0.375	+0.3%
12	1.123	1.397	+24.4%
13	1.833	1.924	+5.0%
14	0.911	0.922	+1.2%
15	0.949	0.981	+3.4%
16	0.493	0.487	-1.2%
17	6.232	6.252	+0.3%
18	12.370	12.531	+1.3%
19	1.586	1.604	+1.1%
20	1.983	1.996	+0.7%
21	5.843	5.854	+0.2%
22	0.587	0.600	+2.2%
Total	56.548	56.295	-0.4%

The per-query differences are noise. Q2 to Q4 look faster only because main ran first and its first iterations were slower (Q2: 1.140 s, then 0.648 s). Q12 differed by +24% in that run. Rerun alternating main, PR, main, PR it averaged 1.219, 1.158, 1.162 and 1.169 s, so that was noise. Both builds logged the same 36 sort-shuffle spills (one memory-pool spill event per task, with matching byte counts), so on this workload encoding never spilled.

[andygrove]

Thanks for this, it's a really clear write-up and the per-partition handling is a nice touch. I traced the reservation through the drain, the spill path and write_task_consolidated, and the rows, stats and pool accounting all line up. The byte-equality test for spill_stream is a great way to pin the format.

Claude's review recommends approving, with just the two small nits below. I want to go through this one manually first to make sure I fully understand it, so I'm holding off on the approval for now.

The red Labeler trigger check looks like a workflow quoting issue unrelated to this change. A rebase should clear it.

[andygrove]

Nice that this reserves the real allocation rather than len(). One thought: buf grows by doubling inside write_ipc_stream, so its capacity can be close to twice the bytes written. That means a partition that would fit can still get spilled. Would a buf.shrink_to_fit() before the try_grow be worth it? It costs a copy, but the writer would then hold less memory for the rest of the task and ask the pool for less. Happy either way.

[NoahKusaba]

Left it out: under mimalloc, realloc keeps the same block when shrinking to at least half its size, so shrink_to_fit frees nothing but lowers capacity(), and we'd under-reserve. Reserving capacity() over-counts slightly, the safe side. Added a comment.

[andygrove]

This test depends on the buffered rows coming to just under 193 KiB. If the BufferedBatches accounting ever grows a little, the input loop spills instead. The test would still pass but would no longer cover the encode path. encoding_spills_the_partitions_the_pool_refuses covers that path directly, so this is minor. Could the test also check that the spill came from encoding, for example that the input loop recorded no spill event? Or a comment saying the margin is tight would help the next person who changes the accounting.

[NoahKusaba]

Both spill paths write identical bytes, so the test now asserts its precondition: the buffered bytes, computed with the writer's own accounting, fit the pool. If the accounting grows, the test fails instead of silently covering the input loop.