fix: leave memory for others when a spilled sort merges its runs

[MisterRaindrop]

Which issue does this PR close?

• Part of ExternalSorter memory accounting: fixes missing from branch-55 and remaining gaps on main #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?

#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 fix: preserve external sort workspace across spilling #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.

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.

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").
• Existing tests pass unchanged: datafusion-physical-plan (2530),
  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 fix: preserve external sort workspace across spilling #24740
  are unaffected.

Reproduction

Cargo.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>:

//! 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");
}

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.

[codecov-commenter]

Codecov Report

❌ Patch coverage is 96.07843% with 2 lines in your changes missing coverage. Please review.
✅ Project coverage is 82.83%. Comparing base (2c568b4) to head (5c59c07).

Files with missing lines	Patch %	Lines
...usion/physical-plan/src/sorts/multi_level_merge.rs	95.45%	1 Missing and 1 partial ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main   #26264      +/-   ##
==========================================
- Coverage   82.83%   82.83%   -0.01%     
==========================================
  Files        1149     1149              
  Lines      452798   452847      +49     
  Branches   452798   452847      +49     
==========================================
+ Hits       375082   375121      +39     
- Misses      54893    54899       +6     
- Partials    22823    22827       +4     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.