This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-24809-f8a4ac82cb1fef35a71068302fa50cbc4b62a75b in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit 40488988ad596c9b093ad60e1453430d803ce33c Author: xudong.w <[email protected]> AuthorDate: Wed Sep 9 06:55:19 2026 +0000 fix: preserve fetch across distribution reoptimization (#24809) ## Which issue does this PR close? - Closes #24807. ## Rationale for this change Reoptimizing an already optimized physical plan could silently remove a pushed-down `LIMIT` stored on `SortPreservingMergeExec` or `CoalescePartitionsExec`. This could make a query return more rows than requested. ## What changes are included in this PR? - Consume a removed `fetch` when a replacement merge operator is inserted. - Remember the outermost removed fetch-capable distribution operator and rebuild it around the optimized child when no replacement consumes its limit. - Preserve the minimum effective fetch across nested distribution operators. - Carry a fetched ordered merge's limit to a replacement sort when order-preserving variants are removed. ## What is the testing strategy for this PR? Added targeted physical optimizer regression tests covering: - reoptimizing a fetched `SortPreservingMergeExec`; - reoptimizing a fetched `CoalescePartitionsExec`; - moving a fetched ordered merge's limit to a replacement sort; - updating an existing fetched single-partition merge snapshot that previously encoded the incorrect removal of its limit. Validated with: ```text # Fails on e4cf35cbc (current main before this patch): both regression tests fail cargo test -p datafusion --test core_integration physical_optimizer::enforce_distribution::preserve_fetch_when_reoptimizing # Passes with this patch cargo test -p datafusion --test core_integration physical_optimizer::enforce_distribution::preserve_fetch_when_reoptimizing cargo test -p datafusion --test core_integration physical_optimizer::enforce_distribution::move_fetch_to_replacement_sort cargo test -p datafusion --test core_integration physical_optimizer::enforce_distribution::test_replace_order_preserving_variants_with_fetch cargo fmt --all -- --check cargo clippy --all-targets --all-features -- -D warnings RUST_BACKTRACE=1 cargo test --profile ci --exclude datafusion-examples --exclude datafusion-benchmarks --exclude datafusion-cli --workspace --lib --tests --bins --features avro,json,backtrace,extended_tests,recursive_protection,parquet_encryption ``` ## Are there any user-facing changes? Queries preserve their requested global limit when physical distribution requirements are optimized more than once. There are no changes to SQL behavior other than fixing the incorrect result, and no changes to documented public APIs. --- .../physical_optimizer/enforce_distribution.rs | 325 ++++++++++++++++++++- .../tests/physical_optimizer/enforce_sorting.rs | 4 +- .../replace_with_order_preserving_variants.rs | 53 +--- .../ensure_requirements/enforce_distribution.rs | 116 +++----- .../src/ensure_requirements/enforce_sorting/mod.rs | 14 +- .../replace_with_order_preserving_variants.rs | 30 +- 6 files changed, 404 insertions(+), 138 deletions(-) diff --git a/datafusion/core/tests/physical_optimizer/enforce_distribution.rs b/datafusion/core/tests/physical_optimizer/enforce_distribution.rs index 4aeee42239..6b8e38d150 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; @@ -47,6 +47,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; @@ -76,11 +77,12 @@ 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::repartition::RepartitionExec; +use datafusion_physical_plan::sorts::sort::SortExec; use datafusion_physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec; use datafusion_physical_plan::union::{InterleaveExec, UnionExec}; use datafusion_physical_plan::{ ChildrenPropertiesMode, DisplayAs, DisplayFormatType, ExecutionPlanProperties, - PlanProperties, ReplaceChildrenOptions, displayable, + PlanProperties, ReplaceChildrenOptions, collect, displayable, }; use insta::Settings; @@ -4844,11 +4846,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::<CoalescePartitionsExec>() - .expect("Expected CoalescePartitionsExec"); + .downcast_ref::<SortExec>() + .expect("Expected a TopK SortExec"); // Verify fetch was preserved assert_eq!( @@ -4860,6 +4863,318 @@ 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<dyn ExecutionPlan> = + 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<dyn ExecutionPlan> = + 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::<CoalescePartitionsExec>() + .expect("expected CoalescePartitionsExec"); + + Ok(()) +} + +#[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<dyn ExecutionPlan> = 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::<SortExec>() + .expect("expected a replacement sort"); + assert_eq!(sort.fetch(), Some(5)); + assert_eq!(fetch_test_values(replacement).await?, expected); + } + Ok(()) +} + +#[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<dyn ExecutionPlan> = + Arc::new(SortPreservingMergeExec::new(sort_key, input).with_fetch(Some(5))); + let plan: Arc<dyn ExecutionPlan> = + Arc::new(CoalescePartitionsExec::new(merge).with_fetch(Some(outer_fetch))); + let expected = (0..outer_fetch.min(5)) + .map(|value| Some(value as i64)) + .collect::<Vec<_>>(); + assert_reoptimized_fetch_values(plan, &expected).await?; + } + Ok(()) +} + +#[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<dyn ExecutionPlan> = + Arc::new(SortPreservingMergeExec::new(sort_key, input).with_fetch(Some(5))); + let plan: Arc<dyn ExecutionPlan> = 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<dyn ExecutionPlan> = + Arc::new(CoalescePartitionsExec::new(input).with_fetch(Some(5))); + let filter: Arc<dyn ExecutionPlan> = + Arc::new(FilterExec::try_new(predicate, coalesce)?); + let mut plan: Arc<dyn ExecutionPlan> = 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<Option<i64>>; 2], + options: SortOptions, +) -> Result<(Arc<dyn ExecutionPlan>, 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::<std::result::Result<Vec<_>, _>>()?; + 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<dyn ExecutionPlan>) -> Result<Vec<Option<i64>>> { + let batches = collect(plan, SessionContext::new().task_ctx()).await?; + Ok(batches + .iter() + .flat_map(|batch| { + batch + .column(0) + .as_any() + .downcast_ref::<Int64Array>() + .unwrap() + .iter() + }) + .collect()) +} + +async fn assert_reoptimized_fetch_values( + plan: Arc<dyn ExecutionPlan>, + expected: &[Option<i64>], +) -> 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(()) +} + +#[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<i64>; 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::<std::result::Result<Vec<_>, _>>()?; + let source = MemorySourceConfig::try_new(&partitions, Arc::clone(&schema), None)? + .try_with_sort_information(vec![sort_key.clone()])?; + let merge: Arc<dyn ExecutionPlan> = 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<dyn ExecutionPlan> = 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::<Int64Array>() + .unwrap() + .values() + .iter() + .copied() + }) + .collect::<Vec<_>>(); + assert_eq!( + values, + expected, + "iteration {iteration}:\n{}", + displayable(plan.as_ref()).indent(true) + ); + assert!( + plan.children()[0].is::<FilterExec>(), + "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/core/tests/physical_optimizer/enforce_sorting.rs b/datafusion/core/tests/physical_optimizer/enforce_sorting.rs index f298fa1e47..04019feb66 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/core/tests/physical_optimizer/replace_with_order_preserving_variants.rs b/datafusion/core/tests/physical_optimizer/replace_with_order_preserving_variants.rs index 601667ea02..b8adb283c2 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 ae9774c9f8..0d850d4113 100644 --- a/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs +++ b/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs @@ -63,6 +63,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; @@ -813,10 +814,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: Option<usize>, -) -> 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 @@ -825,21 +823,16 @@ 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<dyn ExecutionPlan> = 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<dyn ExecutionPlan> = + if let Some(req) = input.plan.output_ordering() { + 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))) + }; DistributionContext::new(new_plan, true, vec![input]) } else { @@ -864,41 +857,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<usize>, -} - +/// 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<RemovedDistOps> { - let mut removed_fetch = 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() { - removed_fetch = Some( - removed_fetch - .map(|existing: usize| existing.min(fetch)) - .unwrap_or(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, - }) + distribution_context } /// Updates the [`DistributionContext`] if preserving ordering while changing partitioning is not helpful or desirable. @@ -934,11 +907,21 @@ pub fn replace_order_preserving_variants( }) .collect::<Result<Vec<_>>>()?; - if is_sort_preserving_merge(&context.plan) { + if let Some(spm) = context.plan.downcast_ref::<SortPreservingMergeExec>() { let child_plan = Arc::clone(&context.children[0].plan); - context.plan = Arc::new( - CoalescePartitionsExec::new(child_plan).with_fetch(context.plan.fetch()), - ); + let fetch = spm.fetch(); + 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. + let ordering = spm.expr().clone(); + context.plan = Arc::new(CoalescePartitionsExec::new(child_plan)); + 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)); return Ok(context); } else if let Some(repartition) = context.plan.downcast_ref::<RepartitionExec>() && repartition.preserve_order() @@ -1383,17 +1366,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, - }, - removed_fetch, - } = 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::<WindowAggExec>() { if let Some(updated_window) = get_best_fitting_window( @@ -1548,7 +1527,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); } Distribution::HashPartitioned(exprs) | Distribution::KeyPartitioned(exprs) => { @@ -1666,12 +1645,13 @@ pub fn ensure_distribution( // make sure ordering requirements are still satisfied after. if ordering_satisfied { // Make sure to satisfy ordering requirement: + let output_fetch = plan + .downcast_ref::<OutputRequirementExec>() + .and_then(|output| output.fetch()); context = add_sort_above_with_check( context, sort_req, - plan.downcast_ref::<OutputRequirementExec>() - .map(|output| output.fetch()) - .unwrap_or(None), + output_fetch, )?; } } 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 e9467d508a..cd2b5bb2f4 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 d7b3f16a8f..2f9d19569b 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<usize>, ) -> Result<OrderPreservationContext> { + 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)) --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
