Addressed in latest PR

Author: zhuqi-lucas

datafusion/physical-optimizer/src/enforce_distribution.rs

@@ -951,7 +951,7 @@ fn add_merge_on_top(
         let new_plan = if let Some(req) = input.plan.output_ordering() {
             Arc::new(
                 SortPreservingMergeExec::new(req.clone(), Arc::clone(&input.plan))
-                    .with_fetch(*fetch),
+                    .with_fetch(fetch.take()),
             ) as _
         } else {
             // If there is no input order, we can simply coalesce partitions:
@@ -1223,7 +1223,7 @@ pub fn ensure_distribution(
             children,
         },
         mut fetch,
-        _spm,
+        spm,
     ) = remove_dist_changing_operators(dist_context)?;
 
     if let Some(exec) = plan.as_any().downcast_ref::<WindowAggExec>() {
@@ -1408,13 +1408,8 @@ pub fn ensure_distribution(
     // It was removed by `remove_dist_changing_operators`
     // and we need to add it back.
     if fetch.is_some() {
-        // We can make sure that `plan` has an ordering because
-        // `SortPreservingMergeExec` requires ordering to be constructed.
-        // If there is no ordering, `SortPreservingMergeExec::new` will panic
-        let ordering = plan.output_ordering().cloned().unwrap();
-        let plan = Arc::new(
-            SortPreservingMergeExec::new(ordering, plan).with_fetch(fetch.take()),
-        );
+        // It's safe to unwrap because `spm` is set only if `fetch` is set.
+        let plan = spm.unwrap().with_fetch(fetch.take()).unwrap();
         optimized_distribution_ctx =
             DistributionContext::new(plan, data, vec![optimized_distribution_ctx]);
     }