andygrove commented on code in PR #2196:
URL:
https://github.com/apache/datafusion-ballista/pull/2196#discussion_r3696800384
##########
ballista/scheduler/src/state/aqe/optimizer_rule/distributed_exchange.rs:
##########
@@ -90,11 +93,36 @@ impl DistributedExchangeRule {
);
return Ok(Transformed::yes(Arc::new(exchange_exec)));
}
+ } else if !execution_plan.is::<ExchangeExec>()
+ && let [child] = execution_plan.children().as_slice()
Review Comment:
Thanks for the fixes on the root-level case. I re-ran my repro and it is
closed, and the fail-loud `set_repartition_routing` is exactly the right call.
One thing the simplification took with it though. Going from the multi-child
walk to `let [child] = execution_plan.children().as_slice()` narrows this arm
to single-child parents. A `UnionExec` no longer matches, and neither would a
sort merge join range-partitioning both of its sides, which I would have
guessed is one of the main consumers you have in mind for this machinery.
I ran the same probe against both commits, a `UnionExec` over two
`RuntimeStatsExec -> URRE` chains:
| commit | exchanges inserted |
|---|---|
| `a92f0617` (before the fix) | 2 |
| `ea450346` (head) | 0 |
At head the plan comes out with no stage boundary anywhere:
```
AdaptiveDatafusionExec: is_final=false, plan_id=0, stage_id=pending,
stage_resolved=false
UnionExec
RuntimeStatsExec: rows + sketch(routing=v@0 asc)
UnorderedRangeRepartitionExec: routing=v@0 asc -> 4 partitions
StatisticsExec: col_count=1, row_count=Absent
RuntimeStatsExec: rows + sketch(routing=v@0 asc)
UnorderedRangeRepartitionExec: routing=v@0 asc -> 4 partitions
StatisticsExec: col_count=1, row_count=Absent
```
The probe, dropped into this file's test module next to your existing ones:
```rust
#[test]
fn probe_range_repartition_under_multi_child_parent() {
let rule = DistributedExchangeRule::default();
let parent: Arc<dyn ExecutionPlan> = Arc::new(
datafusion::physical_plan::union::UnionExec::new(vec![
stats_over_urre_over_leaf(),
stats_over_urre_over_leaf(),
]),
);
let result = rule.optimize(parent, &config()).unwrap();
println!("{}", display_plan(&result));
assert_eq!(count_exchanges(result.as_ref()), 2);
}
```
The good news is this does not go silently wrong the way the root-level case
did. `collect_reachable_stats` only walks single-child chains, so the union
stage collects no reports, and `repartition_routing` then errors with
`range-repartition stage N: no runtime-stats reports`. The job fails rather
than duplicating rows.
The catch is that the message points somewhere else entirely. A rule author
who splices under a join input would read it as a stats collection bug rather
than "your splice position is not supported yet". So either restore the
multi-child walk here, or if single-child-only is a deliberate scope call for
now, would you mind saying so on `is_range_repartitioned` so the contract is
explicit? A test pinning whichever behavior you pick would be worth having too.
Investigation and repro here were LLM assisted, with the plan dumps above
coming from actual runs against both commits.
--
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]
