zhuqi-lucas commented on code in PR #23948:
URL: https://github.com/apache/datafusion/pull/23948#discussion_r3671288691
##########
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() {
Review Comment:
I verified this locally and I'm afraid these two unit tests don't actually
cover the regression: with the fix commit reverted (and `cargo clean -p
datafusion-physical-optimizer` to force a rebuild), **both tests still pass**.
I then instrumented the pre-fix `children[0]`-coalesce-removal branch with an
`eprintln!` probe — it never fires for these mock plans, so they never reach
the buggy path at all. The snapshots are identical pre/post fix.
(The sqllogictest is fine — on the reverted code it fails with exactly the
reported `does not satisfy distribution requirements: SinglePartition` error,
and passes with the fix.)
My guess at the cause: by the time `parallelize_sorts` runs, the
distribution phase has already restructured the probe side, so the coalesce
link that makes the traversal descend into the join in the real reproducer
isn't there in this mock shape. Could you either make these construct the
pre-phase-3a plan shape directly (so the removal branch demonstrably fires — an
`eprintln` probe on `main` is a quick way to confirm), or drop them and let the
sqllogictest carry the regression? As-is they'd stay green if the bug were
reintroduced.
##########
datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/mod.rs:
##########
@@ -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) {
Review Comment:
Worth a short comment that this is deliberately conservative: skipping the
descent entirely also skips legitimate cleanups *below* the protected child
(e.g. a redundant second coalesce underneath the load-bearing one).
Correctness-first is the right call for a fix — just flagging that the gate
could later be narrowed to "descend, but protect only the top-level coalesce"
if anyone cares about the missed optimization.
##########
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:
Question (non-blocking): should this also protect
`Distribution::HashPartitioned` children? A single-partition input trivially
satisfies a hash requirement, so a coalesce sitting under a hash-requiring
child is load-bearing in the same way — removing it exposes whatever
multi-partitioning is below, which may not be hashed on the right keys. I don't
think `EnsureRequirements` itself ever places a coalesce there (hash
requirements get a `RepartitionExec`), so it's probably unreachable via this
rule's own output — but a plan built by other rules or by hand could hit it.
Fine to leave as-is if you agree it's unreachable; a one-line comment saying so
would help the next reader.
--
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]
