refactor: replace SPM prefetch with BufferExec for sort elimination

When PushdownSort eliminates SortExec under SortPreservingMergeExec,
insert BufferExec(8MB) to replace SortExec's buffering role. This is
strictly better: same I/O pipelining with bounded memory, no sort
computation, and no impact on other SPM usage.

Revert SPM prefetch field — SPM stays at spawn_buffered(1) for all
cases. BufferExec only inserted by PushdownSort when needed.

Author: zhuqi-lucas

datafusion/physical-optimizer/src/pushdown_sort.rs

@@ -60,13 +60,19 @@ use datafusion_common::config::ConfigOptions;
 use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode};
 use datafusion_physical_plan::ExecutionPlan;
 use datafusion_physical_plan::SortOrderPushdownResult;
+use datafusion_physical_plan::buffer::BufferExec;
 use datafusion_physical_plan::sorts::sort::SortExec;
 use datafusion_physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec;
 use std::sync::Arc;
 
-/// Prefetch buffer size for SortPreservingMergeExec when sort elimination
-/// removes the buffering SortExec between SPM and DataSourceExec.
-const SPM_PREFETCH_AFTER_SORT_ELIMINATION: usize = 16;
+/// Buffer capacity (in bytes) inserted between SPM and DataSourceExec when
+/// sort elimination removes the buffering SortExec.
+///
+/// SortExec buffers all input data in memory (with spill support) before
+/// outputting sorted results. When we eliminate SortExec, SPM reads directly
+/// from I/O-bound sources. BufferExec compensates by buffering batches in
+/// the background, allowing I/O to pipeline with merge computation.
+const BUFFER_CAPACITY_AFTER_SORT_ELIMINATION: usize = 8 * 1024 * 1024; // 8 MB
 
 /// A PhysicalOptimizerRule that attempts to push down sort requirements to data sources.
 ///
@@ -92,11 +98,11 @@ impl PhysicalOptimizerRule for PushdownSort {
         }
 
         // Use transform_down to find and optimize all SortExec nodes (including nested ones)
-        // Also handles SPM → SortExec pattern to set prefetch when sort is eliminated
+        // Also handles SPM → SortExec pattern to insert BufferExec when sort is eliminated
         plan.transform_down(|plan: Arc<dyn ExecutionPlan>| {
             // Pattern 1: SPM → SortExec(preserve_partitioning)
             // When we eliminate the SortExec, SPM loses its memory buffer and reads
-            // directly from I/O-bound sources. Set a larger prefetch to pipeline I/O.
+            // directly from I/O-bound sources. Insert a BufferExec to compensate.
             if let Some(spm) = plan.downcast_ref::<SortPreservingMergeExec>()
                 && let Some(sort_child) = spm.input().downcast_ref::<SortExec>()
                 && sort_child.preserve_partitioning()
@@ -110,13 +116,16 @@ impl PhysicalOptimizerRule for PushdownSort {
                         } else {
                             inner
                         };
+                        // Insert BufferExec to replace SortExec's buffering role.
+                        // SortExec buffered all data in memory; BufferExec provides
+                        // bounded buffering so SPM doesn't stall on I/O.
+                        let buffered: Arc<dyn ExecutionPlan> = Arc::new(BufferExec::new(
+                            inner,
+                            BUFFER_CAPACITY_AFTER_SORT_ELIMINATION,
+                        ));
                         let new_spm =
-                            SortPreservingMergeExec::new(spm.expr().clone(), inner)
-                                .with_fetch(spm.fetch())
-                                .with_round_robin_repartition(
-                                    spm.enable_round_robin_repartition(),
-                                )
-                                .with_prefetch(SPM_PREFETCH_AFTER_SORT_ELIMINATION);
+                            SortPreservingMergeExec::new(spm.expr().clone(), buffered)
+                                .with_fetch(spm.fetch());
                         return Ok(Transformed::yes(Arc::new(new_spm)));
                     }
                     SortOrderPushdownResult::Inexact { inner } => {
@@ -127,10 +136,7 @@ impl PhysicalOptimizerRule for PushdownSort {
                             spm.expr().clone(),
                             Arc::new(new_sort),
                         )
-                        .with_fetch(spm.fetch())
-                        .with_round_robin_repartition(
-                            spm.enable_round_robin_repartition(),
-                        );
+                        .with_fetch(spm.fetch());
                         return Ok(Transformed::yes(Arc::new(new_spm)));
                     }
                     SortOrderPushdownResult::Unsupported => {

datafusion/physical-plan/src/sorts/sort_preserving_merge.rs

@@ -100,12 +100,6 @@ pub struct SortPreservingMergeExec {
     ///
     /// See [`Self::with_round_robin_repartition`] for more information.
     enable_round_robin_repartition: bool,
-    /// Number of batches to prefetch from each input partition.
-    ///
-    /// When SPM reads directly from I/O-bound sources (e.g., after sort
-    /// elimination removes a buffering SortExec), a larger prefetch allows
-    /// pipelining I/O with merge computation. Defaults to 1.
-    prefetch: usize,
 }
 
 impl SortPreservingMergeExec {
@@ -119,7 +113,6 @@ impl SortPreservingMergeExec {
             fetch: None,
             cache: Arc::new(cache),
             enable_round_robin_repartition: true,
-            prefetch: 1,
         }
     }
 
@@ -146,21 +139,6 @@ impl SortPreservingMergeExec {
         self
     }
 
-    /// Sets the number of batches to prefetch from each input partition.
-    ///
-    /// A larger value allows pipelining I/O with merge computation, which
-    /// helps when inputs are I/O-bound (e.g., reading directly from
-    /// DataSourceExec without a buffering SortExec in between).
-    pub fn with_prefetch(mut self, prefetch: usize) -> Self {
-        self.prefetch = prefetch;
-        self
-    }
-
-    /// Returns the prefetch buffer size
-    pub fn prefetch(&self) -> usize {
-        self.prefetch
-    }
-
     /// Input schema
     pub fn input(&self) -> &Arc<dyn ExecutionPlan> {
         &self.input
@@ -176,11 +154,6 @@ impl SortPreservingMergeExec {
         self.fetch
     }
 
-    /// Whether round-robin repartition is enabled
-    pub fn enable_round_robin_repartition(&self) -> bool {
-        self.enable_round_robin_repartition
-    }
-
     /// Creates the cache object that stores the plan properties
     /// such as schema, equivalence properties, ordering, partitioning, etc.
     fn compute_properties(
@@ -278,7 +251,6 @@ impl ExecutionPlan for SortPreservingMergeExec {
             fetch: limit,
             cache: Arc::clone(&self.cache),
             enable_round_robin_repartition: self.enable_round_robin_repartition,
-            prefetch: self.prefetch,
         }))
     }
 
@@ -389,7 +361,7 @@ impl ExecutionPlan for SortPreservingMergeExec {
                     .map(|partition| {
                         let stream =
                             self.input.execute(partition, Arc::clone(&context))?;
-                        Ok(spawn_buffered(stream, self.prefetch))
+                        Ok(spawn_buffered(stream, 1))
                     })
                     .collect::<Result<_>>()?;
 

datafusion/sqllogictest/test_files/sort_pushdown.slt

@@ -2161,6 +2161,116 @@ DROP TABLE null_src_b;
 statement ok
 DROP TABLE tf_nulls_last;
 
+# ===========================================================
+# Test G: BufferExec insertion when sort elimination removes
+# SortExec under SortPreservingMergeExec.
+#
+# When PushdownSort eliminates SortExec(preserve_partitioning=true),
+# SPM loses SortExec's memory buffer. A BufferExec is inserted to
+# compensate, allowing I/O pipelining with merge computation.
+# ===========================================================
+
+# Create files with reversed naming: c_low has smallest values,
+# a_high has largest — alphabetical order ≠ sort key order.
+statement ok
+CREATE TABLE tg_src_low(id INT, value INT) AS VALUES (1, 100), (2, 200), (3, 300);
+
+statement ok
+CREATE TABLE tg_src_mid(id INT, value INT) AS VALUES (4, 400), (5, 500), (6, 600);
+
+statement ok
+CREATE TABLE tg_src_high(id INT, value INT) AS VALUES (7, 700), (8, 800), (9, 900), (10, 1000);
+
+query I
+COPY (SELECT * FROM tg_src_low ORDER BY id ASC)
+TO 'test_files/scratch/sort_pushdown/tg_buffer/c_low.parquet';
+----
+3
+
+query I
+COPY (SELECT * FROM tg_src_mid ORDER BY id ASC)
+TO 'test_files/scratch/sort_pushdown/tg_buffer/b_mid.parquet';
+----
+3
+
+query I
+COPY (SELECT * FROM tg_src_high ORDER BY id ASC)
+TO 'test_files/scratch/sort_pushdown/tg_buffer/a_high.parquet';
+----
+4
+
+# Use target_partitions=2 so files are split across 2 groups.
+# Files are in wrong alphabetical order → validated_output_ordering strips ordering
+# → EnforceSorting adds SortExec(preserve_partitioning) + SPM
+# → PushdownSort eliminates SortExec and inserts BufferExec
+statement ok
+SET datafusion.execution.target_partitions = 2;
+
+statement ok
+CREATE EXTERNAL TABLE tg_buffer(id INT, value INT)
+STORED AS PARQUET
+LOCATION 'test_files/scratch/sort_pushdown/tg_buffer/'
+WITH ORDER (id ASC);
+
+# Test G.1: BufferExec appears between SPM and DataSourceExec
+query TT
+EXPLAIN SELECT * FROM tg_buffer ORDER BY id ASC;
+----
+logical_plan
+01)Sort: tg_buffer.id ASC NULLS LAST
+02)--TableScan: tg_buffer projection=[id, value]
+physical_plan
+01)SortPreservingMergeExec: [id@0 ASC NULLS LAST]
+02)--BufferExec: capacity=8388608
+03)----DataSourceExec: file_groups={2 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/sort_pushdown/tg_buffer/b_mid.parquet, WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/sort_pushdown/tg_buffer/a_high.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/sort_pushdown/tg_buffer/c_low.parquet]]}, projection=[id, value], output_ordering=[id@0 ASC NULLS LAST], file_type=parquet
+
+# Verify correctness
+query II
+SELECT * FROM tg_buffer ORDER BY id ASC;
+----
+1 100
+2 200
+3 300
+4 400
+5 500
+6 600
+7 700
+8 800
+9 900
+10 1000
+
+# Test G.2: LIMIT query with BufferExec
+query TT
+EXPLAIN SELECT * FROM tg_buffer ORDER BY id ASC LIMIT 3;
+----
+logical_plan
+01)Sort: tg_buffer.id ASC NULLS LAST, fetch=3
+02)--TableScan: tg_buffer projection=[id, value]
+physical_plan
+01)SortPreservingMergeExec: [id@0 ASC NULLS LAST], fetch=3
+02)--BufferExec: capacity=8388608
+03)----DataSourceExec: file_groups={2 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/sort_pushdown/tg_buffer/b_mid.parquet, WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/sort_pushdown/tg_buffer/a_high.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/sort_pushdown/tg_buffer/c_low.parquet]]}, projection=[id, value], limit=3, output_ordering=[id@0 ASC NULLS LAST], file_type=parquet
+
+query II
+SELECT * FROM tg_buffer ORDER BY id ASC LIMIT 3;
+----
+1 100
+2 200
+3 300
+
+# Cleanup Test G
+statement ok
+DROP TABLE tg_src_low;
+
+statement ok
+DROP TABLE tg_src_mid;
+
+statement ok
+DROP TABLE tg_src_high;
+
+statement ok
+DROP TABLE tg_buffer;
+
 # Reset settings (SLT runner uses target_partitions=4, not system default)
 statement ok
 SET datafusion.execution.target_partitions = 4;