Copilot commented on code in PR #23948:
URL: https://github.com/apache/datafusion/pull/23948#discussion_r3668209530
##########
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);
Review Comment:
`requirements.children[0]` will panic if `remove_bottleneck_in_subplan_impl`
is ever called on a node with zero children. Consider adding an early return
for `requirements.children.is_empty()` or rewriting this to use
`requirements.children.get(0)` so the traversal is robust even if `node.data`
is set on a leaf.
##########
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)
+ )
+ };
Review Comment:
`plan.input_distribution_requirements()` is invoked each time
`removable(idx)` runs (including inside the child iteration). If this accessor
does non-trivial work or allocations, it can add overhead to plan rewrites.
Consider storing `let dist_reqs = plan.input_distribution_requirements();` once
and referencing `dist_reqs.child_distribution(idx)`.
##########
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
Review Comment:
This test writes Parquet files under `test_files/scratch/...` but does not
remove them afterwards. If the sqllogictest harness doesn’t reliably clean
`test_files/scratch` between runs (or if `COPY` doesn’t overwrite on all
platforms/filesystems), this can cause disk growth or rerun flakiness. If
supported by the harness, add an explicit cleanup step for the directory/files;
otherwise consider using a run-unique scratch subdirectory to avoid collisions.
##########
datafusion/core/tests/physical_optimizer/ensure_requirements.rs:
##########
@@ -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
+ ");
Review Comment:
These snapshots hard-code `RoundRobinBatch(8)`, which makes the tests
brittle if `TEST_TARGET_PARTITIONS` / the test config’s target partition count
changes (even though the behavior being tested is the presence of the
build-side `CoalescePartitionsExec`). Consider asserting the key structural
property (build side retains `CoalescePartitionsExec` under a `CollectLeft`
join) without pinning the entire formatted plan to a specific partition count.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
