diff --git a/datafusion/core/tests/physical_optimizer/enforce_distribution.rs b/datafusion/core/tests/physical_optimizer/enforce_distribution.rs index 5df634c70bcbb..231325dc51b89 100644 --- a/datafusion/core/tests/physical_optimizer/enforce_distribution.rs +++ b/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; + let optimized = EnforceDistribution::new().optimize(coalesce, &config)?; + let coalesce = optimized + .as_any() + .downcast_ref::() + .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; + let optimized = EnforceDistribution::new().optimize(spm, &config)?; + let spm = optimized + .as_any() + .downcast_ref::() + .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 diff --git a/datafusion/core/tests/physical_optimizer/limited_distinct_aggregation.rs b/datafusion/core/tests/physical_optimizer/limited_distinct_aggregation.rs index c523b4a752a82..8213476f039e3 100644 --- a/datafusion/core/tests/physical_optimizer/limited_distinct_aggregation.rs +++ b/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 = 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(); diff --git a/datafusion/physical-optimizer/src/enforce_distribution.rs b/datafusion/physical-optimizer/src/enforce_distribution.rs index d23a699f715de..6837d73510cc1 100644 --- a/datafusion/physical-optimizer/src/enforce_distribution.rs +++ b/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);