fix: preserve fetch-carrying operators in remove_dist_changing_operators

datafusion/core/tests/physical_optimizer/enforce_distribution.rs

@@ -22,7 +22,7 @@ use std::sync::Arc;
 use crate::physical_optimizer::test_utils::{
     check_integrity, coalesce_partitions_exec, parquet_exec_with_sort,
     parquet_exec_with_stats, repartition_exec, schema, sort_exec,
-    sort_exec_with_preserve_partitioning, sort_merge_join_exec,
+    sort_exec_with_preserve_partitioning, sort_expr, sort_merge_join_exec,
     sort_preserving_merge_exec, union_exec,
 };
 
@@ -3626,6 +3626,43 @@ fn get_schema() -> SchemaRef {
         Field::new("bank_account", DataType::UInt64, true),
     ]))
 }
+#[test]
+fn keep_fetch_carrying_dist_changing_operators() -> Result<()> {
+    let config = ConfigOptions::new();
+
+    // A CoalescePartitionsExec carrying a fetch is not a pure
+    // distribution-changing operator: it also provides the global limit, so
+    // EnforceDistribution must not strip it. This can be hit when the
+    // physical optimizer pipeline runs on an already-optimized plan.
+    let coalesce =
+        Arc::new(CoalescePartitionsExec::new(parquet_exec()).with_fetch(Some(1)))
+            as Arc<dyn ExecutionPlan>;
+    let optimized = EnforceDistribution::new().optimize(coalesce, &config)?;
+    let coalesce = optimized
+        .as_any()
+        .downcast_ref::<CoalescePartitionsExec>()
+        .expect("fetch-carrying CoalescePartitionsExec should be preserved");
+    assert_eq!(coalesce.fetch(), Some(1));
+
+    // Same for SortPreservingMergeExec with a fetch.
+    let ordering = LexOrdering::new(vec![sort_expr("a", &schema())]).unwrap();
+    let spm = Arc::new(
+        SortPreservingMergeExec::new(
+            ordering.clone(),
+            parquet_exec_with_sort(schema(), vec![ordering]),
+        )
+        .with_fetch(Some(1)),
+    ) as Arc<dyn ExecutionPlan>;
+    let optimized = EnforceDistribution::new().optimize(spm, &config)?;
+    let spm = optimized
+        .as_any()
+        .downcast_ref::<SortPreservingMergeExec>()
+        .expect("fetch-carrying SortPreservingMergeExec should be preserved");
+    assert_eq!(spm.fetch(), Some(1));
+
+    Ok(())
+}
+
 #[test]
 fn test_replace_order_preserving_variants_with_fetch() -> Result<()> {
     // Create a base plan

datafusion/core/tests/physical_optimizer/limited_distinct_aggregation.rs

@@ -25,8 +25,11 @@ use crate::physical_optimizer::test_utils::{
     schema,
 };
 
+use arrow::array::{BooleanArray, Int32Array, Int64Array};
 use arrow::datatypes::DataType;
+use arrow::record_batch::RecordBatch;
 use arrow::{compute::SortOptions, util::pretty::pretty_format_batches};
+use datafusion::datasource::MemTable;
 use datafusion::prelude::SessionContext;
 use datafusion_common::Result;
 use datafusion_execution::config::SessionConfig;
@@ -36,7 +39,7 @@ use datafusion_physical_expr_common::sort_expr::PhysicalSortExpr;
 use datafusion_physical_plan::{
     ExecutionPlan,
     aggregates::{AggregateExec, AggregateMode},
-    collect,
+    collect, displayable,
     limit::{GlobalLimitExec, LocalLimitExec},
 };
 
@@ -254,6 +257,52 @@ async fn test_distinct_cols_different_than_group_by_cols() -> Result<()> {
     Ok(())
 }
 
+#[tokio::test]
+async fn test_global_limit_survives_second_optimizer_pass() -> Result<()> {
+    // Regression: EnforceDistribution::remove_dist_changing_operators must not
+    // strip a fetch-carrying CoalescePartitionsExec when the physical
+    // optimizer pipeline runs on an already-optimized plan (e.g. a nested
+    // subplan that was planned and optimized separately).
+    let cfg = SessionConfig::new().with_target_partitions(4);
+    let ctx = SessionContext::new_with_config(cfg);
+    let schema = schema();
+    let batch = RecordBatch::try_new(
+        schema.clone(),
+        vec![
+            Arc::new(Int64Array::from(vec![1, 2, 3])),
+            Arc::new(Int64Array::from(vec![4, 5, 6])),
+            Arc::new(Int64Array::from(vec![7, 8, 9])),
+            Arc::new(Int32Array::from(vec![1, 2, 3])),
+            Arc::new(BooleanArray::from(vec![true, false, true])),
+        ],
+    )?;
+    let table = MemTable::try_new(schema, vec![vec![batch]])?;
+    ctx.register_table("t", Arc::new(table))?;
+
+    let plan = ctx
+        .sql("SELECT DISTINCT a FROM t LIMIT 1")
+        .await?
+        .create_physical_plan()
+        .await?;
+
+    // Run the physical optimizer pipeline a second time over the
+    // already-optimized plan, using the session's own rules and config.
+    let state = ctx.state();
+    let mut optimized: Arc<dyn ExecutionPlan> = plan;
+    for rule in state.physical_optimizers() {
+        optimized = rule.optimize(optimized, state.config_options())?;
+    }
+
+    let display = displayable(optimized.as_ref()).indent(true).to_string();
+    assert!(
+        display.contains("CoalescePartitionsExec: fetch=1")
+            || display.contains("GlobalLimitExec: skip=0, fetch=1"),
+        "global limit must survive a second optimizer pass:\n{display}"
+    );
+
+    Ok(())
+}
+
 #[test]
 fn test_has_order_by() -> Result<()> {
     let schema = schema();

datafusion/physical-optimizer/src/enforce_distribution.rs

@@ -981,6 +981,11 @@ fn remove_dist_changing_operators(
         || is_coalesce_partitions(&distribution_context.plan)
         || is_sort_preserving_merge(&distribution_context.plan)
     {
+        // A fetch carries global limit semantics, not just a distribution
+        // change, so fetch-carrying operators must be kept.
+        if distribution_context.plan.fetch().is_some() {
+            break;
+        }
         // 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);