fix: keep a CoalescePartitionsExec required by a SinglePartition child

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
+///       <build>
+///     RepartitionExec: RoundRobinBatch
+///       CoalescePartitionsExec           <- links the join into the coalesce cascade
+///         <probe multi-partition source>
+/// ```
+///
+/// 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<dyn ExecutionPlan>,
+    join_type: JoinType,
+) -> Result<Arc<dyn ExecutionPlan>> {
+    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<dyn ExecutionPlan> = 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<dyn ExecutionPlan> = 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(())
+}

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<PlanWithCorrespondingCoalescePartitions> {
+    // 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<PlanWithCorrespondingCoalescePartitions> {
     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)
                 }

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;