fix: leave memory for aggregate spill replay - #25383

## Which issue does this PR close?

Prerequisite for apache#25172, which fixes `FairSpillPool` accounting across
reservations belonging to the same memory consumer. This PR prepares
aggregate spill replay to work within that shared allowance and should
merge first.

## Rationale for this change

**Spilling an aggregation to disk only helps if there is enough memory
to read it back and finish the aggregation.** During spill replay,
DataFusion merges sorted spill files and feeds their rows into an
aggregate that combines the intermediate states into final results. The
merge needs buffers for its inputs and output, while the aggregate needs
memory for group keys and accumulators. Those allocations coexist:
producing the next batch and processing it are both part of the same
operator's memory budget.

For example, consider:

```sql
SELECT customer_id, ARRAY_AGG(event_id)
FROM events
GROUP BY customer_id;
```

With many customers and a small memory limit, the aggregate can spill
several times. A customer's events may then be spread across several
spill files. Merging those files brings the customer's intermediate
states together, but the aggregate still has to build the final array.
Even if the merge delivers one row at a time, that array grows until the
customer's group is complete. Small input batches therefore do not
eliminate the need for replay memory.

The current merge tries to merge as many spill files as its reservation
permits, without leaving room for this downstream work. As a simplified
example, suppose the operator has a 100 MiB allowance and no other
reservations. An 80 MiB merge fits by itself, but if the aggregate then
needs 24 MiB to process its output, their combined 104 MiB does not fit.
A pool enforcing the shared allowance must reject that growth, even
though using a smaller merge could let the same aggregation make
progress.

This is particularly relevant to apache#25172. The merge and aggregate use
separate reservations under the same memory consumer. Today,
`FairSpillPool` checks those reservations individually, which can hide
their combined usage exceeding the consumer's allowance. Once apache#25172
enforces that allowance across both reservations, replay must already
leave space for the aggregate. This PR addresses that requirement before
the accounting fix lands.

## What changes are included in this PR?

**Aggregate spill merges now leave room for the aggregate consuming
their output.** When choosing how many spill files to merge and how much
to buffer, the merge asks the configured memory pool to admit its
buffers plus an equal amount of spare capacity for replay. Once
admitted, it releases the spare reservation before replay starts,
retaining only the reservation for the merge buffers. This uses the
pool's existing allocation checks, so the decision follows its admission
policy, including any fair-share restrictions, without adding a public
memory-pool API.

In the example above, the 80 MiB merge would first need approval for 160
MiB and would be rejected. A smaller 40 MiB merge would need approval
for 80 MiB; after returning the 40 MiB of temporary headroom, it would
leave 60 MiB available, enough for the illustrative 24 MiB of aggregate
state. The equal-sized headroom is a practical budgeting policy, not an
estimate of the exact accumulator size or a guarantee that every
aggregation will fit.

The existing multi-level merge adapts by merging fewer files at once,
reducing read-ahead, or splitting oversized spill batches into smaller
batches. This PR makes those adjustments account for replay as well. It
also handles the point where a batch cannot shrink further: a single
wide row may still fit the real pool even when equal headroom does not.
In that case, only the minimum merge may retry without the extra
headroom, with read-ahead disabled and all merge memory still subject to
the pool's checks. Inspecting decoded batches before rewriting a spill
file avoids unnecessary writes when the rows already fit or cannot be
split, including when the original spill files fill the disk quota.

All supported aggregate replay implementations use this policy,
including hash and ordered aggregation and the legacy implementation
selected by `enable_migration_aggregate=false`. The legacy path also
releases unused initial grouping capacity before replay. The headroom
policy is private to aggregate replay; ordinary sort callers retain
their existing admission behavior.

## Are these changes tested?

Regression tests exercise replay across the supported aggregate
implementations; the migrated hash and ordered paths also run with
another registered spilling consumer. In particular, the legacy
`ARRAY_AGG` test uses 64 groups with 64 values per group, one-row
batches, and an 8 KiB pool to check that accumulator state can grow
across replay batches. The tests verify exact aggregate results,
reservation bounds, and release of memory and spill files.

Merge-level coverage checks oversized batches, short and odd-sized
batches, indivisible rows, decoded string-view sizes, and spill files
that already fill the disk quota. It also checks that temporary headroom
is released after a candidate merge is rejected and that the fallback
still fails when the actual merge cannot fit.

<details>
<summary>Validation results and environment</summary>

Validated replay alone, with the existing `FairSpillPool`
implementation:

- 2,263 physical-plan tests and 2,281 core/CLI tests passed.
- The existing permanent-pressure regression passed with its original
success expectation.
- All 520 SQL logic test files passed.
- `cargo fmt --all` and strict all-target/all-feature Clippy passed.
- Full `./dev/rust_lint.sh` passed, including private Rust documentation
and local Markdown links.

Local validation uses Rust 1.98.1 and upstream revision `22651d24` with
its unchanged dependency lockfile; newer main's dependency versions are
unavailable in the local registry. Merge compatibility with current main
is checked separately, and GitHub CI validates its merged revision.

</details>

## Are there any user-facing changes?

Aggregations that spill leave memory available for processing the merged
rows, reducing avoidable replay failures under constrained memory.
Achieving this can require smaller batches or additional merge passes.
SQL semantics and public APIs are unchanged.

The temporary headroom reservations can increase recorded reservation
peaks without allocating additional data buffers. That headroom is
released before replay, so other concurrent allocations or aggregate
state that outgrows the available memory can still cause
`ResourcesExhausted`.

## Which issue does this PR close?

Follow-up to [the review of
apache#25383](apache#25383 (review)),
which suggested reducing intermediate aggregate spill work.

## Rationale for this change

An intermediate aggregate merge can rewrite rows that could remain on
disk until final replay. With eight uniform full-batch runs and an
admitted fan-in of seven, merging six inputs leaves three for the final
pass. The ordinary selection merges seven and leaves two, unnecessarily
rewriting one run.

The final pass must still fit if another partition consumes available
memory during the intermediate write. The retained seven-input buffer
reservation can cover three final inputs plus their replay headroom; it
cannot guarantee seven final inputs with that headroom.

## What changes are included in this PR?

Eligible spill-only intermediate merges select only enough inputs to
leave a final pass whose expected buffers and replay headroom fit the
existing buffer reservation. An actually trimmed pass retains that
reservation in the merge builder through intermediate EOF and spill
completion, transferring it to the next ordinary admission. Sizing
acquires no additional pool memory and preserves read-ahead and output
batching. Passes that are not trimmed keep ordinary reservation
ownership.

Only the newer aggregate implementation (`AggregateSpill`, enabled by
default) opts in. Legacy aggregation keeps its original selection.
Sizing requires equal recorded maximum batch memory and batch-size
limits across admitted and pending runs. Every original run must contain
a full target-size batch; short original runs disable sizing for their
replay. Intermediate and split spill writes also check actual
largest-batch row counts against their output limit. These checks use
private metadata; public APIs are unchanged.

The sizing budget assumes the produced run is no wider than the guarded
inputs. Variable-width output can still require further intermediate
work. Final admission and its normal headroom release remain unchanged.

## What is the testing strategy for this PR?

Head: `66569072850d91476b042010906ac6461ebbbd08`. Comparison base:
`991fd23dd0be0046af5945b8c3905612859b657a`.

The regression tests cover competing memory consumers at replay-headroom
release and intermediate EOF, both beneficial and unchanged merge
selections, short batches, and heterogeneous input budgets. They check
exact ordered results, intermediate rewrite counts, actual competing
reservations, memory and spill-file cleanup, and reservation ownership
when the intermediate stream and builder are dropped.

CI runs for this head: [push Rust
run](https://github.com/apache/datafusion/actions/runs/36461591867) and
[pull-request Rust
run](https://github.com/apache/datafusion/actions/runs/36461597627).

### Performance validation

The updated head has not been rebenchmarked against
`991fd23dd0be0046af5945b8c3905612859b657a`.

## Are there any user-facing changes?

Eligible spilling aggregations can perform fewer intermediate row
writes. Results and configuration remain unchanged. Short runs and
selections without sufficient retained replay capacity use the ordinary
selection.

… memory share (apache#6544)

* fix: let a spilled final aggregate read its spill files back past its memory share

DataFusion 55's FinalHashAggregateStream replays its merged spill files through a stream that
can't spill, so a refused memory request there fails the task. The merge's read buffers are a
sibling reservation of the same consumer and take as many spill files as fit, so the replay often
finds the consumer's share already taken (apache#6254).

Both Comet pools now record that request as overcommit instead of refusing it. They recognize it
as a request from a FinalHashAggregateStream consumer while another of its reservations holds
memory, which in DataFusion 55.1 happens only during the replay. The replay emits its finished
groups after every batch, so the overcommit stays around one batch of groups, and releases repay
it first. Refusals while the aggregate reads its input, and while the merge picks its files, are
unchanged.

Remove this once Comet's DataFusion includes apache/datafusion#25383.

* fix: let an ordered final aggregate read its spill files back too

DataFusion runs a final aggregate whose input is sorted on some of its
grouping keys as an OrderedFinalAggregateStream. It merges and replays its
spill files the same way FinalHashAggregateStream does, so its replay
failed the task the same way. Treat its consumer as a final aggregate too.

Explain in spill_replay.rs why recording the replay's request is safe: the
replay asks for memory only after it has aggregated a batch, so the memory
already exists, as it does for a grow. Describe the exception in the memory
management guide, which said the fair pool always refuses a request that
fails its local checks.

* docs: tighten the spill replay comments and guide

The fair pool is the only one with a share and a pool total to skip, so say
so. The greedy pool takes its tracking lock for other consumers when Spark
refuses them, not never. Point whoever removes the workaround at the tests
that show whether the replay still needs it.

* test: share the spill replay tests between the pools

fair_pool.rs and unified_pool.rs each had a copy of the same two spill
replay scenarios. Run them once, in spill_replay.rs, against both pool
types as createPlan builds them, which also checks that the wrappers pass
register through to the greedy pool's tracking. Keep a greedy pool test for
its per-consumer map, and drop a test import that the pool's own imports
now cover.

Share the final aggregate spill checks between the two apache#6254 tests in
CometAggregateSuite, and use checkCometAnswer, which collects once and
labels the Comet answer as Comet's.

* refactor: move the spill replay workaround into a pool wrapper

The workaround for apache#6254 lived in both Comet pools. The fair pool sent its
three refusals through a check, and the greedy pool tracked each final
aggregate across its reservations. apache#5613 reworks the fair pool's try_grow,
and apache#6583 has to take the workaround out again, so both would have had to
rework the pools.

SpillReplayPool, in spill_replay.rs, now wraps the tracked Comet pool
instead. It keeps a total for each final aggregate's consumer. When the pool
refuses a request from one while another of its reservations holds memory,
it calls the pool's grow, which skips the fair pool's limits and carries
what Spark doesn't grant as overcommit. fair_pool.rs and unified_pool.rs are
back to main's versions, and overcommit() looks through the wrapper.

The Rust tests run against both pool types as createPlan builds them. They
now reach the fair pool's pool-limit refusal too, and check that a failed
Spark call during the replay is returned rather than recorded.

The guide adds the wrapper to the pool stack and keeps a short section. The
DataFusion 55.1 facts that an upgrade has to re-check stay in
spill_replay.rs, which names apache#6583 next to the removal condition.

… memory share (#6544) (#6713)

* fix: let a spilled final aggregate read its spill files back past its memory share

DataFusion 55's FinalHashAggregateStream replays its merged spill files through a stream that
can't spill, so a refused memory request there fails the task. The merge's read buffers are a
sibling reservation of the same consumer and take as many spill files as fit, so the replay often
finds the consumer's share already taken (#6254).

Both Comet pools now record that request as overcommit instead of refusing it. They recognize it
as a request from a FinalHashAggregateStream consumer while another of its reservations holds
memory, which in DataFusion 55.1 happens only during the replay. The replay emits its finished
groups after every batch, so the overcommit stays around one batch of groups, and releases repay
it first. Refusals while the aggregate reads its input, and while the merge picks its files, are
unchanged.

Remove this once Comet's DataFusion includes apache/datafusion#25383.

* fix: let an ordered final aggregate read its spill files back too

DataFusion runs a final aggregate whose input is sorted on some of its
grouping keys as an OrderedFinalAggregateStream. It merges and replays its
spill files the same way FinalHashAggregateStream does, so its replay
failed the task the same way. Treat its consumer as a final aggregate too.

Explain in spill_replay.rs why recording the replay's request is safe: the
replay asks for memory only after it has aggregated a batch, so the memory
already exists, as it does for a grow. Describe the exception in the memory
management guide, which said the fair pool always refuses a request that
fails its local checks.

* docs: tighten the spill replay comments and guide

The fair pool is the only one with a share and a pool total to skip, so say
so. The greedy pool takes its tracking lock for other consumers when Spark
refuses them, not never. Point whoever removes the workaround at the tests
that show whether the replay still needs it.

* test: share the spill replay tests between the pools

fair_pool.rs and unified_pool.rs each had a copy of the same two spill
replay scenarios. Run them once, in spill_replay.rs, against both pool
types as createPlan builds them, which also checks that the wrappers pass
register through to the greedy pool's tracking. Keep a greedy pool test for
its per-consumer map, and drop a test import that the pool's own imports
now cover.

Share the final aggregate spill checks between the two #6254 tests in
CometAggregateSuite, and use checkCometAnswer, which collects once and
labels the Comet answer as Comet's.

* refactor: move the spill replay workaround into a pool wrapper

The workaround for #6254 lived in both Comet pools. The fair pool sent its
three refusals through a check, and the greedy pool tracked each final
aggregate across its reservations. #5613 reworks the fair pool's try_grow,
and #6583 has to take the workaround out again, so both would have had to
rework the pools.

SpillReplayPool, in spill_replay.rs, now wraps the tracked Comet pool
instead. It keeps a total for each final aggregate's consumer. When the pool
refuses a request from one while another of its reservations holds memory,
it calls the pool's grow, which skips the fair pool's limits and carries
what Spark doesn't grant as overcommit. fair_pool.rs and unified_pool.rs are
back to main's versions, and overcommit() looks through the wrapper.

The Rust tests run against both pool types as createPlan builds them. They
now reach the fair pool's pool-limit refusal too, and check that a failed
Spark call during the replay is returned rather than recorded.

The guide adds the wrapper to the pool stack and keeps a short section. The
DataFusion 55.1 facts that an upgrade has to re-check stay in
spill_replay.rs, which names #6583 next to the removal condition.

(cherry picked from commit ba7c892)

Adapted for branch-1.1:
- The tests build the pools over a fake Spark through create_pool, a seam
  that #6271 added to memory_pools. Only that seam is ported, not #6271's
  memory usage log change, which is not on branch-1.1: create_memory_pool
  builds the pools through create_pool, with_spark is pub(super), and the
  pools' new constructors, which nothing calls any more, are removed.
- overcommit() and unwrap_task_shared(), which only that log reads, are not
  ported, so nothing looks through the wrapper. spill_replay.rs drops
  unwrap_spill_replay, and its tests drop their overcommit(&pool)
  assertions, which restate pool.reserved() minus what the fake Spark
  holds. The test that shrinks the replay checks pool.reserved() instead.
- memory_management.md: the paragraph about the task's shared
  CometTaskMemoryManager, from #6261, is not on branch-1.1, so the new
  section follows the task-shared pools section directly.
- CometAggregateSuite: only its imports differ. This branch does not import
  Column, CometBaseAggregateExec or CometSortAggregateExec.

## Which issue does this PR close?

Prerequisite for apache#25172, which fixes `FairSpillPool` accounting across
reservations belonging to the same memory consumer. This PR prepares
aggregate spill replay to work within that shared allowance and should
merge first.

## Rationale for this change

**Spilling an aggregation to disk only helps if there is enough memory
to read it back and finish the aggregation.** During spill replay,
DataFusion merges sorted spill files and feeds their rows into an
aggregate that combines the intermediate states into final results. The
merge needs buffers for its inputs and output, while the aggregate needs
memory for group keys and accumulators. Those allocations coexist:
producing the next batch and processing it are both part of the same
operator's memory budget.

For example, consider:

```sql
SELECT customer_id, ARRAY_AGG(event_id)
FROM events
GROUP BY customer_id;
```

With many customers and a small memory limit, the aggregate can spill
several times. A customer's events may then be spread across several
spill files. Merging those files brings the customer's intermediate
states together, but the aggregate still has to build the final array.
Even if the merge delivers one row at a time, that array grows until the
customer's group is complete. Small input batches therefore do not
eliminate the need for replay memory.

The current merge tries to merge as many spill files as its reservation
permits, without leaving room for this downstream work. As a simplified
example, suppose the operator has a 100 MiB allowance and no other
reservations. An 80 MiB merge fits by itself, but if the aggregate then
needs 24 MiB to process its output, their combined 104 MiB does not fit.
A pool enforcing the shared allowance must reject that growth, even
though using a smaller merge could let the same aggregation make
progress.

This is particularly relevant to apache#25172. The merge and aggregate use
separate reservations under the same memory consumer. Today,
`FairSpillPool` checks those reservations individually, which can hide
their combined usage exceeding the consumer's allowance. Once apache#25172
enforces that allowance across both reservations, replay must already
leave space for the aggregate. This PR addresses that requirement before
the accounting fix lands.

## What changes are included in this PR?

**Aggregate spill merges now leave room for the aggregate consuming
their output.** When choosing how many spill files to merge and how much
to buffer, the merge asks the configured memory pool to admit its
buffers plus an equal amount of spare capacity for replay. Once
admitted, it releases the spare reservation before replay starts,
retaining only the reservation for the merge buffers. This uses the
pool's existing allocation checks, so the decision follows its admission
policy, including any fair-share restrictions, without adding a public
memory-pool API.

In the example above, the 80 MiB merge would first need approval for 160
MiB and would be rejected. A smaller 40 MiB merge would need approval
for 80 MiB; after returning the 40 MiB of temporary headroom, it would
leave 60 MiB available, enough for the illustrative 24 MiB of aggregate
state. The equal-sized headroom is a practical budgeting policy, not an
estimate of the exact accumulator size or a guarantee that every
aggregation will fit.

The existing multi-level merge adapts by merging fewer files at once,
reducing read-ahead, or splitting oversized spill batches into smaller
batches. This PR makes those adjustments account for replay as well. It
also handles the point where a batch cannot shrink further: a single
wide row may still fit the real pool even when equal headroom does not.
In that case, only the minimum merge may retry without the extra
headroom, with read-ahead disabled and all merge memory still subject to
the pool's checks. Inspecting decoded batches before rewriting a spill
file avoids unnecessary writes when the rows already fit or cannot be
split, including when the original spill files fill the disk quota.

All supported aggregate replay implementations use this policy,
including hash and ordered aggregation and the legacy implementation
selected by `enable_migration_aggregate=false`. The legacy path also
releases unused initial grouping capacity before replay. The headroom
policy is private to aggregate replay; ordinary sort callers retain
their existing admission behavior.

## Are these changes tested?

Regression tests exercise replay across the supported aggregate
implementations; the migrated hash and ordered paths also run with
another registered spilling consumer. In particular, the legacy
`ARRAY_AGG` test uses 64 groups with 64 values per group, one-row
batches, and an 8 KiB pool to check that accumulator state can grow
across replay batches. The tests verify exact aggregate results,
reservation bounds, and release of memory and spill files.

Merge-level coverage checks oversized batches, short and odd-sized
batches, indivisible rows, decoded string-view sizes, and spill files
that already fill the disk quota. It also checks that temporary headroom
is released after a candidate merge is rejected and that the fallback
still fails when the actual merge cannot fit.

<details>
<summary>Validation results and environment</summary>

Validated replay alone, with the existing `FairSpillPool`
implementation:

- 2,263 physical-plan tests and 2,281 core/CLI tests passed.
- The existing permanent-pressure regression passed with its original
success expectation.
- All 520 SQL logic test files passed.
- `cargo fmt --all` and strict all-target/all-feature Clippy passed.
- Full `./dev/rust_lint.sh` passed, including private Rust documentation
and local Markdown links.

Local validation uses Rust 1.98.1 and upstream revision `22651d24` with
its unchanged dependency lockfile; newer main's dependency versions are
unavailable in the local registry. Merge compatibility with current main
is checked separately, and GitHub CI validates its merged revision.

</details>

## Are there any user-facing changes?

Aggregations that spill leave memory available for processing the merged
rows, reducing avoidable replay failures under constrained memory.
Achieving this can require smaller batches or additional merge passes.
SQL semantics and public APIs are unchanged.

The temporary headroom reservations can increase recorded reservation
peaks without allocating additional data buffers. That headroom is
released before replay, so other concurrent allocations or aggregate
state that outgrows the available memory can still cause
`ResourcesExhausted`.

## Which issue does this PR close?

Follow-up to [the review of
apache#25383](apache#25383 (review)),
which suggested reducing intermediate aggregate spill work.

## Rationale for this change

An intermediate aggregate merge can rewrite rows that could remain on
disk until final replay. With eight uniform full-batch runs and an
admitted fan-in of seven, merging six inputs leaves three for the final
pass. The ordinary selection merges seven and leaves two, unnecessarily
rewriting one run.

The final pass must still fit if another partition consumes available
memory during the intermediate write. The retained seven-input buffer
reservation can cover three final inputs plus their replay headroom; it
cannot guarantee seven final inputs with that headroom.

## What changes are included in this PR?

Eligible spill-only intermediate merges select only enough inputs to
leave a final pass whose expected buffers and replay headroom fit the
existing buffer reservation. An actually trimmed pass retains that
reservation in the merge builder through intermediate EOF and spill
completion, transferring it to the next ordinary admission. Sizing
acquires no additional pool memory and preserves read-ahead and output
batching. Passes that are not trimmed keep ordinary reservation
ownership.

Only the newer aggregate implementation (`AggregateSpill`, enabled by
default) opts in. Legacy aggregation keeps its original selection.
Sizing requires equal recorded maximum batch memory and batch-size
limits across admitted and pending runs. Every original run must contain
a full target-size batch; short original runs disable sizing for their
replay. Intermediate and split spill writes also check actual
largest-batch row counts against their output limit. These checks use
private metadata; public APIs are unchanged.

The sizing budget assumes the produced run is no wider than the guarded
inputs. Variable-width output can still require further intermediate
work. Final admission and its normal headroom release remain unchanged.

## What is the testing strategy for this PR?

Head: `66569072850d91476b042010906ac6461ebbbd08`. Comparison base:
`991fd23dd0be0046af5945b8c3905612859b657a`.

The regression tests cover competing memory consumers at replay-headroom
release and intermediate EOF, both beneficial and unchanged merge
selections, short batches, and heterogeneous input budgets. They check
exact ordered results, intermediate rewrite counts, actual competing
reservations, memory and spill-file cleanup, and reservation ownership
when the intermediate stream and builder are dropped.

CI runs for this head: [push Rust
run](https://github.com/apache/datafusion/actions/runs/36461591867) and
[pull-request Rust
run](https://github.com/apache/datafusion/actions/runs/36461597627).

### Performance validation

The updated head has not been rebenchmarked against
`991fd23dd0be0046af5945b8c3905612859b657a`.

## Are there any user-facing changes?

Eligible spilling aggregations can perform fewer intermediate row
writes. Results and configuration remain unchanged. Short runs and
selections without sufficient retained replay capacity use the ordinary
selection.

DataFusion 56 includes apache/datafusion#25383, which leaves room for a
final aggregate to read its spill files back, so SpillReplayPool is no
longer needed. It is also now harmful: #25383 holds replay headroom on a
sibling reservation while the spill merge picks its files, so the pool
took every refused read-buffer request for a replay and granted it. The
merge then opened every spill file and starved the next aggregate, which
failed the #6880 PartialMerge test with a 2 MiB pool.

The #6254 and #6857 read-back tests, which the pool was added for, pass
without it. The sort-merge join #6880 test now uses an 8 MiB pool,
because with DataFusion 56 its PartialMerge fits in 12 MiB or more
without spilling.

Closes #6583.

…che#26264)

## Which issue does this PR close?

- Part of apache#25804 (finding 9: "The final spill merge grows its
reservation until the pool refuses").

## Rationale for this change

A sort that has spilled finishes with a merge of its spilled runs. That
merge
seats as many runs as the memory pool will admit (the fan-in is
unlimited by
default), and it holds those buffers until its output is fully consumed.

In a query with several partitions, partitions that finish their input
early
start this final merge while other partitions are still buffering input,
and
while the operator above the sort is still running. Each early merge
takes
whatever the pool admits, and the pool can be exhausted. The partition
that is
still sorting then fails to grow its reservation and the query fails
with
`ResourcesExhausted`, even though only memory for read-ahead buffers was
taken.

A standalone reproduction (below): 6.67M rows,
`SUM(p) OVER (PARTITION BY g ORDER BY k ROWS UNBOUNDED PRECEDING)`
followed by
`MAX`, 10 target partitions, a `FairSpillPool` of 62.5 MB, 40 runs per
setting:

| `sort_spill_reservation_bytes` | main (2c568b4) | this PR |
|---|---|---|
| 1 MB | 25 / 40 succeed | 40 / 40 |
| 1.6 MB | 34 / 40 | 40 / 40 |
| 3 MB | 7 / 40 | 35 / 40 |

The failures are the partitions that are still sorting:
`Failed to allocate additional 3.0 MB for ExternalSorterMerge[6] with
0.0 B
already allocated`, `... 598.8 KB for ExternalSorter[8] with 0.0 B ...`,
or
`... 144.1 KB for ExternalSorterMerge[7]` in `BatchBuilder::push_batch`
while
merging in-memory batches to spill. With a pool instrumented to record
each
consumer's reservation, five partitions in their final merge held 12-29
MB
each when the refusal happened. Capping `max_spill_merge_fan_in` at 4
made
the same runs succeed, which confirms that the merge fan-in causes the
failures.

## What changes are included in this PR?

apache#25383 lets aggregate spill replay leave memory for the aggregate: the
merge
asks the pool to admit its buffers plus as much again, then returns the
extra.
The sort uses a narrower form of that policy:

- The minimum merge of two runs is admitted as before, against the pool
alone. A sort under tight memory keeps its current behavior: it does not
  re-spill runs smaller or emit shorter batches only to leave headroom.
- Each run beyond the minimum is seated only if the pool can admit what
the
  merge takes beyond the workspace it already holds
  (`sort_spill_reservation_bytes`, retained since apache#24740), plus as much
again. The extra is returned before the merge starts. Seating runs
within
  the merge's own workspace needs no headroom, because it takes nothing
  from others.

When the pool refuses, the merge seats fewer runs and merges in more
passes,
as it already does under memory pressure.

The policy is a new `with_fan_in_headroom` option on
`StreamingMergeBuilder`
and `MultiLevelMergeBuilder`, enabled by `ExternalSorter`. The
aggregate's
replay headroom is unchanged.

Headroom is kept only while someone may need it: while another partition
of
the same `SortExec` is still reading its input, or while another
consumer
holds memory in the pool. Partitions that have already started their own
final merge are covered by the second check. A partition leaves the
count
when it has read all its input, or when it is dropped before that (its
stream
dropped early, or its input failed). A sort with the pool to itself
seats every run that fits, as before, and so does the last partition to
finish reading: keeping headroom there cost an extra merge pass (up to
1.8x
the spilled rows) whether or not anyone needed the memory. With all 2M
rows
in one of 4 partitions, most pools between 16 and 64 MB spilled 1.5-1.8x
the
rows with headroom on; with the count of partitions still reading, every
pool
from 16 to 128 MB spills the same rows as without headroom.

On balanced multi-partition sorts (2M rows, greedy pool, 20 runs each,
same
build with headroom off and on; measured by @jayzhan211 before headroom
followed the count of partitions still reading, which moved these
configs
from 54 to 49 of 60 successes, against 32 without headroom):

| partitions / pool / reservation | success, off → on | spilled rows,
off → on |
|---|---|---|
| 4 / 96 MB / 1 MB | 8 → 17 of 20 | 2.76M → 3.35M |
| 8 / 192 MB / 10 MB | 12 → 17 of 20 | 2.00M → 2.16M |
| 2 / 64 MB / 1 MB | 12 → 20 of 20 | 3.00M → 3.82M |

Where another consumer holds memory, headroom is kept and the merge
seats
fewer runs, so it rewrites rows in more passes. That is the price of the
fix:
with a skewed 4-partition sort under `SortPreservingMergeExec`, or two
one-partition sorts sharing a pool, @jayzhan211 measured up to 1.9x the
spilled rows at 24297ff. `sort_tpch` SF1 Q3/Q4/Q11 showed no
wall-clock
change above noise.

<details>
<summary>Measured cost at 24297ff, same build with headroom off and on
(by @jayzhan211)</summary>

| Plan | Pool | Success, off → on | Spilled rows, off → on |
|---|---|---|---|
| 4 partitions, 97/1/1/1 of 2M rows, under `SortPreservingMergeExec` |
24 MB | 9/9 → 9/9 | 7.66M → 11.01M |
| same | 32 MB | 9/9 → 9/9 | 4.99M → 9.38M |
| same | 48 MB | 1/9 → 9/9 | 2.99M → 4.40M |
| same | 64 MB | 9/9 → 9/9 | 1.94M → 3.12M |
| Two one-partition `SortExec`s, 1.5M rows each | 32 MB | 20/20 → 20/20
| 7.83M → 14.28M |
| same | 48 MB | 17/20 → 19/20 | 6.10M → 8.24M |

Greedy pool, 1 MB `sort_spill_reservation_bytes`.

</details>

### What this does not fix

Each merge leaves as much free as it takes, but merges that start one
after
another still compound: in the remaining 3 MB failures, nine partitions
held
3-18 MB each (17.9, 11.5, 7.7, 5.1, ... MB), and the last partition to
start
could not reserve its own 3 MB `sort_spill_reservation_bytes`. With 10
partitions, the workspaces alone are half of this pool. A guarantee
would
need each merge to know how many partitions share the pool, for example
a
per-partition share of the limit. I left that out of this PR, because it
is
a larger policy change and I would like to hear whether it is wanted.

A single-partition sort, or the last partition to finish reading, still
leaves nothing for the operator above it: that operator reserves nothing
before its first batch, so the merge cannot see it. A `SortExec` whose
partitions are not all executed (for example under a `LIMIT` that stops
early)
counts the partitions never started as still reading, so it keeps
headroom.

## What is the testing strategy for this PR?

- `test_final_spilled_merge_leaves_headroom` (sort/spill_tests.rs)
spills a
sort and leaves the pool room for the final merge to seat every run, but
not to seat every run and still leave as much free. It checks that
whatever
the merge takes beyond its workspace, it leaves as much free, and that
the
  output is correct and every reservation is released. It fails on main
("the final merge takes 168128 bytes beyond its workspace and leaves
109728
  free").
- `test_sort_exec_keeps_headroom_only_while_siblings_read_input` runs a
`SortExec` partition with no sibling, an idle sibling, a sibling that
has
finished reading, and a sibling whose stream was dropped unpolled; it
fails
if the condition is replaced by `true` or `false`, or if a dropped
partition
  stays counted.
- `test_sort_exec_rerun_without_reset_state_keeps_no_headroom` executes
a lone
partition three times without `reset_state`; it fails if the count
wraps.
- Existing tests pass unchanged: `datafusion-physical-plan` (2970),
`datafusion-execution`, the `memory_limit` tests in `core_integration`,
and
  the sort, merge, spilling and aggregate fuzz tests (`extended_tests`).
  In particular, the batch-size test under a tight limit
(`..._when_having_to_spill`) and the workspace retention tests from
apache#24740
  are unaffected.

<details>
<summary>Reproduction</summary>

`Cargo.toml`:

```toml
[package]
name = "sortrepro"
version = "0.1.0"
edition = "2024"

[dependencies]
datafusion = { git = "https://github.com/apache/datafusion", default-features = false }
tokio = { version = "1", features = ["rt-multi-thread"] }
futures = "0.3"
```

`src/main.rs`, run as `sortrepro <sort_spill_reservation_bytes> <runs>`:

```rust
//! A window over a spilling sort, 10 partitions, one FairSpillPool.
//! Usage: sortrepro <sort_spill_reservation_bytes> <runs>
use std::sync::Arc;

use datafusion::arrow::array::{Decimal256Array, Int32Array, RecordBatch};
use datafusion::arrow::datatypes::{DataType, Field, Schema, SchemaRef, i256};
use datafusion::catalog::streaming::StreamingTable;
use datafusion::common::ScalarValue;
use datafusion::functions_aggregate::expr_fn::max;
use datafusion::functions_aggregate::sum::sum_udaf;
use datafusion::logical_expr::expr::WindowFunction;
use datafusion::logical_expr::{
    Expr, ExprFunctionExt, WindowFrame, WindowFrameBound, WindowFrameUnits,
    WindowFunctionDefinition, col,
};
use datafusion::execution::memory_pool::FairSpillPool;
use datafusion::execution::runtime_env::RuntimeEnvBuilder;
use datafusion::execution::{SendableRecordBatchStream, TaskContext};
use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
use datafusion::physical_plan::streaming::PartitionStream;
use datafusion::prelude::{SessionConfig, SessionContext};

const ROWS: i32 = 6_666_815;
const BATCH: i32 = 8192;
const POOL: usize = 65_484_800; // 62.5 MB

/// (g, k, p): 1000 groups of increasing k, with a decimal to sum.
#[derive(Debug)]
struct Gen(SchemaRef);

impl PartitionStream for Gen {
    fn schema(&self) -> &SchemaRef {
        &self.0
    }

    fn execute(&self, _: Arc<TaskContext>) -> SendableRecordBatchStream {
        let schema = Arc::clone(&self.0);
        let s = Arc::clone(&schema);
        let batches = (0..(ROWS + BATCH - 1) / BATCH).map(move |b| {
            let (lo, hi) = (b * BATCH, ((b + 1) * BATCH).min(ROWS));
            let k: Vec<i32> = (lo..hi).map(|i| i * 3).collect();
            let g: Vec<i32> = k.iter().map(|k| k % 1000).collect();
            let p: Vec<i256> = k
                .iter()
                .map(|k| i256::from_i128((*k as i128 * 7919) % 10_000_000))
                .collect();
            let p = Decimal256Array::from(p).with_precision_and_scale(76, 2)?;
            Ok(RecordBatch::try_new(
                Arc::clone(&s),
                vec![
                    Arc::new(Int32Array::from(g)),
                    Arc::new(Int32Array::from(k)),
                    Arc::new(p),
                ],
            )?)
        });
        Box::pin(RecordBatchStreamAdapter::new(
            schema,
            futures::stream::iter(batches),
        ))
    }
}

async fn run(sort_spill_reservation_bytes: usize) -> datafusion::error::Result<()> {
    let mut config = SessionConfig::new()
        .with_target_partitions(10)
        .with_batch_size(BATCH as usize);
    config.options_mut().execution.sort_spill_reservation_bytes =
        sort_spill_reservation_bytes;
    let runtime = RuntimeEnvBuilder::new()
        .with_memory_pool(Arc::new(FairSpillPool::new(POOL)))
        .build_arc()?;
    let ctx = SessionContext::new_with_config_rt(config, runtime);
    let schema = Arc::new(Schema::new(vec![
        Field::new("g", DataType::Int32, false),
        Field::new("k", DataType::Int32, false),
        Field::new("p", DataType::Decimal256(76, 2), false),
    ]));
    let table = StreamingTable::try_new(
        Arc::clone(&schema),
        vec![Arc::new(Gen(schema))],
    )?;
    // SELECT max(s) FROM (SELECT sum(p) OVER (PARTITION BY g ORDER BY k
    //   ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW) AS s FROM t)
    let s = Expr::from(WindowFunction::new(
        WindowFunctionDefinition::AggregateUDF(sum_udaf()),
        vec![col("p")],
    ))
    .partition_by(vec![col("g")])
    .order_by(vec![col("k").sort(true, false)])
    .window_frame(WindowFrame::new_bounds(
        WindowFrameUnits::Rows,
        WindowFrameBound::Preceding(ScalarValue::UInt64(None)),
        WindowFrameBound::CurrentRow,
    ))
    .build()?
    .alias("s");
    ctx.read_table(Arc::new(table))?
        .window(vec![s])?
        .aggregate(vec![], vec![max(col("s"))])?
        .collect()
        .await?;
    Ok(())
}

fn main() {
    let args: Vec<String> = std::env::args().collect();
    let reservation: usize = args[1].parse().unwrap();
    let runs: usize = args[2].parse().unwrap();
    let rt = tokio::runtime::Builder::new_multi_thread()
        .worker_threads(10)
        .enable_all()
        .build()
        .unwrap();
    let mut ok = 0;
    for _ in 0..runs {
        match rt.block_on(run(reservation)) {
            Ok(()) => ok += 1,
            Err(e) => println!("{e}"),
        }
    }
    println!("sort_spill_reservation_bytes={reservation}: {ok} / {runs} succeeded");
}
```

</details>

## Are there any user-facing changes?

No API changes. A sort that spills to many runs may merge in more passes
when
the pool is shared, instead of failing other consumers of the pool.