adriangb commented on code in PR #23948:
URL: https://github.com/apache/datafusion/pull/23948#discussion_r3688385115
##########
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:
Agreed on both counts — same hazard in principle, but ensure_distribution
satisfies a hash requirement with a RepartitionExec, never a coalesce, so
widening the match would be dead code. Left as-is with a comment saying so.
##########
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:
Yes, that's the intended trade — the missed cleanup is a redundant coalesce
under a load-bearing one, which I'd rather leave than risk narrowing wrong in a
fix. Added a comment recording that, with "descend, but protect only the
topmost coalesce" as the follow-up.
--
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]
