diff --git a/datafusion/core/tests/physical_optimizer/ensure_requirements.rs b/datafusion/core/tests/physical_optimizer/ensure_requirements.rs index 2c6c46c82985a..580d0b658496a 100644 --- a/datafusion/core/tests/physical_optimizer/ensure_requirements.rs +++ b/datafusion/core/tests/physical_optimizer/ensure_requirements.rs @@ -21,6 +21,8 @@ //! so the tests live alongside the rest of the `physical_optimizer/` integration //! suite and can use real `ExecutionPlan`s where convenient. +use insta::assert_snapshot; + use datafusion_common::config::ConfigOptions; use datafusion_physical_optimizer::PhysicalOptimizerRule; use datafusion_physical_optimizer::ensure_requirements::EnsureRequirements; @@ -65,6 +67,19 @@ struct MockMultiPartitionExec { impl MockMultiPartitionExec { fn new(partition_count: usize) -> Self { + Self::with_partitioning(Partitioning::UnknownPartitioning(partition_count)) + } + + /// A source that is already partitioned on `a`, as an aggregate or a partitioned + /// join below the node under test would be. + fn hash_partitioned_on_a(partition_count: usize) -> Self { + Self::with_partitioning(Partitioning::Hash( + vec![Arc::new(Column::new("a", 0))], + partition_count, + )) + } + + fn with_partitioning(partitioning: Partitioning) -> Self { let schema = Arc::new(Schema::new(vec![ Field::new("a", DataType::Int64, false), Field::new("b", DataType::Int64, false), @@ -81,7 +96,7 @@ impl MockMultiPartitionExec { } let properties = PlanProperties::new( eq, - Partitioning::UnknownPartitioning(partition_count), + partitioning, EmissionType::Incremental, Boundedness::Bounded, ); @@ -1252,3 +1267,112 @@ fn test_idempotent_union_projection_sort() { assert_idempotent(plan); } + +/// A `CollectLeft` `HashJoinExec` requires `Distribution::SinglePartition` on its build +/// (left) child, so `EnsureRequirements` puts a `CoalescePartitionsExec` on top of a +/// multi-partition build side. Its sort-parallelization phase must not take that coalesce +/// back out again. +/// +/// The phase descends into a node when *any* of its children is linked to a +/// `CoalescePartitionsExec` below (`update_coalesce_ctx_children`), so a connected probe +/// side is enough to reach the join, and the removal itself used to look only at +/// `children[0]` without consulting the join's own distribution requirement. The result was +/// a build side left multi-partition with nothing to re-enforce distribution afterwards, +/// which `SanityCheckPlan` then rejected with "does not satisfy distribution requirements: +/// SinglePartition". +#[test] +fn test_collect_left_join_keeps_build_side_coalesce() { + let build: Arc = Arc::new(MockMultiPartitionExec::new(4)); + // The probe side carries the `CoalescePartitionsExec` link that makes the traversal + // descend into the join in the first place. + let probe: Arc = Arc::new(CoalescePartitionsExec::new(Arc::new( + MockMultiPartitionExec::new(4), + ))); + + let on = vec![( + Arc::new(Column::new("a", 0)) as Arc, + Arc::new(Column::new("a", 0)) as Arc, + )]; + let join: Arc = Arc::new( + HashJoinExec::try_new( + build, + probe, + on, + None, + &JoinType::Left, + None, + PartitionMode::CollectLeft, + NullEquality::NullEqualsNothing, + false, + ) + .expect("HashJoinExec creation failed"), + ); + + // A global sort on top is what triggers the sort-parallelization phase. + let plan: Arc = + Arc::new(SortExec::new(sort_expr_on("a", 0, false, false), join)); + + let optimized = optimize_and_sanity_check(plan).expect("plan failed SanityCheckPlan"); + + assert_snapshot!(plan_string(&optimized), @r" + SortPreservingMergeExec: [a@0 ASC NULLS LAST] + SortExec: expr=[a@0 ASC NULLS LAST], preserve_partitioning=[true] + HashJoinExec: mode=CollectLeft, join_type=Left, on=[(a@0, a@0)] + CoalescePartitionsExec + MockMultiPartitionExec + RepartitionExec: partitioning=RoundRobinBatch(8), input_partitions=4 + MockMultiPartitionExec + "); + + assert_idempotent(optimized); +} + +/// The same build-side removal, with a build side that is already hash-partitioned on the +/// join key rather than `UnknownPartitioning`. This is the shape a `JoinSelection` input +/// swap leaves behind (a `CollectLeft` join reported as `join_type=Right`) when the build +/// subtree is the output of an aggregate or a partitioned join. +#[test] +fn test_collect_left_join_keeps_hash_partitioned_build_side_coalesce() { + let build: Arc = Arc::new( + MockMultiPartitionExec::hash_partitioned_on_a(TEST_TARGET_PARTITIONS), + ); + let probe: Arc = Arc::new(CoalescePartitionsExec::new(Arc::new( + MockMultiPartitionExec::new(4), + ))); + + let on = vec![( + Arc::new(Column::new("a", 0)) as Arc, + Arc::new(Column::new("a", 0)) as Arc, + )]; + let join: Arc = Arc::new( + HashJoinExec::try_new( + build, + probe, + on, + None, + &JoinType::Right, + None, + PartitionMode::CollectLeft, + NullEquality::NullEqualsNothing, + false, + ) + .expect("HashJoinExec creation failed"), + ); + + let plan: Arc = + Arc::new(SortExec::new(sort_expr_on("a", 0, false, false), join)); + + let optimized = optimize_and_sanity_check(plan).expect("plan failed SanityCheckPlan"); + + assert_snapshot!(plan_string(&optimized), @r" + SortPreservingMergeExec: [a@0 ASC NULLS LAST] + SortExec: expr=[a@0 ASC NULLS LAST], preserve_partitioning=[true] + HashJoinExec: mode=CollectLeft, join_type=Right, on=[(a@0, a@0)] + CoalescePartitionsExec + MockMultiPartitionExec + RepartitionExec: partitioning=RoundRobinBatch(8), input_partitions=4 + MockMultiPartitionExec + "); + + assert_idempotent(optimized); +} 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 4dce4691f0963..e78f6d245a311 100644 --- a/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/mod.rs +++ b/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/mod.rs @@ -609,11 +609,40 @@ fn adjust_window_sort_removal( /// the plan, some of the remaining `RepartitionExec`s might become unnecessary. /// Removes such `RepartitionExec`s from the plan as well. fn remove_bottleneck_in_subplan( + requirements: PlanWithCorrespondingCoalescePartitions, +) -> Result { + // The root is the node `parallelize_sorts` is rewriting (a `SortExec`, + // `SortPreservingMergeExec` or `CoalescePartitionsExec`). Its own distribution + // requirement does not constrain the removal, because the caller drops the node and + // rebuilds the cascade around the result. + remove_bottleneck_in_subplan_impl(requirements, true) +} + +fn remove_bottleneck_in_subplan_impl( mut requirements: PlanWithCorrespondingCoalescePartitions, + is_root: bool, ) -> Result { let plan = &requirements.plan; + // Below the root, a `CoalescePartitionsExec` feeding a child that requires + // `Distribution::SinglePartition` is not an avoidable bottleneck: it is what satisfies + // that requirement. Removing it leaves the parent with a multi-partition input it cannot + // accept, and nothing re-runs distribution enforcement afterwards, so the plan reaches + // `SanityCheckPlan` invalid. The traversal reaches such a node because + // `update_coalesce_ctx_children` marks a node as connected when *any* child qualifies: + // a `CollectLeft` `HashJoinExec` whose probe side is connected is descended into even + // though its build side must stay single-partition. + let removable = |idx: usize| { + is_root + || !matches!( + plan.input_distribution_requirements() + .child_distribution(idx), + Some(Distribution::SinglePartition) + ) + }; + let remove_from_first_child = + is_coalesce_partitions(&requirements.children[0].plan) && removable(0); let children = &mut requirements.children; - if is_coalesce_partitions(&children[0].plan) { + if remove_from_first_child { // We can safely use the 0th index since we have a `CoalescePartitionsExec`. let mut new_child_node = children[0].children.swap_remove(0); while new_child_node.plan.output_partitioning() == plan.output_partitioning() @@ -627,9 +656,10 @@ fn remove_bottleneck_in_subplan( requirements.children = requirements .children .into_iter() - .map(|node| { - if node.data { - remove_bottleneck_in_subplan(node) + .enumerate() + .map(|(idx, node)| { + if node.data && removable(idx) { + remove_bottleneck_in_subplan_impl(node, false) } else { Ok(node) } diff --git a/datafusion/sqllogictest/test_files/joins.slt b/datafusion/sqllogictest/test_files/joins.slt index 3b8f66def3c34..55efcc3874fce 100644 --- a/datafusion/sqllogictest/test_files/joins.slt +++ b/datafusion/sqllogictest/test_files/joins.slt @@ -5608,3 +5608,77 @@ set datafusion.execution.target_partitions = 4; statement ok reset datafusion.execution.batch_size; + +# Regression test: a `CollectLeft` `HashJoinExec` requires `SinglePartition` on its build +# (left) child, and the `CoalescePartitionsExec` that satisfies it must survive the +# sort-parallelization phase of `EnsureRequirements`. It used to be removed positionally +# (the traversal descends into the join because the *probe* side is linked to a coalesce), +# leaving a multi-partition build side that `SanityCheckPlan` rejects with +# "does not satisfy distribution requirements: SinglePartition". + +statement ok +set datafusion.execution.target_partitions = 8; + +# Keep the scan multi-partition as written, i.e. one partition per file. +statement ok +set datafusion.optimizer.repartition_file_scans = false; + +statement ok +CREATE TABLE collect_left_src (id INT, ts INT) AS VALUES (1, 10), (2, 20), (3, 30); + +query I +COPY (SELECT * FROM collect_left_src) TO 'test_files/scratch/joins/collect_left/0.parquet' STORED AS PARQUET; +---- +3 + +query I +COPY (SELECT * FROM collect_left_src) TO 'test_files/scratch/joins/collect_left/1.parquet' STORED AS PARQUET; +---- +3 + +query I +COPY (SELECT * FROM collect_left_src) TO 'test_files/scratch/joins/collect_left/2.parquet' STORED AS PARQUET; +---- +3 + +query I +COPY (SELECT * FROM collect_left_src) TO 'test_files/scratch/joins/collect_left/3.parquet' STORED AS PARQUET; +---- +3 + +statement ok +CREATE EXTERNAL TABLE collect_left STORED AS PARQUET LOCATION 'test_files/scratch/joins/collect_left/'; + +# The build side is the 4-partition scan; the probe side is the `DISTINCT ON` aggregate, +# whose `CoalescePartitionsExec` is what makes the traversal reach the join. +query I +SELECT a.id +FROM collect_left a +LEFT JOIN (SELECT DISTINCT ON (id) id, ts FROM collect_left ORDER BY id, ts) f + ON a.id = f.id +ORDER BY a.id; +---- +1 +1 +1 +1 +2 +2 +2 +2 +3 +3 +3 +3 + +statement ok +DROP TABLE collect_left; + +statement ok +DROP TABLE collect_left_src; + +statement ok +reset datafusion.optimizer.repartition_file_scans; + +statement ok +set datafusion.execution.target_partitions = 4;