refactor: scope SPM prefetch buffer to sort elimination path only

Add prefetch field to SortPreservingMergeExec (default 1) instead of
hardcoding buffer=16 globally. PushdownSort sets prefetch=16 only when
eliminating SortExec under SPM, avoiding impact on unrelated queries.

Also detect SPM → SortExec(preserve_partitioning) pattern directly in
PushdownSort to handle sort elimination and prefetch in one pass.

Author: zhuqi-lucas

datafusion/physical-optimizer/src/pushdown_sort.rs

@@ -61,8 +61,13 @@ use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode};
 use datafusion_physical_plan::ExecutionPlan;
 use datafusion_physical_plan::SortOrderPushdownResult;
 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;
+
 /// A PhysicalOptimizerRule that attempts to push down sort requirements to data sources.
 ///
 /// See module-level documentation for details.
@@ -87,8 +92,59 @@ 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
         plan.transform_down(|plan: Arc<dyn ExecutionPlan>| {
-            // Check if this is a SortExec
+            // 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.
+            if let Some(spm) = plan.downcast_ref::<SortPreservingMergeExec>()
+                && let Some(sort_child) = spm.input().downcast_ref::<SortExec>()
+                && sort_child.preserve_partitioning()
+            {
+                let sort_input = Arc::clone(sort_child.input());
+                let required_ordering = sort_child.expr();
+                match sort_input.try_pushdown_sort(required_ordering)? {
+                    SortOrderPushdownResult::Exact { inner } => {
+                        let inner = if let Some(fetch) = sort_child.fetch() {
+                            inner.with_fetch(Some(fetch)).unwrap_or(inner)
+                        } else {
+                            inner
+                        };
+                        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);
+                        return Ok(Transformed::yes(Arc::new(new_spm)));
+                    }
+                    SortOrderPushdownResult::Inexact { inner } => {
+                        let new_sort = SortExec::new(
+                            required_ordering.clone(),
+                            inner,
+                        )
+                        .with_fetch(sort_child.fetch())
+                        .with_preserve_partitioning(true);
+                        let new_spm = SortPreservingMergeExec::new(
+                            spm.expr().clone(),
+                            Arc::new(new_sort),
+                        )
+                        .with_fetch(spm.fetch())
+                        .with_round_robin_repartition(
+                            spm.enable_round_robin_repartition(),
+                        );
+                        return Ok(Transformed::yes(Arc::new(new_spm)));
+                    }
+                    SortOrderPushdownResult::Unsupported => {
+                        return Ok(Transformed::no(plan));
+                    }
+                }
+            }
+
+            // Pattern 2: Standalone SortExec (no SPM parent)
             let Some(sort_exec) = plan.downcast_ref::<SortExec>() else {
                 return Ok(Transformed::no(plan));
             };

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

@@ -100,6 +100,12 @@ 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 {
@@ -113,6 +119,7 @@ impl SortPreservingMergeExec {
             fetch: None,
             cache: Arc::new(cache),
             enable_round_robin_repartition: true,
+            prefetch: 1,
         }
     }
 
@@ -139,6 +146,21 @@ 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
@@ -154,6 +176,11 @@ 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(
@@ -250,7 +277,8 @@ impl ExecutionPlan for SortPreservingMergeExec {
             metrics: self.metrics.clone(),
             fetch: limit,
             cache: Arc::clone(&self.cache),
-            enable_round_robin_repartition: true,
+            enable_round_robin_repartition: self.enable_round_robin_repartition,
+            prefetch: self.prefetch,
         }))
     }
 
@@ -361,7 +389,7 @@ impl ExecutionPlan for SortPreservingMergeExec {
                     .map(|partition| {
                         let stream =
                             self.input.execute(partition, Arc::clone(&context))?;
-                        Ok(spawn_buffered(stream, 16))
+                        Ok(spawn_buffered(stream, self.prefetch))
                     })
                     .collect::<Result<_>>()?;