avantgardnerio opened a new pull request, #2196: URL: https://github.com/apache/datafusion-ballista/pull/2196
## Summary Batteries-included AQE machinery so anyone can write a `UnorderedRangeRepartitionExec` / `OrderedRangeRepartitionExec`-inserting optimizer rule and get correct end-to-end routing for free. A range-partitioning rule now needs to do only *one thing*: splice a `RuntimeStatsExec` + `URRE`/`ORRE` above whatever it wants range-repartitioned. Everything downstream is handled: - Executor already ships per-sub-part quantile sketches back to the scheduler (#2094 / #2175). - Scheduler recognizes the range repartition at stage-completion time, merges sketches into `K − 1` global cuts, and rewrites the downstream location map so partition `k` pulls from every producer file whose sketched `[min, max]` overlaps `k`'s assigned cut range. - Cuts + routing expression are parked on the boundary `ExchangeExec`; at task-specialization time the adapter wraps the `ShuffleReader` in a `PerPartitionFilterExec` (#2195) that trims straddling sub-parts to their assigned slice — correctness closes without the rule ever knowing about sketches, filters, or the AQE hook. ## What a rule writer gets | Primitive | Location | Contract | |---|---|---| | `plan_contains_range_repartition` | runtime_stats.rs | detect URRE/ORRE anywhere in a plan subtree | | `find_range_repartition_routing_expr` | runtime_stats.rs | recover `order_by[0].expr` from the range repartition | | `compute_overlapping_locations` | runtime_stats.rs | `(reports, cuts)` → per-downstream-partition sub-part list | | `overlap_remap_partitions` | runtime_stats.rs | full `Vec<Vec<PartitionLocation>>` remap | | `range_partition_predicates` | per_partition_filter.rs (from #2195) | `(routing_expr, cuts)` → K half-open range predicates | | `ExchangeExec.range_repartition_routing` slot | exchange.rs | scheduler → adapter handoff | | DER classifier arm | distributed_exchange.rs | recognizes `RuntimeStatsExec → URRE/ORRE` as a stage boundary | All of this fires only when a rule has actually inserted a range repartition. On any plan without one, `plan_contains_range_repartition` returns `false` and the machinery no-ops — zero overhead for the 99%+ non-range-repartition case. ## Design note: the `ExchangeExec` routing slot `range_repartition_routing: Arc<Mutex<Option<RangeRepartitionRouting>>>` is a direct copy of the existing `coalesce: Arc<Mutex<Option<Arc<CoalescePlan>>>>` slot pattern that landed with `CoalescePartitionsRule` — same lifecycle (AQE rule parks decision, adapter consumes at plan-transform time), same `with_new_children` carry-through, same idempotent-overwrite setter. That slot went through review iteration on main (`b839f036 refactor(scheduler): tighten the AQE coalesce rule after review`, `232d7611 let the AQE coalesce slot be cleared`); this PR follows the pattern rather than inventing a new one. ## Example consumers Two range-repartitioning rules are being prepared as follow-ups against this infrastructure: - **`AdaptiveRangeShuffleRule`** — rewrites Hash exchanges as URRE at join-side boundaries. Proof-of-concept for the machinery. - **`ParallelWindowDetectRule`** (port from `brent/h2o-window-benchmarks`) — wraps `BoundedWindowAggExec` in URRE + `RuntimeStatsExec` for parallel `RANGE`-frame windows. Motivating consumer for the h2o single-partition OOM case. Neither rule is in this PR — they'd conflate two decisions (the machinery vs. the specific consumer). Both are ready to send in short order once this lands. ## Test plan 15 new unit tests covering the pure functions and the exchange slot: - `plan_walker_tests` (6): both URRE and ORRE recognized at root and nested, bare source returns `None`/`false`. - `overlap_remap_tests` (5): disjoint routing, straddling replication, missing-file_id error, default/report mismatch error, empty-sketch passthrough. - `range_repartition_routing_tests` (4): unresolved slot returns `None`, resolve→get roundtrip, second resolve overwrites, `with_new_children` preserves the slot. Full end-to-end wiring (`maybe_range_repartition_overlap_remap` fires at stage completion → cuts parked → adapter injects filter) is exercised in the consumer rule PRs where a live URRE-inserting fixture is available; deferring the ~300 LOC of AQE-graph scaffolding here in favor of pure-function coverage. Existing suites: ballista-core (222), ballista-scheduler (281) — all pass. Clippy clean, fmt clean. ## Caveats - Machinery is inert until a rule inserts a range repartition; the primitives table above is the API surface. - `RangeRepartitionRouting` is scheduler-only state — no wire/codec changes. - `plan_contains_range_repartition` stops at exchange boundaries by design (a nested range repartition lives in a different stage anyway). 🤖 Generated with [Claude Code](https://claude.com/claude-code) -- 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]
