diff --git a/datafusion/core/tests/physical_optimizer/ensure_requirements.rs b/datafusion/core/tests/physical_optimizer/ensure_requirements.rs index 2c6c46c82985a..c04ccd2f3c2ec 100644 --- a/datafusion/core/tests/physical_optimizer/ensure_requirements.rs +++ b/datafusion/core/tests/physical_optimizer/ensure_requirements.rs @@ -21,9 +21,15 @@ //! 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_common::tree_node::{TransformedResult, TreeNode}; use datafusion_physical_optimizer::PhysicalOptimizerRule; use datafusion_physical_optimizer::ensure_requirements::EnsureRequirements; +use datafusion_physical_optimizer::ensure_requirements::enforce_sorting::{ + PlanWithCorrespondingCoalescePartitions, parallelize_sorts, +}; use std::sync::Arc; @@ -65,6 +71,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 +100,7 @@ impl MockMultiPartitionExec { } let properties = PlanProperties::new( eq, - Partitioning::UnknownPartitioning(partition_count), + partitioning, EmissionType::Incremental, Boundedness::Bounded, ); @@ -1252,3 +1271,127 @@ fn test_idempotent_union_projection_sort() { assert_idempotent(plan); } + +/// Builds the plan shape that phase 3a (`parallelize_sorts`) sees in the reproducer, +/// i.e. the output of the distribution + sorting phases, not a freshly planned tree: +/// +/// ```text +/// CoalescePartitionsExec <- the node `parallelize_sorts` rewrites +/// HashJoinExec: mode=CollectLeft +/// CoalescePartitionsExec <- satisfies `SinglePartition` on the build side +/// +/// RepartitionExec: RoundRobinBatch +/// CoalescePartitionsExec <- links the join into the coalesce cascade +/// MockMultiPartitionExec +/// ``` +/// +/// Both coalesces below the join matter. The probe-side one is what makes +/// `update_coalesce_ctx_children` mark the join as connected — it only skips children that +/// require `SinglePartition`, and the probe side does not — so the walk descends into the +/// join. The build-side one is the one that must survive. +fn collect_left_plan_before_parallelize_sorts( + build: Arc, + join_type: JoinType, +) -> Result> { + let build: Arc = Arc::new(CoalescePartitionsExec::new(build)); + let probe: Arc = Arc::new(RepartitionExec::try_new( + Arc::new(CoalescePartitionsExec::new(Arc::new( + MockMultiPartitionExec::new(4), + ))), + Partitioning::RoundRobinBatch(TEST_TARGET_PARTITIONS), + )?); + + 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, + &join_type, + None, + PartitionMode::CollectLeft, + NullEquality::NullEqualsNothing, + false, + )?); + + Ok(Arc::new(CoalescePartitionsExec::new(join))) +} + +/// Runs phase 3a of `EnsureRequirements` (`parallelize_sorts`) on its own, the same way +/// the rule drives it, and checks the result with `SanityCheckPlan`. +/// +/// The phase is driven directly rather than through `EnsureRequirements::optimize` because +/// the earlier phases would rebuild the plan shape above into something that never reaches +/// the code path under test. +fn parallelize_sorts_and_sanity_check( + plan: Arc, +) -> Result> { + let ctx = PlanWithCorrespondingCoalescePartitions::new_default(plan); + let rewritten = ctx.transform_up(parallelize_sorts).data()?.plan; + SanityCheckPlan::new().optimize(Arc::clone(&rewritten), &test_config())?; + Ok(rewritten) +} + +/// A `CollectLeft` `HashJoinExec` requires `Distribution::SinglePartition` on its build +/// (left) child, so the distribution phase puts a `CoalescePartitionsExec` on top of a +/// multi-partition build side. The sort-parallelization phase must not take that coalesce +/// back out again. +/// +/// It used to, because `remove_bottleneck_in_subplan` removed a coalesce found at +/// `children[0]` positionally, without consulting the parent's distribution requirement for +/// that child. The result was a build side left multi-partition with nothing to re-enforce +/// distribution afterwards, which `SanityCheckPlan` rejected with "does not satisfy +/// distribution requirements: SinglePartition". +#[test] +fn test_collect_left_join_keeps_build_side_coalesce() -> Result<()> { + let plan = collect_left_plan_before_parallelize_sorts( + Arc::new(MockMultiPartitionExec::new(4)), + JoinType::Left, + )?; + + let rewritten = parallelize_sorts_and_sanity_check(plan)?; + + // The build-side coalesce is retained; the probe-side one is still removed, which is + // the parallelization this phase exists for. + assert_snapshot!(plan_string(&rewritten), @r" + CoalescePartitionsExec + HashJoinExec: mode=CollectLeft, join_type=Left, on=[(a@0, a@0)] + CoalescePartitionsExec + MockMultiPartitionExec + RepartitionExec: partitioning=RoundRobinBatch(8), input_partitions=4 + MockMultiPartitionExec + "); + + Ok(()) +} + +/// The same 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: the build side satisfies the join's *hash* +/// requirement but still not `SinglePartition`, so the coalesce is just as load-bearing. +#[test] +fn test_collect_left_join_keeps_hash_partitioned_build_side_coalesce() -> Result<()> { + let plan = collect_left_plan_before_parallelize_sorts( + Arc::new(MockMultiPartitionExec::hash_partitioned_on_a( + TEST_TARGET_PARTITIONS, + )), + JoinType::Right, + )?; + + let rewritten = parallelize_sorts_and_sanity_check(plan)?; + + assert_snapshot!(plan_string(&rewritten), @r" + CoalescePartitionsExec + HashJoinExec: mode=CollectLeft, join_type=Right, on=[(a@0, a@0)] + CoalescePartitionsExec + MockMultiPartitionExec + RepartitionExec: partitioning=RoundRobinBatch(8), input_partitions=4 + MockMultiPartitionExec + "); + + Ok(()) +} 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..6efaf76457919 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,49 @@ 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. + // + // Only `SinglePartition` is protected. A `HashPartitioned` child is in principle in the + // same position — a single-partition input trivially satisfies a hash requirement, so a + // coalesce below one is also load-bearing — but nothing puts a coalesce there: + // `ensure_distribution` satisfies a hash requirement with a `RepartitionExec`, never a + // `CoalescePartitionsExec`. Widening the check would be dead code today. + let dist_reqs = plan.input_distribution_requirements(); + let removable = |idx: usize| { + is_root + || !matches!( + dist_reqs.child_distribution(idx), + Some(Distribution::SinglePartition) + ) + }; + let remove_from_first_child = requirements + .children + .first() + .is_some_and(|child| is_coalesce_partitions(&child.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 +665,14 @@ 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)| { + // Deliberately conservative: not descending at all also skips legitimate + // cleanups *below* a protected child (a redundant second coalesce under the + // load-bearing one, say). This could later be narrowed to "descend, but + // protect only the topmost coalesce" if that turns out to matter. + 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;