diff --git a/datafusion/core/tests/physical_optimizer/enforce_sorting.rs b/datafusion/core/tests/physical_optimizer/enforce_sorting.rs index 40bcdbbd6efef..3336caa94b604 100644 --- a/datafusion/core/tests/physical_optimizer/enforce_sorting.rs +++ b/datafusion/core/tests/physical_optimizer/enforce_sorting.rs @@ -2850,3 +2850,120 @@ async fn test_sort_with_streaming_table() -> Result<()> { Ok(()) } + +/// Builds the plan shape that `parallelize_sorts` sees for a `CollectLeft` `HashJoinExec` +/// whose build side is required to be `SinglePartition`, 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 +/// +/// ``` +/// +/// 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> { + use datafusion_common::NullEquality; + use datafusion_physical_plan::joins::{HashJoinExec, PartitionMode}; + + let schema = create_test_schema2()?; + let build = coalesce_partitions_exec(build); + let probe = repartition_exec(coalesce_partitions_exec(parquet_exec(schema))); + + let on = vec![( + Arc::new(Column::new("col_a", 0)) as _, + Arc::new(Column::new("col_a", 0)) as _, + )]; + let join: Arc = Arc::new(HashJoinExec::try_new( + build, + probe, + on, + None, + &join_type, + None, + PartitionMode::CollectLeft, + NullEquality::NullEqualsNothing, + false, + )?); + + Ok(coalesce_partitions_exec(join)) +} + +/// 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 (`parallelize_sorts`) 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". +/// +/// Cherry-picked from apache/datafusion#23948. +#[test] +fn test_collect_left_join_keeps_build_side_coalesce() -> Result<()> { + let schema = create_test_schema2()?; + let build = repartition_exec(parquet_exec(schema)); + let plan = collect_left_plan_before_parallelize_sorts(build, JoinType::Left)?; + + let ctx = PlanWithCorrespondingCoalescePartitions::new_default(plan); + let rewritten = ctx.transform_up(parallelize_sorts).data()?.plan; + + // The build-side coalesce is retained; the probe-side one is still removed, which is + // the parallelization this phase exists for. + assert_snapshot!(displayable(rewritten.as_ref()).indent(true).to_string(), @r" + CoalescePartitionsExec + HashJoinExec: mode=CollectLeft, join_type=Left, on=[(col_a@0, col_a@0)] + CoalescePartitionsExec + RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1 + DataSourceExec: file_groups={1 group: [[x]]}, projection=[col_a, col_b], file_type=parquet + RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1 + DataSourceExec: file_groups={1 group: [[x]]}, projection=[col_a, col_b], file_type=parquet + "); + + Ok(()) +} + +/// The same removal, with a build side that is already hash-partitioned on the join key +/// rather than round-robin partitioned. 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. +/// +/// Cherry-picked from apache/datafusion#23948. +#[test] +fn test_collect_left_join_keeps_hash_partitioned_build_side_coalesce() -> Result<()> { + let schema = create_test_schema2()?; + let build: Arc = Arc::new(RepartitionExec::try_new( + parquet_exec(schema), + Partitioning::Hash(vec![Arc::new(Column::new("col_a", 0))], 10), + )?); + let plan = collect_left_plan_before_parallelize_sorts(build, JoinType::Right)?; + + let ctx = PlanWithCorrespondingCoalescePartitions::new_default(plan); + let rewritten = ctx.transform_up(parallelize_sorts).data()?.plan; + + assert_snapshot!(displayable(rewritten.as_ref()).indent(true).to_string(), @r" + CoalescePartitionsExec + HashJoinExec: mode=CollectLeft, join_type=Right, on=[(col_a@0, col_a@0)] + CoalescePartitionsExec + RepartitionExec: partitioning=Hash([col_a@0], 10), input_partitions=1 + DataSourceExec: file_groups={1 group: [[x]]}, projection=[col_a, col_b], file_type=parquet + RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1 + DataSourceExec: file_groups={1 group: [[x]]}, projection=[col_a, col_b], file_type=parquet + "); + + Ok(()) +} diff --git a/datafusion/physical-optimizer/src/enforce_sorting/mod.rs b/datafusion/physical-optimizer/src/enforce_sorting/mod.rs index 241c843556e1c..2adc702e6b1c1 100644 --- a/datafusion/physical-optimizer/src/enforce_sorting/mod.rs +++ b/datafusion/physical-optimizer/src/enforce_sorting/mod.rs @@ -676,11 +676,45 @@ 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.required_input_distribution(); + let removable = |idx: usize| { + is_root || !matches!(dist_reqs.get(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() @@ -694,9 +728,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 e0be63fe71525..ce71d0038f747 100644 --- a/datafusion/sqllogictest/test_files/joins.slt +++ b/datafusion/sqllogictest/test_files/joins.slt @@ -5527,3 +5527,121 @@ DROP TABLE t1; statement ok DROP TABLE t2; + +# Regression test for a LEFT JOIN with a non-equijoin predicate (forces +# NestedLoopJoinExec) and a multi-partition probe side. Previously the unmatched +# left rows could be emitted before all partitions finished probing, adding +# spurious NULL-padded rows. The result must include every left row exactly once. +statement ok +set datafusion.execution.target_partitions = 4; + +statement ok +set datafusion.execution.batch_size = 2; + +statement ok +CREATE TABLE nlj_left(id INT, v INT) AS VALUES (1, 4), (2, 72), (3, 41), (4, 98), (5, 91); + +statement ok +CREATE TABLE nlj_right(w INT) AS VALUES (49), (58), (83), (3), (76); + +query III +SELECT id, v, w FROM nlj_left LEFT JOIN nlj_right ON nlj_left.v < nlj_right.w ORDER BY id, w; +---- +1 4 49 +1 4 58 +1 4 76 +1 4 83 +2 72 76 +2 72 83 +3 41 49 +3 41 58 +3 41 76 +3 41 83 +4 98 NULL +5 91 NULL + +statement ok +DROP TABLE nlj_left; + +statement ok +DROP TABLE nlj_right; + +statement ok +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;