andygrove commented on PR #2434: URL: https://github.com/apache/datafusion-ballista/pull/2434#issuecomment-5641487705
Thanks for tracing the safety argument rather than taking it on faith — that was the part I most wanted someone to actually walk, and your reading of `find_runnable_exchanges` and the two `Repartitioned` entry points matches mine. Both blocking items are addressed, plus the three inline ones. Pushed. ### 1. Join-type guard — fixed You're right, and it's a real waste rather than a cosmetic one. Applied exactly as you wrote it. One thing worth adding, because it was the only way I could see the guard being *too* strict: the swap decision can't rescue those cases either. `supports_swap_join_order` is `left_bytes > right_bytes`, so `swap_inputs == false` means the build side is already the smaller input — and measuring it only shrinks it further, since the whole premise is that the estimate overshoots. So the swap can never flip `false → true` and turn a `Left` into a broadcastable `Right` after staging. The guard loses nothing. Both directions are pinned now, across all ten join types: `does_not_stage_a_join_type_that_can_never_broadcast` and `stages_the_join_types_that_can_broadcast`. ### 2. Multi-pass coverage — added New file `ballista/scheduler/src/state/aqe/test/stage_build_side.rs`, using the `AdaptivePlanner` + `StatsTable` + `assert_plan!` harness you pointed at. It needed one new fixture helper (`estimated_statistics`, for an `Inexact` `total_byte_size` — `sized_statistics` is `Exact`, which never stages) and a parameterised `mock_partitions_with_size`, with `mock_partitions_with_statistics` delegating to it. Three tests: - `stages_the_build_side_then_broadcasts_the_measured_result` — the mechanism end to end. Pass one asserts an `ExchangeExec` on the build side and a **bare scan** on the probe; `actionable_stages()` returns exactly one stage and that stage is the build side alone; then after the shuffle reports its measured size, pass two resolves to `HashJoinExec: mode=CollectLeft` over `broadcast=true`, with the fact table still unshuffled. - `reuses_the_staged_exchange_when_the_build_side_measures_too_large` — the fallthrough. Asserts one exchange per side, the build's being the resolved staged one (`stage_id=0, stage_resolved=true`) and the probe's fresh, rather than a nested pair. - `does_not_stage_a_left_join` — pins the thing you noted nothing currently pins. You were right that `test_left_join_not_collected_left` passes only incidentally; its MemTable stats don't match the staging shape. These ones do. 418 passed, 0 failed. Clippy and rustfmt clean. ### 3. Splitting the row-threshold change — done Now #2444. Agreed on the reasoning: it isn't what fixes q8, and it deliberately replaces a test that asserted the old semantics, so the two should be revertable independently. ### Smaller **Exchange guard depth.** Acknowledged, and left as is deliberately — I'd rather not widen it without a shape that actually motivates it. The fire-at-most-once bound is unaffected, because staging always produces a *direct* `ExchangeExec` child, which the guard does catch. And in the `FilterExec(ExchangeExec(..))` case the staged measurement is of the filter's output, not the inner shuffle's, so it isn't redundant work — just an extra boundary. A subtree-wide search would also decline legitimately useful cases, like a build side that is itself a join over an already-resolved broadcast exchange. Noted as a follow-up in the description; happy to tighten if you think the edge case is likelier than I'm reading it. **`CoalescePartitionsRule` with a one-exchange join.** I traced it. Short answer: the leaf grouping handles it fine, and the risk you're pointing at isn't reachable — but it's safe by scheduling order rather than by construction, which I think is worth recording. The rule is entirely join-agnostic: it collects leaf `ExchangeExec`s with no notion of which leg they feed, and the only cardinality guard is `leaves.is_empty()`. A one-leaf group is already routine (any single-input stage reading one shuffle), and `start_indices_to_partition_groups` covers `[0, M)` exactly, so coalescing one leg changes task count and never row count. The window you're describing — build leaf carrying a plan from a one-leaf pass while the probe leaf carries none — can't occur, because the rule only runs *after* a full replan. `finalise_stage_internal` calls `replan_stages()`, which runs `SelectJoinRule` and re-decides the join; `CoalescePartitionsRule` runs later, inside `actionable_stages()`. So the instant the staged exchange resolves, the join is already re-decided, and by the time the rule looks the plan is either broadcast (it bails on `any leaf .broadcast`) or two leaves with the probe unresolved (it bails on the unresolved leaf). `m` is also read from `properties().partitioning`, which `set_coalesce` never rewrites, so re-running on an already-coalesced leaf re-derives the same `M` and the heterogeneous-`M` guard doesn't misfire. The thing I'd raise separately, which this PR doesn't introduce but does sit close to: **nothing ever clears a `CoalescePlan`.** `set_coalesce` is the only writer, and both bail paths return without clearing — including the degenerate-`K` path that says so explicitly in its own debug line. If a leaf ever acquires a plan and the group's byte profile later changes such that the rule bails, the stale `K` persists on one leg while its partner stays at `M`. Only the replan-before-rule ordering keeps that from being reachable here. Clearing the slot on the two bail paths, or asserting the group agrees on `coalesce()` before returning, would make it structural. I'd rather do that in its own PR than widen this one — let me know if you'd prefer it bundled. **Executor loss and a resolved exchange's `shuffle_partitions`.** Agreed it looks pre-existing for all exchanges rather than something staging introduces, and agreed that staging makes one more stage's output load-bearing. I haven't checked it either; not planning to in this PR. ### Defaults and benchmark numbers Keeping `stage_build_side` at `true`, per your "fix #1 first, then yes". Both caveats are now stated inline in the description rather than in a footnote: one iteration per query, and the totals were measured **before** the join-type guard, so they include stages that were staged and thrown away on anti and outer joins. The post-fix total should land slightly under 696.09s. I haven't re-measured on a cluster yet. Your partial explanation for the +3.1s band holds up and I've written it into the description: q16 (`NOT IN` → anti join via `not_in_subquery_rewrite`) and q22 (`NOT EXISTS`) are both anti joins, and the guard stops both from staging at all now. It doesn't cover q5/q7/q12, which are plain inner joins — those look like the serialisation cost. Re-running the six with `stage_build_side=false` per query to separate the two is the next measurement I plan to take. -- 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]
