feat(scheduler): range-repartition rule primitives (batteries-included)

[avantgardnerio]

Summary

I'm a little torn on this one.

pros:

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 (feat(core): RuntimeStatsExec — passthrough tap for row counts + quantile sketches #2094 / feat(core,scheduler): transport RuntimeStatsExec reports to scheduler; log merged quantile cuts per stage #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 (feat(core): add PerPartitionFilterExec — per-input-partition boolean predicate #2195) that trims straddling sub-parts to their assigned slice — correctness closes without the rule ever knowing about sketches, filters, or the AQE hook.

cons:

This is lots of plumbing code to achieve that. I hope I've done it in the right places. Honestly piggy backing range_repartition_routing on ExchangeExec feels the worst. However this is following the precedent set by the coalesce field, so maybe it's okay?

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

[avantgardnerio]

@phillipleblanc this one is the base upon which anyone can build a range-repartitioned optimizer rule. I intend to use it for parallel windows in the next PR, but I was also playing with it to prevent spilling in SortShuffle on TPC-H q20. It's interesting because now we can play around and try things out.

[phillipleblanc]

I won't have time to review this in depth for a couple of days, but I will look at it sometime this week.