fix: address review feedback — avoid stale SPM and decouple spm capture from fetch

- Capture only the LexOrdering from the stripped SortPreservingMergeExec
  instead of the entire plan node, then reconstruct a fresh SPM with the
  current (possibly rewritten) child plan in the fallback path. This
  avoids reusing a stale SPM that references an outdated subtree.
- Decouple spm_ordering capture from fetch — use `spm_ordering.is_none()`
  instead of gating on `fetch.is_none()`, so the SPM ordering is recorded
  even when an outer operator already set fetch.
- Remove unwrap() on with_fetch by constructing SPM directly.
- Fix test comment: fallback is CoalescePartitionsExec/SPM, not GlobalLimitExec.
- Add regression test: nested CoalescePartitionsExec(fetch=5) over
  SortPreservingMergeExec(fetch=3) now correctly preserves SPM ordering.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: zhuqi-lucas

datafusion/core/tests/physical_optimizer/enforce_distribution.rs

@@ -3934,7 +3934,7 @@ fn coalesce_partitions_fetch_preserved_by_enforce_distribution() -> Result<()> {
 
     // The fetch=1 must survive. It can appear either as:
     // - CoalescePartitionsExec: fetch=1  (re-inserted with fetch), or
-    // - GlobalLimitExec: skip=0, fetch=1 (fallback when merge wasn't re-added)
+    // - SortPreservingMergeExec: fetch=1 (when a merge is re-added with fetch)
     let plan_str = displayable(result.as_ref()).indent(true).to_string();
     assert!(
         plan_str.contains("fetch=1"),
@@ -3994,6 +3994,51 @@ fn spm_fetch_preserved_by_enforce_distribution() -> Result<()> {
     Ok(())
 }
 
+/// When both a `CoalescePartitionsExec(fetch=N)` and an inner
+/// `SortPreservingMergeExec(fetch=M)` are stripped, the SPM ordering must
+/// still be captured even though fetch was already set when the SPM is
+/// encountered. The reconstructed fallback should use a
+/// `SortPreservingMergeExec` (not a plain `CoalescePartitionsExec`) to
+/// preserve sort semantics, and must wrap the *rewritten* child plan.
+///
+/// Regression test for: the `spm` capture was gated on `fetch.is_none()`,
+/// causing it to be skipped when an outer operator already set `fetch`.
+#[test]
+fn nested_coalesce_over_spm_preserves_spm_ordering() -> Result<()> {
+    let schema = schema();
+    let sort_key: LexOrdering = [PhysicalSortExpr {
+        expr: col("c", &schema)?,
+        options: SortOptions::default(),
+    }]
+    .into();
+
+    // Build: CoalescePartitionsExec(fetch=5)
+    //          -> SortPreservingMergeExec(fetch=3, [c ASC])
+    //               -> sorted multi-partition parquet
+    let parquet = parquet_exec_multiple_sorted(vec![sort_key.clone()]);
+    let spm: Arc<dyn ExecutionPlan> = Arc::new(
+        SortPreservingMergeExec::new(sort_key, parquet).with_fetch(Some(3)),
+    );
+    let coalesce_over_spm: Arc<dyn ExecutionPlan> =
+        Arc::new(CoalescePartitionsExec::new(spm).with_fetch(Some(5)));
+
+    let result = ensure_distribution_helper(coalesce_over_spm, 10, false)?;
+
+    let plan_str = displayable(result.as_ref()).indent(true).to_string();
+    // The minimum fetch (min(5,3)=3) must survive.
+    assert!(
+        plan_str.contains("fetch=3"),
+        "fetch=3 was lost after EnforceDistribution!\nPlan:\n{plan_str}"
+    );
+    // The result should use SortPreservingMergeExec (not CoalescePartitionsExec)
+    // to preserve the ordering semantics.
+    assert!(
+        plan_str.contains("SortPreservingMergeExec"),
+        "Expected SortPreservingMergeExec to preserve ordering, but got:\n{plan_str}"
+    );
+    Ok(())
+}
+
 /// When a parent requires SinglePartition and maintains input order, order-preserving
 /// variants (e.g. SortPreservingMergeExec) should be kept so that ordering can
 /// propagate to ancestors. Replacing them with CoalescePartitionsExec would destroy

datafusion/physical-optimizer/src/enforce_distribution.rs

@@ -40,7 +40,8 @@ use datafusion_expr::logical_plan::{Aggregate, JoinType};
 use datafusion_physical_expr::expressions::{Column, NoOp};
 use datafusion_physical_expr::utils::map_columns_before_projection;
 use datafusion_physical_expr::{
-    EquivalenceProperties, PhysicalExpr, PhysicalExprRef, physical_exprs_equal,
+    EquivalenceProperties, LexOrdering, PhysicalExpr, PhysicalExprRef,
+    physical_exprs_equal,
 };
 use datafusion_physical_plan::ExecutionPlanProperties;
 use datafusion_physical_plan::aggregates::{
@@ -1023,22 +1024,32 @@ fn remove_dist_changing_operators(
 ) -> Result<(
     DistributionContext,
     Option<usize>,
-    Option<Arc<dyn ExecutionPlan>>,
+    Option<LexOrdering>,
 )> {
     let mut fetch = None;
-    let mut spm: Option<Arc<dyn ExecutionPlan>> = None;
+    let mut spm_ordering: Option<LexOrdering> = None;
     while is_repartition(&distribution_context.plan)
         || is_coalesce_partitions(&distribution_context.plan)
         || is_sort_preserving_merge(&distribution_context.plan)
     {
+        // Track whether the stripped operator was a SortPreservingMergeExec,
+        // independently of whether it carries a fetch. We only need the
+        // ordering so we can reconstruct a fresh SPM later if needed.
+        if is_sort_preserving_merge(&distribution_context.plan) && spm_ordering.is_none()
+        {
+            if let Some(spm) = distribution_context
+                .plan
+                .as_any()
+                .downcast_ref::<SortPreservingMergeExec>()
+            {
+                spm_ordering = Some(spm.expr().clone());
+            }
+        }
         // Preserve any `fetch` (limit) that was pushed into a
         // `SortPreservingMergeExec` or `CoalescePartitionsExec` by
         // `LimitPushdown`. Without this, the limit would be lost when
         // the operator is stripped.
         if let Some(child_fetch) = distribution_context.plan.fetch() {
-            if is_sort_preserving_merge(&distribution_context.plan) && fetch.is_none() {
-                spm = Some(Arc::clone(&distribution_context.plan));
-            }
             fetch = Some(fetch.map_or(child_fetch, |f: usize| f.min(child_fetch)));
         }
         // All of above operators have a single child. First child is only child.
@@ -1047,7 +1058,7 @@ fn remove_dist_changing_operators(
         // Note that they will be re-inserted later on if necessary or helpful.
     }
 
-    Ok((distribution_context, fetch, spm))
+    Ok((distribution_context, fetch, spm_ordering))
 }
 
 /// Updates the [`DistributionContext`] if preserving ordering while changing partitioning is not helpful or desirable.
@@ -1250,7 +1261,7 @@ pub fn ensure_distribution(
             children,
         },
         mut fetch,
-        spm,
+        spm_ordering,
     ) = remove_dist_changing_operators(dist_context)?;
 
     if let Some(exec) = plan.as_any().downcast_ref::<WindowAggExec>() {
@@ -1518,17 +1529,27 @@ pub fn ensure_distribution(
     // changing operator would be silently lost. Re-introduce it so the
     // query still returns the correct number of rows.
     if let Some(fetch_val) = fetch.take() {
-        let limit_plan: Arc<dyn ExecutionPlan> = if let Some(spm) = spm {
-            // Re-insert the original SortPreservingMergeExec with fetch.
-            spm.with_fetch(Some(fetch_val)).unwrap()
-        } else {
-            // The fetch came from a CoalescePartitionsExec. Re-introduce
-            // it as a CoalescePartitionsExec(fetch=N) wrapping the output.
-            Arc::new(
-                CoalescePartitionsExec::new(Arc::clone(&dist_context.plan))
+        let limit_plan: Arc<dyn ExecutionPlan> =
+            if let Some(ordering) = spm_ordering {
+                // Reconstruct a fresh SortPreservingMergeExec using the
+                // captured ordering and the *current* (possibly rewritten)
+                // child plan, rather than reusing the stale pre-optimization
+                // SPM which may reference an outdated subtree.
+                Arc::new(
+                    SortPreservingMergeExec::new(
+                        ordering,
+                        Arc::clone(&dist_context.plan),
+                    )
                     .with_fetch(Some(fetch_val)),
-            )
-        };
+                )
+            } else {
+                // The fetch came from a CoalescePartitionsExec. Re-introduce
+                // it as a CoalescePartitionsExec(fetch=N) wrapping the output.
+                Arc::new(
+                    CoalescePartitionsExec::new(Arc::clone(&dist_context.plan))
+                        .with_fetch(Some(fetch_val)),
+                )
+            };
         dist_context = DistributionContext::new(limit_plan, data, vec![dist_context]);
     }