fix: keep a CoalescePartitionsExec required by a SinglePartition child

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<dyn ExecutionPlan> = 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<dyn ExecutionPlan> = Arc::new(CoalescePartitionsExec::new(Arc::new(
+        MockMultiPartitionExec::new(4),
+    )));
+
+    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,
+            &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<dyn ExecutionPlan> =
+        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<dyn ExecutionPlan> = Arc::new(
+        MockMultiPartitionExec::hash_partitioned_on_a(TEST_TARGET_PARTITIONS),
+    );
+    let probe: Arc<dyn ExecutionPlan> = Arc::new(CoalescePartitionsExec::new(Arc::new(
+        MockMultiPartitionExec::new(4),
+    )));
+
+    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,
+            &JoinType::Right,
+            None,
+            PartitionMode::CollectLeft,
+            NullEquality::NullEqualsNothing,
+            false,
+        )
+        .expect("HashJoinExec creation failed"),
+    );
+
+    let plan: Arc<dyn ExecutionPlan> =
+        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);
+}

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

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;