From a78d43395c94b7392e13f1a5bdb185c2fd722284 Mon Sep 17 00:00:00 2001
From: MisterRaindrop <278811821@qq.com>
Date: Tue, 13 Oct 2026 19:39:01 +0800
Subject: [PATCH 1/4] fix: leave memory for others when a spilled sort merges
 its runs

The final merge of a spilled sort seated every run the memory pool
would admit, and held those buffers until its output was consumed.
Partitions that finished their input early could take the whole pool
this way, and a partition still sorting then failed with
ResourcesExhausted (#25804, finding 9).

Seat runs beyond the minimum two only while the pool can admit what
the merge takes beyond its own workspace plus as much again, and
return the extra before merging, as #25383 does for aggregate replay.
The minimum merge is admitted as before, so a sort under tight memory
keeps its batch sizes and does not re-spill runs to leave headroom.
---
 .../src/sorts/multi_level_merge.rs            | 37 +++++++++-
 datafusion/physical-plan/src/sorts/sort.rs    |  5 ++
 .../src/sorts/sort/spill_tests.rs             | 74 ++++++++++++++++++-
 .../src/sorts/streaming_merge.rs              | 12 +++
 4 files changed, 124 insertions(+), 4 deletions(-)

diff --git a/datafusion/physical-plan/src/sorts/multi_level_merge.rs b/datafusion/physical-plan/src/sorts/multi_level_merge.rs
index 28a8000265c7b..643bd9b580d47 100644
--- a/datafusion/physical-plan/src/sorts/multi_level_merge.rs
+++ b/datafusion/physical-plan/src/sorts/multi_level_merge.rs
@@ -28,7 +28,7 @@ use std::task::{Context, Poll};
 
 use arrow::datatypes::SchemaRef;
 use datafusion_common::{Result, internal_err, resources_err};
-use datafusion_execution::memory_pool::{MemoryReservation, MergeMemoryPool};
+use datafusion_execution::memory_pool::{MemoryPool, MemoryReservation, MergeMemoryPool};
 
 use crate::sorts::builder::try_grow_reservation_to_at_least;
 use crate::sorts::sort::get_reserved_bytes_for_record_batch_size;
@@ -159,6 +159,10 @@ pub(crate) struct MultiLevelMergeBuilder {
     merge_pool: Option<Arc<MergeMemoryPool>>,
     /// Leave memory for the aggregate consuming the merged spill rows.
     reserve_replay_headroom: bool,
+    /// Seat runs beyond the minimum two only while the pool can admit their
+    /// buffers plus equal headroom, so a merge does not take memory that the
+    /// operator consuming it and other partitions sharing the pool need.
+    reserve_fan_in_headroom: bool,
     /// Avoid merging more runs than needed before the final replay pass.
     size_intermediate_merges: bool,
     fetch: Option<usize>,
@@ -201,6 +205,7 @@ impl MultiLevelMergeBuilder {
             reservation,
             merge_pool: None,
             reserve_replay_headroom: false,
+            reserve_fan_in_headroom: false,
             size_intermediate_merges: false,
             enable_round_robin_tie_breaker,
             fetch,
@@ -219,6 +224,13 @@ impl MultiLevelMergeBuilder {
         self
     }
 
+    /// Leave headroom when seating runs beyond the minimum two. The minimum
+    /// merge is admitted against the pool alone, as without headroom.
+    pub(super) fn with_fan_in_headroom(mut self, reserve: bool) -> Self {
+        self.reserve_fan_in_headroom = reserve;
+        self
+    }
+
     pub(super) fn with_intermediate_merge_sizing(
         mut self,
         min_spill_batch_rows: Option<usize>,
@@ -586,6 +598,12 @@ impl MultiLevelMergeBuilder {
         // allocation, preventing starvation under memory pressure.
         let mut total_needed: usize = 0;
         let mut accepted_memory: usize = 0;
+        // What this merge already holds from the pool: its retained workspace
+        // when it has one. Seating runs within it takes nothing from others.
+        let owned = self
+            .merge_pool
+            .as_ref()
+            .map_or_else(|| reservation.size(), |pool| pool.reserved());
 
         for (spill, _) in &self.sorted_spill_files {
             if number_of_spills_to_read_for_current_phase >= max_spill_files
@@ -611,7 +629,20 @@ impl MultiLevelMergeBuilder {
                 && number_of_spills_to_read_for_current_phase
                     < minimum_number_of_required_streams;
             let check_headroom = self.reserve_replay_headroom && !skip_headroom;
-            let admission = if check_headroom {
+            let check_fan_in_headroom = !check_headroom
+                && self.reserve_fan_in_headroom
+                && number_of_spills_to_read_for_current_phase
+                    >= minimum_number_of_required_streams;
+            let admission = if check_fan_in_headroom {
+                // Ask the pool for what this merge takes beyond its own
+                // workspace, plus as much again, then return the spare bytes.
+                match total_needed.checked_add(total_needed.saturating_sub(owned)) {
+                    Some(with_headroom) => {
+                        try_grow_reservation_to_at_least(reservation, with_headroom)
+                    }
+                    None => resources_err!("Spill merge headroom exceeds usize::MAX"),
+                }
+            } else if check_headroom {
                 // Ask the pool for merge buffers plus equal replay space, then
                 // return the spare bytes before exposing the merge stream.
                 match total_needed.checked_mul(2) {
@@ -679,7 +710,7 @@ impl MultiLevelMergeBuilder {
             }
         }
 
-        if self.reserve_replay_headroom {
+        if self.reserve_replay_headroom || self.reserve_fan_in_headroom {
             // `total_needed` may include a rejected candidate. Keep only the
             // buffers that were admitted, releasing temporary replay headroom.
             reservation.shrink(reservation.size() - accepted_memory);
diff --git a/datafusion/physical-plan/src/sorts/sort.rs b/datafusion/physical-plan/src/sorts/sort.rs
index 7c34545b77337..909e70fc65679 100644
--- a/datafusion/physical-plan/src/sorts/sort.rs
+++ b/datafusion/physical-plan/src/sorts/sort.rs
@@ -383,6 +383,11 @@ impl ExternalSorter {
                 .with_fetch(None)
                 .with_reservation(self.merge_reservation.take())
                 .with_merge_pool(Arc::clone(&self.merge_pool))
+                // The operator consuming the sorted output, and the other
+                // partitions of this sort still buffering input, share the
+                // pool with this merge. Seating every run the pool admits
+                // would leave them nothing.
+                .with_fan_in_headroom()
                 .build()
         } else {
             // Final output needs no reserve for future spills. Return unused
diff --git a/datafusion/physical-plan/src/sorts/sort/spill_tests.rs b/datafusion/physical-plan/src/sorts/sort/spill_tests.rs
index f41ca4507c25a..21242e3bafef8 100644
--- a/datafusion/physical-plan/src/sorts/sort/spill_tests.rs
+++ b/datafusion/physical-plan/src/sorts/sort/spill_tests.rs
@@ -16,7 +16,8 @@
 // under the License.
 
 use super::{
-    ExternalSorter, get_reserved_bytes_for_record_batch, sort_batch, sort_batch_chunked,
+    ExternalSorter, get_reserved_bytes_for_record_batch,
+    get_reserved_bytes_for_record_batch_size, sort_batch, sort_batch_chunked,
 };
 use crate::metrics::ExecutionPlanMetricsSet;
 use crate::spill::get_record_batch_memory_size;
@@ -1045,3 +1046,74 @@ async fn test_single_batch_spill_returns_live_workspace_loan_on_drop() -> Result
     assert_released(&pool, &runtime).await;
     Ok(())
 }
+
+/// The final merge of a spilled sort shares the pool with the operator that
+/// consumes its output and with sibling partitions still sorting their input.
+/// Seating every run the pool admits leaves them nothing, so whatever the
+/// merge takes beyond its own workspace, it must leave as much free.
+#[tokio::test]
+async fn test_final_spilled_merge_leaves_headroom() -> Result<()> {
+    let fixture = fixture()?;
+    let schema = fixture.parent.schema();
+    let capacity = 16 * fixture.capacity;
+    let pool = AdjustablePool::new(capacity);
+    // Spill under the fixture's limit, then merge with room for every run.
+    pool.set_limit(fixture.capacity);
+    let (mut sorter, runtime) = sorter(&fixture, Arc::clone(&pool))?;
+    for batch in &fixture.batches {
+        sorter.insert_batch(batch.clone()).await?;
+    }
+    if !sorter.in_mem_batches.is_empty() {
+        sorter.sort_and_spill_in_mem_batches().await?;
+    }
+    // What the merge reserves to read each run with two buffers.
+    let needs: Vec<usize> = sorter
+        .finished_spill_files
+        .iter()
+        .map(|spill| {
+            let size = spill.max_record_batch_memory;
+            2 * get_reserved_bytes_for_record_batch_size(size, size)
+        })
+        .collect();
+    assert!(needs.len() >= 3, "need three runs, got {}", needs.len());
+    let all: usize = needs.iter().sum();
+    let smallest = *needs.iter().min().unwrap();
+    pool.set_limit(capacity);
+
+    let mut stream = sorter.sort().await?;
+    drop(sorter);
+    // The merge already holds its workspace. Leave it room to seat every
+    // run, but not to seat them all and leave as much free.
+    let workspace = pool.reserved();
+    let free = (2 * all)
+        .checked_sub(2 * workspace + smallest)
+        .expect("the runs must need more than the workspace");
+    assert!(all <= free + workspace);
+    let contender =
+        MemoryConsumer::new("sibling partition").register(&runtime.memory_pool);
+    contender.try_grow(capacity - workspace - free)?;
+
+    let first = stream
+        .try_next()
+        .await?
+        .expect("the merge must produce output");
+    let taken = (pool.reserved() - contender.size()).saturating_sub(workspace);
+    assert!(
+        capacity - pool.reserved() >= taken,
+        "the final merge takes {taken} bytes beyond its workspace and leaves {} free",
+        capacity - pool.reserved()
+    );
+    drop(contender);
+
+    let mut output = vec![first];
+    while let Some(batch) = stream.try_next().await? {
+        output.push(batch);
+    }
+    assert_eq!(
+        concat_batches(&schema, &output)?,
+        sort_batch(&fixture.parent, &fixture.ordering, None)?
+    );
+    drop(stream);
+    assert_released(&pool, &runtime).await;
+    Ok(())
+}
diff --git a/datafusion/physical-plan/src/sorts/streaming_merge.rs b/datafusion/physical-plan/src/sorts/streaming_merge.rs
index 249b20cf49f38..d90ca86c44be8 100644
--- a/datafusion/physical-plan/src/sorts/streaming_merge.rs
+++ b/datafusion/physical-plan/src/sorts/streaming_merge.rs
@@ -98,6 +98,9 @@ pub struct StreamingMergeBuilder<'a> {
     merge_pool: Option<Arc<MergeMemoryPool>>,
     /// Leave memory for the aggregate consuming the merged spill rows.
     reserve_replay_headroom: bool,
+    /// Leave memory for others sharing the pool when seating more than the
+    /// minimum two runs.
+    reserve_fan_in_headroom: bool,
     min_spill_batch_rows: Option<usize>,
     enable_round_robin_tie_breaker: bool,
 }
@@ -171,6 +174,13 @@ impl<'a> StreamingMergeBuilder<'a> {
         self
     }
 
+    /// Seat spill runs beyond the minimum two only while the pool can admit
+    /// their buffers plus equal headroom. Release the headroom before merging.
+    pub(crate) fn with_fan_in_headroom(mut self) -> Self {
+        self.reserve_fan_in_headroom = true;
+        self
+    }
+
     /// Limit intermediate rewrites when the retained buffer reservation can
     /// also admit the final replay pass with its headroom.
     /// `min_spill_batch_rows` is the minimum largest-batch row count across input
@@ -218,6 +228,7 @@ impl<'a> StreamingMergeBuilder<'a> {
             reservation,
             merge_pool,
             reserve_replay_headroom,
+            reserve_fan_in_headroom,
             min_spill_batch_rows,
             fetch,
             expressions,
@@ -261,6 +272,7 @@ impl<'a> StreamingMergeBuilder<'a> {
             )
             .with_merge_pool(merge_pool)
             .with_replay_headroom(reserve_replay_headroom)
+            .with_fan_in_headroom(reserve_fan_in_headroom)
             .with_intermediate_merge_sizing(min_spill_batch_rows)
             .create_spillable_merge_stream());
         }

From 5c59c075f2b158f12be02d80380887855dba4653 Mon Sep 17 00:00:00 2001
From: MisterRaindrop <278811821@qq.com>
Date: Wed, 14 Oct 2026 09:12:34 +0800
Subject: [PATCH 2/4] test: cover fan-in headroom when seating spill runs

Check at the builder that runs beyond the minimum two are seated only
with equal headroom, and that a headroom request beyond usize::MAX
stops seating runs instead of failing the merge.
---
 .../src/sorts/multi_level_merge.rs            | 67 ++++++++++++++++++-
 1 file changed, 66 insertions(+), 1 deletion(-)

diff --git a/datafusion/physical-plan/src/sorts/multi_level_merge.rs b/datafusion/physical-plan/src/sorts/multi_level_merge.rs
index 643bd9b580d47..eb5c061078668 100644
--- a/datafusion/physical-plan/src/sorts/multi_level_merge.rs
+++ b/datafusion/physical-plan/src/sorts/multi_level_merge.rs
@@ -1010,7 +1010,7 @@ mod tests {
     use arrow::compute::concat_batches;
     use arrow::datatypes::{DataType, Field, Int64Type, Schema};
     use datafusion_execution::memory_pool::{
-        GreedyMemoryPool, MemoryConsumer, MemoryPool,
+        GreedyMemoryPool, MemoryConsumer, MemoryPool, UnboundedMemoryPool,
     };
     use datafusion_execution::runtime_env::{RuntimeEnv, RuntimeEnvBuilder};
     use datafusion_physical_expr::expressions::{Column, col};
@@ -1507,6 +1507,71 @@ mod tests {
         Ok(())
     }
 
+    #[rstest::rstest]
+    #[case::without_headroom(false, 4, 8)]
+    #[case::with_headroom(true, 3, 6)]
+    fn fan_in_headroom_limits_runs_beyond_the_minimum(
+        #[case] headroom: bool,
+        #[case] seated: usize,
+        #[case] retained_batches: usize,
+    ) -> Result<()> {
+        let env = Arc::new(RuntimeEnv::default());
+        let schema = test_schema();
+        let spill_manager = build_spill_manager(&env, &schema);
+        let spills = (0..4)
+            .map(|value| make_sorted_spill_file(&spill_manager, &schema, vec![value]))
+            .collect::<Vec<_>>();
+        let batch_memory = spills[0].max_record_batch_memory;
+        let pool: Arc<dyn MemoryPool> =
+            Arc::new(GreedyMemoryPool::new(13 * batch_memory));
+        let mut builder = build_merge_builder(spill_manager, schema, spills, &pool, 1)
+            .with_fan_in_headroom(headroom);
+        let mut reservation = builder.reservation.new_empty();
+        let SpillFilesToMerge::Ready(spills, buffer_len) =
+            builder.get_sorted_spill_files_to_merge(1, 2, &mut reservation, false)?
+        else {
+            panic!("the minimum merge should fit the pool");
+        };
+        assert_eq!(buffer_len, 1);
+        // Each run takes two batches. With headroom, the third run needs 12
+        // batches and is seated; the fourth needs 16 and is not.
+        assert_eq!(spills.len(), seated);
+        assert_eq!(reservation.size(), retained_batches * batch_memory);
+        drop((reservation, spills, builder));
+        assert_eq!(pool.reserved(), 0);
+        Ok(())
+    }
+
+    #[test]
+    fn fan_in_headroom_beyond_usize_stops_seating_runs() -> Result<()> {
+        let env = Arc::new(RuntimeEnv::default());
+        let schema = test_schema();
+        let spill_manager = build_spill_manager(&env, &schema);
+        // Three runs whose buffers fit in a usize, but whose headroom for the
+        // third does not.
+        let batch_memory = usize::MAX / 10;
+        let spills = (0..3)
+            .map(|value| SortedSpillFile {
+                max_record_batch_memory: batch_memory,
+                ..make_sorted_spill_file(&spill_manager, &schema, vec![value])
+            })
+            .collect::<Vec<_>>();
+        let pool: Arc<dyn MemoryPool> = Arc::new(UnboundedMemoryPool::default());
+        let mut builder = build_merge_builder(spill_manager, schema, spills, &pool, 1)
+            .with_fan_in_headroom(true);
+        let mut reservation = builder.reservation.new_empty();
+        let SpillFilesToMerge::Ready(spills, _) =
+            builder.get_sorted_spill_files_to_merge(1, 2, &mut reservation, false)?
+        else {
+            panic!("the minimum merge should fit an unbounded pool");
+        };
+        assert_eq!(spills.len(), 2);
+        assert_eq!(reservation.size(), 4 * batch_memory);
+        drop((reservation, spills, builder));
+        assert_eq!(pool.reserved(), 0);
+        Ok(())
+    }
+
     #[tokio::test]
     async fn replay_headroom_still_enforces_the_actual_pool() -> Result<()> {
         let env = Arc::new(RuntimeEnv::default());

From c7c5b196b1f564a4521a5107e19ef722bff4ae1b Mon Sep 17 00:00:00 2001
From: MisterRaindrop <278811821@qq.com>
Date: Thu, 5 Nov 2026 21:56:24 +0800
Subject: [PATCH 3/4] fix: keep spill merge headroom only when the pool is
 shared

The fan-in headroom made a lone sort partition that has the pool to
itself rewrite runs in an extra merge pass, for no one: with 60 x 8192
rows of Int64 plus a 120-byte string, 1 MB sort_spill_reservation_bytes
and a 16 MB GreedyMemoryPool, it wrote 1,589,248 rows to spill instead
of 983,040.

Keep the headroom only while the sort has other partitions sharing the
pool, or another consumer holds memory in it. Sibling partitions count
even when they hold nothing at the moment the merge seats its runs.
The operator consuming the output of a lone partition is not counted:
it reserves nothing before its first batch, and keeping headroom for it
would cost the extra pass whether or not it ever needs the memory.

With one partition, the case above writes the same rows as without
headroom in all four configurations measured. With four partitions in
a 32 MB pool and 1 MB reservations, the sort fails with "Not enough
memory to continue external sort" without headroom and succeeds with
it.

Also cover the minimum two runs under a budget below their headroom,
and seating against a retained merge workspace on both sides of
2 * (T - W).
---
 .../src/sorts/multi_level_merge.rs            | 185 ++++++++++++++++--
 datafusion/physical-plan/src/sorts/sort.rs    |  33 +++-
 .../src/sorts/sort/spill_tests.rs             |  31 ++-
 .../src/sorts/streaming_merge.rs              |  30 ++-
 4 files changed, 235 insertions(+), 44 deletions(-)

diff --git a/datafusion/physical-plan/src/sorts/multi_level_merge.rs b/datafusion/physical-plan/src/sorts/multi_level_merge.rs
index eb5c061078668..f41960e1b8b5b 100644
--- a/datafusion/physical-plan/src/sorts/multi_level_merge.rs
+++ b/datafusion/physical-plan/src/sorts/multi_level_merge.rs
@@ -41,6 +41,24 @@ use datafusion_physical_expr_common::sort_expr::LexOrdering;
 use futures::TryStreamExt;
 use futures::{Stream, StreamExt};
 
+/// Leave memory for others sharing the execution pool when a merge seats
+/// more than the minimum two runs.
+#[derive(Debug, Clone)]
+pub(crate) struct FanInHeadroom {
+    /// The execution pool: the parent of the merge pool when there is one.
+    pub(crate) pool: Arc<dyn MemoryPool>,
+    /// Other partitions share `pool`, even if they hold nothing right now.
+    pub(crate) siblings: bool,
+}
+
+impl FanInHeadroom {
+    /// Whether another consumer may need the pool: sibling partitions exist,
+    /// or the pool holds more than the `owned` bytes this merge holds in it.
+    fn is_contended(&self, owned: usize) -> bool {
+        self.siblings || self.pool.reserved() > owned
+    }
+}
+
 /// Merges a stream of sorted cursors and record batches into a single sorted stream
 ///
 /// This is a wrapper around [`SortPreservingMergeStream`](crate::sorts::merge::SortPreservingMergeStream)
@@ -159,10 +177,11 @@ pub(crate) struct MultiLevelMergeBuilder {
     merge_pool: Option<Arc<MergeMemoryPool>>,
     /// Leave memory for the aggregate consuming the merged spill rows.
     reserve_replay_headroom: bool,
-    /// Seat runs beyond the minimum two only while the pool can admit their
-    /// buffers plus equal headroom, so a merge does not take memory that the
-    /// operator consuming it and other partitions sharing the pool need.
-    reserve_fan_in_headroom: bool,
+    /// When set and the pool is shared, seat runs beyond the minimum two only
+    /// if the pool can admit their buffers plus as much again as they exceed
+    /// the retained workspace, so a merge does not take the memory that other
+    /// partitions sharing the pool need.
+    fan_in_headroom: Option<FanInHeadroom>,
     /// Avoid merging more runs than needed before the final replay pass.
     size_intermediate_merges: bool,
     fetch: Option<usize>,
@@ -205,7 +224,7 @@ impl MultiLevelMergeBuilder {
             reservation,
             merge_pool: None,
             reserve_replay_headroom: false,
-            reserve_fan_in_headroom: false,
+            fan_in_headroom: None,
             size_intermediate_merges: false,
             enable_round_robin_tie_breaker,
             fetch,
@@ -224,10 +243,14 @@ impl MultiLevelMergeBuilder {
         self
     }
 
-    /// Leave headroom when seating runs beyond the minimum two. The minimum
-    /// merge is admitted against the pool alone, as without headroom.
-    pub(super) fn with_fan_in_headroom(mut self, reserve: bool) -> Self {
-        self.reserve_fan_in_headroom = reserve;
+    /// Leave headroom when seating runs beyond the minimum two while the pool
+    /// is shared. The minimum merge is admitted against the pool alone, as
+    /// without headroom.
+    pub(super) fn with_fan_in_headroom(
+        mut self,
+        headroom: Option<FanInHeadroom>,
+    ) -> Self {
+        self.fan_in_headroom = headroom;
         self
     }
 
@@ -604,6 +627,14 @@ impl MultiLevelMergeBuilder {
             .merge_pool
             .as_ref()
             .map_or_else(|| reservation.size(), |pool| pool.reserved());
+        // Headroom is for others sharing the pool. With no sibling partitions
+        // and nobody else holding memory in it, seat every run that fits, as
+        // without headroom, rather than rewrite runs in an extra merge pass
+        // for no one.
+        let fan_in_contended = self
+            .fan_in_headroom
+            .as_ref()
+            .is_some_and(|headroom| headroom.is_contended(owned));
 
         for (spill, _) in &self.sorted_spill_files {
             if number_of_spills_to_read_for_current_phase >= max_spill_files
@@ -630,7 +661,7 @@ impl MultiLevelMergeBuilder {
                     < minimum_number_of_required_streams;
             let check_headroom = self.reserve_replay_headroom && !skip_headroom;
             let check_fan_in_headroom = !check_headroom
-                && self.reserve_fan_in_headroom
+                && fan_in_contended
                 && number_of_spills_to_read_for_current_phase
                     >= minimum_number_of_required_streams;
             let admission = if check_fan_in_headroom {
@@ -710,7 +741,7 @@ impl MultiLevelMergeBuilder {
             }
         }
 
-        if self.reserve_replay_headroom || self.reserve_fan_in_headroom {
+        if self.reserve_replay_headroom || fan_in_contended {
             // `total_needed` may include a rejected candidate. Keep only the
             // buffers that were admitted, releasing temporary replay headroom.
             reservation.shrink(reservation.size() - accepted_memory);
@@ -1507,11 +1538,31 @@ mod tests {
         Ok(())
     }
 
+    /// Registers a consumer holding `size` bytes of `pool`, so a merge sharing
+    /// it sees another user and leaves headroom.
+    fn hold(pool: &Arc<dyn MemoryPool>, size: usize) -> Result<MemoryReservation> {
+        let other = MemoryConsumer::new("other consumer").register(pool);
+        other.try_grow(size)?;
+        Ok(other)
+    }
+
+    /// Headroom for others sharing `pool`, with or without sibling partitions.
+    fn headroom(pool: &Arc<dyn MemoryPool>, siblings: bool) -> Option<FanInHeadroom> {
+        Some(FanInHeadroom {
+            pool: Arc::clone(pool),
+            siblings,
+        })
+    }
+
     #[rstest::rstest]
-    #[case::without_headroom(false, 4, 8)]
-    #[case::with_headroom(true, 3, 6)]
+    #[case::without_headroom(false, true, false, 4, 8)]
+    #[case::headroom_without_others(true, false, false, 4, 8)]
+    #[case::headroom_with_others(true, true, false, 3, 6)]
+    #[case::headroom_with_idle_siblings(true, false, true, 3, 6)]
     fn fan_in_headroom_limits_runs_beyond_the_minimum(
-        #[case] headroom: bool,
+        #[case] with_headroom: bool,
+        #[case] others: bool,
+        #[case] siblings: bool,
         #[case] seated: usize,
         #[case] retained_batches: usize,
     ) -> Result<()> {
@@ -1524,8 +1575,11 @@ mod tests {
         let batch_memory = spills[0].max_record_batch_memory;
         let pool: Arc<dyn MemoryPool> =
             Arc::new(GreedyMemoryPool::new(13 * batch_memory));
+        let other = others.then(|| hold(&pool, 1)).transpose()?;
         let mut builder = build_merge_builder(spill_manager, schema, spills, &pool, 1)
-            .with_fan_in_headroom(headroom);
+            .with_fan_in_headroom(
+                with_headroom.then(|| headroom(&pool, siblings).unwrap()),
+            );
         let mut reservation = builder.reservation.new_empty();
         let SpillFilesToMerge::Ready(spills, buffer_len) =
             builder.get_sorted_spill_files_to_merge(1, 2, &mut reservation, false)?
@@ -1534,14 +1588,106 @@ mod tests {
         };
         assert_eq!(buffer_len, 1);
         // Each run takes two batches. With headroom, the third run needs 12
-        // batches and is seated; the fourth needs 16 and is not.
+        // batches and is seated; the fourth needs 16 and is not. When nobody
+        // else holds memory and there are no sibling partitions, all four are
+        // seated, as without headroom.
         assert_eq!(spills.len(), seated);
         assert_eq!(reservation.size(), retained_batches * batch_memory);
-        drop((reservation, spills, builder));
+        drop((reservation, spills, builder, other));
         assert_eq!(pool.reserved(), 0);
         Ok(())
     }
 
+    /// The minimum two runs are admitted against the pool alone, even when
+    /// their buffers plus headroom would not fit.
+    #[test]
+    fn fan_in_headroom_admits_the_minimum_merge_without_headroom() -> Result<()> {
+        let env = Arc::new(RuntimeEnv::default());
+        let schema = test_schema();
+        let spill_manager = build_spill_manager(&env, &schema);
+        let spills = (0..3)
+            .map(|value| make_sorted_spill_file(&spill_manager, &schema, vec![value]))
+            .collect::<Vec<_>>();
+        let batch_memory = spills[0].max_record_batch_memory;
+        let pool: Arc<dyn MemoryPool> = Arc::new(GreedyMemoryPool::new(5 * batch_memory));
+        let other = hold(&pool, 1)?;
+        let mut builder = build_merge_builder(spill_manager, schema, spills, &pool, 1)
+            .with_fan_in_headroom(headroom(&pool, false));
+        let mut reservation = builder.reservation.new_empty();
+        let SpillFilesToMerge::Ready(spills, _) =
+            builder.get_sorted_spill_files_to_merge(1, 2, &mut reservation, false)?
+        else {
+            panic!("the minimum merge should fit the pool");
+        };
+        // Two runs take four batches; with headroom they would take eight. The
+        // third would take twelve.
+        assert_eq!(spills.len(), 2);
+        assert_eq!(reservation.size(), 4 * batch_memory);
+        drop((reservation, spills, builder, other));
+        assert_eq!(pool.reserved(), 0);
+        Ok(())
+    }
+
+    /// With a retained workspace `W`, seating the third run, for buffers `T`
+    /// in all, asks the parent for `2 * (T - W)` more, and nothing when
+    /// `T <= W`. Afterwards the child holds `T` and the parent is charged
+    /// `max(W, T)`.
+    #[rstest::rstest]
+    #[case::runs_within_workspace(6, 0, 0, 3)]
+    #[case::runs_below_workspace(8, 0, 0, 3)]
+    #[case::exactly_the_headroom(2, 8, 0, 3)]
+    #[case::just_below_the_headroom(2, 8, 1, 2)]
+    fn fan_in_headroom_counts_the_retained_workspace(
+        #[case] workspace_batches: usize,
+        #[case] free_batches: usize,
+        #[case] short_by_bytes: usize,
+        #[case] seated: usize,
+    ) -> Result<()> {
+        let env = Arc::new(RuntimeEnv::default());
+        let schema = test_schema();
+        let spill_manager = build_spill_manager(&env, &schema);
+        let spills = (0..3)
+            .map(|value| make_sorted_spill_file(&spill_manager, &schema, vec![value]))
+            .collect::<Vec<_>>();
+        let batch_memory = spills[0].max_record_batch_memory;
+        let per_run =
+            get_reserved_bytes_for_record_batch_size(batch_memory, batch_memory);
+        assert_eq!(per_run, 2 * batch_memory);
+        let workspace = workspace_batches * batch_memory;
+        let free = free_batches * batch_memory - short_by_bytes;
+        let capacity = 32 * batch_memory;
+        let parent: Arc<dyn MemoryPool> = Arc::new(GreedyMemoryPool::new(capacity));
+        let merge_pool = Arc::new(MergeMemoryPool::new(
+            Arc::clone(&parent),
+            MemoryConsumer::new("merge workspace"),
+        ));
+        let pool: Arc<dyn MemoryPool> = Arc::clone(&merge_pool) as Arc<dyn MemoryPool>;
+        let mut builder = build_merge_builder(spill_manager, schema, spills, &pool, 1)
+            .with_merge_pool(Some(Arc::clone(&merge_pool)))
+            .with_fan_in_headroom(headroom(&parent, false));
+        builder.reservation.try_grow(workspace)?;
+        merge_pool.retain(workspace);
+        builder.reservation.free();
+        assert_eq!(merge_pool.reserved(), workspace);
+        let other = hold(&parent, capacity - workspace - free)?;
+
+        let mut reservation = builder.reservation.new_empty();
+        let SpillFilesToMerge::Ready(spills, _) =
+            builder.get_sorted_spill_files_to_merge(1, 2, &mut reservation, false)?
+        else {
+            panic!("the minimum merge should fit the pool");
+        };
+        assert_eq!(spills.len(), seated);
+        let total = seated * per_run;
+        assert_eq!(reservation.size(), total);
+        assert_eq!(merge_pool.reserved(), workspace.max(total));
+        assert_eq!(parent.reserved(), other.size() + workspace.max(total));
+        drop((reservation, spills, builder, other));
+        merge_pool.release_unused();
+        assert_eq!(parent.reserved(), 0);
+        Ok(())
+    }
+
     #[test]
     fn fan_in_headroom_beyond_usize_stops_seating_runs() -> Result<()> {
         let env = Arc::new(RuntimeEnv::default());
@@ -1557,8 +1703,9 @@ mod tests {
             })
             .collect::<Vec<_>>();
         let pool: Arc<dyn MemoryPool> = Arc::new(UnboundedMemoryPool::default());
+        let other = hold(&pool, 1)?;
         let mut builder = build_merge_builder(spill_manager, schema, spills, &pool, 1)
-            .with_fan_in_headroom(true);
+            .with_fan_in_headroom(headroom(&pool, false));
         let mut reservation = builder.reservation.new_empty();
         let SpillFilesToMerge::Ready(spills, _) =
             builder.get_sorted_spill_files_to_merge(1, 2, &mut reservation, false)?
@@ -1567,7 +1714,7 @@ mod tests {
         };
         assert_eq!(spills.len(), 2);
         assert_eq!(reservation.size(), 4 * batch_memory);
-        drop((reservation, spills, builder));
+        drop((reservation, spills, builder, other));
         assert_eq!(pool.reserved(), 0);
         Ok(())
     }
diff --git a/datafusion/physical-plan/src/sorts/sort.rs b/datafusion/physical-plan/src/sorts/sort.rs
index 909e70fc65679..82c57896b2d54 100644
--- a/datafusion/physical-plan/src/sorts/sort.rs
+++ b/datafusion/physical-plan/src/sorts/sort.rs
@@ -272,6 +272,9 @@ struct ExternalSorter {
     /// How much memory to reserve for performing in-memory sort/merges
     /// prior to spilling.
     sort_spill_reservation_bytes: usize,
+    /// Other partitions of this sort run alongside this one and share its
+    /// memory pool.
+    has_sibling_partitions: bool,
 }
 
 impl ExternalSorter {
@@ -325,9 +328,16 @@ impl ExternalSorter {
             batch_size,
             sort_spill_reservation_bytes,
             sort_in_place_threshold_bytes,
+            has_sibling_partitions: false,
         })
     }
 
+    /// Record whether other partitions of this sort share its memory pool.
+    fn with_sibling_partitions(mut self, has_sibling_partitions: bool) -> Self {
+        self.has_sibling_partitions = has_sibling_partitions;
+        self
+    }
+
     /// Appends an unsorted [`RecordBatch`] to `in_mem_batches`
     ///
     /// Updates memory usage metrics, and possibly triggers spilling to disk
@@ -383,11 +393,19 @@ impl ExternalSorter {
                 .with_fetch(None)
                 .with_reservation(self.merge_reservation.take())
                 .with_merge_pool(Arc::clone(&self.merge_pool))
-                // The operator consuming the sorted output, and the other
-                // partitions of this sort still buffering input, share the
-                // pool with this merge. Seating every run the pool admits
-                // would leave them nothing.
-                .with_fan_in_headroom()
+                // Other partitions of this sort, still buffering input, share
+                // the pool with this merge, and may hold nothing at the moment
+                // it seats its runs. Seating every run the pool admits would
+                // leave them nothing, so keep headroom while they exist or
+                // anyone else holds memory in the pool. A lone partition with
+                // the pool to itself seats every run that fits: the operator
+                // consuming its output reserves nothing before its first
+                // batch, and keeping headroom for it would cost an extra merge
+                // pass whether or not it ever needs the memory.
+                .with_fan_in_headroom(
+                    Arc::clone(&self.runtime.memory_pool),
+                    self.has_sibling_partitions,
+                )
                 .build()
         } else {
             // Final output needs no reserve for future spills. Return unused
@@ -1522,7 +1540,10 @@ impl ExecutionPlan for SortExec {
                     context.session_config().spill_compression(),
                     &self.metrics_set,
                     context.runtime_env(),
-                )?;
+                )?
+                .with_sibling_partitions(
+                    self.input.output_partitioning().partition_count() > 1,
+                );
                 Ok(Box::pin(RecordBatchStreamAdapter::new(
                     self.schema(),
                     futures::stream::once(async move {
diff --git a/datafusion/physical-plan/src/sorts/sort/spill_tests.rs b/datafusion/physical-plan/src/sorts/sort/spill_tests.rs
index 21242e3bafef8..2b291043c8256 100644
--- a/datafusion/physical-plan/src/sorts/sort/spill_tests.rs
+++ b/datafusion/physical-plan/src/sorts/sort/spill_tests.rs
@@ -1047,19 +1047,24 @@ async fn test_single_batch_spill_returns_live_workspace_loan_on_drop() -> Result
     Ok(())
 }
 
-/// The final merge of a spilled sort shares the pool with the operator that
-/// consumes its output and with sibling partitions still sorting their input.
-/// Seating every run the pool admits leaves them nothing, so whatever the
-/// merge takes beyond its own workspace, it must leave as much free.
+/// The final merge of a spilled sort shares the pool with sibling partitions
+/// still sorting their input. Seating every run the pool admits leaves them
+/// nothing, so whatever the merge takes beyond its own workspace, it must leave
+/// as much free: while another consumer holds memory, and while sibling
+/// partitions exist even if they hold nothing at the moment.
+#[rstest::rstest]
+#[case::other_consumer_holds_memory(false)]
+#[case::idle_sibling_partitions(true)]
 #[tokio::test]
-async fn test_final_spilled_merge_leaves_headroom() -> Result<()> {
+async fn test_final_spilled_merge_leaves_headroom(#[case] siblings: bool) -> Result<()> {
     let fixture = fixture()?;
     let schema = fixture.parent.schema();
     let capacity = 16 * fixture.capacity;
     let pool = AdjustablePool::new(capacity);
     // Spill under the fixture's limit, then merge with room for every run.
     pool.set_limit(fixture.capacity);
-    let (mut sorter, runtime) = sorter(&fixture, Arc::clone(&pool))?;
+    let (sorter, runtime) = sorter(&fixture, Arc::clone(&pool))?;
+    let mut sorter = sorter.with_sibling_partitions(siblings);
     for batch in &fixture.batches {
         sorter.insert_batch(batch.clone()).await?;
     }
@@ -1089,9 +1094,17 @@ async fn test_final_spilled_merge_leaves_headroom() -> Result<()> {
         .checked_sub(2 * workspace + smallest)
         .expect("the runs must need more than the workspace");
     assert!(all <= free + workspace);
+    // Either another consumer holds the rest of the pool, or idle siblings
+    // hold nothing and the pool is only as large as the merge may use.
     let contender =
         MemoryConsumer::new("sibling partition").register(&runtime.memory_pool);
-    contender.try_grow(capacity - workspace - free)?;
+    let limit = if siblings {
+        pool.set_limit(workspace + free);
+        workspace + free
+    } else {
+        contender.try_grow(capacity - workspace - free)?;
+        capacity
+    };
 
     let first = stream
         .try_next()
@@ -1099,9 +1112,9 @@ async fn test_final_spilled_merge_leaves_headroom() -> Result<()> {
         .expect("the merge must produce output");
     let taken = (pool.reserved() - contender.size()).saturating_sub(workspace);
     assert!(
-        capacity - pool.reserved() >= taken,
+        limit - pool.reserved() >= taken,
         "the final merge takes {taken} bytes beyond its workspace and leaves {} free",
-        capacity - pool.reserved()
+        limit - pool.reserved()
     );
     drop(contender);
 
diff --git a/datafusion/physical-plan/src/sorts/streaming_merge.rs b/datafusion/physical-plan/src/sorts/streaming_merge.rs
index d90ca86c44be8..e45a7fb2b3812 100644
--- a/datafusion/physical-plan/src/sorts/streaming_merge.rs
+++ b/datafusion/physical-plan/src/sorts/streaming_merge.rs
@@ -19,7 +19,7 @@
 //! This is an order-preserving merge.
 
 use crate::metrics::BaselineMetrics;
-use crate::sorts::multi_level_merge::MultiLevelMergeBuilder;
+use crate::sorts::multi_level_merge::{FanInHeadroom, MultiLevelMergeBuilder};
 use crate::sorts::{
     merge::SortPreservingMergeStream,
     stream::{FieldCursorStream, RowCursorStream},
@@ -98,9 +98,9 @@ pub struct StreamingMergeBuilder<'a> {
     merge_pool: Option<Arc<MergeMemoryPool>>,
     /// Leave memory for the aggregate consuming the merged spill rows.
     reserve_replay_headroom: bool,
-    /// Leave memory for others sharing the pool when seating more than the
-    /// minimum two runs.
-    reserve_fan_in_headroom: bool,
+    /// Leave memory for others sharing the execution pool when seating more
+    /// than the minimum two runs.
+    fan_in_headroom: Option<FanInHeadroom>,
     min_spill_batch_rows: Option<usize>,
     enable_round_robin_tie_breaker: bool,
 }
@@ -174,10 +174,20 @@ impl<'a> StreamingMergeBuilder<'a> {
         self
     }
 
-    /// Seat spill runs beyond the minimum two only while the pool can admit
-    /// their buffers plus equal headroom. Release the headroom before merging.
-    pub(crate) fn with_fan_in_headroom(mut self) -> Self {
-        self.reserve_fan_in_headroom = true;
+    /// While `siblings` is true or another consumer holds memory in `pool`,
+    /// seat spill runs beyond the minimum two only if the pool can admit their
+    /// buffers plus as much again as they exceed the retained workspace.
+    /// Release the headroom before merging.
+    ///
+    /// `pool` is the execution pool: the parent of the merge pool when one is
+    /// set, not the merge pool itself. `siblings` says other partitions share
+    /// that pool even when they hold nothing at the moment.
+    pub(crate) fn with_fan_in_headroom(
+        mut self,
+        pool: Arc<dyn MemoryPool>,
+        siblings: bool,
+    ) -> Self {
+        self.fan_in_headroom = Some(FanInHeadroom { pool, siblings });
         self
     }
 
@@ -228,7 +238,7 @@ impl<'a> StreamingMergeBuilder<'a> {
             reservation,
             merge_pool,
             reserve_replay_headroom,
-            reserve_fan_in_headroom,
+            fan_in_headroom,
             min_spill_batch_rows,
             fetch,
             expressions,
@@ -272,7 +282,7 @@ impl<'a> StreamingMergeBuilder<'a> {
             )
             .with_merge_pool(merge_pool)
             .with_replay_headroom(reserve_replay_headroom)
-            .with_fan_in_headroom(reserve_fan_in_headroom)
+            .with_fan_in_headroom(fan_in_headroom)
             .with_intermediate_merge_sizing(min_spill_batch_rows)
             .create_spillable_merge_stream());
         }

From 24297ff6c44d7fbaff7120631f98800a1b1284f9 Mon Sep 17 00:00:00 2001
From: MisterRaindrop <278811821@qq.com>
Date: Fri, 6 Nov 2026 23:03:42 +0800
Subject: [PATCH 4/4] fix: keep sort merge headroom only while sibling
 partitions read input

A partition that finishes last kept fan-in headroom for siblings that had
already finished, and rewrote rows in an extra merge pass. Count the
partitions of a SortExec still reading their input instead, and keep
headroom only while another one is. Merges already in progress hold pool
memory, so the existing pool check still covers them.

Also add a SortExec-level test for the condition, and return the
headroom test helper's struct directly.
---
 .../src/sorts/multi_level_merge.rs            | 16 ++--
 datafusion/physical-plan/src/sorts/sort.rs    | 62 +++++++++----
 .../src/sorts/sort/spill_tests.rs             | 87 ++++++++++++++++++-
 3 files changed, 134 insertions(+), 31 deletions(-)

diff --git a/datafusion/physical-plan/src/sorts/multi_level_merge.rs b/datafusion/physical-plan/src/sorts/multi_level_merge.rs
index a2d1b30ee05f7..39d06121bf111 100644
--- a/datafusion/physical-plan/src/sorts/multi_level_merge.rs
+++ b/datafusion/physical-plan/src/sorts/multi_level_merge.rs
@@ -1534,11 +1534,11 @@ mod tests {
     }
 
     /// Headroom for others sharing `pool`, with or without sibling partitions.
-    fn headroom(pool: &Arc<dyn MemoryPool>, siblings: bool) -> Option<FanInHeadroom> {
-        Some(FanInHeadroom {
+    fn headroom(pool: &Arc<dyn MemoryPool>, siblings: bool) -> FanInHeadroom {
+        FanInHeadroom {
             pool: Arc::clone(pool),
             siblings,
-        })
+        }
     }
 
     #[rstest::rstest]
@@ -1564,9 +1564,7 @@ mod tests {
             Arc::new(GreedyMemoryPool::new(13 * batch_memory));
         let other = others.then(|| hold(&pool, 1)).transpose()?;
         let mut builder = build_merge_builder(spill_manager, schema, spills, &pool, 1)
-            .with_fan_in_headroom(
-                with_headroom.then(|| headroom(&pool, siblings).unwrap()),
-            );
+            .with_fan_in_headroom(with_headroom.then(|| headroom(&pool, siblings)));
         let mut reservation = builder.reservation.new_empty();
         let SpillFilesToMerge::Ready(spills, buffer_len) =
             builder.get_sorted_spill_files_to_merge(1, 2, &mut reservation, false)?
@@ -1599,7 +1597,7 @@ mod tests {
         let pool: Arc<dyn MemoryPool> = Arc::new(GreedyMemoryPool::new(5 * batch_memory));
         let other = hold(&pool, 1)?;
         let mut builder = build_merge_builder(spill_manager, schema, spills, &pool, 1)
-            .with_fan_in_headroom(headroom(&pool, false));
+            .with_fan_in_headroom(Some(headroom(&pool, false)));
         let mut reservation = builder.reservation.new_empty();
         let SpillFilesToMerge::Ready(spills, _) =
             builder.get_sorted_spill_files_to_merge(1, 2, &mut reservation, false)?
@@ -1651,7 +1649,7 @@ mod tests {
         let pool: Arc<dyn MemoryPool> = Arc::clone(&merge_pool) as Arc<dyn MemoryPool>;
         let mut builder = build_merge_builder(spill_manager, schema, spills, &pool, 1)
             .with_merge_pool(Some(Arc::clone(&merge_pool)))
-            .with_fan_in_headroom(headroom(&parent, false));
+            .with_fan_in_headroom(Some(headroom(&parent, false)));
         builder.reservation.try_grow(workspace)?;
         merge_pool.retain(workspace);
         builder.reservation.free();
@@ -1692,7 +1690,7 @@ mod tests {
         let pool: Arc<dyn MemoryPool> = Arc::new(UnboundedMemoryPool::default());
         let other = hold(&pool, 1)?;
         let mut builder = build_merge_builder(spill_manager, schema, spills, &pool, 1)
-            .with_fan_in_headroom(headroom(&pool, false));
+            .with_fan_in_headroom(Some(headroom(&pool, false)));
         let mut reservation = builder.reservation.new_empty();
         let SpillFilesToMerge::Ready(spills, _) =
             builder.get_sorted_spill_files_to_merge(1, 2, &mut reservation, false)?
diff --git a/datafusion/physical-plan/src/sorts/sort.rs b/datafusion/physical-plan/src/sorts/sort.rs
index 7ab2224865f30..b1fd5a05e5efe 100644
--- a/datafusion/physical-plan/src/sorts/sort.rs
+++ b/datafusion/physical-plan/src/sorts/sort.rs
@@ -9,6 +9,7 @@
 use std::fmt;
 use std::fmt::{Debug, Formatter};
 use std::future::Future;
+use std::sync::atomic::{AtomicUsize, Ordering};
 use std::sync::{Arc, OnceLock};
 
 use parking_lot::RwLock;
@@ -309,9 +310,9 @@ struct ExternalSorter {
     /// How much memory to reserve for performing in-memory sort/merges
     /// prior to spilling.
     sort_spill_reservation_bytes: usize,
-    /// Other partitions of this sort run alongside this one and share its
-    /// memory pool.
-    has_sibling_partitions: bool,
+    /// Partitions of the owning `SortExec`, this one included, still reading
+    /// their input. Each one will need memory from the pool this sort shares.
+    partitions_reading_input: Arc<AtomicUsize>,
     /// Shared by the owning SortExec's partitions. Initialized on the runtime
     /// that polls the sort, which may differ from the one calling `execute`.
     sort_concurrency: Arc<OnceLock<SortConcurrency>>,
@@ -375,17 +376,25 @@ impl ExternalSorter {
             batch_size,
             sort_spill_reservation_bytes,
             sort_in_place_threshold_bytes,
-            has_sibling_partitions: false,
+            partitions_reading_input: Arc::new(AtomicUsize::new(1)),
             sort_concurrency: Arc::new(OnceLock::new()),
         })
     }
 
-    /// Record whether other partitions of this sort share its memory pool.
-    fn with_sibling_partitions(mut self, has_sibling_partitions: bool) -> Self {
-        self.has_sibling_partitions = has_sibling_partitions;
+    /// Share the count of partitions still reading their input with the
+    /// other partitions of this sort. It must include this partition.
+    fn with_partitions_reading_input(mut self, reading: Arc<AtomicUsize>) -> Self {
+        self.partitions_reading_input = reading;
         self
     }
 
+    /// Records that this partition has read all its input, and returns whether
+    /// another partition still reads its own. Executing partitions again
+    /// without `reset_state` can wrap the count, which only keeps headroom.
+    fn finish_reading_input(&self) -> bool {
+        self.partitions_reading_input.fetch_sub(1, Ordering::AcqRel) > 1
+    }
+
     /// Appends an unsorted [`RecordBatch`] to `in_mem_batches`
     ///
     /// Updates memory usage metrics, and possibly triggers spilling to disk
@@ -416,6 +425,7 @@ impl ExternalSorter {
     /// 2. A combined streaming merge incorporating both in-memory
     ///    batches and data from spill files on disk.
     async fn sort(&mut self) -> Result<SendableRecordBatchStream> {
+        let other_partitions_reading = self.finish_reading_input();
         if self.spilled_before() {
             // Sort `in_mem_batches` and spill it first. If there are many
             // `in_mem_batches` and the memory limit is almost reached, merging
@@ -444,15 +454,17 @@ impl ExternalSorter {
                 // Other partitions of this sort, still buffering input, share
                 // the pool with this merge, and may hold nothing at the moment
                 // it seats its runs. Seating every run the pool admits would
-                // leave them nothing, so keep headroom while they exist or
-                // anyone else holds memory in the pool. A lone partition with
-                // the pool to itself seats every run that fits: the operator
-                // consuming its output reserves nothing before its first
-                // batch, and keeping headroom for it would cost an extra merge
-                // pass whether or not it ever needs the memory.
+                // leave them nothing, so keep headroom while any of them is
+                // still reading or anyone else holds memory in the pool;
+                // partitions already merging hold memory, so the second check
+                // covers them. A lone partition, or the last one to finish
+                // reading, with the pool to itself seats every run that fits:
+                // the operator consuming its output reserves nothing before
+                // its first batch, and keeping headroom for it would cost an
+                // extra merge pass whether or not it ever needs the memory.
                 .with_fan_in_headroom(
                     Arc::clone(&self.runtime.memory_pool),
-                    self.has_sibling_partitions,
+                    other_partitions_reading,
                 )
                 .build()
         } else {
@@ -1349,6 +1361,10 @@ pub struct SortExec {
     /// Batch-sort limit shared by partitions and clones.
     /// `reset_state` returns a new plan with a fresh limit.
     sort_concurrency: Arc<OnceLock<SortConcurrency>>,
+    /// Partitions still reading their input, shared with their sorters. Set
+    /// on the first `execute`, once the input can no longer change.
+    /// `reset_state` returns a new plan with a fresh count.
+    partitions_reading_input: Arc<OnceLock<Arc<AtomicUsize>>>,
 }
 
 impl SortExec {
@@ -1369,6 +1385,7 @@ impl SortExec {
             cache: Arc::new(cache),
             filter: None,
             sort_concurrency: Arc::new(OnceLock::new()),
+            partitions_reading_input: Arc::new(OnceLock::new()),
         }
     }
 
@@ -1446,6 +1463,7 @@ impl SortExec {
             cache: Arc::clone(&self.cache),
             filter: self.filter.clone(),
             sort_concurrency: Arc::clone(&self.sort_concurrency),
+            partitions_reading_input: Arc::clone(&self.partitions_reading_input),
         }
     }
 
@@ -1776,14 +1794,15 @@ impl ExecutionPlan for SortExec {
             .downcast_ref::<SortExec>()
             .expect("rebuilt SortExec with new children")
             .clone();
-        // Our dynamic filter, execution metrics and sort concurrency are the state
-        // we need to reset.
+        // Our dynamic filter, execution metrics, sort concurrency and count of
+        // partitions reading input are the state we need to reset.
         if let Some(filter) = &new_sort.filter {
             filter.read().expr().reset();
             new_sort.filter = Some(new_sort.create_filter());
         }
         new_sort.metrics_set = ExecutionPlanMetricsSet::new();
         new_sort.sort_concurrency = Arc::new(OnceLock::new());
+        new_sort.partitions_reading_input = Arc::new(OnceLock::new());
 
         Ok(Arc::new(new_sort))
     }
@@ -1865,9 +1884,13 @@ impl ExecutionPlan for SortExec {
                     &self.metrics_set,
                     context.runtime_env(),
                 )?
-                .with_sibling_partitions(
-                    self.input.output_partitioning().partition_count() > 1,
-                );
+                .with_partitions_reading_input(Arc::clone(
+                    self.partitions_reading_input.get_or_init(|| {
+                        Arc::new(AtomicUsize::new(
+                            self.input.output_partitioning().partition_count(),
+                        ))
+                    }),
+                ));
                 sorter.sort_concurrency = Arc::clone(&self.sort_concurrency);
                 Ok(Box::pin(RecordBatchStreamAdapter::new(
                     self.schema(),
@@ -2070,6 +2093,7 @@ impl ExecutionPlan for SortExec {
             filter,
             // Execution state, initialized on the runtime polling the decoded plan.
             sort_concurrency: _,
+            partitions_reading_input: _,
         } = self;
         let input = ctx.encode_child(input)?;
         let expr = expr
diff --git a/datafusion/physical-plan/src/sorts/sort/spill_tests.rs b/datafusion/physical-plan/src/sorts/sort/spill_tests.rs
index ca996ea3283cd..0b6f97ff4281c 100644
--- a/datafusion/physical-plan/src/sorts/sort/spill_tests.rs
+++ b/datafusion/physical-plan/src/sorts/sort/spill_tests.rs
@@ -3,12 +3,14 @@
 // SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements
 
 use super::{
-    ExternalSorter, get_reserved_bytes_for_record_batch,
+    ExternalSorter, SortExec, get_reserved_bytes_for_record_batch,
     get_reserved_bytes_for_record_batch_size, sort_batch, sort_batch_chunked,
 };
+use crate::ExecutionPlan;
 use crate::metrics::ExecutionPlanMetricsSet;
 use crate::spill::get_record_batch_memory_size;
 use crate::spill::spill_manager::GetSlicedSize;
+use crate::test::TestMemoryExec;
 use arrow::array::{
     ArrayRef, Decimal128Array, DictionaryArray, Int8Array, Int64Array, StringArray,
     StringViewArray,
@@ -18,6 +20,8 @@ use arrow::datatypes::{DataType, Field, Int8Type, Schema};
 use arrow::record_batch::RecordBatch;
 use datafusion_common::config::{ExecutionOptions, SpillCompression};
 use datafusion_common::{DataFusionError, Result};
+use datafusion_execution::TaskContext;
+use datafusion_execution::config::SessionConfig;
 use datafusion_execution::disk_manager::{DiskManagerBuilder, DiskManagerMode};
 use datafusion_execution::memory_pool::{
     GreedyMemoryPool, MemoryConsumer, MemoryLimit, MemoryPool, MemoryReservation,
@@ -27,6 +31,7 @@ use datafusion_physical_expr::expressions::Column;
 use datafusion_physical_expr::{LexOrdering, PhysicalSortExpr};
 use futures::TryStreamExt;
 use std::fmt::{Display, Formatter};
+use std::sync::atomic::AtomicUsize;
 use std::sync::{Arc, Mutex};
 use std::time::Duration;
 
@@ -1208,7 +1213,7 @@ async fn test_single_batch_spill_returns_live_workspace_loan_on_drop() -> Result
 /// still sorting their input. Seating every run the pool admits leaves them
 /// nothing, so whatever the merge takes beyond its own workspace, it must leave
 /// as much free: while another consumer holds memory, and while sibling
-/// partitions exist even if they hold nothing at the moment.
+/// partitions still read their input even if they hold nothing at the moment.
 #[rstest::rstest]
 #[case::other_consumer_holds_memory(false)]
 #[case::idle_sibling_partitions(true)]
@@ -1221,7 +1226,9 @@ async fn test_final_spilled_merge_leaves_headroom(#[case] siblings: bool) -> Res
     // Spill under the fixture's limit, then merge with room for every run.
     pool.set_limit(fixture.capacity);
     let (sorter, runtime) = sorter(&fixture, Arc::clone(&pool))?;
-    let mut sorter = sorter.with_sibling_partitions(siblings);
+    let mut sorter = sorter.with_partitions_reading_input(Arc::new(AtomicUsize::new(
+        1 + usize::from(siblings),
+    )));
     for batch in &fixture.batches {
         sorter.insert_batch(batch.clone()).await?;
     }
@@ -1287,3 +1294,77 @@ async fn test_final_spilled_merge_leaves_headroom(#[case] siblings: bool) -> Res
     assert_released(&pool, &runtime).await;
     Ok(())
 }
+
+/// Whether a `SortExec` keeps fan-in headroom follows its partitions still
+/// reading input. The pool fits every run's buffers but not the headroom on
+/// top, so with headroom the final merge seats fewer runs and rewrites rows in
+/// an extra pass. Partition 0 runs alone, so nothing else ever holds pool
+/// memory and the outcome is deterministic. In the last case its sibling has
+/// already read all its input, so there is nobody to keep headroom for.
+#[rstest::rstest]
+#[case::lone_partition(1, false, false)]
+#[case::idle_sibling(2, false, true)]
+#[case::finished_sibling(2, true, false)]
+#[tokio::test]
+async fn test_sort_exec_keeps_headroom_only_while_siblings_read_input(
+    #[case] input_partitions: usize,
+    #[case] finish_sibling_first: bool,
+    #[case] rewrites_rows: bool,
+) -> Result<()> {
+    let schema = Arc::new(Schema::new(vec![
+        Field::new("key", DataType::Int64, false),
+        Field::new("payload", DataType::Utf8, false),
+    ]));
+    let (batches, rows) = (64, 1024);
+    let batches = (0..batches)
+        .map(|batch| {
+            let keys = (0..rows).map(|row| ((batch * rows + row) * 2654435761) % 1000003);
+            let payload = (0..rows).map(|row| format!("{row:0120}"));
+            let columns: Vec<ArrayRef> = vec![
+                Arc::new(Int64Array::from_iter_values(keys)),
+                Arc::new(StringArray::from_iter_values(payload)),
+            ];
+            Ok(RecordBatch::try_new(Arc::clone(&schema), columns)?)
+        })
+        .collect::<Result<Vec<_>>>()?;
+    let total_rows: usize = batches.iter().map(RecordBatch::num_rows).sum();
+    let mut partitions = vec![batches];
+    partitions.resize(input_partitions, vec![]);
+    let input = TestMemoryExec::try_new_exec(&partitions, Arc::clone(&schema), None)?;
+    let ordering: LexOrdering = [PhysicalSortExpr::new_default(Arc::new(Column::new(
+        "key", 0,
+    )))]
+    .into();
+    let sort = SortExec::new(ordering, input).with_preserve_partitioning(true);
+    let runtime = RuntimeEnvBuilder::new()
+        .with_memory_pool(Arc::new(GreedyMemoryPool::new(3500 * 1024)))
+        .build_arc()?;
+    let context = Arc::new(
+        TaskContext::default()
+            .with_session_config(
+                SessionConfig::new()
+                    .with_batch_size(1024)
+                    .with_sort_spill_reservation_bytes(128 * 1024),
+            )
+            .with_runtime(runtime),
+    );
+
+    if finish_sibling_first {
+        let sibling: Vec<RecordBatch> =
+            sort.execute(1, Arc::clone(&context))?.try_collect().await?;
+        assert!(sibling.is_empty());
+    }
+    let output: Vec<RecordBatch> = sort.execute(0, context)?.try_collect().await?;
+
+    assert_eq!(
+        output.iter().map(RecordBatch::num_rows).sum::<usize>(),
+        total_rows
+    );
+    let spilled_rows = sort.metrics().unwrap().spilled_rows().unwrap();
+    if rewrites_rows {
+        assert!(spilled_rows > total_rows, "{spilled_rows} <= {total_rows}");
+    } else {
+        assert_eq!(spilled_rows, total_rows);
+    }
+    Ok(())
+}