From 3109bdb300b3c056ea6f14ddea781c6a5960566a Mon Sep 17 00:00:00 2001
From: Ariel Miculas <ariel.miculas@gmail.com>
Date: Wed, 15 Jul 2026 15:36:25 +0300
Subject: [PATCH 01/15] fix: reduce peak memory usage when round robin
 tiebreaker is disabled

---
 ...spilling_fuzz_in_memory_constrained_env.rs | 165 ++++++++++++++++++
 datafusion/physical-plan/src/sorts/merge.rs   |   8 +-
 2 files changed, 171 insertions(+), 2 deletions(-)

diff --git a/datafusion/core/tests/fuzz_cases/spilling_fuzz_in_memory_constrained_env.rs b/datafusion/core/tests/fuzz_cases/spilling_fuzz_in_memory_constrained_env.rs
index 103c3e03c06df..5e6eaa467f8f4 100644
--- a/datafusion/core/tests/fuzz_cases/spilling_fuzz_in_memory_constrained_env.rs
+++ b/datafusion/core/tests/fuzz_cases/spilling_fuzz_in_memory_constrained_env.rs
@@ -48,6 +48,16 @@ use datafusion_physical_plan::metrics::MetricValue;
 use datafusion_physical_plan::stream::RecordBatchStreamAdapter;
 use futures::StreamExt;
 
+use arrow::array::Int32Array;
+use datafusion::datasource::memory::MemorySourceConfig;
+use datafusion::physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec;
+use datafusion_execution::memory_pool::{
+    MemoryPool, TrackConsumersPool, UnboundedMemoryPool,
+};
+use datafusion_physical_plan::metrics::{ExecutionPlanMetricsSet, SpillMetrics};
+use datafusion_physical_plan::spill::SpillManager;
+use std::num::NonZeroUsize;
+
 #[tokio::test]
 async fn test_sort_with_limited_memory() -> Result<()> {
     let record_batch_size = 8192;
@@ -290,6 +300,161 @@ async fn test_sort_with_limited_memory_and_oversized_record_batch() -> Result<()
     Ok(())
 }
 
+#[tokio::test]
+async fn test_sort_preserving_merge_peak_memory_with_spilled_input_round_robin()
+-> Result<()> {
+    run_sort_preserving_merge_peak_memory_with_spilled_input(true).await
+}
+
+#[tokio::test]
+async fn test_sort_preserving_merge_peak_memory_with_spilled_input_no_round_robin()
+-> Result<()> {
+    run_sort_preserving_merge_peak_memory_with_spilled_input(false).await
+}
+
+/// Intended to measure the maximum number of record batches held in memory by
+/// the SortPreservingMergeStream in a convoluted way by measuring the peak
+/// memory reservation. Relevant for merging spilled streams, where the produced
+/// record batches suffer from the following issue:
+/// https://github.com/apache/arrow-rs/issues/6363
+///
+/// After an IPC roundtrip, all columns in a [`RecordBatch`] share a single
+/// parent buffer. It causes the memory reservation to be inflated, but the
+/// bigger issue is the increase in the peak allocated memory caused by
+/// prev_cursors in SortPreservingMergeExec. The increase is caused by the fact
+/// that the FieldCursor inside prev_cursors holds a reference for the entire
+/// Buffer allocated for the input record batch, preventing it from being
+/// dropped and thus increasing the number of concomitent input record batches
+/// living during the merging phase
+async fn run_sort_preserving_merge_peak_memory_with_spilled_input(
+    round_robin: bool,
+) -> Result<()> {
+    let num_batches = 10usize;
+    let num_rows_per_batch = 100usize;
+    // payload is ~100x larger than the sort key (i32 = 4 bytes, string ≈ 400 bytes)
+    let large_string = "x".repeat(400);
+
+    let schema = Arc::new(Schema::new(vec![
+        Field::new("sort_key", DataType::Int32, false),
+        Field::new("payload", DataType::Utf8, false),
+    ]));
+
+    // Unbounded env used only for spilling the input; the merge runs under its
+    // own pool below.
+    let spill_env = Arc::new(RuntimeEnvBuilder::new().build()?);
+
+    let mut partition_batches: Vec<Vec<RecordBatch>> = Vec::new();
+
+    for stream_idx in 0..2usize {
+        // Each stream covers a non-overlapping key range so both are individually
+        // sorted: stream 0 → [0, 1000), stream 1 → [1000, 2000).
+        let batches: Vec<RecordBatch> = (0..num_batches)
+            .map(|b| {
+                // Interleave streams: stream 0 → even slots [0,200,400,...],
+                // stream 1 → odd slots [100,300,500,...] so the merge
+                // alternates between them on every batch.
+                let base = ((b * 2 + stream_idx) * num_rows_per_batch) as i32;
+                let sort_col: Int32Array =
+                    (base..base + num_rows_per_batch as i32).collect();
+                let payload_col: StringArray =
+                    std::iter::repeat_n(large_string.as_str(), num_rows_per_batch)
+                        .map(Some)
+                        .collect();
+                RecordBatch::try_new(
+                    Arc::clone(&schema),
+                    vec![Arc::new(sort_col), Arc::new(payload_col)],
+                )
+                .unwrap()
+            })
+            .collect();
+
+        // Spill to disk then read back: each RecordBatch is now IPC-backed,
+        // meaning all columns share a single parent buffer.  As a result,
+        // get_buffer_memory_size() on the sort_key column returns the full
+        // parent-buffer capacity (≈ batch size of both columns combined) rather
+        // than just the key data (num_rows * 4 bytes).
+        let metrics = SpillMetrics::new(&ExecutionPlanMetricsSet::new(), 0);
+        let manager =
+            SpillManager::new(Arc::clone(&spill_env), metrics, Arc::clone(&schema));
+        let spill_file = manager
+            .spill_record_batch_and_finish(&batches, "stream")?
+            .expect("non-empty input should produce a spill file");
+
+        let mut stream = manager.read_spill_as_stream(spill_file, None)?;
+        let mut ipc_batches: Vec<RecordBatch> = Vec::new();
+        while let Some(batch) = stream.next().await {
+            ipc_batches.push(batch?);
+        }
+        if stream_idx == 0 {
+            use datafusion_physical_plan::spill::get_record_batch_memory_size;
+        }
+        partition_batches.push(ipc_batches);
+    }
+
+    use datafusion_physical_plan::spill::get_record_batch_memory_size;
+    let ipc_batch_size = get_record_batch_memory_size(&partition_batches[0][0]);
+
+    // Build a 2-partition plan from the IPC-recovered batches.
+    let input =
+        MemorySourceConfig::try_new_exec(&partition_batches, Arc::clone(&schema), None)?;
+
+    let sort_expr = PhysicalSortExpr {
+        expr: col("sort_key", &schema)?,
+        options: SortOptions {
+            descending: false,
+            nulls_first: true,
+        },
+    };
+    let merge = Arc::new(
+        SortPreservingMergeExec::new(LexOrdering::new(vec![sort_expr]).unwrap(), input)
+            .with_round_robin_repartition(round_robin),
+    );
+
+    // TrackConsumersPool wraps an unbounded pool so the merge never OOMs;
+    // we keep the typed Arc to call .metrics() after the run.
+    let tracking_pool = Arc::new(TrackConsumersPool::new(
+        UnboundedMemoryPool::default(),
+        NonZeroUsize::new(10).unwrap(),
+    ));
+    let runtime = RuntimeEnvBuilder::new()
+        .with_memory_pool(Arc::clone(&tracking_pool) as Arc<dyn MemoryPool>)
+        .build()?;
+    let task_ctx = Arc::new(
+        TaskContext::default()
+            .with_session_config(SessionConfig::new().with_batch_size(num_rows_per_batch))
+            .with_runtime(Arc::new(runtime)),
+    );
+
+    let mut output = merge.execute(0, task_ctx)?;
+    let mut total_rows = 0usize;
+    while let Some(batch) = output.next().await {
+        total_rows += batch?.num_rows();
+    }
+    assert_eq!(total_rows, 2 * num_batches * num_rows_per_batch);
+
+    let mut metrics = tracking_pool.metrics();
+    metrics.sort_by_key(|m| std::cmp::Reverse(m.peak));
+    let peak_bytes: usize = metrics.iter().map(|m| m.peak).sum();
+
+    // with round robin enabled, 2 extra record batches live in memory
+    // see https://github.com/apache/datafusion/issues/23604
+    // 5 comes from:
+    // BatchBuilder needs to hold 3 Record batches simultaneously to merge two
+    // streams (because a stream can cross a record batch boundary)
+    // the `cursors` also need 1 record batch worth of memory each (because of
+    // the IPC issue)
+    let max_batches = if round_robin { 7 } else { 5 };
+    let max_peak = max_batches * ipc_batch_size;
+
+    assert!(
+        peak_bytes <= max_peak,
+        "peak reservation {peak_bytes} bytes exceeds {max_batches}x IPC batch size \
+         ({max_peak} bytes); round_robin={round_robin}",
+    );
+
+    Ok(())
+}
+
 struct RunTestWithLimitedMemoryArgs {
     pool_size: usize,
     task_ctx: Arc<TaskContext>,
diff --git a/datafusion/physical-plan/src/sorts/merge.rs b/datafusion/physical-plan/src/sorts/merge.rs
index 310416c22d982..e52189f378694 100644
--- a/datafusion/physical-plan/src/sorts/merge.rs
+++ b/datafusion/physical-plan/src/sorts/merge.rs
@@ -420,8 +420,12 @@ impl<C: CursorValues> SortPreservingMergeStream<C> {
             let _ = cursor.advance();
             let finished = cursor.is_finished();
             if finished {
-                // Take the current cursor, leaving `None` in its place
-                self.prev_cursors[stream_idx] = self.cursors[stream_idx].take();
+                if self.enable_round_robin_tie_breaker {
+                    // Take the current cursor, leaving `None` in its place
+                    self.prev_cursors[stream_idx] = self.cursors[stream_idx].take();
+                } else {
+                    self.cursors[stream_idx].take();
+                }
             }
             return finished;
         }

From d0518a7405193c44b971c0336b2995b70a1a8099 Mon Sep 17 00:00:00 2001
From: Ariel Miculas <ariel.miculas@gmail.com>
Date: Wed, 15 Jul 2026 17:21:06 +0300
Subject: [PATCH 02/15] chore: remove useless code

---
 .../fuzz_cases/spilling_fuzz_in_memory_constrained_env.rs      | 3 ---
 1 file changed, 3 deletions(-)

diff --git a/datafusion/core/tests/fuzz_cases/spilling_fuzz_in_memory_constrained_env.rs b/datafusion/core/tests/fuzz_cases/spilling_fuzz_in_memory_constrained_env.rs
index 5e6eaa467f8f4..e139aa62101bc 100644
--- a/datafusion/core/tests/fuzz_cases/spilling_fuzz_in_memory_constrained_env.rs
+++ b/datafusion/core/tests/fuzz_cases/spilling_fuzz_in_memory_constrained_env.rs
@@ -385,9 +385,6 @@ async fn run_sort_preserving_merge_peak_memory_with_spilled_input(
         while let Some(batch) = stream.next().await {
             ipc_batches.push(batch?);
         }
-        if stream_idx == 0 {
-            use datafusion_physical_plan::spill::get_record_batch_memory_size;
-        }
         partition_batches.push(ipc_batches);
     }
 

From 4ec66180e235a47417e5005328dc5e52ae86cd49 Mon Sep 17 00:00:00 2001
From: Ariel Miculas <ariel.miculas@gmail.com>
Date: Wed, 15 Jul 2026 18:48:13 +0300
Subject: [PATCH 03/15] refactor: simplify the expression

---
 datafusion/physical-plan/src/sorts/merge.rs | 7 +++----
 1 file changed, 3 insertions(+), 4 deletions(-)

diff --git a/datafusion/physical-plan/src/sorts/merge.rs b/datafusion/physical-plan/src/sorts/merge.rs
index e52189f378694..1daa969465db4 100644
--- a/datafusion/physical-plan/src/sorts/merge.rs
+++ b/datafusion/physical-plan/src/sorts/merge.rs
@@ -420,11 +420,10 @@ impl<C: CursorValues> SortPreservingMergeStream<C> {
             let _ = cursor.advance();
             let finished = cursor.is_finished();
             if finished {
+                // Take the current cursor, leaving `None` in its place
+                let taken = self.cursors[stream_idx].take();
                 if self.enable_round_robin_tie_breaker {
-                    // Take the current cursor, leaving `None` in its place
-                    self.prev_cursors[stream_idx] = self.cursors[stream_idx].take();
-                } else {
-                    self.cursors[stream_idx].take();
+                    self.prev_cursors[stream_idx] = taken;
                 }
             }
             return finished;

From c05609d6c8083d4b2e1d1c5277f1118a4274b6c3 Mon Sep 17 00:00:00 2001
From: Ariel Miculas <ariel.miculas@gmail.com>
Date: Thu, 16 Jul 2026 15:02:31 +0300
Subject: [PATCH 04/15] feat: add test for multi-column sort

---
 ...spilling_fuzz_in_memory_constrained_env.rs | 98 ++++++++++++++++---
 1 file changed, 84 insertions(+), 14 deletions(-)

diff --git a/datafusion/core/tests/fuzz_cases/spilling_fuzz_in_memory_constrained_env.rs b/datafusion/core/tests/fuzz_cases/spilling_fuzz_in_memory_constrained_env.rs
index e139aa62101bc..b387e8800aca3 100644
--- a/datafusion/core/tests/fuzz_cases/spilling_fuzz_in_memory_constrained_env.rs
+++ b/datafusion/core/tests/fuzz_cases/spilling_fuzz_in_memory_constrained_env.rs
@@ -23,6 +23,7 @@ use std::sync::Arc;
 use crate::fuzz_cases::aggregate_fuzz::assert_spill_count_metric;
 use crate::fuzz_cases::once_exec::OnceExec;
 use arrow::array::UInt64Array;
+use arrow::row::{RowConverter, SortField};
 use arrow::{array::StringArray, compute::SortOptions, record_batch::RecordBatch};
 use arrow_schema::{DataType, Field, Schema};
 use datafusion::common::Result;
@@ -45,6 +46,7 @@ use datafusion_physical_plan::aggregates::{
     AggregateExec, AggregateMode, PhysicalGroupBy,
 };
 use datafusion_physical_plan::metrics::MetricValue;
+use datafusion_physical_plan::spill::get_record_batch_memory_size;
 use datafusion_physical_plan::stream::RecordBatchStreamAdapter;
 use futures::StreamExt;
 
@@ -303,13 +305,25 @@ async fn test_sort_with_limited_memory_and_oversized_record_batch() -> Result<()
 #[tokio::test]
 async fn test_sort_preserving_merge_peak_memory_with_spilled_input_round_robin()
 -> Result<()> {
-    run_sort_preserving_merge_peak_memory_with_spilled_input(true).await
+    run_sort_preserving_merge_peak_memory_with_spilled_input(true, false).await
 }
 
 #[tokio::test]
 async fn test_sort_preserving_merge_peak_memory_with_spilled_input_no_round_robin()
 -> Result<()> {
-    run_sort_preserving_merge_peak_memory_with_spilled_input(false).await
+    run_sort_preserving_merge_peak_memory_with_spilled_input(false, false).await
+}
+
+#[tokio::test]
+async fn test_sort_preserving_merge_peak_memory_with_spilled_input_round_robin_multi_column()
+-> Result<()> {
+    run_sort_preserving_merge_peak_memory_with_spilled_input(true, true).await
+}
+
+#[tokio::test]
+async fn test_sort_preserving_merge_peak_memory_with_spilled_input_no_round_robin_multi_column()
+-> Result<()> {
+    run_sort_preserving_merge_peak_memory_with_spilled_input(false, true).await
 }
 
 /// Intended to measure the maximum number of record batches held in memory by
@@ -328,6 +342,7 @@ async fn test_sort_preserving_merge_peak_memory_with_spilled_input_no_round_robi
 /// living during the merging phase
 async fn run_sort_preserving_merge_peak_memory_with_spilled_input(
     round_robin: bool,
+    multi_column_sort: bool,
 ) -> Result<()> {
     let num_batches = 10usize;
     let num_rows_per_batch = 100usize;
@@ -388,22 +403,67 @@ async fn run_sort_preserving_merge_peak_memory_with_spilled_input(
         partition_batches.push(ipc_batches);
     }
 
-    use datafusion_physical_plan::spill::get_record_batch_memory_size;
     let ipc_batch_size = get_record_batch_memory_size(&partition_batches[0][0]);
 
     // Build a 2-partition plan from the IPC-recovered batches.
     let input =
         MemorySourceConfig::try_new_exec(&partition_batches, Arc::clone(&schema), None)?;
 
-    let sort_expr = PhysicalSortExpr {
+    let sort_key_expr = PhysicalSortExpr {
         expr: col("sort_key", &schema)?,
         options: SortOptions {
             descending: false,
             nulls_first: true,
         },
     };
+    // `payload` has the same value in every row, so adding it as a secondary
+    // sort key doesn't change the resulting order — it only forces the merge
+    // onto the row-oriented (`RowValues`/`RowCursorStream`) comparison path
+    // used whenever more than one sort expression is present.
+    let mut sort_exprs = vec![sort_key_expr];
+    if multi_column_sort {
+        sort_exprs.push(PhysicalSortExpr {
+            expr: col("payload", &schema)?,
+            options: SortOptions {
+                descending: false,
+                nulls_first: true,
+            },
+        });
+    }
+
+    // When sorting by more than one column, the merge switches to the
+    // row-oriented `RowValues`/`RowCursorStream` path
+    //
+    // `RowCursorStream` also tracks one *shared* (not per-partition)
+    // reservation sized to `converter.size()` (`stream.rs`:
+    // `self.reservation.try_resize(self.converter.size())`) — the
+    // `RowConverter`'s own fixed internal state, separate from the `Rows`
+    // it produces per batch.
+    let (row_batch_size, converter_size) = if multi_column_sort {
+        let sort_fields = sort_exprs
+            .iter()
+            .map(|s| {
+                let data_type = s.expr.data_type(&schema)?;
+                Ok(SortField::new_with_options(data_type, s.options))
+            })
+            .collect::<Result<Vec<_>>>()?;
+        let converter = RowConverter::new(sort_fields)?;
+        let cols = sort_exprs
+            .iter()
+            .map(|s| {
+                s.expr
+                    .evaluate(&partition_batches[0][0])?
+                    .into_array(partition_batches[0][0].num_rows())
+            })
+            .collect::<Result<Vec<_>>>()?;
+        let rows = converter.convert_columns(&cols)?;
+        (rows.size(), converter.size())
+    } else {
+        (0, 0)
+    };
+
     let merge = Arc::new(
-        SortPreservingMergeExec::new(LexOrdering::new(vec![sort_expr]).unwrap(), input)
+        SortPreservingMergeExec::new(LexOrdering::new(sort_exprs).unwrap(), input)
             .with_round_robin_repartition(round_robin),
     );
 
@@ -433,20 +493,30 @@ async fn run_sort_preserving_merge_peak_memory_with_spilled_input(
     metrics.sort_by_key(|m| std::cmp::Reverse(m.peak));
     let peak_bytes: usize = metrics.iter().map(|m| m.peak).sum();
 
-    // with round robin enabled, 2 extra record batches live in memory
-    // see https://github.com/apache/datafusion/issues/23604
-    // 5 comes from:
+    // in the single column case, the cursor takes up an ipc_batch_size worth of memory due to the
+    // IPC roundtrip issue
+    // for the multi-column case, we've calculated row_batch_size above
+    let cursor_unit = if multi_column_sort {
+        row_batch_size
+    } else {
+        ipc_batch_size
+    };
+
     // BatchBuilder needs to hold 3 Record batches simultaneously to merge two
     // streams (because a stream can cross a record batch boundary)
-    // the `cursors` also need 1 record batch worth of memory each (because of
-    // the IPC issue)
-    let max_batches = if round_robin { 7 } else { 5 };
-    let max_peak = max_batches * ipc_batch_size;
+    // there is also one cursor needed per stream
+    let mut max_peak = 3 * ipc_batch_size + 2 * cursor_unit + converter_size;
+
+    // with round robin enabled, 2 extra cursors live in memory
+    // see https://github.com/apache/datafusion/issues/23604
+    if round_robin {
+        max_peak += 2 * cursor_unit;
+    };
 
     assert!(
         peak_bytes <= max_peak,
-        "peak reservation {peak_bytes} bytes exceeds {max_batches}x IPC batch size \
-         ({max_peak} bytes); round_robin={round_robin}",
+        "peak reservation {peak_bytes} bytes exceeds max_peak ({max_peak} bytes); \
+         round_robin={round_robin}, multi_column_sort={multi_column_sort}",
     );
 
     Ok(())

From ed55e6573402a801e10eb4eed73d024827e5ccaf Mon Sep 17 00:00:00 2001
From: Ariel Miculas <ariel.miculas@gmail.com>
Date: Thu, 16 Jul 2026 15:32:08 +0300
Subject: [PATCH 05/15] refactor: make prev_cursors optional

---
 datafusion/physical-plan/src/sorts/merge.rs | 21 ++++++++++++++++-----
 1 file changed, 16 insertions(+), 5 deletions(-)

diff --git a/datafusion/physical-plan/src/sorts/merge.rs b/datafusion/physical-plan/src/sorts/merge.rs
index 1daa969465db4..7713e4d42f850 100644
--- a/datafusion/physical-plan/src/sorts/merge.rs
+++ b/datafusion/physical-plan/src/sorts/merge.rs
@@ -126,8 +126,9 @@ pub(crate) struct SortPreservingMergeStream<C: CursorValues> {
     /// Current reset count
     current_reset_epoch: usize,
 
-    /// Stores the previous value of each partitions for tracking the poll counts on the same value.
-    prev_cursors: Vec<Option<Cursor<C>>>,
+    /// Stores the previous value of each partitions for tracking the poll counts on the same value
+    /// when round_robin_tie_breaker is enabled
+    prev_cursors: Option<Vec<Option<Cursor<C>>>>,
 
     /// Optional number of rows to fetch
     fetch: Option<usize>,
@@ -156,7 +157,11 @@ impl<C: CursorValues> SortPreservingMergeStream<C> {
             streams,
             metrics,
             cursors: (0..stream_count).map(|_| None).collect(),
-            prev_cursors: (0..stream_count).map(|_| None).collect(),
+            prev_cursors: if enable_round_robin_tie_breaker {
+                Some((0..stream_count).map(|_| None).collect())
+            } else {
+                None
+            },
             round_robin_tie_breaker_mode: false,
             num_of_polled_with_same_value: vec![0; stream_count],
             current_reset_epoch: 0,
@@ -396,7 +401,13 @@ impl<C: CursorValues> SortPreservingMergeStream<C> {
 
         if let Some(c) = cursor.as_mut() {
             // Compare with the last row in the previous batch
-            let prev_cursor = &self.prev_cursors[partition_idx];
+            let prev_cursor = self
+                .prev_cursors
+                .as_ref()
+                .map(|v| &v[partition_idx])
+                .expect(
+                    "prev_cursor should be set when round robin tie breaker is enabled",
+                );
             if c.is_eq_to_prev_one(prev_cursor.as_ref()) {
                 self.num_of_polled_with_same_value[partition_idx] += 1;
             } else {
@@ -423,7 +434,7 @@ impl<C: CursorValues> SortPreservingMergeStream<C> {
                 // Take the current cursor, leaving `None` in its place
                 let taken = self.cursors[stream_idx].take();
                 if self.enable_round_robin_tie_breaker {
-                    self.prev_cursors[stream_idx] = taken;
+                    self.prev_cursors.as_mut().expect("prev_cursor should be set when round robin tie breaker is enabled")[stream_idx] = taken;
                 }
             }
             return finished;

From 72eb210ce4e4f45bc50fc6651a4c7576762e82df Mon Sep 17 00:00:00 2001
From: Ariel Miculas <ariel.miculas@gmail.com>
Date: Mon, 20 Jul 2026 19:34:31 +0300
Subject: [PATCH 06/15] refactor: use prev_cursors as source of truth for
 tie-breaker config

---
 datafusion/physical-plan/src/sorts/merge.rs | 57 +++++++++++----------
 1 file changed, 29 insertions(+), 28 deletions(-)

diff --git a/datafusion/physical-plan/src/sorts/merge.rs b/datafusion/physical-plan/src/sorts/merge.rs
index 7713e4d42f850..647649038766d 100644
--- a/datafusion/physical-plan/src/sorts/merge.rs
+++ b/datafusion/physical-plan/src/sorts/merge.rs
@@ -89,29 +89,6 @@ pub(crate) struct SortPreservingMergeStream<C: CursorValues> {
     /// Cursors for each input partition. `None` means the input is exhausted
     cursors: Vec<Option<Cursor<C>>>,
 
-    /// Configuration parameter to enable round-robin selection of tied winners of loser tree.
-    ///
-    /// This option controls the tie-breaker strategy and attempts to avoid the
-    /// issue of unbalanced polling between partitions
-    ///
-    /// If `true`, when multiple partitions have the same value, the partition
-    /// that has the fewest poll counts is selected. This strategy ensures that
-    /// multiple partitions with the same value are chosen equally, distributing
-    /// the polling load in a round-robin fashion. This approach balances the
-    /// workload more effectively across partitions and avoids excessive buffer
-    /// growth.
-    ///
-    /// if `false`, partitions with smaller indices are consistently chosen as
-    /// the winners, which can lead to an uneven distribution of polling and potentially
-    /// causing upstream operator buffers for the other partitions to grow
-    /// excessively, as they continued receiving data without consuming it.
-    ///
-    /// For example, an upstream operator like `RepartitionExec` execution would
-    /// keep sending data to certain partitions, but those partitions wouldn't
-    /// consume the data if they weren't selected as winners. This resulted in
-    /// inefficient buffer usage.
-    enable_round_robin_tie_breaker: bool,
-
     /// Flag indicating whether we are in the mode of round-robin
     /// tie breaker for the loser tree winners.
     round_robin_tie_breaker_mode: bool,
@@ -127,7 +104,7 @@ pub(crate) struct SortPreservingMergeStream<C: CursorValues> {
     current_reset_epoch: usize,
 
     /// Stores the previous value of each partitions for tracking the poll counts on the same value
-    /// when round_robin_tie_breaker is enabled
+    /// Used if and only if round robin tie breaker is enabled, otherwise None
     prev_cursors: Option<Vec<Option<Cursor<C>>>>,
 
     /// Optional number of rows to fetch
@@ -170,7 +147,6 @@ impl<C: CursorValues> SortPreservingMergeStream<C> {
             batch_size,
             fetch,
             produced: 0,
-            enable_round_robin_tie_breaker,
         }
     }
 
@@ -416,6 +392,31 @@ impl<C: CursorValues> SortPreservingMergeStream<C> {
         }
     }
 
+    /// Whether round-robin selection of tied winners of loser tree is enabled.
+    ///
+    /// This option controls the tie-breaker strategy and attempts to avoid the
+    /// issue of unbalanced polling between partitions
+    ///
+    /// If `true`, when multiple partitions have the same value, the partition
+    /// that has the fewest poll counts is selected. This strategy ensures that
+    /// multiple partitions with the same value are chosen equally, distributing
+    /// the polling load in a round-robin fashion. This approach balances the
+    /// workload more effectively across partitions and avoids excessive buffer
+    /// growth.
+    ///
+    /// if `false`, partitions with smaller indices are consistently chosen as
+    /// the winners, which can lead to an uneven distribution of polling and potentially
+    /// causing upstream operator buffers for the other partitions to grow
+    /// excessively, as they continued receiving data without consuming it.
+    ///
+    /// For example, an upstream operator like `RepartitionExec` execution would
+    /// keep sending data to certain partitions, but those partitions wouldn't
+    /// consume the data if they weren't selected as winners. This resulted in
+    /// inefficient buffer usage.
+    fn round_robin_tie_breaker_enabled(&self) -> bool {
+        self.prev_cursors.is_some()
+    }
+
     fn fetch_reached(&mut self) -> bool {
         self.fetch
             .map(|fetch| self.produced + self.in_progress.len() >= fetch)
@@ -433,8 +434,8 @@ impl<C: CursorValues> SortPreservingMergeStream<C> {
             if finished {
                 // Take the current cursor, leaving `None` in its place
                 let taken = self.cursors[stream_idx].take();
-                if self.enable_round_robin_tie_breaker {
-                    self.prev_cursors.as_mut().expect("prev_cursor should be set when round robin tie breaker is enabled")[stream_idx] = taken;
+                if let Some(prev_cursors) = &mut self.prev_cursors {
+                    prev_cursors[stream_idx] = taken;
                 }
             }
             return finished;
@@ -602,7 +603,7 @@ impl<C: CursorValues> SortPreservingMergeStream<C> {
         if cmp_node == 1 {
             let challenger = self.loser_tree[1];
             // If round-robin tie-breaker is enabled and we're at the final comparison (cmp_node == 1)
-            if self.enable_round_robin_tie_breaker {
+            if self.round_robin_tie_breaker_enabled() {
                 match (&self.cursors[winner], &self.cursors[challenger]) {
                     (Some(ac), Some(bc)) => match ac.cmp(bc) {
                         std::cmp::Ordering::Equal => {

From 80bec049679ef5d5a8a0374555b81a1fa40d11b9 Mon Sep 17 00:00:00 2001
From: Ariel Miculas <ariel.miculas@gmail.com>
Date: Tue, 4 Aug 2026 19:23:54 +0300
Subject: [PATCH 07/15] feat: add test case where all sort values are equal
 across all batches

---
 ...spilling_fuzz_in_memory_constrained_env.rs | 33 +++++++++++++++----
 1 file changed, 26 insertions(+), 7 deletions(-)

diff --git a/datafusion/core/tests/fuzz_cases/spilling_fuzz_in_memory_constrained_env.rs b/datafusion/core/tests/fuzz_cases/spilling_fuzz_in_memory_constrained_env.rs
index b387e8800aca3..708dcac26d884 100644
--- a/datafusion/core/tests/fuzz_cases/spilling_fuzz_in_memory_constrained_env.rs
+++ b/datafusion/core/tests/fuzz_cases/spilling_fuzz_in_memory_constrained_env.rs
@@ -305,25 +305,37 @@ async fn test_sort_with_limited_memory_and_oversized_record_batch() -> Result<()
 #[tokio::test]
 async fn test_sort_preserving_merge_peak_memory_with_spilled_input_round_robin()
 -> Result<()> {
-    run_sort_preserving_merge_peak_memory_with_spilled_input(true, false).await
+    run_sort_preserving_merge_peak_memory_with_spilled_input(true, false, false).await
 }
 
 #[tokio::test]
 async fn test_sort_preserving_merge_peak_memory_with_spilled_input_no_round_robin()
 -> Result<()> {
-    run_sort_preserving_merge_peak_memory_with_spilled_input(false, false).await
+    run_sort_preserving_merge_peak_memory_with_spilled_input(false, false, false).await
 }
 
 #[tokio::test]
 async fn test_sort_preserving_merge_peak_memory_with_spilled_input_round_robin_multi_column()
 -> Result<()> {
-    run_sort_preserving_merge_peak_memory_with_spilled_input(true, true).await
+    run_sort_preserving_merge_peak_memory_with_spilled_input(true, true, false).await
 }
 
 #[tokio::test]
 async fn test_sort_preserving_merge_peak_memory_with_spilled_input_no_round_robin_multi_column()
 -> Result<()> {
-    run_sort_preserving_merge_peak_memory_with_spilled_input(false, true).await
+    run_sort_preserving_merge_peak_memory_with_spilled_input(false, true, false).await
+}
+
+#[tokio::test]
+async fn test_sort_preserving_merge_peak_memory_with_spilled_input_round_robin_tied_values()
+-> Result<()> {
+    run_sort_preserving_merge_peak_memory_with_spilled_input(true, false, true).await
+}
+
+#[tokio::test]
+async fn test_sort_preserving_merge_peak_memory_with_spilled_input_no_round_robin_tied_values()
+-> Result<()> {
+    run_sort_preserving_merge_peak_memory_with_spilled_input(false, false, true).await
 }
 
 /// Intended to measure the maximum number of record batches held in memory by
@@ -343,6 +355,7 @@ async fn test_sort_preserving_merge_peak_memory_with_spilled_input_no_round_robi
 async fn run_sort_preserving_merge_peak_memory_with_spilled_input(
     round_robin: bool,
     multi_column_sort: bool,
+    tied_values: bool,
 ) -> Result<()> {
     let num_batches = 10usize;
     let num_rows_per_batch = 100usize;
@@ -362,15 +375,21 @@ async fn run_sort_preserving_merge_peak_memory_with_spilled_input(
 
     for stream_idx in 0..2usize {
         // Each stream covers a non-overlapping key range so both are individually
-        // sorted: stream 0 → [0, 1000), stream 1 → [1000, 2000).
+        // sorted: stream 0 → [0, 1000), stream 1 → [1000, 2000). When
+        // `tied_values` is set, every row of every batch in both streams
+        // instead carries the same sort key, so every comparison between the
+        // two streams is a tie.
         let batches: Vec<RecordBatch> = (0..num_batches)
             .map(|b| {
                 // Interleave streams: stream 0 → even slots [0,200,400,...],
                 // stream 1 → odd slots [100,300,500,...] so the merge
                 // alternates between them on every batch.
                 let base = ((b * 2 + stream_idx) * num_rows_per_batch) as i32;
-                let sort_col: Int32Array =
-                    (base..base + num_rows_per_batch as i32).collect();
+                let sort_col: Int32Array = if tied_values {
+                    std::iter::repeat_n(0, num_rows_per_batch).collect()
+                } else {
+                    (base..base + num_rows_per_batch as i32).collect()
+                };
                 let payload_col: StringArray =
                     std::iter::repeat_n(large_string.as_str(), num_rows_per_batch)
                         .map(Some)

From 4bf99eccfefd46b381d0acc54e623fb09dd7d694 Mon Sep 17 00:00:00 2001
From: Ariel Miculas <ariel.miculas@gmail.com>
Date: Tue, 4 Aug 2026 22:39:46 +0300
Subject: [PATCH 08/15] fix: use PeakRecordingPool to avoid missing
 unregistered memory consumers

---
 Cargo.lock                                    |  1 +
 datafusion/core/Cargo.toml                    |  1 +
 ...spilling_fuzz_in_memory_constrained_env.rs | 27 ++++++++++---------
 3 files changed, 17 insertions(+), 12 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 619fee7603629..0861eb8cb335d 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1715,6 +1715,7 @@ dependencies = [
  "criterion",
  "ctor",
  "dashmap",
+ "datafusion-benchmarks",
  "datafusion-catalog",
  "datafusion-catalog-listing",
  "datafusion-common",
diff --git a/datafusion/core/Cargo.toml b/datafusion/core/Cargo.toml
index 8679dad9f9a32..bfa257830e458 100644
--- a/datafusion/core/Cargo.toml
+++ b/datafusion/core/Cargo.toml
@@ -163,6 +163,7 @@ zstd = { workspace = true, optional = true }
 async-trait = { workspace = true }
 criterion = { workspace = true, features = ["async_tokio", "async_futures"] }
 ctor = { workspace = true }
+datafusion-benchmarks = { path = "../../benchmarks" }
 dashmap = "6.2.1"
 datafusion-doc = { workspace = true }
 datafusion-functions-window-common = { workspace = true }
diff --git a/datafusion/core/tests/fuzz_cases/spilling_fuzz_in_memory_constrained_env.rs b/datafusion/core/tests/fuzz_cases/spilling_fuzz_in_memory_constrained_env.rs
index 708dcac26d884..93d2dc7e5b4a5 100644
--- a/datafusion/core/tests/fuzz_cases/spilling_fuzz_in_memory_constrained_env.rs
+++ b/datafusion/core/tests/fuzz_cases/spilling_fuzz_in_memory_constrained_env.rs
@@ -53,12 +53,10 @@ use futures::StreamExt;
 use arrow::array::Int32Array;
 use datafusion::datasource::memory::MemorySourceConfig;
 use datafusion::physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec;
-use datafusion_execution::memory_pool::{
-    MemoryPool, TrackConsumersPool, UnboundedMemoryPool,
-};
+use datafusion_benchmarks::util::PeakRecordingPool;
+use datafusion_execution::memory_pool::{MemoryPool, UnboundedMemoryPool};
 use datafusion_physical_plan::metrics::{ExecutionPlanMetricsSet, SpillMetrics};
 use datafusion_physical_plan::spill::SpillManager;
-use std::num::NonZeroUsize;
 
 #[tokio::test]
 async fn test_sort_with_limited_memory() -> Result<()> {
@@ -486,12 +484,15 @@ async fn run_sort_preserving_merge_peak_memory_with_spilled_input(
             .with_round_robin_repartition(round_robin),
     );
 
-    // TrackConsumersPool wraps an unbounded pool so the merge never OOMs;
-    // we keep the typed Arc to call .metrics() after the run.
-    let tracking_pool = Arc::new(TrackConsumersPool::new(
+    // PeakRecordingPool records peak reserved bytes as a running high-water mark
+    // (via grow/shrink deltas), independent of any per-consumer registration
+    // bookkeeping - unlike TrackConsumersPool, whose tracked-consumer entry (and
+    // its peak) gets discarded the moment the consumer unregisters, which now
+    // happens mid-poll (inside the drain loop below) rather than when the
+    // caller eventually drops the returned stream.
+    let tracking_pool = Arc::new(PeakRecordingPool::new(Arc::new(
         UnboundedMemoryPool::default(),
-        NonZeroUsize::new(10).unwrap(),
-    ));
+    )));
     let runtime = RuntimeEnvBuilder::new()
         .with_memory_pool(Arc::clone(&tracking_pool) as Arc<dyn MemoryPool>)
         .build()?;
@@ -508,9 +509,7 @@ async fn run_sort_preserving_merge_peak_memory_with_spilled_input(
     }
     assert_eq!(total_rows, 2 * num_batches * num_rows_per_batch);
 
-    let mut metrics = tracking_pool.metrics();
-    metrics.sort_by_key(|m| std::cmp::Reverse(m.peak));
-    let peak_bytes: usize = metrics.iter().map(|m| m.peak).sum();
+    let peak_bytes = tracking_pool.peak_reserved();
 
     // in the single column case, the cursor takes up an ipc_batch_size worth of memory due to the
     // IPC roundtrip issue
@@ -532,6 +531,10 @@ async fn run_sort_preserving_merge_peak_memory_with_spilled_input(
         max_peak += 2 * cursor_unit;
     };
 
+    assert!(
+        peak_bytes > 0,
+        "peak reservation {peak_bytes} should be greater than 0"
+    );
     assert!(
         peak_bytes <= max_peak,
         "peak reservation {peak_bytes} bytes exceeds max_peak ({max_peak} bytes); \

From fa82f5a63f6fa79ce2682da2f4b143895105c216 Mon Sep 17 00:00:00 2001
From: Ariel Miculas <ariel.miculas@gmail.com>
Date: Tue, 4 Aug 2026 23:09:12 +0300
Subject: [PATCH 09/15] refactor: move PeakRecordingPool to
 datafusion-execution

---
 .../execution/src/memory_pool/peak_recording.rs                   | 0
 1 file changed, 0 insertions(+), 0 deletions(-)
 rename benchmarks/src/util/memory_pool.rs => datafusion/execution/src/memory_pool/peak_recording.rs (100%)

diff --git a/benchmarks/src/util/memory_pool.rs b/datafusion/execution/src/memory_pool/peak_recording.rs
similarity index 100%
rename from benchmarks/src/util/memory_pool.rs
rename to datafusion/execution/src/memory_pool/peak_recording.rs

From 4aaf8dd65c4c16c11a3cd349e1e4bfab26198487 Mon Sep 17 00:00:00 2001
From: Ariel Miculas <ariel.miculas@gmail.com>
Date: Tue, 4 Aug 2026 23:23:17 +0300
Subject: [PATCH 10/15] refactor: avoid circular dependency

---
 Cargo.lock                                                | 1 -
 benchmarks/Cargo.toml                                     | 1 +
 benchmarks/src/bin/external_aggr.rs                       | 6 ++----
 benchmarks/src/util/memory.rs                             | 4 +---
 benchmarks/src/util/mod.rs                                | 2 --
 benchmarks/src/util/options.rs                            | 7 +++++--
 benchmarks/src/util/run.rs                                | 3 +--
 datafusion/core/Cargo.toml                                | 1 -
 .../fuzz_cases/spilling_fuzz_in_memory_constrained_env.rs | 5 +++--
 datafusion/execution/src/memory_pool/mod.rs               | 2 ++
 datafusion/execution/src/memory_pool/peak_recording.rs    | 8 +++-----
 11 files changed, 18 insertions(+), 22 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 0861eb8cb335d..619fee7603629 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1715,7 +1715,6 @@ dependencies = [
  "criterion",
  "ctor",
  "dashmap",
- "datafusion-benchmarks",
  "datafusion-catalog",
  "datafusion-catalog-listing",
  "datafusion-common",
diff --git a/benchmarks/Cargo.toml b/benchmarks/Cargo.toml
index 282b27e48101d..9a27310143a50 100644
--- a/benchmarks/Cargo.toml
+++ b/benchmarks/Cargo.toml
@@ -48,6 +48,7 @@ criterion = { workspace = true, features = ["html_reports"] }
 datafusion = { workspace = true, default-features = true }
 datafusion-common = { workspace = true, default-features = true }
 datafusion-common-runtime = { workspace = true }
+datafusion-execution = { workspace = true }
 env_logger = { workspace = true }
 futures = { workspace = true }
 libmimalloc-sys = { version = "0.1", optional = true }
diff --git a/benchmarks/src/bin/external_aggr.rs b/benchmarks/src/bin/external_aggr.rs
index 226a619192ac9..c19554eb33583 100644
--- a/benchmarks/src/bin/external_aggr.rs
+++ b/benchmarks/src/bin/external_aggr.rs
@@ -34,14 +34,12 @@ use datafusion::datasource::listing::{
 use datafusion::datasource::{MemTable, TableProvider};
 use datafusion::error::Result;
 use datafusion::execution::SessionStateBuilder;
-use datafusion::execution::memory_pool::FairSpillPool;
+use datafusion::execution::memory_pool::{FairSpillPool, PeakRecordingPool};
 use datafusion::execution::runtime_env::RuntimeEnvBuilder;
 use datafusion::physical_plan::display::DisplayableExecutionPlan;
 use datafusion::physical_plan::{collect, displayable};
 use datafusion::prelude::*;
-use datafusion_benchmarks::util::{
-    BenchmarkRun, CommonOpt, PeakRecordingPool, QueryResult,
-};
+use datafusion_benchmarks::util::{BenchmarkRun, CommonOpt, QueryResult};
 use datafusion_common::instant::Instant;
 use datafusion_common::utils::get_available_parallelism;
 use datafusion_common::{DEFAULT_PARQUET_EXTENSION, exec_err};
diff --git a/benchmarks/src/util/memory.rs b/benchmarks/src/util/memory.rs
index 2b186c79c3516..f0339c9cf0c95 100644
--- a/benchmarks/src/util/memory.rs
+++ b/benchmarks/src/util/memory.rs
@@ -15,9 +15,7 @@
 // specific language governing permissions and limitations
 // under the License.
 
-use datafusion::execution::memory_pool::MemoryPool;
-
-use super::PeakRecordingPool;
+use datafusion::execution::memory_pool::{MemoryPool, PeakRecordingPool};
 
 /// Print Peak RSS, Peak Commit, Page Faults based on mimalloc api, followed by
 /// the peak reservation of `memory_pool` when a memory limit was configured.
diff --git a/benchmarks/src/util/mod.rs b/benchmarks/src/util/mod.rs
index 43855ea468ef5..6dc11c0f425bd 100644
--- a/benchmarks/src/util/mod.rs
+++ b/benchmarks/src/util/mod.rs
@@ -18,11 +18,9 @@
 //! Shared benchmark utilities
 pub mod latency_object_store;
 mod memory;
-mod memory_pool;
 mod options;
 mod run;
 
 pub use memory::print_memory_stats;
-pub use memory_pool::PeakRecordingPool;
 pub use options::CommonOpt;
 pub use run::{BenchQuery, BenchmarkRun, QueryResult};
diff --git a/benchmarks/src/util/options.rs b/benchmarks/src/util/options.rs
index c744d0bf31c7f..4a1c14674a1d0 100644
--- a/benchmarks/src/util/options.rs
+++ b/benchmarks/src/util/options.rs
@@ -21,7 +21,10 @@ use clap::Args;
 use datafusion::{
     execution::{
         disk_manager::DiskManagerBuilder,
-        memory_pool::{FairSpillPool, GreedyMemoryPool, MemoryPool, TrackConsumersPool},
+        memory_pool::{
+            FairSpillPool, GreedyMemoryPool, MemoryPool, PeakRecordingPool,
+            TrackConsumersPool,
+        },
         object_store::ObjectStoreUrl,
         runtime_env::{RuntimeEnv, RuntimeEnvBuilder},
     },
@@ -30,7 +33,7 @@ use datafusion::{
 use datafusion_common::{DataFusionError, Result};
 use object_store::local::LocalFileSystem;
 
-use super::{latency_object_store::LatencyObjectStore, memory_pool::PeakRecordingPool};
+use super::latency_object_store::LatencyObjectStore;
 
 // Common benchmark options (don't use doc comments otherwise this doc
 // shows up in help files)
diff --git a/benchmarks/src/util/run.rs b/benchmarks/src/util/run.rs
index 6c63ceec6423c..772d421bc7bf4 100644
--- a/benchmarks/src/util/run.rs
+++ b/benchmarks/src/util/run.rs
@@ -15,8 +15,7 @@
 // specific language governing permissions and limitations
 // under the License.
 
-use super::memory_pool::PeakRecordingPool;
-use datafusion::execution::memory_pool::MemoryPool;
+use datafusion::execution::memory_pool::{MemoryPool, PeakRecordingPool};
 use datafusion::{DATAFUSION_VERSION, error::Result};
 use datafusion_common::utils::get_available_parallelism;
 use serde::{Serialize, Serializer};
diff --git a/datafusion/core/Cargo.toml b/datafusion/core/Cargo.toml
index bfa257830e458..8679dad9f9a32 100644
--- a/datafusion/core/Cargo.toml
+++ b/datafusion/core/Cargo.toml
@@ -163,7 +163,6 @@ zstd = { workspace = true, optional = true }
 async-trait = { workspace = true }
 criterion = { workspace = true, features = ["async_tokio", "async_futures"] }
 ctor = { workspace = true }
-datafusion-benchmarks = { path = "../../benchmarks" }
 dashmap = "6.2.1"
 datafusion-doc = { workspace = true }
 datafusion-functions-window-common = { workspace = true }
diff --git a/datafusion/core/tests/fuzz_cases/spilling_fuzz_in_memory_constrained_env.rs b/datafusion/core/tests/fuzz_cases/spilling_fuzz_in_memory_constrained_env.rs
index 93d2dc7e5b4a5..f82d0165f2fdb 100644
--- a/datafusion/core/tests/fuzz_cases/spilling_fuzz_in_memory_constrained_env.rs
+++ b/datafusion/core/tests/fuzz_cases/spilling_fuzz_in_memory_constrained_env.rs
@@ -53,8 +53,9 @@ use futures::StreamExt;
 use arrow::array::Int32Array;
 use datafusion::datasource::memory::MemorySourceConfig;
 use datafusion::physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec;
-use datafusion_benchmarks::util::PeakRecordingPool;
-use datafusion_execution::memory_pool::{MemoryPool, UnboundedMemoryPool};
+use datafusion_execution::memory_pool::{
+    MemoryPool, PeakRecordingPool, UnboundedMemoryPool,
+};
 use datafusion_physical_plan::metrics::{ExecutionPlanMetricsSet, SpillMetrics};
 use datafusion_physical_plan::spill::SpillManager;
 
diff --git a/datafusion/execution/src/memory_pool/mod.rs b/datafusion/execution/src/memory_pool/mod.rs
index 2b36ee7f40add..40a79d136b84e 100644
--- a/datafusion/execution/src/memory_pool/mod.rs
+++ b/datafusion/execution/src/memory_pool/mod.rs
@@ -24,6 +24,7 @@ use std::fmt::Display;
 use std::hash::{Hash, Hasher};
 use std::{cmp::Ordering, sync::Arc, sync::atomic};
 
+mod peak_recording;
 mod pool;
 
 #[cfg(feature = "arrow_buffer_pool")]
@@ -36,6 +37,7 @@ pub mod proxy {
 pub use datafusion_common::{
     human_readable_count, human_readable_duration, human_readable_size, units,
 };
+pub use peak_recording::*;
 pub use pool::*;
 
 /// Tracks and potentially limits memory use across operators during execution.
diff --git a/datafusion/execution/src/memory_pool/peak_recording.rs b/datafusion/execution/src/memory_pool/peak_recording.rs
index a3606ca0a7b7a..5da8787787da6 100644
--- a/datafusion/execution/src/memory_pool/peak_recording.rs
+++ b/datafusion/execution/src/memory_pool/peak_recording.rs
@@ -49,9 +49,7 @@ use std::{
     },
 };
 
-use datafusion::execution::memory_pool::{
-    MemoryConsumer, MemoryLimit, MemoryPool, MemoryReservation,
-};
+use super::{MemoryConsumer, MemoryLimit, MemoryPool, MemoryReservation};
 use datafusion_common::Result;
 
 /// Wraps a [`MemoryPool`], recording the high-water mark of
@@ -227,7 +225,7 @@ impl MemoryPool for PeakRecordingPool {
 
 #[cfg(test)]
 mod tests {
-    use datafusion::execution::memory_pool::GreedyMemoryPool;
+    use crate::memory_pool::GreedyMemoryPool;
 
     use super::*;
 
@@ -360,8 +358,8 @@ mod tests {
     /// This test pins that.
     #[test]
     fn records_reservations_arriving_through_the_arrow_adapter() {
+        use crate::memory_pool::arrow::ArrowMemoryPool;
         use arrow_buffer::MemoryPool as ArrowMemoryPoolTrait;
-        use datafusion_execution::memory_pool::arrow::ArrowMemoryPool;
 
         let (recording, pool) = pool(4096);
 

From a5f0136ad06bfde3695a7f366dca7e30647a1f49 Mon Sep 17 00:00:00 2001
From: Ariel Miculas <ariel.miculas@gmail.com>
Date: Tue, 4 Aug 2026 23:34:49 +0300
Subject: [PATCH 11/15] refactor: remove unused dependencies

---
 benchmarks/Cargo.toml | 5 -----
 1 file changed, 5 deletions(-)

diff --git a/benchmarks/Cargo.toml b/benchmarks/Cargo.toml
index 9a27310143a50..49a4dc882f4a9 100644
--- a/benchmarks/Cargo.toml
+++ b/benchmarks/Cargo.toml
@@ -65,11 +65,6 @@ tokio = { workspace = true, features = ["rt-multi-thread", "parking_lot"] }
 tokio-util = { version = "0.7.17" }
 
 [dev-dependencies]
-# `pool`/`arrow_buffer_pool` are enabled only for tests, so the benchmark
-# binaries are built exactly as before. They let `memory_pool`'s tests cover
-# Arrow-side reservations reaching the pool via `ArrowMemoryPool`.
-arrow-buffer = { workspace = true, features = ["pool"] }
-datafusion-execution = { workspace = true, features = ["arrow_buffer_pool"] }
 datafusion-proto = { workspace = true, features = ["parquet"] }
 tempfile = { workspace = true }
 

From c1d3bf59517e5fe77580df4c767ce7bbda984272 Mon Sep 17 00:00:00 2001
From: Ariel Miculas <ariel.miculas@gmail.com>
Date: Tue, 4 Aug 2026 23:45:57 +0300
Subject: [PATCH 12/15] refactor: remove unused dependency and regenerate
 cargo.lock

---
 Cargo.lock            | 2 --
 benchmarks/Cargo.toml | 1 -
 2 files changed, 3 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 619fee7603629..9666352539111 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1780,7 +1780,6 @@ name = "datafusion-benchmarks"
 version = "54.1.0"
 dependencies = [
  "arrow",
- "arrow-buffer",
  "async-trait",
  "bytes",
  "clap",
@@ -1788,7 +1787,6 @@ dependencies = [
  "datafusion",
  "datafusion-common",
  "datafusion-common-runtime",
- "datafusion-execution",
  "datafusion-proto",
  "env_logger",
  "futures",
diff --git a/benchmarks/Cargo.toml b/benchmarks/Cargo.toml
index 49a4dc882f4a9..5dae70761f9a7 100644
--- a/benchmarks/Cargo.toml
+++ b/benchmarks/Cargo.toml
@@ -48,7 +48,6 @@ criterion = { workspace = true, features = ["html_reports"] }
 datafusion = { workspace = true, default-features = true }
 datafusion-common = { workspace = true, default-features = true }
 datafusion-common-runtime = { workspace = true }
-datafusion-execution = { workspace = true }
 env_logger = { workspace = true }
 futures = { workspace = true }
 libmimalloc-sys = { version = "0.1", optional = true }

From 79bda6795005b3d8c52e497ce45248d2d38b8eb8 Mon Sep 17 00:00:00 2001
From: Ariel Miculas <ariel.miculas@gmail.com>
Date: Tue, 4 Aug 2026 23:49:37 +0300
Subject: [PATCH 13/15] refactor: enable arrow_buffer_pool feature for the
 arrow memory pool test

---
 Cargo.lock                      | 1 +
 datafusion/execution/Cargo.toml | 4 ++++
 2 files changed, 5 insertions(+)

diff --git a/Cargo.lock b/Cargo.lock
index 9666352539111..75f785af4516e 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -2134,6 +2134,7 @@ dependencies = [
  "chrono",
  "dashmap",
  "datafusion-common",
+ "datafusion-execution",
  "datafusion-expr",
  "datafusion-physical-expr-common",
  "futures",
diff --git a/datafusion/execution/Cargo.toml b/datafusion/execution/Cargo.toml
index c9d4acd3644ba..0d2dff8d4f990 100644
--- a/datafusion/execution/Cargo.toml
+++ b/datafusion/execution/Cargo.toml
@@ -75,5 +75,9 @@ url = { workspace = true }
 tokio = { workspace = true, features = ["fs"] }
 
 [dev-dependencies]
+# Enables `arrow_buffer_pool` for this crate's own tests only (e.g. the
+# ArrowMemoryPool-adapter test in peak_recording.rs), without turning it on
+# as a default feature for consumers of this crate.
 chrono = { workspace = true }
+datafusion-execution = { path = ".", features = ["arrow_buffer_pool"] }
 insta = { workspace = true }

From 4e5ba442f3bce8a4dcf879b5d88820f6521781a1 Mon Sep 17 00:00:00 2001
From: Ariel Miculas <ariel.miculas@gmail.com>
Date: Wed, 5 Aug 2026 00:02:58 +0300
Subject: [PATCH 14/15] refactor: cannot have circular dependencies so gate via
 cfg instead

---
 Cargo.lock                                             | 1 -
 datafusion/execution/Cargo.toml                        | 4 ----
 datafusion/execution/src/memory_pool/peak_recording.rs | 5 +++++
 3 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 75f785af4516e..9666352539111 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -2134,7 +2134,6 @@ dependencies = [
  "chrono",
  "dashmap",
  "datafusion-common",
- "datafusion-execution",
  "datafusion-expr",
  "datafusion-physical-expr-common",
  "futures",
diff --git a/datafusion/execution/Cargo.toml b/datafusion/execution/Cargo.toml
index 0d2dff8d4f990..c9d4acd3644ba 100644
--- a/datafusion/execution/Cargo.toml
+++ b/datafusion/execution/Cargo.toml
@@ -75,9 +75,5 @@ url = { workspace = true }
 tokio = { workspace = true, features = ["fs"] }
 
 [dev-dependencies]
-# Enables `arrow_buffer_pool` for this crate's own tests only (e.g. the
-# ArrowMemoryPool-adapter test in peak_recording.rs), without turning it on
-# as a default feature for consumers of this crate.
 chrono = { workspace = true }
-datafusion-execution = { path = ".", features = ["arrow_buffer_pool"] }
 insta = { workspace = true }
diff --git a/datafusion/execution/src/memory_pool/peak_recording.rs b/datafusion/execution/src/memory_pool/peak_recording.rs
index 5da8787787da6..fd02d2d13eba6 100644
--- a/datafusion/execution/src/memory_pool/peak_recording.rs
+++ b/datafusion/execution/src/memory_pool/peak_recording.rs
@@ -356,6 +356,11 @@ mod tests {
     /// bytes show up in this peak without further changes — as long as the
     /// adapter is built from the `RuntimeEnv`'s pool, which is the wrapped one.
     /// This test pins that.
+    ///
+    /// Only compiled with `--features arrow_buffer_pool`, since that's what
+    /// gates `crate::memory_pool::arrow` and `arrow_buffer::MemoryPool` in the
+    /// first place; not part of this crate's default feature set.
+    #[cfg(feature = "arrow_buffer_pool")]
     #[test]
     fn records_reservations_arriving_through_the_arrow_adapter() {
         use crate::memory_pool::arrow::ArrowMemoryPool;

From 468f0da063a2c34262f01e89ebcab44afaa95da5 Mon Sep 17 00:00:00 2001
From: Ariel Miculas <ariel.miculas@gmail.com>
Date: Wed, 5 Aug 2026 11:01:05 +0300
Subject: [PATCH 15/15] fix: rust doctests

---
 .../execution/src/memory_pool/peak_recording.rs   | 15 ++++-----------
 1 file changed, 4 insertions(+), 11 deletions(-)

diff --git a/datafusion/execution/src/memory_pool/peak_recording.rs b/datafusion/execution/src/memory_pool/peak_recording.rs
index fd02d2d13eba6..b407cc0eaf36b 100644
--- a/datafusion/execution/src/memory_pool/peak_recording.rs
+++ b/datafusion/execution/src/memory_pool/peak_recording.rs
@@ -27,7 +27,7 @@
 //! itself is never recorded — [`MemoryPool::reserved`] is a live value that has
 //! usually fallen back to zero by the time a query finishes. This module records
 //! the high-water mark so benchmarks can emit it alongside the peak RSS that
-//! [`print_memory_stats`] already prints, making the gap between the two
+//! `print_memory_stats` already prints, making the gap between the two
 //! measurable.
 //!
 //! This is measurement only: nothing here enforces a relationship between the
@@ -38,8 +38,6 @@
 //! through `ArrowMemoryPool` are included, because that adapter grows a
 //! DataFusion reservation against the pool it wraps; nothing claims buffers
 //! today, but the peak picks it up when something does.
-//!
-//! [`print_memory_stats`]: super::print_memory_stats
 
 use std::{
     fmt::{Debug, Display, Formatter},
@@ -69,8 +67,7 @@ use datafusion_common::Result;
 ///
 /// ```
 /// # use std::sync::Arc;
-/// # use datafusion::execution::memory_pool::{GreedyMemoryPool, MemoryConsumer, MemoryPool};
-/// # use datafusion_benchmarks::util::PeakRecordingPool;
+/// # use datafusion_execution::memory_pool::{GreedyMemoryPool, MemoryConsumer, MemoryPool, PeakRecordingPool};
 /// let recording = Arc::new(PeakRecordingPool::new(Arc::new(GreedyMemoryPool::new(1024))));
 /// let pool: Arc<dyn MemoryPool> = Arc::clone(&recording) as _;
 ///
@@ -114,10 +111,8 @@ impl PeakRecordingPool {
     /// The recorder installed as `pool`, if there is one.
     ///
     /// Returns `None` whenever a benchmark runs without a memory limit, since
-    /// [`CommonOpt::runtime_env_builder`] only installs the wrapper alongside a
+    /// `CommonOpt::runtime_env_builder` only installs the wrapper alongside a
     /// pool it has a limit for.
-    ///
-    /// [`CommonOpt::runtime_env_builder`]: super::CommonOpt::runtime_env_builder
     pub fn from_pool(pool: &dyn MemoryPool) -> Option<&Self> {
         pool.downcast_ref::<Self>()
     }
@@ -138,12 +133,10 @@ impl PeakRecordingPool {
     /// Reset the value returned by [`Self::peak_reserved`] to what is reserved
     /// right now, so the next reading covers only what follows.
     ///
-    /// [`BenchmarkRun::start_new_case`] calls this, giving each benchmark query
+    /// `BenchmarkRun::start_new_case` calls this, giving each benchmark query
     /// its own reading. Anything still held when a query starts — data the
     /// benchmark loaded up front, say — stays in the reading, since the query
     /// runs with those bytes reserved.
-    ///
-    /// [`BenchmarkRun::start_new_case`]: super::BenchmarkRun::start_new_case
     pub fn reset_peak(&self) {
         self.peak
             .store(self.reserved.load(Ordering::Relaxed), Ordering::Relaxed);