From b8aad02b7583f077a1ed2e4d2b405644bf688e88 Mon Sep 17 00:00:00 2001 From: "xudong.w" Date: Thu, 27 Aug 2026 17:04:16 +0800 Subject: [PATCH 1/2] Preserve fetch across distribution reoptimization --- .../enforce_distribution.rs | 49 +++++- .../enforce_distribution.rs | 145 ++++++++++++------ 2 files changed, 143 insertions(+), 51 deletions(-) diff --git a/datafusion/core/tests/physical_optimizer/enforce_distribution.rs b/datafusion/core/tests/physical_optimizer/enforce_distribution.rs index 2dbacf1d898ac..c805be44552cc 100644 --- a/datafusion/core/tests/physical_optimizer/enforce_distribution.rs +++ b/datafusion/core/tests/physical_optimizer/enforce_distribution.rs @@ -4442,7 +4442,7 @@ fn test_replace_order_preserving_variants_with_fetch() -> Result<()> { ); // Apply the function - let result = replace_order_preserving_variants(dist_context)?; + let result = replace_order_preserving_variants(dist_context, false)?.0; // Verify the plan was transformed to CoalescePartitionsExec result @@ -4460,6 +4460,53 @@ fn test_replace_order_preserving_variants_with_fetch() -> Result<()> { Ok(()) } +#[test] +fn preserve_fetch_when_reoptimizing_ordered_merge() -> Result<()> { + let schema = schema(); + let sort_key: LexOrdering = + [PhysicalSortExpr::new_default(col("c", &schema)?)].into(); + let input = parquet_exec_multiple_sorted(vec![sort_key.clone()]); + let plan: Arc = + Arc::new(SortPreservingMergeExec::new(sort_key, input).with_fetch(Some(5))); + + let optimized = + EnsureRequirements::new().optimize(plan, &test_suite_default_config_options())?; + let plan = displayable(optimized.as_ref()).indent(true).to_string(); + + assert!( + plan.contains("SortPreservingMergeExec: [c@2 ASC], fetch=5"), + "expected the optimizer to preserve fetch:\n{plan}" + ); + + Ok(()) +} + +#[test] +fn move_fetch_to_replacement_sort() -> Result<()> { + let schema = schema(); + let sort_key: LexOrdering = + [PhysicalSortExpr::new_default(col("c", &schema)?)].into(); + let input = parquet_exec_multiple_sorted(vec![sort_key.clone()]); + let merge: Arc = Arc::new( + SortPreservingMergeExec::new(sort_key, input.clone()).with_fetch(Some(5)), + ); + let context = DistributionContext::new( + merge, + true, + vec![DistributionContext::new(input, false, vec![])], + ); + let (context, fetch) = replace_order_preserving_variants(context, true)?; + + context + .plan + .downcast_ref::() + .expect("expected CoalescePartitionsExec"); + assert_eq!(context.plan.fetch(), None); + assert_eq!(fetch, Some(5)); + + 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 diff --git a/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs b/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs index 07bc98b2db798..5d8d5bac4a385 100644 --- a/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs +++ b/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs @@ -41,6 +41,7 @@ use crate::utils::{ use arrow::compute::SortOptions; use datafusion_common::config::ConfigOptions; use datafusion_common::error::Result; +use datafusion_common::internal_err; use datafusion_common::stats::Precision; use datafusion_common::tree_node::Transformed; use datafusion_expr::logical_plan::{Aggregate, JoinType}; @@ -784,7 +785,7 @@ fn preserving_order_enables_streaming( /// requirement is satisfied. fn add_merge_on_top( input: DistributionContext, - fetch: Option, + fetch: &mut Option, ) -> DistributionContext { // Apply only when the partition count is larger than one. if input.plan.output_partitioning().partition_count() > 1 { @@ -794,21 +795,19 @@ fn add_merge_on_top( // - 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: Arc = 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)).with_fetch(fetch), - ) - }; + let new_plan: Arc = + if let Some(req) = input.plan.output_ordering() { + let mut spm = + SortPreservingMergeExec::new(req.clone(), Arc::clone(&input.plan)); + spm = spm.with_fetch(fetch.take()); + Arc::new(spm) + } else { + // If there is no input order, we can simply coalesce partitions: + Arc::new( + CoalescePartitionsExec::new(Arc::clone(&input.plan)) + .with_fetch(fetch.take()), + ) + }; DistributionContext::new(new_plan, true, vec![input]) } else { @@ -840,23 +839,33 @@ struct RemovedDistOps { /// The fetch value from the removed SPM/Coalesce, if any. /// Must be re-applied when distribution operators are re-inserted. removed_fetch: Option, + /// The outermost removed operator carrying a fetch, used to restore the + /// limit when no replacement distribution operator consumes it. + fetch_plan: Option>, +} + +fn min_fetch(left: Option, right: Option) -> Option { + match (left, right) { + (Some(left), Some(right)) => Some(left.min(right)), + (left, right) => left.or(right), + } } fn remove_dist_changing_operators( mut distribution_context: DistributionContext, ) -> Result { let mut removed_fetch = None; + let mut fetch_plan = 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), - ); + if fetch_plan.is_none() { + fetch_plan = Some(Arc::clone(&distribution_context.plan)); + } + removed_fetch = min_fetch(removed_fetch, Some(fetch)); } // All of above operators have a single child. First child is only child. // Remove any distribution changing operators at the beginning: @@ -867,6 +876,7 @@ fn remove_dist_changing_operators( Ok(RemovedDistOps { context: distribution_context, removed_fetch, + fetch_plan, }) } @@ -890,25 +900,32 @@ fn remove_dist_changing_operators( /// ``` pub fn replace_order_preserving_variants( mut context: DistributionContext, -) -> Result { - context.children = context - .children - .into_iter() - .map(|child| { - if child.data { - replace_order_preserving_variants(child) - } else { - Ok(child) - } - }) - .collect::>>()?; + ordering_satisfied: bool, +) -> Result<(DistributionContext, Option)> { + let mut children = Vec::with_capacity(context.children.len()); + let mut fetch = None; + for child in context.children { + if child.data { + let (child, child_fetch) = + replace_order_preserving_variants(child, ordering_satisfied)?; + children.push(child); + fetch = min_fetch(fetch, child_fetch); + } else { + children.push(child); + } + } + context.children = children; if is_sort_preserving_merge(&context.plan) { + let fetch = min_fetch(fetch, context.plan.fetch()); let child_plan = Arc::clone(&context.children[0].plan); - context.plan = Arc::new( - CoalescePartitionsExec::new(child_plan).with_fetch(context.plan.fetch()), - ); - return Ok(context); + if ordering_satisfied { + context.plan = Arc::new(CoalescePartitionsExec::new(child_plan)); + return Ok((context, fetch)); + } + context.plan = + Arc::new(CoalescePartitionsExec::new(child_plan).with_fetch(fetch)); + return Ok((context, None)); } else if let Some(repartition) = context.plan.downcast_ref::() && repartition.preserve_order() { @@ -916,10 +933,12 @@ pub fn replace_order_preserving_variants( Arc::clone(&context.children[0].plan), repartition.partitioning().clone(), )?); - return Ok(context); + return Ok((context, fetch)); } - context.update_plan_from_children() + context + .update_plan_from_children() + .map(|context| (context, fetch)) } /// A struct to keep track of repartition requirements for each child node. @@ -1175,7 +1194,8 @@ pub fn ensure_distribution( data, children, }, - removed_fetch, + mut removed_fetch, + fetch_plan, } = remove_dist_changing_operators(dist_context)?; if let Some(exec) = plan.downcast_ref::() { @@ -1331,7 +1351,7 @@ pub fn ensure_distribution( // Satisfy the distribution requirement if it is unmet. match &requirement { Distribution::SinglePartition => { - child = add_merge_on_top(child, removed_fetch); + child = add_merge_on_top(child, &mut removed_fetch); } Distribution::HashPartitioned(exprs) | Distribution::KeyPartitioned(exprs) => { @@ -1444,17 +1464,23 @@ pub fn ensure_distribution( && !streaming_benefit && context.data { - context = replace_order_preserving_variants(context)?; + let (replaced_context, preserved_fetch) = + replace_order_preserving_variants( + context, + ordering_satisfied, + )?; + context = replaced_context; // If ordering requirements were satisfied before repartitioning, // make sure ordering requirements are still satisfied after. if ordering_satisfied { // Make sure to satisfy ordering requirement: + let output_fetch = plan + .downcast_ref::() + .and_then(|output| output.fetch()); context = add_sort_above_with_check( context, sort_req, - plan.downcast_ref::() - .map(|output| output.fetch()) - .unwrap_or(None), + min_fetch(preserved_fetch, output_fetch), )?; } } @@ -1472,7 +1498,8 @@ pub fn ensure_distribution( // input order, we keep order-preserving variants so // ordering can flow through to ancestors that need it. if !maintains_input_order && !streaming_benefit { - context = replace_order_preserving_variants(context)?; + context = + replace_order_preserving_variants(context, false)?.0; } } Distribution::UnspecifiedDistribution => { @@ -1480,7 +1507,8 @@ pub fn ensure_distribution( if !maintains_input_order || plan.is::() { - context = replace_order_preserving_variants(context)?; + context = + replace_order_preserving_variants(context, false)?.0; } } } @@ -1536,9 +1564,26 @@ pub fn ensure_distribution( replace_children_if_necessary(plan, children_plans)? }; - Ok(Transformed::yes(DistributionContext::new( - plan, data, children, - ))) + let mut optimized_context = DistributionContext::new(plan, data, children); + + // A removed fetch must survive even when this node does not need a new + // distribution operator. Otherwise a second optimizer pass can silently + // remove the query's LIMIT. + if let Some(fetch) = removed_fetch { + let Some(fetch_plan) = fetch_plan else { + return internal_err!("removed distribution fetch has no source plan"); + }; + let fetch_plan = replace_children_if_necessary( + fetch_plan, + vec![Arc::clone(&optimized_context.plan)], + )?; + let Some(plan) = fetch_plan.with_fetch(Some(fetch)) else { + return internal_err!("removed distribution operator cannot restore fetch"); + }; + optimized_context = DistributionContext::new(plan, data, vec![optimized_context]); + } + + Ok(Transformed::yes(optimized_context)) } /// Keeps track of distribution changing operators (like `RepartitionExec`, From 414ee99a559bbdcb0b6bd1152778be53dbb98149 Mon Sep 17 00:00:00 2001 From: "xudong.w" Date: Thu, 27 Aug 2026 17:39:59 +0800 Subject: [PATCH 2/2] Update fetch preservation snapshot --- datafusion/core/tests/physical_optimizer/enforce_sorting.rs | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/datafusion/core/tests/physical_optimizer/enforce_sorting.rs b/datafusion/core/tests/physical_optimizer/enforce_sorting.rs index a8162f137ed0a..bf811c8758735 100644 --- a/datafusion/core/tests/physical_optimizer/enforce_sorting.rs +++ b/datafusion/core/tests/physical_optimizer/enforce_sorting.rs @@ -2293,7 +2293,9 @@ async fn test_remove_unnecessary_spm2() -> Result<()> { DataSourceExec: partitions=1, partition_sizes=[0] Optimized Plan: - DataSourceExec: partitions=1, partition_sizes=[0] + LocalLimitExec: fetch=100 + SortExec: expr=[non_nullable_col@1 ASC], preserve_partitioning=[false] + DataSourceExec: partitions=1, partition_sizes=[0] "); Ok(())