diff --git a/datafusion/core/tests/physical_optimizer/enforce_distribution.rs b/datafusion/core/tests/physical_optimizer/enforce_distribution.rs index 4aeee422394ca..49f435c328de0 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::{Int32Array, RecordBatch, UInt8Array, UInt64Array}; use arrow::compute::SortOptions; use arrow_schema::{DataType, Field, Schema, SchemaRef}; use datafusion::config::ConfigOptions; @@ -4821,42 +4821,376 @@ fn get_schema() -> SchemaRef { Field::new("bank_account", DataType::UInt64, true), ])) } +async fn collect_i32( + plan: Arc, + task_ctx: Arc, +) -> Result> { + Ok(datafusion_physical_plan::collect(plan, task_ctx) + .await? + .iter() + .flat_map(|batch| { + batch + .column(0) + .as_any() + .downcast_ref::() + .expect("test plan produces Int32") + .values() + .to_vec() + }) + .collect()) +} + #[test] -fn test_replace_order_preserving_variants_with_fetch() -> Result<()> { - // Create a base plan - let parquet_exec = parquet_exec(); +fn replace_order_preserving_variants_preserves_fetched_merge_before_recursing() +-> Result<()> { + for fetch in [None, Some(0), Some(5)] { + let ordering: LexOrdering = + [PhysicalSortExpr::new_default(Arc::new(Column::new("a", 0)))].into(); + let input: Arc = + parquet_exec_multiple_sorted(vec![ordering.clone()]); + let repartition = Arc::new( + RepartitionExec::try_new( + Arc::clone(&input), + Partitioning::RoundRobinBatch(2), + )? + .with_preserve_order(), + ); + let child = DistributionContext::new( + repartition, + true, + vec![DistributionContext::new(input, false, vec![])], + ); + let context = DistributionContext::new( + Arc::new( + SortPreservingMergeExec::new(ordering, Arc::clone(&child.plan)) + .with_fetch(fetch), + ), + true, + vec![child], + ); + let result = replace_order_preserving_variants(context)?; + assert_eq!(result.plan.fetch(), fetch); + assert_eq!(result.plan.is::(), fetch.is_some()); + assert_eq!(result.plan.is::(), fetch.is_none()); + assert_eq!( + result.children[0] + .plan + .downcast_ref::() + .unwrap() + .preserve_order(), + fetch.is_some() + ); + } + Ok(()) +} - let sort_expr = PhysicalSortExpr::new_default(Arc::new(Column::new("id", 0))); +/// Each case executes unoptimized, then two complete physical-optimizer passes. +#[tokio::test] +async fn fetched_coalesce_survives_distribution_rewrites() -> Result<()> { + use datafusion::physical_planner::DefaultPhysicalPlanner; + use datafusion_physical_plan::test::TestMemoryExec; + + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int32Array::from(vec![1, 1, 1]))], + )?; + let state = + SessionContext::new_with_config(SessionConfig::new().with_target_partitions(3)) + .state(); + let planner = DefaultPhysicalPlanner::default(); + let mut failures = vec![]; + + for partitions in [1, 3] { + for fetch in [0, 1] { + let input = Arc::new(TestMemoryExec::try_new( + &vec![vec![batch.clone()]; partitions], + Arc::clone(&schema), + None, + )?); + let unoptimized: Arc = + Arc::new(CoalescePartitionsExec::new(input).with_fetch(Some(fetch))); + assert_eq!( + collect_i32(Arc::clone(&unoptimized), state.task_ctx()).await?, + vec![1; fetch] + ); + let first = planner.optimize_physical_plan(unoptimized, &state, |_, _| {})?; + let second = + planner.optimize_physical_plan(Arc::clone(&first), &state, |_, _| {})?; + let first_rows = collect_i32(first, state.task_ctx()).await?; + let second_rows = collect_i32(second, state.task_ctx()).await?; + eprintln!( + "fetched_coalesce partitions={partitions}, fetch={fetch}, first={first_rows:?}, second={second_rows:?}" + ); + if first_rows.len() != fetch || second_rows.len() != fetch { + failures.push(format!("partitions={partitions}, fetch={fetch}, first={first_rows:?}, second={second_rows:?}")); + } + } + } + assert!(failures.is_empty(), "{failures:#?}"); + Ok(()) +} - // Create a SortPreservingMergeExec with fetch=5 - let spm_exec = Arc::new( - SortPreservingMergeExec::new([sort_expr].into(), parquet_exec.clone()) - .with_fetch(Some(5)), - ); +/// SQL planning must retain a fetched TopK below an unordered aggregate when +/// a physical plan is optimized again. +#[tokio::test] +async fn sql_fetched_topk_survives_reoptimization_under_unordered_parent() -> Result<()> { + use datafusion::physical_planner::DefaultPhysicalPlanner; + + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, false), + ])); + let partitions = [vec![(1, 10), (100, 99)], vec![(2, 20), (200, 98)]] + .into_iter() + .map(|rows| { + Ok(RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from( + rows.iter().map(|(a, _)| *a).collect::>(), + )), + Arc::new(Int32Array::from( + rows.iter().map(|(_, b)| *b).collect::>(), + )), + ], + )?) + }) + .collect::>>()?; + let table = MemTable::try_new( + schema, + vec![vec![partitions[0].clone()], vec![partitions[1].clone()]], + )?; + let ctx = + SessionContext::new_with_config(SessionConfig::new().with_target_partitions(2)); + ctx.register_table("t", Arc::new(table))?; + let state = ctx.state(); + let planner = DefaultPhysicalPlanner::default(); + + for (sql, unordered_parent) in [ + ("SELECT b FROM t ORDER BY a LIMIT 2", false), + ( + "SELECT b FROM (SELECT a, b FROM t ORDER BY a LIMIT 2) GROUP BY b", + true, + ), + ] { + let plan = ctx.sql(sql).await?.create_physical_plan().await?; + let mut sql_rows = collect_i32(Arc::clone(&plan), state.task_ctx()).await?; + let reoptimized = planner.optimize_physical_plan(plan, &state, |_, _| {})?; + let mut reoptimized_rows = + collect_i32(Arc::clone(&reoptimized), state.task_ctx()).await?; + if unordered_parent { + sql_rows.sort_unstable(); + reoptimized_rows.sort_unstable(); + } + assert_eq!(sql_rows, [10, 20], "SQL plan: {sql}"); + assert_eq!(reoptimized_rows, [10, 20], "reoptimized plan: {sql}"); + } + Ok(()) +} - // Create distribution context - let dist_context = DistributionContext::new( - spm_exec, - true, - vec![DistributionContext::new(parquet_exec, false, vec![])], +/// A fetched TopK must remain ordered even when its aggregate parent cannot use that ordering. +#[tokio::test] +async fn fetched_topk_preserves_selected_rows_under_unordered_parent() -> Result<()> { + use datafusion::physical_planner::DefaultPhysicalPlanner; + use datafusion_physical_plan::test::TestMemoryExec; + + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, false), + ])); + let partitions = [vec![(1, 10), (100, 99)], vec![(2, 20), (200, 98)]] + .into_iter() + .map(|rows| { + Ok(vec![RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from( + rows.iter().map(|(a, _)| *a).collect::>(), + )), + Arc::new(Int32Array::from( + rows.iter().map(|(_, b)| *b).collect::>(), + )), + ], + )?]) + }) + .collect::>>()?; + let ordering: LexOrdering = + [PhysicalSortExpr::new_default(col("a", &schema)?)].into(); + let input = Arc::new( + TestMemoryExec::try_new(&partitions, Arc::clone(&schema), None)? + .try_with_sort_information(vec![ordering.clone()])?, + ); + let repartition = Arc::new( + RepartitionExec::try_new(input, Partitioning::RoundRobinBatch(2))? + .with_preserve_order(), ); + let topk = + Arc::new(SortPreservingMergeExec::new(ordering, repartition).with_fetch(Some(2))); + // Grouping by b (rather than the a sort key) makes this a genuinely unordered parent. + let unoptimized: Arc = Arc::new(AggregateExec::try_new( + AggregateMode::Single, + PhysicalGroupBy::new_single(vec![(col("b", &schema)?, "b".to_string())]), + vec![], + vec![], + topk, + schema, + )?); + let state = + SessionContext::new_with_config(SessionConfig::new().with_target_partitions(2)) + .state(); + let planner = DefaultPhysicalPlanner::default(); + let mut original_rows = + collect_i32(Arc::clone(&unoptimized), state.task_ctx()).await?; + original_rows.sort_unstable(); + assert_eq!(original_rows, [10, 20]); + let first = planner.optimize_physical_plan(unoptimized, &state, |_, _| {})?; + let second = planner.optimize_physical_plan(Arc::clone(&first), &state, |_, _| {})?; + let mut first_rows = collect_i32(first, state.task_ctx()).await?; + let mut second_rows = collect_i32(second, state.task_ctx()).await?; + first_rows.sort_unstable(); + second_rows.sort_unstable(); + eprintln!("fetched_topk first={first_rows:?}, second={second_rows:?}"); + assert_eq!(first_rows, [10, 20]); + assert_eq!(second_rows, [10, 20]); + Ok(()) +} + +/// LIMIT, OFFSET + LIMIT, and OFFSET-only retain their global cardinality on reoptimization. +#[tokio::test] +async fn global_limit_over_single_partitioned_aggregate_survives_reoptimization() +-> Result<()> { + use datafusion::physical_planner::DefaultPhysicalPlanner; + use datafusion_physical_plan::test::TestMemoryExec; + + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let partitions = [1, 2, 3] + .into_iter() + .map(|value| { + Ok(vec![RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int32Array::from(vec![value]))], + )?]) + }) + .collect::>>()?; + let state = + SessionContext::new_with_config(SessionConfig::new().with_target_partitions(3)) + .state(); + let planner = DefaultPhysicalPlanner::default(); + for (skip, fetch, expected) in [(0, Some(1), 1), (1, Some(1), 1), (1, None, 2)] { + let input = Arc::new(TestMemoryExec::try_new( + &partitions, + Arc::clone(&schema), + None, + )?); + let aggregate = Arc::new(AggregateExec::try_new( + AggregateMode::SinglePartitioned, + PhysicalGroupBy::new_single(vec![(col("a", &schema)?, "a".to_string())]), + vec![], + vec![], + input, + Arc::clone(&schema), + )?); + let unoptimized: Arc = + Arc::new(GlobalLimitExec::new(aggregate, skip, fetch)); + let first = planner.optimize_physical_plan(unoptimized, &state, |_, _| {})?; + let second = + planner.optimize_physical_plan(Arc::clone(&first), &state, |_, _| {})?; + let first_rows = collect_i32(first, state.task_ctx()).await?.len(); + let second_rows = collect_i32(second, state.task_ctx()).await?.len(); + eprintln!( + "global_limit skip={skip}, fetch={fetch:?}, first={first_rows}, second={second_rows}" + ); + assert_eq!(first_rows, expected); + assert_eq!(second_rows, expected); + } + Ok(()) +} - // Apply the function - let result = replace_order_preserving_variants(dist_context)?; +/// A local limit still caps each partition, rather than the combined stream. +#[tokio::test] +async fn local_limit_remains_per_partition_after_reoptimization() -> Result<()> { + use datafusion::physical_planner::DefaultPhysicalPlanner; + use datafusion_physical_plan::test::TestMemoryExec; - // Verify the plan was transformed to CoalescePartitionsExec - result - .plan - .downcast_ref::() - .expect("Expected CoalescePartitionsExec"); + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let partitions = [vec![1, 10], vec![2, 20], vec![3, 30]] + .into_iter() + .map(|values| { + Ok(vec![RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from(values))], + )?]) + }) + .collect::>>()?; + let input = Arc::new(TestMemoryExec::try_new(&partitions, schema, None)?); + let mut plan: Arc = Arc::new(LocalLimitExec::new(input, 1)); + let state = + SessionContext::new_with_config(SessionConfig::new().with_target_partitions(3)) + .state(); + for pass in 0..=2 { + if pass > 0 { + plan = DefaultPhysicalPlanner::default().optimize_physical_plan( + plan, + &state, + |_, _| {}, + )?; + } + let mut rows = collect_i32(plan.clone(), state.task_ctx()).await?; + rows.sort_unstable(); + assert_eq!(rows, [1, 2, 3], "pass={pass}"); + } + Ok(()) +} - // Verify fetch was preserved - assert_eq!( - result.plan.fetch(), - Some(5), - "Fetch value was not preserved after transformation" - ); +/// OFFSET must discard the same ordered rows exactly once, with or without LIMIT. +#[tokio::test] +async fn ordered_offset_preserves_rows_after_reoptimization() -> Result<()> { + use datafusion::physical_planner::DefaultPhysicalPlanner; + use datafusion_physical_plan::test::TestMemoryExec; + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let partitions = [vec![1, 3], vec![2]] + .into_iter() + .map(|values| { + Ok(vec![RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from(values))], + )?]) + }) + .collect::>>()?; + let ordering: LexOrdering = + [PhysicalSortExpr::new_default(col("a", &schema)?)].into(); + let input = Arc::new( + TestMemoryExec::try_new(&partitions, schema, None)? + .try_with_sort_information(vec![ordering.clone()])?, + ); + let state = + SessionContext::new_with_config(SessionConfig::new().with_target_partitions(2)) + .state(); + for (fetch, expected) in [(Some(1), vec![2]), (None, vec![2, 3])] { + let merge = Arc::new( + SortPreservingMergeExec::new(ordering.clone(), input.clone()) + .with_fetch(fetch.map(|n| n + 1)), + ); + let mut limit = GlobalLimitExec::new(merge, 1, fetch); + limit.set_required_ordering(Some(ordering.clone())); + let mut plan: Arc = Arc::new(limit); + for pass in 0..=2 { + if pass > 0 { + plan = DefaultPhysicalPlanner::default().optimize_physical_plan( + plan, + &state, + |_, _| {}, + )?; + } + assert_eq!( + collect_i32(plan.clone(), state.task_ctx()).await?, + expected, + "fetch={fetch:?}, pass={pass}" + ); + } + } Ok(()) } diff --git a/datafusion/core/tests/physical_optimizer/enforce_sorting.rs b/datafusion/core/tests/physical_optimizer/enforce_sorting.rs index f298fa1e47097..b2d14293e3a53 100644 --- a/datafusion/core/tests/physical_optimizer/enforce_sorting.rs +++ b/datafusion/core/tests/physical_optimizer/enforce_sorting.rs @@ -2321,13 +2321,15 @@ async fn test_remove_unnecessary_spm2() -> Result<()> { ); let test = EnforceSortingTest::new(input.clone()).with_repartition_sorts(true); - assert_snapshot!(test.run(), @r" + assert_snapshot!(test.run(), @" Input Plan: SortPreservingMergeExec: [non_nullable_col@1 ASC], fetch=100 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 ae9774c9f8c2d..9ca42603e7f86 100644 --- a/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs +++ b/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs @@ -813,10 +813,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, -) -> 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 +822,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 = 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() { + 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,30 +856,16 @@ 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, -} - fn remove_dist_changing_operators( mut distribution_context: DistributionContext, -) -> Result { - let mut removed_fetch = None; +) -> Result { 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), - ); + // Fetch-bearing distribution operators also select rows globally. + 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: @@ -895,10 +873,7 @@ fn remove_dist_changing_operators( // Note that they will be re-inserted later on if necessary or helpful. } - Ok(RemovedDistOps { - context: distribution_context, - removed_fetch, - }) + Ok(distribution_context) } /// Updates the [`DistributionContext`] if preserving ordering while changing partitioning is not helpful or desirable. @@ -922,6 +897,12 @@ fn remove_dist_changing_operators( pub fn replace_order_preserving_variants( mut context: DistributionContext, ) -> Result { + // A fetched merge selects globally ordered rows, not merely a distribution. + // Preserve it before rewriting its order-preserving descendants. + if is_sort_preserving_merge(&context.plan) && context.plan.fetch().is_some() { + return Ok(context); + } + context.children = context .children .into_iter() @@ -1384,15 +1365,10 @@ pub fn ensure_distribution( 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, + let DistributionContext { + mut plan, + data, + children, } = remove_dist_changing_operators(dist_context)?; if let Some(exec) = plan.downcast_ref::() { @@ -1548,7 +1524,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) => { 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 e9467d508aa10..54c352e72aa62 100644 --- a/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/mod.rs +++ b/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/mod.rs @@ -742,6 +742,10 @@ fn remove_corresponding_sort_from_sub_plan( mut node: PlanWithCorrespondingSort, requires_single_partition: bool, ) -> Result { + if is_sort_preserving_merge(&node.plan) && node.plan.fetch().is_some() { + return Ok(node); + } + // A `SortExec` is always at the bottom of the tree. if let Some(sort_exec) = node.plan.downcast_ref::() { // Do not remove sorts with fetch: