Add EnsureRequirements: idempotent merged EnforceDistribution + EnforceSorting

## Summary

Replace the separate `EnforceDistribution` and `EnforceSorting` optimizer rules
with a single `EnsureRequirements` rule in the default optimizer chain. This makes
the composition idempotent by fixing distribution-awareness in `pushdown_sorts`
and fetch preservation in `EnforceDistribution`.

## Problem

`EnforceDistribution` and `EnforceSorting` are coupled through
`SortExec.preserve_partitioning` but run as independent rules. This caused:

1. **Production 502 errors**: `pushdown_sorts` set `preserve_partitioning=true`
   without `SortPreservingMergeExec`, violating `SinglePartition` requirements
   from `GlobalLimitExec` → `SanityCheckPlan` failure.

2. **Non-idempotent composition**: Running the rules multiple times produced
   different (sometimes invalid) plans.

3. **Lost fetch values** (#14150): `EnforceDistribution` dropped `fetch` from
   `SortPreservingMergeExec` when stripping and re-adding distribution operators.

DataFusion was the only major engine with separate rules — Spark (`EnsureRequirements`)
and Presto (`AddExchanges`) use a single rule.

## Changes

### `EnsureRequirements` rule (new)
- Composes `EnforceDistribution::optimize()` + `EnforceSorting::optimize()`
- Replaces both rules in the default optimizer chain
- 53 comprehensive tests including idempotency verification

### Distribution-aware `pushdown_sorts` (fix)
- Add `distribution_requirement` field to `ParentRequirements`
- New `add_sort_above_with_distribution()` inserts `SortPreservingMergeExec`
  when parent requires `SinglePartition` and input has multiple partitions
- Propagate distribution through recursion with `stronger_distribution()`
- Reset distribution below partition-merging nodes (SPM, single-partition outputs)

### Fix `EnforceDistribution` fetch preservation (#14150)
- `remove_dist_changing_operators()` now saves fetch from removed SPM/Coalesce
- `add_merge_on_top()` re-applies saved fetch to new operators

## Testing

| Suite | Result |
|-------|--------|
| EnsureRequirements (new) | 53 passed |
| enforce_sorting (existing) | 124 passed, 0 regressions |
| enforce_distribution (existing) | 66 passed, 0 regressions |
| SLT (465 files) | 1 pre-existing failure only |
| **Total** | **243 unit + 464 SLT, 0 new failures** |

Idempotency verified:
- All partition counts 1-64
- Triple + 10x consecutive optimization passes
- SortMergeJoin, HashJoin, Window, Aggregate topologies
- PR #53/#54 regression scenarios
- #14150 fetch preservation across passes

Closes: #14150
Part of: #21973

Author: zhuqi-lucas

datafusion/physical-optimizer/src/enforce_distribution.rs

@@ -970,7 +970,10 @@ fn preserving_order_enables_streaming(
 ///
 /// Updated node with an execution plan, where the desired single distribution
 /// requirement is satisfied.
-fn add_merge_on_top(input: DistributionContext) -> DistributionContext {
+fn add_merge_on_top(
+    input: DistributionContext,
+    fetch: Option<usize>,
+) -> DistributionContext {
     // Apply only when the partition count is larger than one.
     if input.plan.output_partitioning().partition_count() > 1 {
         // When there is an existing ordering, we preserve ordering
@@ -979,14 +982,20 @@ fn add_merge_on_top(input: DistributionContext) -> DistributionContext {
         // - Preserving ordering is not helpful in terms of satisfying ordering requirements
         // - Usage of order preserving variants is not desirable
         // (determined by flag `config.optimizer.prefer_existing_sort`)
-        let new_plan = if let Some(req) = input.plan.output_ordering() {
-            Arc::new(SortPreservingMergeExec::new(
-                req.clone(),
-                Arc::clone(&input.plan),
-            )) as _
+        let new_plan: Arc<dyn ExecutionPlan> = if let Some(req) =
+            input.plan.output_ordering()
+        {
+            let mut spm =
+                SortPreservingMergeExec::new(req.clone(), Arc::clone(&input.plan));
+            if let Some(f) = fetch {
+                spm = spm.with_fetch(Some(f));
+            }
+            Arc::new(spm)
         } else {
             // If there is no input order, we can simply coalesce partitions:
-            Arc::new(CoalescePartitionsExec::new(Arc::clone(&input.plan))) as _
+            Arc::new(
+                CoalescePartitionsExec::new(Arc::clone(&input.plan)).with_fetch(fetch),
+            )
         };
 
         DistributionContext::new(new_plan, true, vec![input])
@@ -1012,20 +1021,41 @@ fn add_merge_on_top(input: DistributionContext) -> DistributionContext {
 /// ```text
 /// "DataSourceExec: file_groups={2 groups: \[\[x], \[y]]}, projection=\[a, b, c, d, e], output_ordering=\[a@0 ASC], file_type=parquet",
 /// ```
+/// Returned by [`remove_dist_changing_operators`] to carry the fetch value
+/// that may have been on a removed `SortPreservingMergeExec` or `CoalescePartitionsExec`.
+struct RemovedDistOps {
+    context: DistributionContext,
+    /// The fetch value from the removed SPM/Coalesce, if any.
+    /// Must be re-applied when distribution operators are re-inserted.
+    removed_fetch: Option<usize>,
+}
+
 fn remove_dist_changing_operators(
     mut distribution_context: DistributionContext,
-) -> Result<DistributionContext> {
+) -> Result<RemovedDistOps> {
+    let mut removed_fetch = None;
     while is_repartition(&distribution_context.plan)
         || is_coalesce_partitions(&distribution_context.plan)
         || is_sort_preserving_merge(&distribution_context.plan)
     {
+        // Preserve fetch from SPM or CoalescePartitions before removing (#14150).
+        if let Some(fetch) = distribution_context.plan.fetch() {
+            removed_fetch = Some(
+                removed_fetch
+                    .map(|existing: usize| existing.min(fetch))
+                    .unwrap_or(fetch),
+            );
+        }
         // All of above operators have a single child. First child is only child.
         // Remove any distribution changing operators at the beginning:
         distribution_context = distribution_context.children.swap_remove(0);
         // Note that they will be re-inserted later on if necessary or helpful.
     }
 
-    Ok(distribution_context)
+    Ok(RemovedDistOps {
+        context: distribution_context,
+        removed_fetch,
+    })
 }
 
 /// Updates the [`DistributionContext`] if preserving ordering while changing partitioning is not helpful or desirable.
@@ -1219,11 +1249,16 @@ pub fn ensure_distribution(
     let order_preserving_variants_desirable =
         unbounded_and_pipeline_friendly || config.optimizer.prefer_existing_sort;
 
-    // Remove unnecessary repartition from the physical plan if any
-    let DistributionContext {
-        mut plan,
-        data,
-        children,
+    // Remove unnecessary repartition from the physical plan if any.
+    // Preserve fetch from removed SPM/Coalesce (#14150).
+    let RemovedDistOps {
+        context:
+            DistributionContext {
+                mut plan,
+                data,
+                children,
+            },
+        removed_fetch,
     } = remove_dist_changing_operators(dist_context)?;
 
     if let Some(exec) = plan.downcast_ref::<WindowAggExec>() {
@@ -1359,7 +1394,7 @@ pub fn ensure_distribution(
             // Satisfy the distribution requirement if it is unmet.
             match &requirement {
                 Distribution::SinglePartition => {
-                    child = add_merge_on_top(child);
+                    child = add_merge_on_top(child, removed_fetch);
                 }
                 Distribution::HashPartitioned(exprs) => {
                     // See https://github.com/apache/datafusion/issues/18341#issuecomment-3503238325 for background

datafusion/physical-optimizer/src/enforce_sorting/mod.rs

@@ -262,7 +262,7 @@ impl PhysicalOptimizerRule for EnforceSorting {
 /// If the plan is not a [`SortExec`] or its child is not unbounded, returns the original plan.
 /// Otherwise, by checking the requirement satisfaction searches for a replacement chance.
 /// If there's one replaces the [`SortExec`] plan with a [`PartialSortExec`]
-fn replace_with_partial_sort(
+pub fn replace_with_partial_sort(
     plan: Arc<dyn ExecutionPlan>,
 ) -> Result<Arc<dyn ExecutionPlan>> {
     let Some(sort_plan) = plan.downcast_ref::<SortExec>() else {

datafusion/physical-optimizer/src/enforce_sorting/sort_pushdown.rs

@@ -19,7 +19,8 @@ use std::fmt::Debug;
 use std::sync::Arc;
 
 use crate::utils::{
-    add_sort_above, is_sort, is_sort_preserving_merge, is_union, is_window,
+    add_sort_above_with_distribution, is_sort, is_sort_preserving_merge, is_union,
+    is_window,
 };
 
 use arrow::datatypes::SchemaRef;
@@ -29,7 +30,7 @@ use datafusion_expr::JoinType;
 use datafusion_physical_expr::expressions::Column;
 use datafusion_physical_expr::utils::collect_columns;
 use datafusion_physical_expr::{
-    EquivalenceProperties, add_offset_to_physical_sort_exprs,
+    Distribution, EquivalenceProperties, add_offset_to_physical_sort_exprs,
 };
 use datafusion_physical_expr_common::sort_expr::{
     LexOrdering, LexRequirement, OrderingRequirements, PhysicalSortExpr,
@@ -55,23 +56,42 @@ use datafusion_physical_plan::{ExecutionPlan, ExecutionPlanProperties};
 /// of the parent node as its data.
 ///
 /// [`EnforceSorting`]: crate::enforce_sorting::EnforceSorting
-#[derive(Default, Clone, Debug)]
+#[derive(Clone, Debug)]
 pub struct ParentRequirements {
     ordering_requirement: Option<OrderingRequirements>,
     fetch: Option<usize>,
+    /// The distribution required by the consumer above any SortExec we insert.
+    /// When this is `SinglePartition` and the input has multiple partitions,
+    /// `add_sort_above_with_distribution` wraps the sort in `SortPreservingMergeExec`.
+    distribution_requirement: Distribution,
+}
+
+impl Default for ParentRequirements {
+    fn default() -> Self {
+        Self {
+            ordering_requirement: None,
+            fetch: None,
+            distribution_requirement: Distribution::UnspecifiedDistribution,
+        }
+    }
 }
 
 pub type SortPushDown = PlanContext<ParentRequirements>;
 
 /// Assigns the ordering requirement of the root node to the its children.
 pub fn assign_initial_requirements(sort_push_down: &mut SortPushDown) {
     let reqs = sort_push_down.plan.required_input_ordering();
-    for (child, requirement) in sort_push_down.children.iter_mut().zip(reqs) {
+    let dists = sort_push_down.plan.required_input_distribution();
+    for (idx, (child, requirement)) in
+        sort_push_down.children.iter_mut().zip(reqs).enumerate()
+    {
         child.data = ParentRequirements {
             ordering_requirement: requirement,
-            // If the parent has a fetch value, assign it to the children
-            // Or use the fetch value of the child.
             fetch: child.plan.fetch(),
+            distribution_requirement: dists
+                .get(idx)
+                .cloned()
+                .unwrap_or(Distribution::UnspecifiedDistribution),
         };
     }
 }
@@ -92,11 +112,25 @@ fn min_fetch(f1: Option<usize>, f2: Option<usize>) -> Option<usize> {
     }
 }
 
+/// Returns the stricter of two distribution requirements.
+/// `SinglePartition` is the strictest.
+fn stronger_distribution(a: &Distribution, b: &Distribution) -> Distribution {
+    match (a, b) {
+        (Distribution::SinglePartition, _) | (_, Distribution::SinglePartition) => {
+            Distribution::SinglePartition
+        }
+        (Distribution::HashPartitioned(_), _) => a.clone(),
+        (_, Distribution::HashPartitioned(_)) => b.clone(),
+        _ => Distribution::UnspecifiedDistribution,
+    }
+}
+
 fn pushdown_sorts_helper(
     mut sort_push_down: SortPushDown,
 ) -> Result<Transformed<SortPushDown>> {
     let plan = sort_push_down.plan;
     let parent_fetch = sort_push_down.data.fetch;
+    let parent_distribution = sort_push_down.data.distribution_requirement.clone();
 
     let Some(parent_requirement) = sort_push_down.data.ordering_requirement.clone()
     else {
@@ -121,6 +155,14 @@ fn pushdown_sorts_helper(
             return pushdown_sorts_helper(sort_push_down);
         }
         sort_push_down.plan = plan;
+        // No ordering is being pushed; use each child's own distribution requirement
+        let dists = sort_push_down.plan.required_input_distribution();
+        for (idx, child) in sort_push_down.children.iter_mut().enumerate() {
+            child.data.distribution_requirement = dists
+                .get(idx)
+                .cloned()
+                .unwrap_or(Distribution::UnspecifiedDistribution);
+        }
         return Ok(Transformed::no(sort_push_down));
     };
 
@@ -149,22 +191,29 @@ fn pushdown_sorts_helper(
             // The sort was imposing a different ordering than the one being
             // pushed down. Replace it with a sort that matches the pushed-down
             // ordering, and continue the pushdown.
-            // Add back the sort:
-            sort_push_down = add_sort_above(
+            // Add back the sort (distribution-aware):
+            sort_push_down = add_sort_above_with_distribution(
                 sort_push_down,
                 parent_requirement.into_single(),
                 parent_fetch,
+                &parent_distribution,
             );
             // Update pushdown requirements:
             sort_push_down.children[0].data = ParentRequirements {
                 ordering_requirement: Some(OrderingRequirements::from(sort_ordering)),
                 fetch: sort_fetch,
+                distribution_requirement: Distribution::UnspecifiedDistribution,
             };
             return Ok(Transformed::yes(sort_push_down));
         } else {
             // Sort was unnecessary, just propagate the stricter fetch and
-            // ordering requirements:
+            // ordering requirements. Reset distribution to Unspecified
+            // because the sort we're removing may have been below a
+            // partition-merging node (like SortPreservingMergeExec) that
+            // already satisfies SinglePartition.
             sort_push_down.data.fetch = min_fetch(sort_fetch, parent_fetch);
+            sort_push_down.data.distribution_requirement =
+                Distribution::UnspecifiedDistribution;
             let current_is_stricter = eqp.requirements_compatible(
                 sort_ordering.clone().into(),
                 parent_requirement.first().clone(),
@@ -184,30 +233,65 @@ fn pushdown_sorts_helper(
     if satisfy_parent {
         // For non-sort operators which satisfy ordering:
         let reqs = sort_push_down.plan.required_input_ordering();
+        let dists = sort_push_down.plan.required_input_distribution();
+
+        // If this node already outputs single partition, don't push SinglePartition
+        // requirement to children (they're below the merge point).
+        let effective_parent_dist =
+            if sort_push_down.plan.output_partitioning().partition_count() == 1 {
+                Distribution::UnspecifiedDistribution
+            } else {
+                parent_distribution.clone()
+            };
 
-        for (child, order) in sort_push_down.children.iter_mut().zip(reqs) {
+        for (idx, (child, order)) in
+            sort_push_down.children.iter_mut().zip(reqs).enumerate()
+        {
             child.data.ordering_requirement = order;
             child.data.fetch = min_fetch(parent_fetch, child.data.fetch);
+            child.data.distribution_requirement = stronger_distribution(
+                &effective_parent_dist,
+                dists
+                    .get(idx)
+                    .unwrap_or(&Distribution::UnspecifiedDistribution),
+            );
         }
     } else if let Some(adjusted) = pushdown_requirement_to_children(
         &sort_push_down.plan,
         parent_requirement.clone(),
         parent_fetch,
     )? {
         // For operators that can take a sort pushdown, continue with updated
-        // requirements:
+        // requirements. If this node already outputs single partition (e.g. SPM),
+        // don't push SinglePartition to children.
         let current_fetch = sort_push_down.plan.fetch();
-        for (child, order) in sort_push_down.children.iter_mut().zip(adjusted) {
+        let dists = sort_push_down.plan.required_input_distribution();
+        let effective_dist =
+            if sort_push_down.plan.output_partitioning().partition_count() == 1 {
+                Distribution::UnspecifiedDistribution
+            } else {
+                parent_distribution.clone()
+            };
+        for (idx, (child, order)) in
+            sort_push_down.children.iter_mut().zip(adjusted).enumerate()
+        {
             child.data.ordering_requirement = order;
             child.data.fetch = min_fetch(current_fetch, parent_fetch);
+            child.data.distribution_requirement = stronger_distribution(
+                &effective_dist,
+                dists
+                    .get(idx)
+                    .unwrap_or(&Distribution::UnspecifiedDistribution),
+            );
         }
         sort_push_down.data.ordering_requirement = None;
     } else {
-        // Can not push down requirements, add new `SortExec`:
-        sort_push_down = add_sort_above(
+        // Can not push down requirements, add new `SortExec` (distribution-aware):
+        sort_push_down = add_sort_above_with_distribution(
             sort_push_down,
             parent_requirement.into_single(),
             parent_fetch,
+            &parent_distribution,
         );
         assign_initial_requirements(&mut sort_push_down);
     }