From cab7e9c907d35a77fdffaabe20fefadca9313e51 Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Tue, 28 Jul 2026 11:48:50 -0500 Subject: [PATCH 1/4] fix: keep a CoalescePartitionsExec required by a SinglePartition child `parallelize_sorts` (phase 3a of `EnsureRequirements`) removes `CoalescePartitionsExec`s that only act as a parallelism bottleneck below a global sort. `remove_bottleneck_in_subplan` decides that positionally: if `children[0]` is a `CoalescePartitionsExec`, it is removed, without consulting the parent's own distribution requirement for that child. The traversal reaches such a parent because `update_coalesce_ctx_children` marks a node as connected when *any* of its children is linked to a coalesce below, and it deliberately excludes children that require `Distribution::SinglePartition` only from *setting* that flag. So a `CollectLeft` `HashJoinExec` whose probe side is connected is descended into, and the coalesce satisfying the build side's `SinglePartition` requirement is then taken out of the plan. Nothing re-runs distribution enforcement afterwards, so the first thing to notice is `SanityCheckPlan`, which rejects the plan with `does not satisfy distribution requirements: SinglePartition`. Consult the per-child distribution requirement before removing, both for `children[0]` and when recursing into other children. The node `parallelize_sorts` is itself rewriting (the root of the call) is exempt: the caller drops that node and rebuilds the sort cascade around the result, so its requirement does not constrain the removal. Co-Authored-By: Claude Opus 5 --- .../physical_optimizer/ensure_requirements.rs | 126 +++++++++++++++++- .../enforce_sorting/mod.rs | 38 +++++- 2 files changed, 159 insertions(+), 5 deletions(-) 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) } From 4b31a7d84e33ad978f333c9b379d6c00580ed00f Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Tue, 28 Jul 2026 12:02:01 -0500 Subject: [PATCH 2/4] test: add sqllogictest reproducing the CollectLeft build-side removal Covers the same regression end to end from SQL: a multi-file (and therefore multi-partition) parquet scan on the build side of a `CollectLeft` join, with a `DISTINCT ON` aggregate on the probe side supplying the coalesce link that makes the traversal descend into the join. Co-Authored-By: Claude Opus 5 --- datafusion/sqllogictest/test_files/joins.slt | 74 ++++++++++++++++++++ 1 file changed, 74 insertions(+) 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; From c752ac44f69ea1530532da1d9aea3a7ef0f79b38 Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Fri, 31 Jul 2026 00:07:57 -0500 Subject: [PATCH 3/4] refactor: review follow-ups in remove_bottleneck_in_subplan - Hoist `input_distribution_requirements()` out of the `removable` closure so it is computed once per node instead of once per child check. - Use `children.first()` instead of indexing, so the traversal is robust if it is ever reached on a node with no children. - Document why only `SinglePartition` is protected (a `HashPartitioned` child is in the same position in principle, but `ensure_distribution` never puts a coalesce there) and why skipping the descent entirely is deliberate. Co-Authored-By: Claude Opus 5 --- .../enforce_sorting/mod.rs | 21 +++++++++++++++---- 1 file changed, 17 insertions(+), 4 deletions(-) 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 e78f6d245a311..6efaf76457919 100644 --- a/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/mod.rs +++ b/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/mod.rs @@ -631,16 +631,25 @@ fn remove_bottleneck_in_subplan_impl( // `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!( - plan.input_distribution_requirements() - .child_distribution(idx), + dist_reqs.child_distribution(idx), Some(Distribution::SinglePartition) ) }; - let remove_from_first_child = - is_coalesce_partitions(&requirements.children[0].plan) && removable(0); + 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 remove_from_first_child { // We can safely use the 0th index since we have a `CoalescePartitionsExec`. @@ -658,6 +667,10 @@ fn remove_bottleneck_in_subplan_impl( .into_iter() .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 { From 8fb3d307445538d334954dbd534b7ad605b38473 Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Fri, 31 Jul 2026 00:07:57 -0500 Subject: [PATCH 4/4] test: make the CollectLeft regression tests actually exercise the removal MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The two optimizer-level tests went through the full `EnsureRequirements` rule from a freshly built plan. By the time phase 3a ran, the earlier phases had already dropped the probe-side `CoalescePartitionsExec`, so the join was never marked as connected and `remove_bottleneck_in_subplan` never descended into it. Both tests passed with the fix reverted. Build the pre-phase-3a plan shape directly instead — the shape the reproducer actually produces, with a coalesce on both the build and probe sides of the join — and drive `parallelize_sorts` on its own, then check the result with `SanityCheckPlan`. Both tests now fail with the fix reverted, with the same "does not satisfy distribution requirements: SinglePartition" error the PR reports. Co-Authored-By: Claude Opus 5 --- .../physical_optimizer/ensure_requirements.rs | 203 ++++++++++-------- 1 file changed, 111 insertions(+), 92 deletions(-) diff --git a/datafusion/core/tests/physical_optimizer/ensure_requirements.rs b/datafusion/core/tests/physical_optimizer/ensure_requirements.rs index 580d0b658496a..c04ccd2f3c2ec 100644 --- a/datafusion/core/tests/physical_optimizer/ensure_requirements.rs +++ b/datafusion/core/tests/physical_optimizer/ensure_requirements.rs @@ -24,8 +24,12 @@ 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; @@ -1268,111 +1272,126 @@ 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. +/// 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: /// -/// 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), - ))); +/// ```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, - &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 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))) +} - let optimized = optimize_and_sanity_check(plan).expect("plan failed SanityCheckPlan"); +/// 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) +} - 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 +/// 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 "); - assert_idempotent(optimized); + Ok(()) } -/// 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. +/// 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() { - 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 +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 "); - assert_idempotent(optimized); + Ok(()) }