fix: reduce peak memory usage when round robin tiebreaker is disabled

benchmarks/Cargo.toml

@@ -64,11 +64,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 }
 

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};

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.

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};

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)

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};

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,9 +46,19 @@ 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;
 
+use arrow::array::Int32Array;
+use datafusion::datasource::memory::MemorySourceConfig;
+use datafusion::physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec;
+use datafusion_execution::memory_pool::{
+    MemoryPool, PeakRecordingPool, UnboundedMemoryPool,
+};
+use datafusion_physical_plan::metrics::{ExecutionPlanMetricsSet, SpillMetrics};
+use datafusion_physical_plan::spill::SpillManager;
+
 #[tokio::test]
 async fn test_sort_with_limited_memory() -> Result<()> {
     let record_batch_size = 8192;
@@ -290,6 +301,250 @@ 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, 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, 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, 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, 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
+/// 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,
+    multi_column_sort: bool,
+    tied_values: 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). 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 = 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)
+                        .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?);
+        }
+        partition_batches.push(ipc_batches);
+    }
+
+    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_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(sort_exprs).unwrap(), input)
+            .with_round_robin_repartition(round_robin),
+    );
+
+    // 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(),
+    )));
+    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 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
+    // 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)
+    // 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 > 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); \
+         round_robin={round_robin}, multi_column_sort={multi_column_sort}",
+    );
+
+    Ok(())
+}
+
 struct RunTestWithLimitedMemoryArgs {
     pool_size: usize,
     task_ctx: Arc<TaskContext>,

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.