fix: EnforceDistribution optimizer preserves fetch (LIMIT) from distribution-changing operators

When `LimitPushdown` merges a `GlobalLimitExec` into a
`CoalescePartitionsExec` or `SortPreservingMergeExec` as a `fetch`
value, `EnforceDistribution` would strip and re-insert these operators
without preserving the fetch. This silently drops the LIMIT for queries
over multi-partition sources.

The fix captures the `fetch` in `remove_dist_changing_operators` and
threads it through `add_merge_on_top`, or re-introduces it as a
fallback when the merge operator is not re-inserted.

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

Author: zhuqi-lucas

datafusion/core/tests/physical_optimizer/enforce_distribution.rs

@@ -3915,6 +3915,85 @@ fn test_replace_order_preserving_variants_with_fetch() -> Result<()> {
     Ok(())
 }
 
+/// When `LimitPushdown` merges a `GlobalLimitExec` into `CoalescePartitionsExec`
+/// as `fetch`, `remove_dist_changing_operators` must preserve that fetch value.
+/// Otherwise, queries with LIMIT over multi-partition sources silently lose
+/// the limit and return duplicate/extra rows.
+///
+/// Regression test for: https://github.com/apache/datafusion/issues/21169
+#[test]
+fn coalesce_partitions_fetch_preserved_by_enforce_distribution() -> Result<()> {
+    // Simulate what LimitPushdown produces:
+    // CoalescePartitionsExec(fetch=1)
+    //   DataSourceExec (2 partitions)
+    let parquet = parquet_exec_multiple();
+    let coalesce_with_fetch: Arc<dyn ExecutionPlan> =
+        Arc::new(CoalescePartitionsExec::new(parquet).with_fetch(Some(1)));
+
+    let result = ensure_distribution_helper(coalesce_with_fetch, 10, false)?;
+
+    // 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)
+    let plan_str = displayable(result.as_ref()).indent(true).to_string();
+    assert!(
+        plan_str.contains("fetch=1"),
+        "fetch=1 was lost after EnforceDistribution!\nPlan:\n{plan_str}"
+    );
+    Ok(())
+}
+
+/// Same as above, but with a sorted multi-partition source.
+/// The fetch should be preserved on `SortPreservingMergeExec`.
+#[test]
+fn coalesce_partitions_fetch_preserved_sorted() -> Result<()> {
+    let schema = schema();
+    let sort_key: LexOrdering = [PhysicalSortExpr {
+        expr: col("c", &schema)?,
+        options: SortOptions::default(),
+    }]
+    .into();
+
+    // CoalescePartitionsExec(fetch=5) over sorted multi-partition source
+    let parquet = parquet_exec_multiple_sorted(vec![sort_key]);
+    let coalesce_with_fetch: Arc<dyn ExecutionPlan> =
+        Arc::new(CoalescePartitionsExec::new(parquet).with_fetch(Some(5)));
+
+    let result = ensure_distribution_helper(coalesce_with_fetch, 10, false)?;
+
+    let plan_str = displayable(result.as_ref()).indent(true).to_string();
+    assert!(
+        plan_str.contains("fetch=5"),
+        "fetch=5 was lost after EnforceDistribution!\nPlan:\n{plan_str}"
+    );
+    Ok(())
+}
+
+/// SortPreservingMergeExec with fetch should also be preserved.
+#[test]
+fn spm_fetch_preserved_by_enforce_distribution() -> Result<()> {
+    let schema = schema();
+    let sort_key: LexOrdering = [PhysicalSortExpr {
+        expr: col("c", &schema)?,
+        options: SortOptions::default(),
+    }]
+    .into();
+
+    // SortPreservingMergeExec(fetch=3) over sorted multi-partition source
+    let parquet = parquet_exec_multiple_sorted(vec![sort_key.clone()]);
+    let spm_with_fetch: Arc<dyn ExecutionPlan> =
+        Arc::new(SortPreservingMergeExec::new(sort_key, parquet).with_fetch(Some(3)));
+
+    let result = ensure_distribution_helper(spm_with_fetch, 10, false)?;
+
+    let plan_str = displayable(result.as_ref()).indent(true).to_string();
+    assert!(
+        plan_str.contains("fetch=3"),
+        "fetch=3 was lost after EnforceDistribution!\nPlan:\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

@@ -967,9 +967,12 @@ fn preserving_order_enables_streaming(
 
 /// # Returns
 ///
-/// Updated node with an execution plan, where the desired single distribution
-/// requirement is satisfied.
-fn add_merge_on_top(input: DistributionContext) -> DistributionContext {
+/// Updated node with an execution plan, where desired single
+/// distribution is satisfied by adding [`SortPreservingMergeExec`].
+fn add_merge_on_top(
+    input: DistributionContext,
+    fetch: &mut 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,13 +982,16 @@ fn add_merge_on_top(input: DistributionContext) -> DistributionContext {
         // - 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 _
+            Arc::new(
+                SortPreservingMergeExec::new(req.clone(), Arc::clone(&input.plan))
+                    .with_fetch(fetch.take()),
+            ) as _
         } 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.take()),
+            ) as _
         };
 
         DistributionContext::new(new_plan, true, vec![input])
@@ -1011,20 +1017,39 @@ 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",
 /// ```
+#[expect(clippy::type_complexity)]
 fn remove_dist_changing_operators(
     mut distribution_context: DistributionContext,
-) -> Result<DistributionContext> {
+) -> Result<(
+    DistributionContext,
+    Option<usize>,
+    Option<Arc<dyn ExecutionPlan>>,
+)> {
+    let mut fetch = None;
+    let mut spm: Option<Arc<dyn ExecutionPlan>> = None;
     while is_repartition(&distribution_context.plan)
         || is_coalesce_partitions(&distribution_context.plan)
         || is_sort_preserving_merge(&distribution_context.plan)
     {
+        // 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.
         // 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((distribution_context, fetch, spm))
 }
 
 /// Updates the [`DistributionContext`] if preserving ordering while changing partitioning is not helpful or desirable.
@@ -1220,11 +1245,15 @@ pub fn ensure_distribution(
         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_dist_changing_operators(dist_context)?;
+    let (
+        DistributionContext {
+            mut plan,
+            data,
+            children,
+        },
+        mut fetch,
+        spm,
+    ) = remove_dist_changing_operators(dist_context)?;
 
     if let Some(exec) = plan.as_any().downcast_ref::<WindowAggExec>() {
         if let Some(updated_window) = get_best_fitting_window(
@@ -1354,7 +1383,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, &mut fetch);
                 }
                 Distribution::HashPartitioned(exprs) => {
                     // See https://github.com/apache/datafusion/issues/18341#issuecomment-3503238325 for background
@@ -1483,9 +1512,29 @@ pub fn ensure_distribution(
         plan.with_new_children(children_plans)?
     };
 
-    Ok(Transformed::yes(DistributionContext::new(
-        plan, data, children,
-    )))
+    let mut dist_context = DistributionContext::new(Arc::clone(&plan), data, children);
+
+    // If `fetch` was not consumed by `add_merge_on_top` (e.g. when the parent
+    // had `UnspecifiedDistribution` or the child already had a single
+    // partition), the limit that was originally embedded in a 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))
+                    .with_fetch(Some(fetch_val)),
+            )
+        };
+        dist_context = DistributionContext::new(limit_plan, data, vec![dist_context]);
+    }
+
+    Ok(Transformed::yes(dist_context))
 }
 
 /// Keeps track of distribution changing operators (like `RepartitionExec`,