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.

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::{
@@ -1017,28 +1018,32 @@ fn add_merge_on_top(
 /// ```text
 /// "DataSourceExec: file_groups={2 groups: \[\[x], \[y]]}, projection=\[a, b, c, d, e], output_ordering=\[a@0 ASC], file_type=parquet",
 /// ```
-#[expect(clippy::type_complexity)]
 fn remove_dist_changing_operators(
     mut distribution_context: DistributionContext,
-) -> Result<(
-    DistributionContext,
-    Option<usize>,
-    Option<Arc<dyn ExecutionPlan>>,
-)> {
+) -> Result<(DistributionContext, Option<usize>, 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()
+            && 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 +1052,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 +1255,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,9 +1523,15 @@ 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()
+        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.