avantgardnerio opened a new pull request, #2175: URL: https://github.com/apache/datafusion-ballista/pull/2175
## Summary Barrier 1.5 groundwork for the parallel-window / range-repartition flow ([design](https://github.com/apache/datafusion-ballista/pull/2094)). 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 #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 arrival`** — `collect_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 cuts`** — `RunningStage` 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 - [x] `cargo test -p ballista-core --lib` — 176 pass (was 166 pre-PR). - [x] `cargo test -p ballista-scheduler --lib` — 269 pass, 2 ignored. - [x] `cargo test -p ballista-executor --lib` — 53 pass. - [x] `cargo clippy --all-targets` — clean. - [x] `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](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]
