fix: keep a CoalescePartitionsExec required by a SinglePartition child

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
+///       <build>
+///     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<dyn ExecutionPlan>,
+    join_type: JoinType,
+) -> Result<Arc<dyn ExecutionPlan>> {
+    let build: Arc<dyn ExecutionPlan> = Arc::new(CoalescePartitionsExec::new(build));
+    let probe: Arc<dyn ExecutionPlan> = 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<dyn PhysicalExpr>,
+        Arc::new(Column::new("a", 0)) as Arc<dyn PhysicalExpr>,
+    )];
+    let join: Arc<dyn ExecutionPlan> = 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<dyn ExecutionPlan>,
+) -> Result<Arc<dyn ExecutionPlan>> {
+    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(())
+}

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<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.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)
                 }

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;