From 1d957a5558960c94975134f90f2ae2da187d2318 Mon Sep 17 00:00:00 2001 From: "xudong.w" Date: Mon, 31 Aug 2026 12:06:10 +0800 Subject: [PATCH 1/3] fix: preserve fetch across distribution reoptimization --- .../enforce_distribution.rs | 66 ++++++++ .../physical_optimizer/enforce_sorting.rs | 4 +- .../enforce_distribution.rs | 153 ++++++++++++------ 3 files changed, 174 insertions(+), 49 deletions(-) diff --git a/datafusion/core/tests/physical_optimizer/enforce_distribution.rs b/datafusion/core/tests/physical_optimizer/enforce_distribution.rs index 489076331bbf5..91cd64e208dea 100644 --- a/datafusion/core/tests/physical_optimizer/enforce_distribution.rs +++ b/datafusion/core/tests/physical_optimizer/enforce_distribution.rs @@ -4558,6 +4558,72 @@ 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 preserve_fetch_when_reoptimizing_coalesce_partitions() -> Result<()> { + let input = parquet_exec_multiple(); + let plan: Arc = + Arc::new(CoalescePartitionsExec::new(input).with_fetch(Some(5))); + + let optimized = + EnsureRequirements::new().optimize(plan, &test_suite_default_config_options())?; + + assert_eq!(optimized.fetch(), Some(5)); + optimized + .downcast_ref::() + .expect("expected CoalescePartitionsExec"); + + 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.clone(), input).with_fetch(Some(5)), + ); + let plan = sort_required_exec_with_req(merge, sort_key); + + let optimized = ensure_distribution_helper(plan, 10, false)?; + let plan = displayable(optimized.as_ref()).indent(true).to_string(); + + assert!( + plan.contains( + "SortExec: TopK(fetch=5), expr=[c@2 ASC], preserve_partitioning=[false]" + ), + "expected the replacement sort to preserve fetch:\n{plan}" + ); + assert!( + !plan.contains("CoalescePartitionsExec: fetch=5"), + "fetch below the replacement sort would change TopK results:\n{plan}" + ); + + 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/core/tests/physical_optimizer/enforce_sorting.rs b/datafusion/core/tests/physical_optimizer/enforce_sorting.rs index 76af9b0c29218..87d8ac3b159e1 100644 --- a/datafusion/core/tests/physical_optimizer/enforce_sorting.rs +++ b/datafusion/core/tests/physical_optimizer/enforce_sorting.rs @@ -2327,7 +2327,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(()) diff --git a/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs b/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs index 0368577f9a24f..fc49c5c9d8fc3 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, }) } @@ -889,26 +899,47 @@ fn remove_dist_changing_operators( /// " DataSourceExec: file_groups={2 groups: \[\[x], \[y]]}, projection=\[a, b, c, d, e], output_ordering=\[a@0 ASC], file_type=parquet", /// ``` pub fn replace_order_preserving_variants( - mut context: DistributionContext, + context: DistributionContext, ) -> Result { - context.children = context - .children - .into_iter() - .map(|child| { - if child.data { - replace_order_preserving_variants(child) - } else { - Ok(child) - } - }) - .collect::>>()?; + let (context, fetch) = replace_order_preserving_variants_with_fetch(context, false)?; + debug_assert!( + fetch.is_none(), + "fetch must stay in the plan when no replacement sort is needed" + ); + Ok(context) +} + +/// Also returns a fetch that must be applied to the replacement sort when +/// removing an ordered merge whose ordering satisfied the requirement. A +/// `None` value means any fetch remains enforced within the returned context. +fn replace_order_preserving_variants_with_fetch( + mut context: DistributionContext, + 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_with_fetch(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 +947,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. @@ -1361,7 +1394,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::() { @@ -1517,7 +1551,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) => { @@ -1630,17 +1664,23 @@ pub fn ensure_distribution( && !streaming_benefit && context.data { - context = replace_order_preserving_variants(context)?; + let (replaced_context, preserved_fetch) = + replace_order_preserving_variants_with_fetch( + 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), )?; } } @@ -1722,9 +1762,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 c90e58bad288d466e78d15c9880138aa440b862d Mon Sep 17 00:00:00 2001 From: "xudong.w" Date: Mon, 7 Sep 2026 14:07:52 +0800 Subject: [PATCH 2/3] fix: preserve TopK position when replacing ordered merges --- .../enforce_distribution.rs | 99 ++++++++++++++++++- .../enforce_distribution.rs | 74 +++++++------- 2 files changed, 131 insertions(+), 42 deletions(-) diff --git a/datafusion/core/tests/physical_optimizer/enforce_distribution.rs b/datafusion/core/tests/physical_optimizer/enforce_distribution.rs index 91cd64e208dea..03847207407d3 100644 --- a/datafusion/core/tests/physical_optimizer/enforce_distribution.rs +++ b/datafusion/core/tests/physical_optimizer/enforce_distribution.rs @@ -26,7 +26,7 @@ use crate::physical_optimizer::test_utils::{ sort_merge_join_exec, sort_preserving_merge_exec, union_exec, }; -use arrow::array::{RecordBatch, UInt8Array, UInt64Array}; +use arrow::array::{Int64Array, RecordBatch, UInt8Array, UInt64Array}; use arrow::compute::SortOptions; use arrow_schema::{DataType, Field, Schema, SchemaRef}; use datafusion::config::ConfigOptions; @@ -45,6 +45,7 @@ use datafusion_common::tree_node::{ }; use datafusion_datasource::file_groups::FileGroup; use datafusion_datasource::file_scan_config::FileScanConfigBuilder; +use datafusion_datasource::memory::MemorySourceConfig; use datafusion_expr::{JoinType, Operator}; use datafusion_functions_aggregate::count::count_udaf; use datafusion_physical_expr::aggregate::AggregateExprBuilder; @@ -75,7 +76,7 @@ use datafusion_physical_plan::sorts::sort_preserving_merge::SortPreservingMergeE use datafusion_physical_plan::union::UnionExec; use datafusion_physical_plan::{ ChildrenPropertiesMode, DisplayAs, DisplayFormatType, ExecutionPlanProperties, - PlanProperties, ReplaceChildrenOptions, displayable, + PlanProperties, ReplaceChildrenOptions, collect, displayable, }; use insta::Settings; @@ -4624,6 +4625,100 @@ fn move_fetch_to_replacement_sort() -> Result<()> { Ok(()) } +#[tokio::test] +async fn preserve_fetch_below_filter_when_reoptimizing() -> Result<()> { + check_fetch_below_filter( + Operator::Gt, + [vec![-2, 0, 2, 4], vec![-1, 1, 3, 5]], + &[1, 2], + ) + .await +} + +#[tokio::test] +async fn preserve_fetch_below_filter_with_constant_ordering() -> Result<()> { + check_fetch_below_filter( + Operator::Eq, + [vec![-2, 0, 0, 0], vec![-1, 0, 0, 0]], + &[0, 0, 0], + ) + .await +} + +async fn check_fetch_below_filter( + op: Operator, + partitions: [Vec; 2], + expected: &[i64], +) -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new("c", DataType::Int64, false)])); + let sort_key: LexOrdering = + [PhysicalSortExpr::new_default(col("c", &schema)?)].into(); + let partitions = partitions + .into_iter() + .map(|values| { + RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int64Array::from(values))], + ) + .map(|batch| vec![batch]) + }) + .collect::, _>>()?; + let source = MemorySourceConfig::try_new(&partitions, Arc::clone(&schema), None)? + .try_with_sort_information(vec![sort_key.clone()])?; + let merge: Arc = Arc::new( + SortPreservingMergeExec::new( + sort_key.clone(), + DataSourceExec::from_data_source(source), + ) + .with_fetch(Some(5)), + ); + let predicate = Arc::new(BinaryExpr::new(col("c", &schema)?, op, lit(0_i64))); + let filter: Arc = Arc::new(FilterExec::try_new(predicate, merge)?); + let mut plan = sort_required_exec_with_req(filter, sort_key); + let mut config = test_suite_default_config_options(); + config.optimizer.enable_round_robin_repartition = false; + let task_context = SessionContext::new().task_ctx(); + + // The test operator only declares ordering requirements. Execute its child + // to compare query results before optimization and after repeated passes. + for iteration in 0..3 { + if iteration > 0 { + let distribution = DistributionContext::new_default(Arc::clone(&plan)) + .transform_up(|context| ensure_distribution(context, &config))? + .data; + check_integrity(distribution)?; + plan = EnsureRequirements::new().optimize(plan, &config)?; + } + let input = Arc::clone(plan.children()[0]); + let batches = collect(input, Arc::clone(&task_context)).await?; + let values = batches + .iter() + .flat_map(|batch| { + batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .iter() + .copied() + }) + .collect::>(); + assert_eq!( + values, + expected, + "iteration {iteration}:\n{}", + displayable(plan.as_ref()).indent(true) + ); + assert!( + plan.children()[0].is::(), + "fetch must stay below the filter:\n{}", + displayable(plan.as_ref()).indent(true) + ); + } + 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 fc49c5c9d8fc3..546225f83ee26 100644 --- a/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs +++ b/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs @@ -64,6 +64,7 @@ use datafusion_physical_plan::joins::{ }; use datafusion_physical_plan::projection::{ProjectionExec, ProjectionExpr}; use datafusion_physical_plan::repartition::RepartitionExec; +use datafusion_physical_plan::sorts::sort::SortExec; use datafusion_physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec; use datafusion_physical_plan::statistics::{StatisticsArgs, StatisticsContext}; use datafusion_physical_plan::tree_node::PlanContext; @@ -901,45 +902,42 @@ fn remove_dist_changing_operators( pub fn replace_order_preserving_variants( context: DistributionContext, ) -> Result { - let (context, fetch) = replace_order_preserving_variants_with_fetch(context, false)?; - debug_assert!( - fetch.is_none(), - "fetch must stay in the plan when no replacement sort is needed" - ); - Ok(context) + replace_order_preserving_variants_impl(context, false) } -/// Also returns a fetch that must be applied to the replacement sort when -/// removing an ordered merge whose ordering satisfied the requirement. A -/// `None` value means any fetch remains enforced within the returned context. -fn replace_order_preserving_variants_with_fetch( +fn replace_order_preserving_variants_impl( mut context: DistributionContext, 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_with_fetch(child, ordering_satisfied)?; - children.push(child); - fetch = min_fetch(fetch, child_fetch); - } else { - children.push(child); - } - } - context.children = children; +) -> Result { + context.children = context + .children + .into_iter() + .map(|child| { + if child.data { + replace_order_preserving_variants_impl(child, ordering_satisfied) + } else { + Ok(child) + } + }) + .collect::>>()?; - if is_sort_preserving_merge(&context.plan) { - let fetch = min_fetch(fetch, context.plan.fetch()); + if let Some(spm) = context.plan.downcast_ref::() { let child_plan = Arc::clone(&context.children[0].plan); - if ordering_satisfied { + let fetch = spm.fetch(); + if ordering_satisfied && fetch.is_some() { + // A fetched merge selects the first rows in its sort order. Moving + // fetch to an ancestor's sort can cross a filter, or lose the limit + // entirely if the ancestor needs no additional sort. + let ordering = spm.expr().clone(); context.plan = Arc::new(CoalescePartitionsExec::new(child_plan)); - return Ok((context, fetch)); + let sort = Arc::new( + SortExec::new(ordering, Arc::clone(&context.plan)).with_fetch(fetch), + ); + return Ok(DistributionContext::new(sort, false, vec![context])); } context.plan = Arc::new(CoalescePartitionsExec::new(child_plan).with_fetch(fetch)); - return Ok((context, None)); + return Ok(context); } else if let Some(repartition) = context.plan.downcast_ref::() && repartition.preserve_order() { @@ -947,12 +945,10 @@ fn replace_order_preserving_variants_with_fetch( Arc::clone(&context.children[0].plan), repartition.partitioning().clone(), )?); - return Ok((context, fetch)); + return Ok(context); } - context - .update_plan_from_children() - .map(|context| (context, fetch)) + context.update_plan_from_children() } /// A struct to keep track of repartition requirements for each child node. @@ -1664,12 +1660,10 @@ pub fn ensure_distribution( && !streaming_benefit && context.data { - let (replaced_context, preserved_fetch) = - replace_order_preserving_variants_with_fetch( - context, - ordering_satisfied, - )?; - context = replaced_context; + context = replace_order_preserving_variants_impl( + context, + ordering_satisfied, + )?; // If ordering requirements were satisfied before repartitioning, // make sure ordering requirements are still satisfied after. if ordering_satisfied { @@ -1680,7 +1674,7 @@ pub fn ensure_distribution( context = add_sort_above_with_check( context, sort_req, - min_fetch(preserved_fetch, output_fetch), + output_fetch, )?; } } From 809870632aac65dbe98db5ee433a87e23fda4e12 Mon Sep 17 00:00:00 2001 From: "xudong.w" Date: Mon, 7 Sep 2026 16:01:26 +0800 Subject: [PATCH 3/3] fix: preserve fetched distribution operator boundaries --- .../enforce_distribution.rs | 204 +++++++++++++++--- .../replace_with_order_preserving_variants.rs | 53 ++--- .../enforce_distribution.rs | 133 +++--------- .../enforce_sorting/mod.rs | 14 +- .../replace_with_order_preserving_variants.rs | 30 +-- 5 files changed, 243 insertions(+), 191 deletions(-) diff --git a/datafusion/core/tests/physical_optimizer/enforce_distribution.rs b/datafusion/core/tests/physical_optimizer/enforce_distribution.rs index 03847207407d3..2837b64c19ef7 100644 --- a/datafusion/core/tests/physical_optimizer/enforce_distribution.rs +++ b/datafusion/core/tests/physical_optimizer/enforce_distribution.rs @@ -72,6 +72,7 @@ use datafusion_physical_plan::filter::FilterExec; use datafusion_physical_plan::joins::utils::JoinOn; use datafusion_physical_plan::limit::{GlobalLimitExec, LocalLimitExec}; use datafusion_physical_plan::projection::{ProjectionExec, ProjectionExpr}; +use datafusion_physical_plan::sorts::sort::SortExec; use datafusion_physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec; use datafusion_physical_plan::union::UnionExec; use datafusion_physical_plan::{ @@ -4543,11 +4544,12 @@ fn test_replace_order_preserving_variants_with_fetch() -> Result<()> { // Apply the function let result = replace_order_preserving_variants(dist_context)?; - // Verify the plan was transformed to CoalescePartitionsExec + // A fetched ordered merge must still select the TopK rows. + let result = check_integrity(result)?; result .plan - .downcast_ref::() - .expect("Expected CoalescePartitionsExec"); + .downcast_ref::() + .expect("Expected a TopK SortExec"); // Verify fetch was preserved assert_eq!( @@ -4597,31 +4599,183 @@ fn preserve_fetch_when_reoptimizing_coalesce_partitions() -> Result<()> { 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.clone(), input).with_fetch(Some(5)), - ); - let plan = sort_required_exec_with_req(merge, sort_key); +#[tokio::test] +async fn move_fetch_to_replacement_sort() -> Result<()> { + for (options, partitions, expected) in [ + ( + SortOptions::default(), + [ + vec![None, Some(1), Some(1), Some(6)], + vec![None, Some(1), Some(2), Some(7)], + ], + vec![None, None, Some(1), Some(1), Some(1)], + ), + ( + SortOptions { + descending: true, + nulls_first: false, + }, + [vec![Some(7), Some(1), None], vec![Some(6), Some(1), None]], + vec![Some(7), Some(6), Some(1), Some(1), None], + ), + ] { + let (input, sort_key) = sorted_memory_input(partitions, options)?; + let merge: Arc = Arc::new( + SortPreservingMergeExec::new(sort_key.clone(), input).with_fetch(Some(5)), + ); + assert_eq!(fetch_test_values(Arc::clone(&merge)).await?, expected); + let plan = sort_required_exec_with_req(merge, sort_key); + let optimized = ensure_distribution_helper(plan, 10, false)?; + let replacement = Arc::clone(optimized.children()[0]); + let sort = replacement + .downcast_ref::() + .expect("expected a replacement sort"); + assert_eq!(sort.fetch(), Some(5)); + assert_eq!(fetch_test_values(replacement).await?, expected); + } + Ok(()) +} - let optimized = ensure_distribution_helper(plan, 10, false)?; - let plan = displayable(optimized.as_ref()).indent(true).to_string(); +#[tokio::test] +async fn preserve_fetch_in_nested_distribution_operators() -> Result<()> { + for outer_fetch in [0, 3, 10] { + let (input, sort_key) = sorted_memory_input( + [0, 1].map(|start| (start..10).step_by(2).map(Some).collect()), + SortOptions::default(), + )?; + let merge: Arc = + Arc::new(SortPreservingMergeExec::new(sort_key, input).with_fetch(Some(5))); + let plan: Arc = + Arc::new(CoalescePartitionsExec::new(merge).with_fetch(Some(outer_fetch))); + let expected = (0..outer_fetch.min(5)) + .map(|value| Some(value as i64)) + .collect::>(); + assert_reoptimized_fetch_values(plan, &expected).await?; + } + Ok(()) +} - assert!( - plan.contains( - "SortExec: TopK(fetch=5), expr=[c@2 ASC], preserve_partitioning=[false]" - ), - "expected the replacement sort to preserve fetch:\n{plan}" - ); - assert!( - !plan.contains("CoalescePartitionsExec: fetch=5"), - "fetch below the replacement sort would change TopK results:\n{plan}" - ); +#[tokio::test] +async fn preserve_topk_when_parent_changes_ordering() -> Result<()> { + let (input, sort_key) = sorted_memory_input( + [0, 1].map(|start| (start..10).step_by(2).map(Some).collect()), + SortOptions::default(), + )?; + let descending = [PhysicalSortExpr::new( + col("c", &input.schema())?, + SortOptions { + descending: true, + nulls_first: false, + }, + )] + .into(); + let merge: Arc = + Arc::new(SortPreservingMergeExec::new(sort_key, input).with_fetch(Some(5))); + let plan: Arc = Arc::new(SortExec::new(descending, merge)); + assert_reoptimized_fetch_values(plan, &[Some(4), Some(3), Some(2), Some(1), Some(0)]) + .await +} + +#[tokio::test] +async fn preserve_fetch_when_parallelizing_sort_above_filter() -> Result<()> { + let (input, sort_key) = sorted_memory_input( + [ + vec![Some(-4), Some(-2), Some(2), Some(4), Some(6)], + vec![Some(-3), Some(-1), Some(3), Some(5), Some(7)], + ], + SortOptions::default(), + )?; + let predicate = Arc::new(BinaryExpr::new( + col("c", &input.schema())?, + Operator::Gt, + lit(0_i64), + )); + let coalesce: Arc = + Arc::new(CoalescePartitionsExec::new(input).with_fetch(Some(5))); + let filter: Arc = + Arc::new(FilterExec::try_new(predicate, coalesce)?); + let mut plan: Arc = Arc::new(SortExec::new(sort_key, filter)); + let mut config = test_suite_default_config_options(); + config.optimizer.enable_round_robin_repartition = false; + config.optimizer.repartition_sorts = true; + for iteration in 0..3 { + if iteration > 0 { + plan = EnsureRequirements::new().optimize(plan, &config)?; + } + // Either input batch can arrive first. Both contain three positive + // rows, so keeping the limit below the filter always returns three. + assert_eq!( + fetch_test_values(Arc::clone(&plan)).await?.len(), + 3, + "iteration {iteration}:\n{}", + displayable(plan.as_ref()).indent(true) + ); + } + Ok(()) +} +fn sorted_memory_input( + partitions: [Vec>; 2], + options: SortOptions, +) -> Result<(Arc, LexOrdering)> { + let schema = Arc::new(Schema::new(vec![Field::new("c", DataType::Int64, true)])); + let order: LexOrdering = [PhysicalSortExpr::new(col("c", &schema)?, options)].into(); + let partitions = partitions + .into_iter() + .map(|values| { + RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int64Array::from(values))], + ) + .map(|batch| vec![batch]) + }) + .collect::, _>>()?; + let source = MemorySourceConfig::try_new(&partitions, schema, None)? + .try_with_sort_information(vec![order.clone()])?; + Ok((DataSourceExec::from_data_source(source), order)) +} + +async fn fetch_test_values(plan: Arc) -> Result>> { + let batches = collect(plan, SessionContext::new().task_ctx()).await?; + Ok(batches + .iter() + .flat_map(|batch| { + batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .iter() + }) + .collect()) +} + +async fn assert_reoptimized_fetch_values( + plan: Arc, + expected: &[Option], +) -> Result<()> { + for repartition_sorts in [false, true] { + let mut optimized = Arc::clone(&plan); + let mut config = test_suite_default_config_options(); + config.optimizer.enable_round_robin_repartition = false; + config.optimizer.repartition_sorts = repartition_sorts; + for iteration in 0..3 { + if iteration > 0 { + let distribution = + DistributionContext::new_default(Arc::clone(&optimized)) + .transform_up(|context| ensure_distribution(context, &config))? + .data; + check_integrity(distribution)?; + optimized = EnsureRequirements::new().optimize(optimized, &config)?; + } + assert_eq!( + fetch_test_values(Arc::clone(&optimized)).await?, + expected, + "iteration {iteration}, repartition_sorts={repartition_sorts}:\n{}", + displayable(optimized.as_ref()).indent(true) + ); + } + } Ok(()) } diff --git a/datafusion/core/tests/physical_optimizer/replace_with_order_preserving_variants.rs b/datafusion/core/tests/physical_optimizer/replace_with_order_preserving_variants.rs index 601667ea02c0d..b8adb283c2d31 100644 --- a/datafusion/core/tests/physical_optimizer/replace_with_order_preserving_variants.rs +++ b/datafusion/core/tests/physical_optimizer/replace_with_order_preserving_variants.rs @@ -31,7 +31,7 @@ use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; use arrow::record_batch::RecordBatch; use insta::{allow_duplicates, assert_snapshot}; use datafusion_common::tree_node::{TransformedResult, TreeNode}; -use datafusion_common::{assert_contains, NullEquality, Result}; +use datafusion_common::{NullEquality, Result}; use datafusion_common::config::ConfigOptions; use datafusion_datasource::source::DataSourceExec; use datafusion_execution::TaskContext; @@ -1199,43 +1199,22 @@ fn test_plan_with_order_preserving_variants_preserves_fetch() -> Result<()> { .with_fetch(Some(10)) .unwrap(); - // Test sort's fetch is greater than coalesce fetch, return error because it's not reasonable - let requirements = OrderPreservationContext::new( - coalesced.clone(), - false, - vec![OrderPreservationContext::new( - parquet_exec.clone(), - false, - vec![], - )], - ); - let res = plan_with_order_preserving_variants(requirements, false, true, Some(15)); - assert_contains!( - res.unwrap_err().to_string(), - "CoalescePartitionsExec fetch [10] should be greater than or equal to SortExec fetch [15]" - ); - - // Test sort is without fetch, expected to get the fetch value from the coalesced - let requirements = OrderPreservationContext::new( - coalesced.clone(), - false, - vec![OrderPreservationContext::new( - parquet_exec.clone(), + // Keep the coalesce's row selection independently of an ancestor's fetch. + for sort_fetch in [Some(15), None, Some(5)] { + let requirements = OrderPreservationContext::new( + Arc::clone(&coalesced), false, - vec![], - )], - ); - let res = plan_with_order_preserving_variants(requirements, false, true, None)?; - assert_eq!(res.plan.fetch(), Some(10),); - - // Test sort's fetch is less than coalesces fetch, expected to get the fetch value from the sort - let requirements = OrderPreservationContext::new( - coalesced, - false, - vec![OrderPreservationContext::new(parquet_exec, false, vec![])], - ); - let res = plan_with_order_preserving_variants(requirements, false, true, Some(5))?; - assert_eq!(res.plan.fetch(), Some(5),); + vec![OrderPreservationContext::new( + parquet_exec.clone(), + false, + vec![], + )], + ); + let res = + plan_with_order_preserving_variants(requirements, false, true, sort_fetch)?; + assert!(Arc::ptr_eq(&res.plan, &coalesced)); + assert_eq!(res.plan.fetch(), Some(10)); + } Ok(()) } diff --git a/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs b/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs index 546225f83ee26..6e3f82ec16bdf 100644 --- a/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs +++ b/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs @@ -41,7 +41,6 @@ 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,10 +783,7 @@ 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, - fetch: &mut Option, -) -> DistributionContext { +fn add_merge_on_top(input: DistributionContext) -> 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 @@ -798,16 +794,13 @@ fn add_merge_on_top( // (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)); - spm = spm.with_fetch(fetch.take()); - Arc::new(spm) + Arc::new(SortPreservingMergeExec::new( + req.clone(), + Arc::clone(&input.plan), + )) } else { // If there is no input order, we can simply coalesce partitions: - Arc::new( - CoalescePartitionsExec::new(Arc::clone(&input.plan)) - .with_fetch(fetch.take()), - ) + Arc::new(CoalescePartitionsExec::new(Arc::clone(&input.plan))) }; DistributionContext::new(new_plan, true, vec![input]) @@ -833,52 +826,21 @@ fn add_merge_on_top( /// ```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, - /// 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), - } -} - +/// A distribution operator with a fetch also selects rows. Stop at that +/// boundary so neither its limit nor an ordered merge's TopK selection is +/// moved across another operator. 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) +) -> DistributionContext { + while distribution_context.plan.fetch().is_none() + && (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() { - 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: + // All of the above operators have a single child. distribution_context = distribution_context.children.swap_remove(0); - // Note that they will be re-inserted later on if necessary or helpful. } - - Ok(RemovedDistOps { - context: distribution_context, - removed_fetch, - fetch_plan, - }) + distribution_context } /// Updates the [`DistributionContext`] if preserving ordering while changing partitioning is not helpful or desirable. @@ -900,21 +862,14 @@ fn remove_dist_changing_operators( /// " DataSourceExec: file_groups={2 groups: \[\[x], \[y]]}, projection=\[a, b, c, d, e], output_ordering=\[a@0 ASC], file_type=parquet", /// ``` pub fn replace_order_preserving_variants( - context: DistributionContext, -) -> Result { - replace_order_preserving_variants_impl(context, false) -} - -fn replace_order_preserving_variants_impl( mut context: DistributionContext, - ordering_satisfied: bool, ) -> Result { context.children = context .children .into_iter() .map(|child| { if child.data { - replace_order_preserving_variants_impl(child, ordering_satisfied) + replace_order_preserving_variants(child) } else { Ok(child) } @@ -924,7 +879,7 @@ fn replace_order_preserving_variants_impl( if let Some(spm) = context.plan.downcast_ref::() { let child_plan = Arc::clone(&context.children[0].plan); let fetch = spm.fetch(); - if ordering_satisfied && fetch.is_some() { + if fetch.is_some() { // A fetched merge selects the first rows in its sort order. Moving // fetch to an ancestor's sort can cross a filter, or lose the limit // entirely if the ancestor needs no additional sort. @@ -935,8 +890,7 @@ fn replace_order_preserving_variants_impl( ); return Ok(DistributionContext::new(sort, false, vec![context])); } - context.plan = - Arc::new(CoalescePartitionsExec::new(child_plan).with_fetch(fetch)); + context.plan = Arc::new(CoalescePartitionsExec::new(child_plan)); return Ok(context); } else if let Some(repartition) = context.plan.downcast_ref::() && repartition.preserve_order() @@ -1381,18 +1335,13 @@ 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. - // Preserve fetch from removed SPM/Coalesce (#14150). - let RemovedDistOps { - context: - DistributionContext { - mut plan, - data, - children, - }, - mut removed_fetch, - fetch_plan, - } = remove_dist_changing_operators(dist_context)?; + // Remove distribution-only operators, retaining any fetched operator as + // a row-selection boundary. + let DistributionContext { + mut plan, + data, + children, + } = remove_dist_changing_operators(dist_context); if let Some(exec) = plan.downcast_ref::() { if let Some(updated_window) = get_best_fitting_window( @@ -1547,7 +1496,7 @@ pub fn ensure_distribution( // Satisfy the distribution requirement if it is unmet. match &requirement { Distribution::SinglePartition => { - child = add_merge_on_top(child, &mut removed_fetch); + child = add_merge_on_top(child); } Distribution::HashPartitioned(exprs) | Distribution::KeyPartitioned(exprs) => { @@ -1660,10 +1609,7 @@ pub fn ensure_distribution( && !streaming_benefit && context.data { - context = replace_order_preserving_variants_impl( - context, - ordering_satisfied, - )?; + context = replace_order_preserving_variants(context)?; // If ordering requirements were satisfied before repartitioning, // make sure ordering requirements are still satisfied after. if ordering_satisfied { @@ -1756,26 +1702,9 @@ pub fn ensure_distribution( replace_children_if_necessary(plan, children_plans)? }; - 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)) + Ok(Transformed::yes(DistributionContext::new( + plan, data, children, + ))) } /// Keeps track of distribution changing operators (like `RepartitionExec`, diff --git a/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/mod.rs b/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/mod.rs index 42a157257341d..697a3f3d686f5 100644 --- a/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/mod.rs +++ b/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/mod.rs @@ -170,8 +170,10 @@ fn update_coalesce_ctx_children( // Plan has no children, it cannot be a `CoalescePartitionsExec`. false } else if is_coalesce_partitions(&coalesce_context.plan) { - // Initiate a connection: - true + // A fetched coalesce selects rows before the sort. Removing it would + // lose that limit, and moving its fetch onto the sort changes which + // rows are selected. + coalesce_context.plan.fetch().is_none() } else { children.iter().enumerate().any(|(idx, node)| { // Only consider operators that don't require a single partition, @@ -668,11 +670,9 @@ fn remove_bottleneck_in_subplan_impl( Some(Distribution::SinglePartition) ) }; - let remove_from_first_child = requirements - .children - .first() - .is_some_and(|child| is_coalesce_partitions(&child.plan)) - && removable(0); + let remove_from_first_child = requirements.children.first().is_some_and(|child| { + is_coalesce_partitions(&child.plan) && child.plan.fetch().is_none() + }) && removable(0); let children = &mut requirements.children; if remove_from_first_child { // We can safely use the 0th index since we have a `CoalescePartitionsExec`. diff --git a/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/replace_with_order_preserving_variants.rs b/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/replace_with_order_preserving_variants.rs index 6ab84dc95eab9..2f9d19569b357 100644 --- a/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/replace_with_order_preserving_variants.rs +++ b/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/replace_with_order_preserving_variants.rs @@ -25,9 +25,9 @@ use crate::utils::{ is_coalesce_partitions, is_repartition, is_sort, is_sort_preserving_merge, }; +use datafusion_common::Result; use datafusion_common::config::ConfigOptions; use datafusion_common::tree_node::Transformed; -use datafusion_common::{Result, assert_or_internal_err}; use datafusion_physical_plan::ExecutionPlanProperties; use datafusion_physical_plan::coalesce_partitions::CoalescePartitionsExec; use datafusion_physical_plan::execution_plan::EmissionType; @@ -64,7 +64,9 @@ pub fn update_order_preservation_ctx_children_data(opc: &mut OrderPreservationCo } let plan_children = plan.children(); - *data = if plan_children.is_empty() { + *data = if plan_children.is_empty() + || (is_coalesce_partitions(plan) && plan.fetch().is_some()) + { false } else if !children[0].data && ((is_repartition(plan) && !maintains_input_order[0]) @@ -102,6 +104,12 @@ pub fn plan_with_order_preserving_variants( is_spm_better: bool, fetch: Option, ) -> Result { + if is_coalesce_partitions(&sort_input.plan) && sort_input.plan.fetch().is_some() { + // A fetched coalesce selects rows in arrival order. An ordered merge + // would select different rows, even with the same fetch value. + sort_input.data = false; + return Ok(sort_input); + } sort_input.children = sort_input .children .into_iter() @@ -137,24 +145,6 @@ pub fn plan_with_order_preserving_variants( } else if is_coalesce_partitions(&sort_input.plan) && is_spm_better { let child = &sort_input.children[0].plan; if let Some(ordering) = child.output_ordering() { - let mut fetch = fetch; - if let Some(coalesce_fetch) = sort_input.plan.fetch() { - fetch = match fetch { - Some(sort_fetch) => { - assert_or_internal_err!( - coalesce_fetch >= sort_fetch, - "CoalescePartitionsExec fetch [{:?}] should be greater than or equal to SortExec fetch [{:?}]", - coalesce_fetch, - sort_fetch - ); - Some(sort_fetch) - } - None => { - // If the sort node does not have a fetch, we need to keep the coalesce node's fetch. - Some(coalesce_fetch) - } - }; - }; // When the input of a `CoalescePartitionsExec` has an ordering, // replace it with a `SortPreservingMergeExec` if appropriate: let spm = SortPreservingMergeExec::new(ordering.clone(), Arc::clone(child))