Skip to content

feat(core,scheduler): transport RuntimeStatsExec reports to scheduler; log merged quantile cuts per stage - #2175

Merged
avantgardnerio merged 4 commits into
apache:mainfrom
avantgardnerio:brent/runtime-stats-transport-merge
Jul 27, 2026
Merged

feat(core,scheduler): transport RuntimeStatsExec reports to scheduler; log merged quantile cuts per stage#2175
avantgardnerio merged 4 commits into
apache:mainfrom
avantgardnerio:brent/runtime-stats-transport-merge

Conversation

@avantgardnerio

@avantgardnerio avantgardnerio commented Jul 24, 2026

Copy link
Copy Markdown
Contributor

Summary

Groundwork for the parallel-window / range-repartition flow (design). Ships RuntimeStatsExec observations (per-partition row counts and T-Digest quantile sketches) from executor back to scheduler, merges them per stage, and logs the K-1 quantile cuts a globally-informed router would use for K output partitions. No planner change yet — reports don't naturally appear in production plans until a follow-up rule inserts the Stage-1 chain under BWAG.

Three commits, each mergeable on its own:

  1. proto: add RuntimeStatsReport wire format on SuccessfulTask — pure wire additions: SuccessfulTask.runtime_stats, RuntimeStatsReport { order_by, partitions }, RuntimeStatsPartitionEntry { partition_id, row_count, optional sketch }. QuantileSketchState and sketch_to_proto / sketch_from_proto already landed with feat(core): RuntimeStatsExec — passthrough tap for row counts + quantile sketches #2094 and are reused unchanged. All existing SuccessfulTask construction sites populated with runtime_stats: vec![]. Behavior unchanged.

  2. executor: ship RuntimeStatsReports to scheduler; log on arrivalcollect_reports(plan) walks through the preserves_distribution whitelist from the shuffle writer's plan root and emits one report per reachable RuntimeStatsExec. Stats-taps below any distribution-changing operator (e.g. UnorderedRangeRepartitionExec) are excluded automatically. QueryStageExecutor::collect_runtime_stats_reports() (default empty) with DefaultQueryStageExec overriding. as_task_status gains the Vec<RuntimeStatsReport> param and threads it through both call-sites. Scheduler-side TaskManager::update_task_statuses logs a per-arrival summary at debug!RUST_LOG promotes for local verification.

  3. scheduler: accumulate runtime-stats reports per stage; log merged cutsRunningStage carries a runtime_stats_reports accumulator; each successful task update appends. When the stage transitions to Successful, log_merged_runtime_stats groups by order_by wire tag (prost-encoded bytes as HashMap key), merges T-Digests with TDigest::merge_digests, and logs the K-1 quantile cuts at debug!. K is read from the reports (report.partitions.len()). Mismatched-K within a group and sketch decode failures both surface as errors from a fallible merge_reports() -> Result<Vec<MergedRuntimeStats>>; the log-side helper catches at warn!. MergedRuntimeStats is public so a future AQE consumer can call merge_reports directly without paying the log-formatting cost.

Plus a small review-fixup commit tightening merge_reports to be honestly fallible (dropped warn!+continue per-group silent fallbacks) and renaming .k.partition_count.

Test plan

  • cargo test -p ballista-core --lib — 176 pass (was 166 pre-PR).
  • cargo test -p ballista-scheduler --lib — 269 pass, 2 ignored.
  • cargo test -p ballista-executor --lib — 53 pass.
  • cargo clippy --all-targets — clean.
  • cargo fmt --all — clean.

Nine new unit tests total: four walker tests exercising the whitelist boundary behavior (positive: BufferExec doesn't block, negative: SortExec with preserve_partitioning=false does), five merge tests covering disjoint-range union, quartile cuts on a uniform sample, row-count-only groups, mismatched-partition-count errors, sketch-decode-error propagation, and empty input.

  • End-to-end cluster verification is intentionally deferred to the follow-up rule PR — no production plan naturally contains a RuntimeStatsExec today. Once the parallel-window detection rule lands, real queries will produce reports and the merged cuts will show up in the scheduler log against actual data.

🤖 Generated with Claude Code

@avantgardnerio

Copy link
Copy Markdown
Contributor Author

@phillipleblanc WDYT?

@andygrove andygrove left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nice, clean plumbing and the test coverage on the walker and merge paths is good. One thing I want to sort out before this merges is the public API break on as_task_status, plus a couple of smaller notes inline. Details on the lines below.

Comment thread ballista/executor/src/lib.rs Outdated
Comment thread ballista/scheduler/src/state/execution_stage.rs Outdated
Comment thread ballista/executor/src/execution_engine.rs
@github-actions github-actions Bot added the documentation Improvements or additions to documentation label Jul 25, 2026
Groundwork for shipping RuntimeStatsExec observations (row counts +
quantile sketches) back to the scheduler for global-cut selection.
Wire format only; behavior unchanged. The executor still emits
`runtime_stats: vec![]` at every callsite. Slice B populates it.

Proto additions (`ballista.proto`):
  * `SuccessfulTask.runtime_stats: repeated RuntimeStatsReport`. One
    entry per RuntimeStatsExec that is still valid at the plan's
    output (i.e., reachable from the top through distribution-
    preserving nodes only, per `preserves_distribution` in
    range_repartition_common). Pre-repartition stats stay local to
    the executor (used only for the repartitioner's own approximate
    routing) — the walker's whitelist naturally excludes them.
  * `RuntimeStatsReport { order_by, partitions }`. `order_by` tags
    the routing expression so the scheduler groups sketches across
    tasks that were sampling the same expression. `partitions`
    carries one entry per observed partition — pre-repartition one
    per input partition, post-repartition one per output sub-
    partition; the operator's position determines the interpretation.
  * `RuntimeStatsPartitionEntry { partition_id, row_count, optional
    sketch }`. A future `optional MinMaxState min_max` slot is
    called out as a TODO — the eventual lighter mode for post-
    repartition bin-pack where full T-Digests are overkill.

Note: `QuantileSketchState` and the `sketch_to_proto` /
`sketch_from_proto` helpers already landed in main with apache#2094 and are
reused here without change.

Constructor plumbing:
  * Five existing `SuccessfulTask { ... }` sites populated with
    `runtime_stats: vec![]` — the real executor path
    (executor/src/lib.rs), one scheduler-server integration-test
    path, and three scheduler test fixtures.

Verification:
  * `cargo build --workspace` — clean.
  * `cargo test -p ballista-core --lib` — 166 pass (existing wire
    tests over `sketch_to_proto` / `sketch_from_proto` cover the
    payload; prost codegen refuses field-number clashes at build
    time so no new roundtrip test is warranted for a purely
    structural addition).
  * `cargo test -p ballista-scheduler --lib` — 269 pass.
  * `cargo test -p ballista-executor --lib` — 53 pass.
  * `cargo clippy --all-targets` — clean.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

executor: ship RuntimeStatsReports to scheduler; log on arrival

End-to-end path for stage-1 quantile sketches — extract on the
executor at task completion, transport in `SuccessfulTask.runtime_stats`
(landed in the previous commit), and log a per-report summary on the
scheduler as proof-of-life. Per-stage accumulation + merged-quantile-cut
logging comes next.

# Executor extraction (runtime_stats.rs)

`collect_reports(plan)` walks the plan through the
`preserves_distribution` whitelist and collects one `RuntimeStatsReport`
per reachable `RuntimeStatsExec`. Stats-taps sitting below any
distribution-changing operator (e.g. `UnorderedRangeRepartitionExec`)
are excluded automatically because the walker stops at that boundary;
their sketches describe data the repartitioner then routed away and are
no longer meaningful at the plan's output.

Each emitted report carries the operator's `order_by` (serialised as a
`PhysicalSortExprNode` — the wire tag the scheduler will group on) and
one `RuntimeStatsPartitionEntry` per partition slot with
`{partition_id, row_count, optional sketch}`. Empty sketches are elided
(`sketch: None`) so bin-pack space is proportional to actual samples,
not slot count.

`sketch_to_proto` / `sketch_from_proto`, `partition_count()`, and the
widened `preserves_distribution` whitelist already exist in main —
introduced with the RuntimeStatsExec / URRE PRs. This commit reuses
them.

# Trait plumbing (execution_engine.rs)

`QueryStageExecutor::collect_runtime_stats_reports()` — default
returns empty. `DefaultQueryStageExec` overrides to call
`collect_reports` against whichever `ShuffleWriterVariant` it holds
(`Passthrough` or `Sort`). Serialisation errors are logged and the
report dropped rather than failing the task; the task's data was
already produced correctly, telemetry loss shouldn't tank the query.

# Transport (executor lib + call sites)

`as_task_status` gains a `runtime_stats: Vec<RuntimeStatsReport>`
parameter (marked `#[allow(clippy::too_many_arguments)]`) and populates
`SuccessfulTask.runtime_stats` with it. Both executor call-sites
(`execution_loop.rs`, `executor_server.rs`) call
`collect_runtime_stats_reports()` on the query-stage executor and
thread the result through. Failure paths pass `vec![]`.

# Scheduler receive (task_manager.rs)

`TaskManager::update_task_statuses` calls a new
`log_runtime_stats_arrival` helper on each status before batching.
Logged at `debug!` — RUST_LOG can promote it locally during
verification. Fields: executor / job / stage / task + per-report
summary (order_by_len, partition count, non-empty partitions,
total rows, sketches).

# Tests

Four new walker tests in `execution_plans::runtime_stats::collect_tests`:
  * `collect_reports_finds_stats_and_ships_sketch` — stats at plan
    root, populated sketch survives `sketch_from_proto`, row_count
    matches drained data.
  * `collect_reports_row_count_only_emits_report_without_sketch` —
    row-count-only mode emits a report but no sketch payload.
  * `collect_reports_descends_through_whitelisted_op` — a `BufferExec`
    intermediary does not block the walker.
  * `collect_reports_stops_at_sort_that_collapses_partitions` —
    N→1 SortExec is explicitly excluded from the whitelist; the
    walker respects that.

# Verification

  * `cargo test -p ballista-core --lib` — 170 pass (was 166).
  * `cargo test -p ballista-scheduler --lib` — 269 pass, 2 ignored.
  * `cargo test -p ballista-executor --lib` — 53 pass.
  * `cargo clippy --all-targets` — clean.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

scheduler: accumulate runtime-stats reports per stage; log merged cuts

Barrier 1.5 prep. On the scheduler, each stage-attempt now accumulates
the `RuntimeStatsReport`s that arrive with successful task statuses.
When the stage transitions to Successful, the reports are grouped by
`order_by` wire tag, sketches merged via `TDigest::merge_digests`, and
one debug! line emitted per group with the `K-1` quantile cuts a
globally-informed router would have used for `K` output partitions.

`K` is read from the reports themselves — every report carries one
partition entry per repartitioner-output slot, so
`K == report.partitions.len()`. Reports with mismatched `K` within a
group are surfaced as a `warn!` and the group is skipped (planner
invariant break).

# ballista_core::execution_plans::runtime_stats

Two new symbols, exported from `execution_plans::mod`:

  * `merge_reports(reports) -> Vec<MergedRuntimeStats>`. Groups by
    `order_by` (prost-encoded bytes as HashMap key), merges T-Digests
    within each group, computes `K-1` cuts at quantiles `i/K`, returns
    `MergedRuntimeStats { order_by_len, k, task_count, total_rows,
    cuts, min, max }`. Empty input → empty output. Sketches with
    `count() == 0` are dropped from the merge input rather than
    corrupting the result.
  * `log_merged_runtime_stats(job_id, stage_id, reports)`. Calls
    `merge_reports` and formats each group at `debug!` — `RUST_LOG`
    can promote when verifying against a live cluster. Splits into
    two log lines depending on whether any sketches were present
    (drops `min` / `max` in row-count-only mode).

`MergedRuntimeStats` is `pub` so slice-D consumers (an AQE rule that
plans stage 2 from the cuts) can call `merge_reports` directly without
paying the log-formatting cost.

# ballista_scheduler::state::execution_stage

`RunningStage` gains a `runtime_stats_reports: Vec<RuntimeStatsReport>`
field, initialized empty in `RunningStage::new` and in the retry path
`FailedStage::to_running` (a fresh attempt discards the previous
attempt's stats — merged-cut logging fires per-attempt). New method
`append_runtime_stats_reports(reports)` for the graph to push into on
each successful task update.

# ballista_scheduler::state::execution_graph

In the `Successful` arm of `update_task_status`, destructure
`SuccessfulTask` to grab both `partitions` and `runtime_stats` and
push the reports onto the running stage. When `is_final_successful`
flips true (stage transitions to Successful), call
`log_merged_runtime_stats` alongside the existing `print_stage_metrics`
hook. Reports are dropped on the floor after logging — slice D
propagates them onto `SuccessfulStage` when the AQE consumer needs
them at plan-time.

# Tests

Five new tests in `execution_plans::runtime_stats::merge_tests`:

  * `merge_reports_combines_disjoint_ranges` — two reports over
    disjoint value ranges; merged sketch spans the union, K=2 midpoint
    cut falls between the two ranges, total_rows sums.
  * `merge_reports_k_of_4_produces_three_quartile_cuts` — uniform
    [0, 100) sample over K=4; the three cuts land in the expected
    quartile bands.
  * `merge_reports_row_count_only_emits_empty_cuts` — reports with
    only row_counts (no sketches) sum correctly with empty `cuts`
    and `None` `min` / `max`.
  * `merge_reports_skips_group_with_mismatched_partition_counts` —
    mismatched `K` within a group drops the whole group.
  * `merge_reports_empty_input_is_empty_output` — degenerate case.

# Verification

  * `cargo test -p ballista-core --lib` — 175 pass (was 170).
  * `cargo test -p ballista-scheduler --lib` — 269 pass, 2 ignored.
  * `cargo test -p ballista-executor --lib` — 53 pass.
  * `cargo clippy --all-targets` — clean.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

review: merge_reports returns Result; drop silent per-group fallbacks

Address review feedback on the slice-C merge helper:

  * merge_reports is fallible now. Two failure modes it used to swallow
    silently — a group with mismatched partition counts, and a
    corrupted sketch that won't decode — both surface as errors. The
    scheduler-side caller (log_merged_runtime_stats) catches at warn!
    and returns; a slice-D consumer that actually plans off the merged
    cuts can propagate or fail the stage as it sees fit. If the
    function's job is to merge, it either merges or errors — no
    "returns empty because the third group was funky" middle ground.

  * The group-loop body moved into its own `merge_group(&[&Report]) ->
    Result<MergedRuntimeStats>` helper. The outer loop is one `?` per
    group; each group's own failure modes live in one place.

  * `MergedRuntimeStats.k` → `.partition_count`. Struct is public;
    slice D will consume it. `k` is a math-notation letter, not a
    field name.

  * Local single-letter names retired: `r1`/`r2` → `low_range` /
    `high_range`, `k`/`vs` inside `sketching_report` → `slot_id` /
    `slot_values`, `make` closure → `make_report`, closure params
    `|r|` → `|report|`. Per the max-specificity naming preference in
    memory.

  * No array accessors in the merge helper or its tests. The empty-
    group branch of `merge_group` uses `let [first, rest @ ..] = group`
    and returns an internal error rather than indexing at `[0]`. Tests
    unpack `merge_reports(...)?` results with `match cuts.as_slice()`
    and `only_group(...)` helper that panics with a real message
    rather than `groups[0]`.

New test: `merge_reports_propagates_sketch_decode_errors` — builds a
wire-shape-corrupt `QuantileSketchState` (3 scalars instead of 6) and
asserts the decode error propagates through `merge_reports`. This
was the case slice C's original `warn!` + `continue` was hiding.

Existing test renamed `merge_reports_skips_group_with_mismatched_...`
→ `merge_reports_errors_on_mismatched_partition_counts` to reflect
the new contract.

Verification:
  * cargo test -p ballista-core --lib — 176 pass (was 175; +1 test).
  * cargo clippy --all-targets — clean.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

fix(docs): drop rustdoc links to items outside the doc'd crate

CI's `cargo doc --document-private-items --no-deps --workspace` runs
with `-D warnings`, which promotes two rustdoc lint classes to errors:

  * `rustdoc::private-intra-doc-links` — a public item's doc can't
    link to a private one. `collect_reports` (and its re-export
    `collect_runtime_stats_reports`) linked to
    `super::range_repartition_common::preserves_distribution`, but
    `range_repartition_common` is `mod`, not `pub mod`.
  * `rustdoc::broken-intra-doc-links` — a link target must resolve
    in the doc'd crate. `RunningStage::append_runtime_stats_reports`
    linked to `log_merged_runtime_stats`, which now lives in
    ballista-core (moved there in the slice-C refactor because
    ballista-scheduler doesn't depend on
    `datafusion_functions_aggregate_common`).

Both docstrings switched from intra-doc `[...]` links to plain
backtick prose that names the target function fully. The prose
still points a reader to the right place; the compiler stops
complaining.

Verified locally with the exact CI invocation:
`RUSTDOCFLAGS='-D warnings' cargo doc --document-private-items \
    --no-deps --workspace` — clean.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

review: reshape as_task_status to take TaskCompletionExtras

Fold operator_metrics + runtime_stats into a TaskCompletionExtras struct
marked #[non_exhaustive] + Default, so future additions are non-breaking
for external callers via ..Default::default(). Drops the too_many_arguments
allow on as_task_status.

Documents the break in docs/source/upgrading/55.0.0.md since the parameter
list still changes shape (execution_times now precedes extras, and
operator_metrics moves from a positional param into the struct).

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
@avantgardnerio
avantgardnerio force-pushed the brent/runtime-stats-transport-merge branch from 0b9d5ac to ca63093 Compare July 25, 2026 15:46
@avantgardnerio

Copy link
Copy Markdown
Contributor Author

@andygrove I really appreciate your review. I think I've addressed your feedback, so PTAL again whenever you have a moment 🙌

@avantgardnerio
avantgardnerio requested a review from andygrove July 25, 2026 16:46

@phillipleblanc phillipleblanc left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks good to me, one minor comment about making this resilient to stage retries - but I like where this is going.

Comment thread ballista/scheduler/src/state/execution_graph.rs Outdated
avantgardnerio and others added 3 commits July 27, 2026 08:52
`AdaptiveExecutionGraph::update_task_status` in `aqe/mod.rs:774` unpacked
`SuccessfulTask.partitions` from an incoming task-success message but
silently dropped the `runtime_stats` field. It also skipped
`log_merged_runtime_stats` at final-success. The RSE-follow-up wire
path was only plumbed on `StaticExecutionGraph`; any tap running under
AQE mode (`ballista.planner.adaptive.enabled=true`) had its sketches
discarded at the scheduler.

Symptom on the adaptive-range-shuffle path: stage-1 tasks correctly
harvested one `RuntimeStatsReport` each (visible in executor logs as
`Task N finished with … 1 runtime-stats report(s)`), the wire message
carried them, but nothing on the scheduler side ever consumed them —
`log_merged_runtime_stats` never fired, so the merged cut points that
downstream stages need weren't visible.

The fix mirrors the static path (`execution_graph.rs:956`): unpack the
`runtime_stats` field, hand it to `running_stage.append_runtime_stats_reports`,
and call `log_merged_runtime_stats` once the stage is final-successful.
Verified on Q20 SF10 — merged output now appears:

```
merged runtime stats: job=… stage=1 order_by_len=1 partition_count=16
    task_count=4 total_rows=9088057
    cuts=[125078.52, 249995.47, …, 1874842.86] min=1 max=2000000
```

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Wrap each incoming `RuntimeStatsReport` in a `TaskRuntimeStats { producer_task_id, report }` when the scheduler accumulates it on `RunningStage.runtime_stats_reports`. Threads the producer `task_id` through both accumulation paths (static- and adaptive-graph `update_task_status`).

Motivation: reports are appended on every successful-task update, but a task's "success" can be invalidated later (executor loss → task marked `Failed(ResultLost)`, partition rescheduled). Without a producer key on the accumulated entries, the merged view at stage-final-success would double-count the reset producer's contribution — the old ghost plus the retry's fresh report. Keying by producer `task_id` gives the follow-up purge helper something to filter on.

No behavior change here (purge lands in the next commit). `log_merged_runtime_stats` extracts the raw reports internally for the existing merge, so the logged output is unchanged.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
`RunningStage.runtime_stats_reports` accumulates per-task sketches on every successful-task update, but a task's "success" is not permanent — an executor loss flips previously-successful tasks to `Failed(ResultLost)` and reschedules their partitions on another executor. Without a purge, the ghost reports stayed in the vec; the retry appended its fresh reports under a new `task_id`; the merged view at stage-final-success double-counted the bounced slices.

Addresses reviewer feedback on this PR (Phillip LeBlanc, apache#2175): "Reports are not associated with task attempts. When successful work is invalidated and rerun (i.e. lost executor, and the scheduler requeues the task), the old report cannot be removed."

They are now — each accumulated entry carries the producer `task_id`, so both reset paths on `RunningStage` filter out stale contributions:
- `reset_task_info(task_id)`: retryable single-task reset (task killed / result lost). Drops entries where `producer_task_id == task_id`.
- `RunningStage::reset_tasks(executor)`: whole-executor loss on a running stage. Collects the set of reset task_ids while flipping their statuses, then filters the reports vec against that set in one pass.

`SuccessfulStage::reset_tasks` doesn't need matching treatment — `SuccessfulStage` doesn't carry the reports (they were consumed by `log_merged_runtime_stats` at stage-final-success), and `to_running` starts the next attempt with a fresh empty vec.

Two unit tests cover both paths: single-task reset purges only the matching entry; executor-loss reset purges every entry from the lost executor while leaving surviving executors' entries alone.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
@avantgardnerio
avantgardnerio merged commit dc6b380 into apache:main Jul 27, 2026
21 checks passed
@avantgardnerio
avantgardnerio deleted the brent/runtime-stats-transport-merge branch July 27, 2026 16:44
avantgardnerio added a commit to avantgardnerio/arrow-ballista that referenced this pull request Jul 27, 2026
Scheduler-side AQE machinery so any optimizer rule that inserts an
`UnorderedRangeRepartitionExec` or `OrderedRangeRepartitionExec` gets
end-to-end correct 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 to the scheduler
  (apache#2094 / apache#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` (apache#2195) that trims straddling sub-parts to
  their assigned slice.

Primitives a rule writer now has:

- `plan_contains_range_repartition` / `find_range_repartition_routing_expr`
  (runtime_stats.rs) — detect URRE or ORRE anywhere in a plan subtree.
- `compute_overlapping_locations` / `overlap_remap_partitions`
  (runtime_stats.rs) — merged sketches + cuts → per-downstream-partition
  location map with correct producer-file assignment.
- `ExchangeExec.range_repartition_routing` slot + `resolve_*` /
  `.range_repartition_routing()` accessors — scheduler → adapter handoff.
  Mirrors the existing `coalesce: Arc<Mutex<Option<Arc<CoalescePlan>>>>`
  slot pattern (same lifecycle, same `with_new_children` carry-through,
  same idempotent-overwrite semantics — see `b839f036` /`232d7611` for
  the coalesce refactors that established the pattern).
- DER classifier arm (distributed_exchange.rs) — recognizes
  `RuntimeStatsExec → URRE/ORRE` as a stage boundary so the sketches
  ship to the scheduler at stage-N completion.

All of this fires only when a rule has actually inserted an RRE. On
plans without one, `plan_contains_range_repartition` returns `false`
and the machinery no-ops — zero overhead for the 99%+ non-RRE case.

Test coverage (15 new unit tests, all pure-function):
- `plan_walker_tests` (6): URRE and ORRE at root and nested, bare
  source returns `None`/`false`.
- `overlap_remap_tests` (5): disjoint, straddling, missing-file_id
  error, default/report mismatch error, empty-sketch passthrough.
- `range_repartition_routing_tests` (4): unresolved returns `None`,
  resolve→get roundtrip, second resolve overwrites, `with_new_children`
  preserves the slot.

Full end-to-end wiring is exercised in the follow-up consumer rule
PRs (AdaptiveRangeShuffleRule for join-side rewrites,
ParallelWindowDetectRule for parallel windows) where a live
URRE-inserting fixture is naturally available; deferring the
~300 LOC of AQE-graph scaffolding here in favor of pure-function
coverage.

Existing tests: 222 in ballista-core, 281 in ballista-scheduler — all
pass. Clippy clean, fmt clean.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

api-change documentation Improvements or additions to documentation

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants