From ca630937b999d0c2b818c49ac1e1f1a415649839 Mon Sep 17 00:00:00 2001 From: Brent Gardner Date: Fri, 24 Jul 2026 12:09:40 -0600 Subject: [PATCH 1/4] proto: add RuntimeStatsReport wire format on SuccessfulTask MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 #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) 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` 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) 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`. 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` 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) 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` 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) 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) 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) --- ballista/core/proto/ballista.proto | 34 + ballista/core/src/execution_plans/mod.rs | 6 +- .../core/src/execution_plans/runtime_stats.rs | 643 ++++++++++++++++++ ballista/core/src/serde/generated/ballista.rs | 43 ++ ballista/executor/src/execution_engine.rs | 37 + ballista/executor/src/execution_loop.rs | 10 +- ballista/executor/src/executor_server.rs | 10 +- ballista/executor/src/lib.rs | 32 +- .../scheduler/src/scheduler_server/mod.rs | 1 + .../scheduler/src/state/execution_graph.rs | 23 +- .../scheduler/src/state/execution_stage.rs | 23 +- ballista/scheduler/src/state/task_manager.rs | 44 ++ ballista/scheduler/src/test_utils.rs | 2 + docs/source/upgrading/55.0.0.md | 44 ++ 14 files changed, 933 insertions(+), 19 deletions(-) diff --git a/ballista/core/proto/ballista.proto b/ballista/core/proto/ballista.proto index c2471c7640..84474dedeb 100644 --- a/ballista/core/proto/ballista.proto +++ b/ballista/core/proto/ballista.proto @@ -592,6 +592,40 @@ message SuccessfulTask { // TODO tasks are currently always shuffle writes but this will not always be the case // so we might want to think about some refactoring of the task definitions repeated ShuffleWritePartition partitions = 2; + // Reports from `RuntimeStatsExec` operators in this task's plan that + // are still valid at the plan's output (walked from the top through + // distribution-preserving nodes only — see + // `range_repartition_common::preserves_distribution`). Empty when the + // plan has no such stats-taps. Currently reports one entry per + // executed `RuntimeStatsExec`; the scheduler groups by `order_by` tag + // to combine reports across tasks/executors. + repeated RuntimeStatsReport runtime_stats = 3; +} + +// One report per `RuntimeStatsExec` in the executed plan. +message RuntimeStatsReport { + // What the operator was sampling. Empty = row-count-only mode. When + // non-empty, the first entry identifies which routing expression this + // sketch describes; the scheduler groups sketches by this tag to + // combine samples across tasks that were sampling the same expression. + repeated datafusion.PhysicalSortExprNode order_by = 1; + // Per-partition observations. Interpretation depends on the operator's + // position in the plan — pre-repartition gets one entry per input + // partition, post-repartition gets one per output sub-partition. The + // scheduler groups by `order_by` tag and aggregates. + repeated RuntimeStatsPartitionEntry partitions = 2; +} + +// One partition's observations from a `RuntimeStatsExec`. +message RuntimeStatsPartitionEntry { + uint32 partition_id = 1; + uint64 row_count = 2; + // Present when the `RuntimeStatsExec` was in sketch mode AND this + // partition observed at least one non-null routing value. + optional QuantileSketchState sketch = 3; + // TODO: `optional MinMaxState min_max` — for a lighter post-repartition + // mode where the bin-packer just needs (min, max, count) per + // sub-partition and a full T-Digest is overkill. } message ExecutionError { diff --git a/ballista/core/src/execution_plans/mod.rs b/ballista/core/src/execution_plans/mod.rs index 7aec3cf47d..a7825ef396 100644 --- a/ballista/core/src/execution_plans/mod.rs +++ b/ballista/core/src/execution_plans/mod.rs @@ -40,7 +40,11 @@ use datafusion::common::exec_err; pub use distributed_explain_analyze::DistributedExplainAnalyzeExec; pub use distributed_query::{DistributedQueryExec, execute_physical_plan}; pub use ordered_range_repartition::OrderedRangeRepartitionExec; -pub use runtime_stats::{RuntimeStatsExec, sketch_from_proto, sketch_to_proto}; +pub use runtime_stats::{ + MergedRuntimeStats, RuntimeStatsExec, + collect_reports as collect_runtime_stats_reports, log_merged_runtime_stats, + merge_reports as merge_runtime_stats_reports, sketch_from_proto, sketch_to_proto, +}; pub use shuffle_reader::{CoalescePlan, PartitionGroup, ShuffleReaderExec}; pub use shuffle_reader::{stats_for_partition, stats_for_partitions}; pub use shuffle_writer::DEFAULT_SHUFFLE_CHANNEL_CAPACITY; diff --git a/ballista/core/src/execution_plans/runtime_stats.rs b/ballista/core/src/execution_plans/runtime_stats.rs index 43760b55e8..8559d9b611 100644 --- a/ballista/core/src/execution_plans/runtime_stats.rs +++ b/ballista/core/src/execution_plans/runtime_stats.rs @@ -506,6 +506,94 @@ impl Drop for StreamState { } } +/// Walk `plan` and collect one [`crate::serde::protobuf::RuntimeStatsReport`] +/// per [`RuntimeStatsExec`] that remains valid at the plan's output. +/// "Valid" means reachable through single-child chains of distribution- +/// preserving operators only — see the `preserves_distribution` +/// whitelist in the range-repartition module. A stats-tap sitting +/// *below* an [`super::UnorderedRangeRepartitionExec`] (or any other +/// distribution-changing operator) is excluded automatically because +/// the walker stops at that boundary; its sketch describes data the +/// repartitioner then routed away and is no longer meaningful at the +/// plan's output. +/// +/// Executors call this once per task at completion to package what to +/// return to the scheduler. +pub fn collect_reports( + plan: &Arc, +) -> Result> { + use datafusion_proto::physical_plan::{ + DefaultPhysicalExtensionCodec, DefaultPhysicalProtoConverter, + }; + let codec = DefaultPhysicalExtensionCodec {}; + let converter = DefaultPhysicalProtoConverter {}; + let mut found: Vec<&RuntimeStatsExec> = Vec::new(); + collect_reachable_stats(plan, &mut found); + found + .into_iter() + .map(|stats| stats_to_report(stats, &codec, &converter)) + .collect() +} + +/// DFS `plan` through single-child chains only, descending through +/// distribution-preserving nodes and past any [`RuntimeStatsExec`] found +/// on the way. Stops at any branch, leaf, or non-whitelisted node. +/// Similar in shape to `range_repartition_common::find_runtime_stats` +/// but collects *all* reachable stats rather than returning the first +/// match keyed to a specific routing expression. +fn collect_reachable_stats<'a>( + plan: &'a Arc, + out: &mut Vec<&'a RuntimeStatsExec>, +) { + if let Some(stats) = plan.downcast_ref::() { + out.push(stats); + // Continue descending — a plan could conceivably chain multiple + // stats-taps; `preserves_distribution` still guards the recursion. + } else if !super::range_repartition_common::preserves_distribution(plan.as_ref()) { + return; + } + let children = plan.children(); + let [only_child] = children.as_slice() else { + return; + }; + collect_reachable_stats(only_child, out); +} + +fn stats_to_report( + stats: &RuntimeStatsExec, + codec: &dyn datafusion_proto::physical_plan::PhysicalExtensionCodec, + converter: &datafusion_proto::physical_plan::DefaultPhysicalProtoConverter, +) -> Result { + use datafusion_proto::physical_plan::to_proto::serialize_physical_sort_exprs; + let order_by = match stats.order_by() { + Some(order_by) => { + serialize_physical_sort_exprs(order_by.iter().cloned(), codec, converter)? + } + None => Vec::new(), + }; + // Iterate every partition slot the operator holds. Slots the task + // didn't touch have row_count = 0 and an empty sketch; we still emit + // them so the scheduler sees a shape-consistent view. + let partition_count = stats.partition_count(); + let mut partitions = Vec::with_capacity(partition_count); + for partition_id in 0..partition_count { + let row_count = stats.row_count(partition_id)? as u64; + let sketch = match stats.quantile_sketch(partition_id)? { + Some(sk) if sk.count() > 0.0 => Some(sketch_to_proto(&sk)?), + _ => None, + }; + partitions.push(crate::serde::protobuf::RuntimeStatsPartitionEntry { + partition_id: partition_id as u32, + row_count, + sketch, + }); + } + Ok(crate::serde::protobuf::RuntimeStatsReport { + order_by, + partitions, + }) +} + /// Serialize a T-Digest to the on-wire /// [`crate::serde::protobuf::QuantileSketchState`]. /// @@ -555,6 +643,200 @@ pub fn sketch_from_proto( Ok(TDigest::from_scalar_state(&scalars)) } +/// One group's merged view: sketches from every report sharing the same +/// `order_by` wire tag combined, plus the `partition_count - 1` quantile +/// cuts a globally-informed router would use for `partition_count` +/// output partitions. +#[derive(Debug, Clone)] +pub struct MergedRuntimeStats { + /// How many `PhysicalSortExprNode`s were in the shared `order_by` tag. + pub order_by_len: usize, + /// Number of output partitions the router used (per-report + /// `partitions.len()`, must agree across reports in the group). + pub partition_count: usize, + /// Number of `RuntimeStatsReport`s contributing to this group. + pub task_count: usize, + /// Sum of `row_count` across every partition entry in the group. + pub total_rows: u64, + /// `partition_count - 1` cut points at quantiles `i/partition_count` + /// on the merged T-Digest. Empty when `partition_count < 2` or no + /// non-empty sketches were merged. + pub cuts: Vec, + /// Merged T-Digest's `min()` if at least one non-empty sketch + /// contributed; `None` in row-count-only mode. + pub min: Option, + /// Merged T-Digest's `max()` if at least one non-empty sketch + /// contributed; `None` in row-count-only mode. + pub max: Option, +} + +/// Group `RuntimeStatsReport`s by `order_by` wire tag, merge the T-Digests +/// within each group, and return one [`MergedRuntimeStats`] per group. +/// +/// Errors surface planner / wire invariants the caller can't quietly +/// paper over: mismatched partition counts within a group (the planner +/// gave two tasks in the same stage different plans), or a sketch that +/// won't decode (wire corruption). The caller decides whether to log +/// and drop, propagate to the user, or fail the stage. +pub fn merge_reports( + reports: &[crate::serde::protobuf::RuntimeStatsReport], +) -> Result> { + use prost::Message; + use std::collections::HashMap; + + if reports.is_empty() { + return Ok(Vec::new()); + } + + // Group by the bytes of the encoded `order_by`. Prost-encoding each + // `PhysicalSortExprNode` and concatenating gives a stable, cheap + // grouping key without needing `Hash` on the generated proto types. + let mut groups: HashMap, Vec<&crate::serde::protobuf::RuntimeStatsReport>> = + HashMap::new(); + for report in reports { + let mut group_key = Vec::new(); + for expr in &report.order_by { + expr.encode(&mut group_key) + .expect("Vec is an infallible sink for prost::Message::encode"); + } + groups.entry(group_key).or_default().push(report); + } + + let mut merged_groups = Vec::with_capacity(groups.len()); + for group in groups.into_values() { + merged_groups.push(merge_group(&group)?); + } + Ok(merged_groups) +} + +/// Merge one group of reports (all sharing the same `order_by` tag). +/// Kept separate from `merge_reports` so the group iteration reads as a +/// single fallible step per group. +fn merge_group( + group: &[&crate::serde::protobuf::RuntimeStatsReport], +) -> Result { + let [first, rest @ ..] = group else { + // `merge_reports` only builds groups from `HashMap::entry().push()`, + // so an empty group is unreachable. Surface as internal error + // rather than panicking. + return internal_err!( + "runtime stats merge: empty group — merge_reports invariant broken" + ); + }; + let partition_count = first.partitions.len(); + let task_count = group.len(); + + // Every task ran the same stage plan, so partition counts must + // agree. Mismatch = internal invariant break. + for report in rest { + if report.partitions.len() != partition_count { + return internal_err!( + "runtime stats merge: order_by_len={} mismatched partition \ + counts across reports ({} vs {})", + first.order_by.len(), + partition_count, + report.partitions.len() + ); + } + } + + let mut total_rows: u64 = 0; + let mut sketches: Vec = Vec::new(); + for report in group { + for entry in &report.partitions { + total_rows = total_rows.saturating_add(entry.row_count); + if let Some(proto_sketch) = entry.sketch.as_ref() { + let sketch = sketch_from_proto(proto_sketch)?; + if sketch.count() > 0.0 { + sketches.push(sketch); + } + } + } + } + + if sketches.is_empty() { + return Ok(MergedRuntimeStats { + order_by_len: first.order_by.len(), + partition_count, + task_count, + total_rows, + cuts: Vec::new(), + min: None, + max: None, + }); + } + + let merged_sketch = TDigest::merge_digests(sketches.iter()); + let cuts: Vec = if partition_count > 1 { + (1..partition_count) + .map(|cut_index| { + merged_sketch.estimate_quantile(cut_index as f64 / partition_count as f64) + }) + .collect() + } else { + Vec::new() + }; + Ok(MergedRuntimeStats { + order_by_len: first.order_by.len(), + partition_count, + task_count, + total_rows, + cuts, + min: Some(merged_sketch.min()), + max: Some(merged_sketch.max()), + }) +} + +/// Merge `reports` and log each group's merged view at `debug!` +/// (`RUST_LOG` promotes when needed). Any merge error is logged at +/// `warn!` — the scheduler doesn't want telemetry loss to tank a query +/// whose data was already produced correctly. The scheduler calls this +/// once per stage-attempt at final-success. +pub fn log_merged_runtime_stats( + job_id: &str, + stage_id: usize, + reports: &[crate::serde::protobuf::RuntimeStatsReport], +) { + let merged_groups = match merge_reports(reports) { + Ok(groups) => groups, + Err(err) => { + log::warn!( + "runtime stats merge failed for job={job_id} stage={stage_id}: {err}" + ); + return; + } + }; + for merged in merged_groups { + match (merged.min, merged.max) { + (Some(min), Some(max)) => log::debug!( + "merged runtime stats: job={} stage={} order_by_len={} \ + partition_count={} task_count={} total_rows={} cuts={:?} \ + min={} max={}", + job_id, + stage_id, + merged.order_by_len, + merged.partition_count, + merged.task_count, + merged.total_rows, + merged.cuts, + min, + max, + ), + _ => log::debug!( + "merged runtime stats: job={} stage={} order_by_len={} \ + partition_count={} task_count={} total_rows={} cuts=[] \ + (no sketches)", + job_id, + stage_id, + merged.order_by_len, + merged.partition_count, + merged.task_count, + merged.total_rows, + ), + } + } +} + #[cfg(test)] mod wire_tests { use super::*; @@ -732,3 +1014,364 @@ mod stream_tests { assert!(stats.merged_quantile_sketch().unwrap().is_none()); } } + +#[cfg(test)] +mod collect_tests { + //! Walker behavior: which `RuntimeStatsExec`s does `collect_reports` + //! see through the whitelist, and what do the emitted reports look + //! like once the plan has been drained? + + use super::*; + use crate::execution_plans::BufferExec; + use crate::execution_plans::buffer::BufferMode; + use datafusion::arrow::array::Float64Array; + use datafusion::arrow::compute::SortOptions; + use datafusion::arrow::datatypes::{DataType, Field, Schema}; + use datafusion::arrow::record_batch::RecordBatch; + use datafusion::datasource::memory::MemorySourceConfig; + use datafusion::datasource::source::DataSourceExec; + use datafusion::physical_expr::{LexOrdering, PhysicalSortExpr}; + use datafusion::physical_plan::common; + use datafusion::physical_plan::expressions::col; + use datafusion::physical_plan::sorts::sort::SortExec; + use datafusion::prelude::SessionContext; + + fn schema_v() -> Arc { + Arc::new(Schema::new(vec![Field::new("v", DataType::Float64, false)])) + } + + fn v_batch(schema: &Arc, v: Vec) -> RecordBatch { + RecordBatch::try_new(schema.clone(), vec![Arc::new(Float64Array::from(v))]) + .unwrap() + } + + fn v_input(schema: Arc) -> Arc { + let b1 = v_batch(&schema, vec![1.0, 3.0, 5.0]); + let b2 = v_batch(&schema, vec![2.0, 4.0]); + let memory = + Arc::new(MemorySourceConfig::try_new(&[vec![b1, b2]], schema, None).unwrap()); + Arc::new(DataSourceExec::new(memory)) + } + + fn sort_expr_on_v(schema: &Arc) -> PhysicalSortExpr { + PhysicalSortExpr { + expr: col("v", schema.as_ref()).unwrap(), + options: SortOptions { + descending: false, + nulls_first: false, + }, + } + } + + /// Stats sit at plan root, sketching mode: `collect_reports` returns + /// exactly one report whose partition entry carries the observed + /// row_count and a populated sketch that survives the on-wire round- + /// trip via `sketch_from_proto`. + #[tokio::test] + async fn collect_reports_finds_stats_and_ships_sketch() { + let schema = schema_v(); + let input = v_input(schema.clone()); + let stats = Arc::new( + RuntimeStatsExec::try_new(input, Some(vec![sort_expr_on_v(&schema)])) + .unwrap(), + ); + + // Drive the stream so counters and the sketch actually fill. + let ctx = SessionContext::new().task_ctx(); + let stream = stats.clone().execute(0, ctx).unwrap(); + let _ = common::collect(stream).await.unwrap(); + + let plan: Arc = stats; + let reports = collect_reports(&plan).expect("collect_reports must succeed"); + let [report] = reports.as_slice() else { + panic!( + "expected exactly one report, got {} (order_by tags: {:?})", + reports.len(), + reports.iter().map(|r| r.order_by.len()).collect::>() + ); + }; + assert_eq!(report.order_by.len(), 1, "one sort expr encoded"); + let [entry] = report.partitions.as_slice() else { + panic!( + "expected one partition entry, got {}", + report.partitions.len() + ); + }; + assert_eq!(entry.partition_id, 0); + assert_eq!(entry.row_count, 5); + let proto_sketch = entry.sketch.as_ref().expect("sketch present in wire"); + let round_tripped = sketch_from_proto(proto_sketch).unwrap(); + assert_eq!(round_tripped.count(), 5.0); + assert_eq!(round_tripped.min(), 1.0); + assert_eq!(round_tripped.max(), 5.0); + } + + /// Row-count-only mode: report emitted, but its partition entry + /// carries no sketch. + #[tokio::test] + async fn collect_reports_row_count_only_emits_report_without_sketch() { + let schema = schema_v(); + let input = v_input(schema.clone()); + let stats = Arc::new(RuntimeStatsExec::try_new(input, None).unwrap()); + + let ctx = SessionContext::new().task_ctx(); + let stream = stats.clone().execute(0, ctx).unwrap(); + let _ = common::collect(stream).await.unwrap(); + + let plan: Arc = stats; + let reports = collect_reports(&plan).unwrap(); + let [report] = reports.as_slice() else { + panic!("expected one report, got {}", reports.len()); + }; + assert!(report.order_by.is_empty()); + assert_eq!(report.partitions.len(), 1); + assert!( + report.partitions[0].sketch.is_none(), + "no sketch in row-count-only mode" + ); + assert_eq!(report.partitions[0].row_count, 5); + } + + /// Whitelisted intermediary (`BufferExec` in Dam mode) between plan + /// root and the stats-tap: walker still descends to it. + #[tokio::test] + async fn collect_reports_descends_through_whitelisted_op() { + let schema = schema_v(); + let input = v_input(schema.clone()); + let stats = Arc::new( + RuntimeStatsExec::try_new(input, Some(vec![sort_expr_on_v(&schema)])) + .unwrap(), + ); + let buffer: Arc = + Arc::new(BufferExec::try_new(stats, BufferMode::Dam).unwrap()); + + // Drain via the outer plan so counters fill. + let ctx = SessionContext::new().task_ctx(); + let stream = buffer.clone().execute(0, ctx).unwrap(); + let _ = common::collect(stream).await.unwrap(); + + let reports = collect_reports(&buffer).unwrap(); + assert_eq!(reports.len(), 1, "buffer must not block the walker"); + assert_eq!(reports[0].partitions[0].row_count, 5); + } + + /// A `SortExec` with `preserve_partitioning=false` collapses N→1; + /// the whitelist excludes that variant explicitly. The walker + /// stops at the collapse and doesn't reach the stats below. + #[tokio::test] + async fn collect_reports_stops_at_sort_that_collapses_partitions() { + let schema = schema_v(); + let input = v_input(schema.clone()); + let stats: Arc = Arc::new( + RuntimeStatsExec::try_new(input, Some(vec![sort_expr_on_v(&schema)])) + .unwrap(), + ); + // Default SortExec has preserve_partitioning=false — the + // whitelist path we're testing rejects it. + let sort = SortExec::new( + LexOrdering::new(vec![sort_expr_on_v(&schema)]).unwrap(), + stats, + ); + assert!( + !sort.preserve_partitioning(), + "test fixture assumes N→1 sort" + ); + let plan: Arc = Arc::new(sort); + let reports = collect_reports(&plan).unwrap(); + assert!( + reports.is_empty(), + "N→1 sort must block the walker; got {} reports", + reports.len() + ); + } +} + +#[cfg(test)] +mod merge_tests { + //! Scheduler-side aggregation: given several `RuntimeStatsReport`s + //! sharing an `order_by` tag, verify the merged view (total rows, + //! cuts, min/max) reflects the union of the underlying samples. + + use super::*; + use crate::serde::protobuf::{RuntimeStatsPartitionEntry, RuntimeStatsReport}; + use datafusion_proto::protobuf::PhysicalSortExprNode; + + /// Build a report whose partition slots each carry a sketch made + /// from that slot's `values`. Slot `slot_id` in the resulting + /// report has `row_count = values[slot_id].len()` and a sketch + /// over those values. + fn sketching_report( + order_by: Vec, + values_per_slot: Vec>, + ) -> RuntimeStatsReport { + let partitions = values_per_slot + .into_iter() + .enumerate() + .map(|(slot_id, slot_values)| { + let row_count = slot_values.len() as u64; + let sketch = if slot_values.is_empty() { + None + } else { + let digest = TDigest::new(100).merge_unsorted_f64(slot_values); + Some(sketch_to_proto(&digest).unwrap()) + }; + RuntimeStatsPartitionEntry { + partition_id: slot_id as u32, + row_count, + sketch, + } + }) + .collect(); + RuntimeStatsReport { + order_by, + partitions, + } + } + + fn only_group(reports: &[RuntimeStatsReport]) -> MergedRuntimeStats { + let mut groups = merge_reports(reports).expect("merge should succeed"); + match groups.as_slice() { + [_] => groups.remove(0), + other => panic!("expected exactly one group, got {}", other.len()), + } + } + + /// Two reports over disjoint value ranges — merged sketch spans the + /// union, total_rows sums, and the partition_count=2 midpoint cut + /// falls between the two ranges. + #[test] + fn merge_reports_combines_disjoint_ranges() { + // Both reports share an empty `order_by` — we just need two + // reports that land in the same group. + let low_range = sketching_report(vec![], vec![vec![1.0, 2.0, 3.0], vec![]]); + let high_range = sketching_report(vec![], vec![vec![], vec![10.0, 11.0, 12.0]]); + + let group = only_group(&[low_range, high_range]); + assert_eq!(group.partition_count, 2); + assert_eq!(group.task_count, 2); + assert_eq!(group.total_rows, 6); + let midpoint = match group.cuts.as_slice() { + [midpoint] => *midpoint, + other => panic!("expected exactly one cut, got {other:?}"), + }; + assert!( + (3.0..=10.0).contains(&midpoint), + "midpoint cut should land between ranges (got {midpoint})" + ); + assert_eq!(group.min, Some(1.0)); + assert_eq!(group.max, Some(12.0)); + } + + /// partition_count=4 cuts on a uniform [0, 100) sample land roughly + /// at quartiles — verifies the quantile indices `i / partition_count` + /// for `i in 1..partition_count`. + #[test] + fn merge_reports_partition_count_of_four_produces_three_quartile_cuts() { + let uniform: Vec = (0..100).map(|value| value as f64).collect(); + // Single report, partition_count=4: each slot gets 25 uniform + // samples. + let values_per_slot = vec![ + uniform[0..25].to_vec(), + uniform[25..50].to_vec(), + uniform[50..75].to_vec(), + uniform[75..100].to_vec(), + ]; + let report = sketching_report(vec![], values_per_slot); + + let group = only_group(&[report]); + assert_eq!(group.partition_count, 4); + let (p25, p50, p75) = match group.cuts.as_slice() { + [p25, p50, p75] => (*p25, *p50, *p75), + other => panic!("expected 3 cuts, got {other:?}"), + }; + // Loose bounds — T-Digest quantile estimates aren't exact, but + // must land in the expected quartile bands. + assert!((10.0..40.0).contains(&p25), "p25 near 25, got {p25}"); + assert!((35.0..65.0).contains(&p50), "p50 near 50, got {p50}"); + assert!((60.0..90.0).contains(&p75), "p75 near 75, got {p75}"); + } + + /// Row-count-only reports (no sketches) still produce a group with + /// summed `total_rows` — just empty `cuts` and `None` min/max. + #[test] + fn merge_reports_row_count_only_emits_empty_cuts() { + let make_report = |row_counts: [u64; 2]| RuntimeStatsReport { + order_by: vec![], + partitions: vec![ + RuntimeStatsPartitionEntry { + partition_id: 0, + row_count: row_counts[0], + sketch: None, + }, + RuntimeStatsPartitionEntry { + partition_id: 1, + row_count: row_counts[1], + sketch: None, + }, + ], + }; + let group = only_group(&[make_report([100, 200]), make_report([300, 400])]); + assert_eq!(group.partition_count, 2); + assert_eq!(group.total_rows, 1000); + assert!(group.cuts.is_empty()); + assert!(group.min.is_none()); + assert!(group.max.is_none()); + } + + /// Mismatched partition counts within a group surface as an error — + /// the caller (scheduler / slice-D consumer) sees the invariant + /// break rather than silently getting a partial merge. + #[test] + fn merge_reports_errors_on_mismatched_partition_counts() { + let two_partitions = sketching_report(vec![], vec![vec![1.0], vec![2.0]]); + let one_partition = sketching_report(vec![], vec![vec![3.0]]); + let err = merge_reports(&[two_partitions, one_partition]) + .expect_err("mismatched partition counts must error"); + let message = err.to_string(); + assert!( + message.contains("mismatched partition counts"), + "expected mismatch error, got: {message}" + ); + } + + /// A wire-corrupted sketch — one whose scalar-state length is wrong + /// — surfaces the underlying `sketch_from_proto` error rather than + /// getting silently dropped. + #[test] + fn merge_reports_propagates_sketch_decode_errors() { + use crate::serde::protobuf::QuantileSketchState; + use datafusion::common::ScalarValue; + + // Six scalars is the valid shape; three is a corrupted wire. + let corrupt_sketch = QuantileSketchState { + state: (0..3) + .map(|_| { + datafusion_proto_common::ScalarValue::try_from(&ScalarValue::Float64( + Some(0.0), + )) + .unwrap() + }) + .collect(), + }; + let report = RuntimeStatsReport { + order_by: vec![], + partitions: vec![RuntimeStatsPartitionEntry { + partition_id: 0, + row_count: 1, + sketch: Some(corrupt_sketch), + }], + }; + let err = merge_reports(&[report]) + .expect_err("corrupt sketch must surface as an error"); + assert!( + err.to_string().contains("expected 6 elements"), + "expected shape-error propagation, got: {err}" + ); + } + + /// Empty input → empty output; ensures no panics or spurious groups. + #[test] + fn merge_reports_empty_input_is_empty_output() { + assert!(merge_reports(&[]).unwrap().is_empty()); + } +} diff --git a/ballista/core/src/serde/generated/ballista.rs b/ballista/core/src/serde/generated/ballista.rs index ca5b589b7d..2a23eaf1a4 100644 --- a/ballista/core/src/serde/generated/ballista.rs +++ b/ballista/core/src/serde/generated/ballista.rs @@ -871,6 +871,49 @@ pub struct SuccessfulTask { /// so we might want to think about some refactoring of the task definitions #[prost(message, repeated, tag = "2")] pub partitions: ::prost::alloc::vec::Vec, + /// Reports from `RuntimeStatsExec` operators in this task's plan that + /// are still valid at the plan's output (walked from the top through + /// distribution-preserving nodes only — see + /// `range_repartition_common::preserves_distribution`). Empty when the + /// plan has no such stats-taps. Currently reports one entry per + /// executed `RuntimeStatsExec`; the scheduler groups by `order_by` tag + /// to combine reports across tasks/executors. + #[prost(message, repeated, tag = "3")] + pub runtime_stats: ::prost::alloc::vec::Vec, +} +/// One report per `RuntimeStatsExec` in the executed plan. +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct RuntimeStatsReport { + /// What the operator was sampling. Empty = row-count-only mode. When + /// non-empty, the first entry identifies which routing expression this + /// sketch describes; the scheduler groups sketches by this tag to + /// combine samples across tasks that were sampling the same expression. + #[prost(message, repeated, tag = "1")] + pub order_by: ::prost::alloc::vec::Vec< + ::datafusion_proto::protobuf::PhysicalSortExprNode, + >, + /// Per-partition observations. Interpretation depends on the operator's + /// position in the plan — pre-repartition gets one entry per input + /// partition, post-repartition gets one per output sub-partition. The + /// scheduler groups by `order_by` tag and aggregates. + #[prost(message, repeated, tag = "2")] + pub partitions: ::prost::alloc::vec::Vec, +} +/// One partition's observations from a `RuntimeStatsExec`. +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct RuntimeStatsPartitionEntry { + #[prost(uint32, tag = "1")] + pub partition_id: u32, + #[prost(uint64, tag = "2")] + pub row_count: u64, + /// Present when the `RuntimeStatsExec` was in sketch mode AND this + /// partition observed at least one non-null routing value. + /// + /// TODO: `optional MinMaxState min_max` — for a lighter post-repartition + /// mode where the bin-packer just needs (min, max, count) per + /// sub-partition and a full T-Digest is overkill. + #[prost(message, optional, tag = "3")] + pub sketch: ::core::option::Option, } #[derive(Clone, Copy, PartialEq, Eq, Hash, ::prost::Message)] pub struct ExecutionError {} diff --git a/ballista/executor/src/execution_engine.rs b/ballista/executor/src/execution_engine.rs index 7f779d2c82..4f8e6663d8 100644 --- a/ballista/executor/src/execution_engine.rs +++ b/ballista/executor/src/execution_engine.rs @@ -86,6 +86,18 @@ pub trait QueryStageExecutor: Sync + Send + Debug + Display { /// Collects execution metrics from all operators in the plan. fn collect_plan_metrics(&self) -> Vec; + + /// Collect runtime-stats reports for every `RuntimeStatsExec` still + /// valid at the plan's output (walked through the distribution- + /// preserving whitelist). Called at task completion; whatever comes + /// back rides along in the task's `SuccessfulTask` message to the + /// scheduler. Default returns empty — implementers with real plans + /// override to walk their operator tree. + fn collect_runtime_stats_reports( + &self, + ) -> Vec { + Vec::new() + } } /// Default execution engine using DataFusion's ShuffleWriterExec. @@ -290,6 +302,31 @@ impl QueryStageExecutor for DefaultQueryStageExec { ShuffleWriterVariant::Sort(writer) => utils::collect_plan_metrics(writer), } } + + fn collect_runtime_stats_reports( + &self, + ) -> Vec { + // Walk from the shuffle writer's plan through the whitelist. If + // no `RuntimeStatsExec` sits within reach, we return an empty + // Vec — the majority of plans (anything not on the parallel- + // window path today). Serialization 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. + let plan: Arc = match &self.shuffle_writer { + ShuffleWriterVariant::Passthrough(writer) => Arc::new(writer.clone()), + ShuffleWriterVariant::Sort(writer) => Arc::new(writer.clone()), + }; + match ballista_core::execution_plans::collect_runtime_stats_reports(&plan) { + Ok(reports) => reports, + Err(e) => { + log::warn!( + "collect_runtime_stats_reports failed, task will report empty stats: {e}" + ); + Vec::new() + } + } + } } /// Spawn K parallel `plan.execute(N, ctx)` calls against a shuffle writer, diff --git a/ballista/executor/src/execution_loop.rs b/ballista/executor/src/execution_loop.rs index e0eea30e41..b84c83b9d4 100644 --- a/ballista/executor/src/execution_loop.rs +++ b/ballista/executor/src/execution_loop.rs @@ -24,7 +24,7 @@ use crate::cpu_bound_executor::DedicatedExecutor; use crate::executor::Executor; use crate::executor_process::remove_job_data; -use crate::{TaskExecutionTimes, as_task_status}; +use crate::{TaskCompletionExtras, TaskExecutionTimes, as_task_status}; use ballista_core::JobId; use ballista_core::error::BallistaError; use ballista_core::extension::SessionConfigHelperExt; @@ -245,8 +245,8 @@ where executor.metadata.id.clone(), task.task_attempt_num as usize, task_key, - None, task_execution_times, + TaskCompletionExtras::default(), )) { warn!("failed to send task status: {error:?}"); }; @@ -392,6 +392,7 @@ async fn run_received_task, BallistaError>>() .ok(); + let runtime_stats = query_stage_exec.collect_runtime_stats_reports(); let end_exec_time = SystemTime::now() .duration_since(UNIX_EPOCH) @@ -409,8 +410,11 @@ async fn run_received_task ExecutorServer, BallistaError>>() .ok(); + let runtime_stats = exec.collect_runtime_stats_reports(); let executor_id = &self.executor.metadata.id; let end_exec_time = SystemTime::now() @@ -468,8 +469,11 @@ impl ExecutorServer ExecutorServer>, + /// Runtime-stats reports harvested from `RuntimeStatsExec` taps in the plan. + pub runtime_stats: Vec, +} + /// Converts a task execution result into a [`TaskStatus`] protobuf message. /// /// This function wraps the outcome of task execution (success or failure) @@ -108,16 +123,22 @@ pub fn as_task_status( executor_id: String, stage_attempt_num: usize, key: TaskKey, - operator_metrics: Option>, execution_times: TaskExecutionTimes, + extras: TaskCompletionExtras, ) -> TaskStatus { + let TaskCompletionExtras { + operator_metrics, + runtime_stats, + } = extras; let metrics = operator_metrics.unwrap_or_default(); let task_id = key.task_id; match execution_result { Ok(partitions) => { debug!( - "Task {task_id} finished with operator_metrics array size {}", - metrics.len() + "Task {task_id} finished with operator_metrics array size {} \ + and {} runtime-stats report(s)", + metrics.len(), + runtime_stats.len(), ); TaskStatus { task_id: task_id as u32, @@ -131,6 +152,7 @@ pub fn as_task_status( status: Some(task_status::Status::Successful(SuccessfulTask { executor_id, partitions, + runtime_stats, })), } } diff --git a/ballista/scheduler/src/scheduler_server/mod.rs b/ballista/scheduler/src/scheduler_server/mod.rs index 66ed01b505..6fac8b0e45 100644 --- a/ballista/scheduler/src/scheduler_server/mod.rs +++ b/ballista/scheduler/src/scheduler_server/mod.rs @@ -569,6 +569,7 @@ mod test { status: Some(task_status::Status::Successful(SuccessfulTask { executor_id: "executor-1".to_owned(), partitions, + runtime_stats: vec![], })), }; diff --git a/ballista/scheduler/src/state/execution_graph.rs b/ballista/scheduler/src/state/execution_graph.rs index dd5704e769..e353945da5 100644 --- a/ballista/scheduler/src/state/execution_graph.rs +++ b/ballista/scheduler/src/state/execution_graph.rs @@ -36,7 +36,8 @@ use ballista_core::serde::protobuf::failed_task::FailedReason; use ballista_core::serde::protobuf::job_status::Status; use ballista_core::serde::protobuf::{FailedJob, ShuffleWritePartition, job_status}; use ballista_core::serde::protobuf::{ - FailedTask, JobStatus, ResultLost, RunningJob, SuccessfulJob, TaskStatus, + FailedTask, JobStatus, ResultLost, RunningJob, SuccessfulJob, SuccessfulTask, + TaskStatus, }; use ballista_core::serde::protobuf::{RunningTask, task_status}; use ballista_core::serde::scheduler::{ @@ -47,6 +48,8 @@ use crate::display::print_stage_metrics; use crate::planner::DistributedPlanner; use crate::scheduler_server::event::QueryStageSchedulerEvent; use crate::scheduler_server::timestamp_millis; +use ballista_core::execution_plans::log_merged_runtime_stats; + use crate::state::execution_stage::RunningStage; pub(crate) use crate::state::execution_stage::{ ExecutionStage, ResolvedStage, StageOutput, TaskInfo, UnresolvedStage, @@ -950,12 +953,15 @@ impl ExecutionGraph for StaticExecutionGraph { running_stage .update_task_metrics(task_id, operator_metrics)?; + let SuccessfulTask { + partitions, + runtime_stats, + .. + } = successful_task; + running_stage.append_runtime_stats_reports(runtime_stats); + locations.append(&mut partition_to_location( - &job_id, - task_id, - stage_id, - executor, - successful_task.partitions, + &job_id, task_id, stage_id, executor, partitions, )); } else { warn!( @@ -978,6 +984,11 @@ impl ExecutionGraph for StaticExecutionGraph { stage_metrics, ); } + log_merged_runtime_stats( + job_id.as_str(), + stage_id, + &running_stage.runtime_stats_reports, + ); } let output_links = running_stage.output_links.clone(); diff --git a/ballista/scheduler/src/state/execution_stage.rs b/ballista/scheduler/src/state/execution_stage.rs index 2b17a1e16d..339e161119 100644 --- a/ballista/scheduler/src/state/execution_stage.rs +++ b/ballista/scheduler/src/state/execution_stage.rs @@ -35,7 +35,8 @@ use ballista_core::error::{BallistaError, Result}; use ballista_core::execution_plans::{ShuffleWriterExec, SortShuffleWriterExec}; use ballista_core::serde::protobuf::failed_task::FailedReason; use ballista_core::serde::protobuf::{ - FailedTask, OperatorMetricsSet, ResultLost, SuccessfulTask, TaskKilled, TaskStatus, + FailedTask, OperatorMetricsSet, ResultLost, RuntimeStatsReport, SuccessfulTask, + TaskKilled, TaskStatus, }; use ballista_core::serde::protobuf::{RunningTask, task_status}; use ballista_core::serde::scheduler::PartitionLocation; @@ -218,6 +219,10 @@ pub struct RunningStage { pub stage_metrics: Option>, /// [SessionConfig] used for this stage pub session_config: Arc, + /// `RuntimeStatsReport`s collected from every successful task in + /// this stage attempt. Merged and logged once the stage finalizes; + /// dropped when the stage transitions to Successful. + pub runtime_stats_reports: Vec, } /// If a stage finishes successfully, its task statuses and metrics will be finalized @@ -630,6 +635,7 @@ impl RunningStage { task_failure_numbers: vec![0; partitions], stage_metrics: None, session_config, + runtime_stats_reports: Vec::new(), } } @@ -825,6 +831,17 @@ impl RunningStage { true } + /// Accumulate the `RuntimeStatsReport`s a successful task shipped + /// back in its `SuccessfulTask` payload. Consumed at final-success + /// by `ballista_core::execution_plans::log_merged_runtime_stats`; + /// each stage attempt starts with an empty accumulator. + pub fn append_runtime_stats_reports(&mut self, reports: Vec) { + if reports.is_empty() { + return; + } + self.runtime_stats_reports.extend(reports); + } + /// update and upsert the task metrics to the stage metrics pub fn update_task_metrics( &mut self, @@ -1132,6 +1149,9 @@ impl SuccessfulStage { .duration_since(UNIX_EPOCH) .unwrap() .as_millis(), + // Fresh attempt: previous attempt's stats are irrelevant. + // Merged-cut logging fires per-attempt on final success. + runtime_stats_reports: Vec::new(), } } @@ -1351,6 +1371,7 @@ mod tests { status: Some(task_status::Status::Successful(SuccessfulTask { executor_id: "executor-1".to_string(), partitions: vec![], + runtime_stats: vec![], })), metrics: vec![], } diff --git a/ballista/scheduler/src/state/task_manager.rs b/ballista/scheduler/src/state/task_manager.rs index fe84ef252b..7d46488c77 100644 --- a/ballista/scheduler/src/state/task_manager.rs +++ b/ballista/scheduler/src/state/task_manager.rs @@ -550,6 +550,7 @@ impl TaskManager let mut job_updates: HashMap> = HashMap::new(); for status in task_status { trace!("Task Update\n{status:?}"); + log_runtime_stats_arrival(executor, &status); let job_id = status.job_id.clone(); let job_task_statuses = job_updates.entry(job_id).or_default(); job_task_statuses.push(status); @@ -999,6 +1000,49 @@ impl From<&ExecutionGraphBox> for JobOverview { } } +/// Log any `RuntimeStatsReport`s that arrived with this task status. For +/// now this is proof-of-life for the wire — the reports flow from +/// executor to scheduler and land here observably. The per-stage +/// accumulator + merged-quantile-cut logging comes in a follow-up commit; +/// this hook is what verifies the plumbing works against a live cluster. +fn log_runtime_stats_arrival( + executor: &ExecutorMetadata, + status: &ballista_core::serde::protobuf::TaskStatus, +) { + use ballista_core::serde::protobuf::task_status::Status; + let Some(Status::Successful(successful)) = status.status.as_ref() else { + return; + }; + if successful.runtime_stats.is_empty() { + return; + } + for (report_idx, report) in successful.runtime_stats.iter().enumerate() { + let non_empty_partitions = + report.partitions.iter().filter(|p| p.row_count > 0).count(); + let total_rows: u64 = report.partitions.iter().map(|p| p.row_count).sum(); + let sketch_count = report + .partitions + .iter() + .filter(|p| p.sketch.is_some()) + .count(); + debug!( + "RuntimeStats arrival: executor={} job={} stage={} task={} \ + report[{}] order_by_len={} partitions={} non_empty={} \ + total_rows={} sketches={}", + executor.id, + status.job_id, + status.stage_id, + status.task_id, + report_idx, + report.order_by.len(), + report.partitions.len(), + non_empty_partitions, + total_rows, + sketch_count, + ); + } +} + #[cfg(test)] mod tests { use super::*; diff --git a/ballista/scheduler/src/test_utils.rs b/ballista/scheduler/src/test_utils.rs index be184cd6e6..c5231ed101 100644 --- a/ballista/scheduler/src/test_utils.rs +++ b/ballista/scheduler/src/test_utils.rs @@ -308,6 +308,7 @@ pub fn default_task_runner() -> impl TaskRunner { status: Some(task_status::Status::Successful(SuccessfulTask { executor_id: executor_id.clone(), partitions: partitions.clone(), + runtime_stats: vec![], })), }); } @@ -1209,6 +1210,7 @@ pub fn mock_completed_task(task: TaskDescription, executor_id: &str) -> TaskStat status: Some(task_status::Status::Successful(protobuf::SuccessfulTask { executor_id: executor_id.to_owned(), partitions, + runtime_stats: vec![], })), } } diff --git a/docs/source/upgrading/55.0.0.md b/docs/source/upgrading/55.0.0.md index ccbabfe89e..8c2432bd77 100644 --- a/docs/source/upgrading/55.0.0.md +++ b/docs/source/upgrading/55.0.0.md @@ -71,3 +71,47 @@ behavior. If you relied on it being `true` and want small joins to broadcast, se `HashJoinExec`, which remains eligible for broadcast promotion. Note this opts you into the non-spilling hash join for _all_ joins in the session, not only the small ones. + +### API changes + +#### `ballista_executor::as_task_status` signature reshaped + +`as_task_status` now takes a `TaskCompletionExtras` struct in place of the +`operator_metrics: Option>` positional parameter. The +struct also carries a new `runtime_stats: Vec` field used to +transport `RuntimeStatsExec` reports back to the scheduler. It is marked +`#[non_exhaustive]` with a `Default` impl, so future additions to the struct +are non-breaking for callers that construct via `..Default::default()`. + +The parameter order also changed: `execution_times` now precedes the extras +struct. + +**Action required:** external callers building a `TaskStatus` via +`as_task_status` should migrate from: + +```rust +as_task_status( + execution_result, + executor_id, + stage_attempt_num, + key, + operator_metrics, + execution_times, +) +``` + +to: + +```rust +as_task_status( + execution_result, + executor_id, + stage_attempt_num, + key, + execution_times, + TaskCompletionExtras { + operator_metrics, + ..Default::default() + }, +) +``` From 1af1e3d09a11300a178348b783c642ddb114d54e Mon Sep 17 00:00:00 2001 From: Brent Gardner Date: Sun, 26 Jul 2026 11:10:03 -0600 Subject: [PATCH 2/4] fix: AQE path forwards SuccessfulTask.runtime_stats to the scheduler MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `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) --- ballista/scheduler/src/state/aqe/mod.rs | 25 ++++++++++++++++++++----- 1 file changed, 20 insertions(+), 5 deletions(-) diff --git a/ballista/scheduler/src/state/aqe/mod.rs b/ballista/scheduler/src/state/aqe/mod.rs index deec100639..8ec8c7abb1 100644 --- a/ballista/scheduler/src/state/aqe/mod.rs +++ b/ballista/scheduler/src/state/aqe/mod.rs @@ -779,13 +779,23 @@ impl ExecutionGraph for AdaptiveExecutionGraph { running_stage .update_task_metrics(task_id, operator_metrics)?; + let ballista_core::serde::protobuf::SuccessfulTask { + partitions, + runtime_stats, + .. + } = successful_task; + debug!( + "append_runtime_stats_reports: job={} stage={} task={} report_count={}", + job_id, + stage_id, + task_id, + runtime_stats.len(), + ); + running_stage.append_runtime_stats_reports(runtime_stats); + locations.append( &mut crate::state::execution_graph::partition_to_location( - &job_id, - task_id, - stage_id, - executor, - successful_task.partitions, + &job_id, task_id, stage_id, executor, partitions, ), ); } else { @@ -815,6 +825,11 @@ impl ExecutionGraph for AdaptiveExecutionGraph { stage_metrics, ); } + ballista_core::execution_plans::log_merged_runtime_stats( + job_id.as_str(), + stage_id, + &running_stage.runtime_stats_reports, + ); } let stages_to_cancel = self.update_stage_progress( stage_id, From 68783140d45b213e07509401d2f5fb37810605d8 Mon Sep 17 00:00:00 2001 From: Brent Gardner Date: Mon, 27 Jul 2026 09:06:02 -0600 Subject: [PATCH 3/4] key runtime-stats reports by producer task_id MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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) --- ballista/core/src/execution_plans/mod.rs | 2 +- .../core/src/execution_plans/runtime_stats.rs | 22 ++++++++++-- ballista/scheduler/src/state/aqe/mod.rs | 3 +- .../scheduler/src/state/execution_graph.rs | 3 +- .../scheduler/src/state/execution_stage.rs | 36 +++++++++++++------ 5 files changed, 51 insertions(+), 15 deletions(-) diff --git a/ballista/core/src/execution_plans/mod.rs b/ballista/core/src/execution_plans/mod.rs index a7825ef396..c1b53e7baf 100644 --- a/ballista/core/src/execution_plans/mod.rs +++ b/ballista/core/src/execution_plans/mod.rs @@ -41,7 +41,7 @@ pub use distributed_explain_analyze::DistributedExplainAnalyzeExec; pub use distributed_query::{DistributedQueryExec, execute_physical_plan}; pub use ordered_range_repartition::OrderedRangeRepartitionExec; pub use runtime_stats::{ - MergedRuntimeStats, RuntimeStatsExec, + MergedRuntimeStats, RuntimeStatsExec, TaskRuntimeStats, collect_reports as collect_runtime_stats_reports, log_merged_runtime_stats, merge_reports as merge_runtime_stats_reports, sketch_from_proto, sketch_to_proto, }; diff --git a/ballista/core/src/execution_plans/runtime_stats.rs b/ballista/core/src/execution_plans/runtime_stats.rs index 8559d9b611..d654f93956 100644 --- a/ballista/core/src/execution_plans/runtime_stats.rs +++ b/ballista/core/src/execution_plans/runtime_stats.rs @@ -787,6 +787,22 @@ fn merge_group( }) } +/// One producer task's runtime-stats report, kept alongside the +/// `producer_task_id` that emitted it. The scheduler stores these on +/// `RunningStage.runtime_stats_reports` so downstream stages can address +/// individual producer files as `(producer_task_id, partition_id)` pairs — +/// the partition_id inside a report is producer-local (0..K URRE sub-parts), +/// so the pair is what uniquely identifies a shuffle file across producers. +#[derive(Debug, Clone)] +pub struct TaskRuntimeStats { + /// Producer task's task_id at the time it emitted the report. Matches + /// the `file_id` stamped on `ShuffleWritePartition` records. + pub producer_task_id: usize, + /// The report itself: per-partition row counts and (in sketch mode) + /// quantile sketches for the routing expression. + pub report: crate::serde::protobuf::RuntimeStatsReport, +} + /// Merge `reports` and log each group's merged view at `debug!` /// (`RUST_LOG` promotes when needed). Any merge error is logged at /// `warn!` — the scheduler doesn't want telemetry loss to tank a query @@ -795,9 +811,11 @@ fn merge_group( pub fn log_merged_runtime_stats( job_id: &str, stage_id: usize, - reports: &[crate::serde::protobuf::RuntimeStatsReport], + reports: &[TaskRuntimeStats], ) { - let merged_groups = match merge_reports(reports) { + let raw: Vec = + reports.iter().map(|t| t.report.clone()).collect(); + let merged_groups = match merge_reports(&raw) { Ok(groups) => groups, Err(err) => { log::warn!( diff --git a/ballista/scheduler/src/state/aqe/mod.rs b/ballista/scheduler/src/state/aqe/mod.rs index 8ec8c7abb1..1c92580e93 100644 --- a/ballista/scheduler/src/state/aqe/mod.rs +++ b/ballista/scheduler/src/state/aqe/mod.rs @@ -791,7 +791,8 @@ impl ExecutionGraph for AdaptiveExecutionGraph { task_id, runtime_stats.len(), ); - running_stage.append_runtime_stats_reports(runtime_stats); + running_stage + .append_runtime_stats_reports(task_id, runtime_stats); locations.append( &mut crate::state::execution_graph::partition_to_location( diff --git a/ballista/scheduler/src/state/execution_graph.rs b/ballista/scheduler/src/state/execution_graph.rs index e353945da5..8bf4fb884b 100644 --- a/ballista/scheduler/src/state/execution_graph.rs +++ b/ballista/scheduler/src/state/execution_graph.rs @@ -958,7 +958,8 @@ impl ExecutionGraph for StaticExecutionGraph { runtime_stats, .. } = successful_task; - running_stage.append_runtime_stats_reports(runtime_stats); + running_stage + .append_runtime_stats_reports(task_id, runtime_stats); locations.append(&mut partition_to_location( &job_id, task_id, stage_id, executor, partitions, diff --git a/ballista/scheduler/src/state/execution_stage.rs b/ballista/scheduler/src/state/execution_stage.rs index 339e161119..68468c5f9a 100644 --- a/ballista/scheduler/src/state/execution_stage.rs +++ b/ballista/scheduler/src/state/execution_stage.rs @@ -32,7 +32,9 @@ use datafusion::prelude::SessionConfig; use log::{debug, warn}; use ballista_core::error::{BallistaError, Result}; -use ballista_core::execution_plans::{ShuffleWriterExec, SortShuffleWriterExec}; +use ballista_core::execution_plans::{ + ShuffleWriterExec, SortShuffleWriterExec, TaskRuntimeStats, +}; use ballista_core::serde::protobuf::failed_task::FailedReason; use ballista_core::serde::protobuf::{ FailedTask, OperatorMetricsSet, ResultLost, RuntimeStatsReport, SuccessfulTask, @@ -219,10 +221,13 @@ pub struct RunningStage { pub stage_metrics: Option>, /// [SessionConfig] used for this stage pub session_config: Arc, - /// `RuntimeStatsReport`s collected from every successful task in - /// this stage attempt. Merged and logged once the stage finalizes; - /// dropped when the stage transitions to Successful. - pub runtime_stats_reports: Vec, + /// Per-task runtime-stats reports collected from every successful task + /// in this stage attempt. Each entry pairs the producer task's `task_id` + /// (== `file_id` on the emitted shuffle files) with its report — the pair + /// is what uniquely addresses a producer file, since the report's own + /// `partition_id` field is producer-local. Merged and logged once the + /// stage finalizes; dropped when the stage transitions to Successful. + pub runtime_stats_reports: Vec, } /// If a stage finishes successfully, its task statuses and metrics will be finalized @@ -832,14 +837,25 @@ impl RunningStage { } /// Accumulate the `RuntimeStatsReport`s a successful task shipped - /// back in its `SuccessfulTask` payload. Consumed at final-success - /// by `ballista_core::execution_plans::log_merged_runtime_stats`; - /// each stage attempt starts with an empty accumulator. - pub fn append_runtime_stats_reports(&mut self, reports: Vec) { + /// back in its `SuccessfulTask` payload, tagging each with the producer + /// `task_id` so downstream stages can address individual producer files. + /// Consumed at final-success by + /// `ballista_core::execution_plans::log_merged_runtime_stats` and by the + /// step-3 overlap router (`compute_overlapping_locations`); each stage + /// attempt starts with an empty accumulator. + pub fn append_runtime_stats_reports( + &mut self, + producer_task_id: usize, + reports: Vec, + ) { if reports.is_empty() { return; } - self.runtime_stats_reports.extend(reports); + self.runtime_stats_reports + .extend(reports.into_iter().map(|report| TaskRuntimeStats { + producer_task_id, + report, + })); } /// update and upsert the task metrics to the stage metrics From 8e9be38a18f92f7935f03348f4949f8427356153 Mon Sep 17 00:00:00 2001 From: Brent Gardner Date: Mon, 27 Jul 2026 09:10:51 -0600 Subject: [PATCH 4/4] purge stale runtime-stats reports when a task is reset for retry MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `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, #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) --- .../scheduler/src/state/execution_stage.rs | 98 ++++++++++++++++++- 1 file changed, 97 insertions(+), 1 deletion(-) diff --git a/ballista/scheduler/src/state/execution_stage.rs b/ballista/scheduler/src/state/execution_stage.rs index 68468c5f9a..4643e3efdd 100644 --- a/ballista/scheduler/src/state/execution_stage.rs +++ b/ballista/scheduler/src/state/execution_stage.rs @@ -1023,6 +1023,11 @@ impl RunningStage { /// Mark the task as lost/killed and push its partitions back to the /// front of `pending` so they are retried on the next bind. Does not /// touch failure counts — those are updated in `update_task_info`. + /// + /// Any runtime-stats reports the task already contributed are also + /// dropped: the retry will produce fresh sketches under a new + /// `task_id`, and leaving the ghost behind would double-count the + /// slice in the stage's merged view. pub fn reset_task_info(&mut self, task_id: usize) { let task = &mut self.task_infos[task_id]; let partitions = task.global_input_partition_ids.clone(); @@ -1033,16 +1038,25 @@ impl RunningStage { failed_reason: Some(FailedReason::TaskKilled(TaskKilled {})), }); self.pending.reschedule(partitions); + self.runtime_stats_reports + .retain(|s| s.producer_task_id != task_id); } /// Reset the running and completed tasks on a given executor by /// marking their `TaskInfo` as `Failed(ResultLost)` and pushing their /// partition slices back to `pending`. Returns the number of tasks /// reset. + /// + /// Runtime-stats reports contributed by any reset producer are + /// dropped — the retries will report fresh sketches under new + /// `task_id`s. Reports are append-only otherwise, so without this + /// purge the stage's merged view would double-count every slice that + /// bounced through a lost executor. pub fn reset_tasks(&mut self, executor: &str) -> usize { let mut reset = 0; let mut to_reschedule: Vec = vec![]; - for task in self.task_infos.iter_mut() { + let mut reset_task_ids: HashSet = HashSet::new(); + for (task_id, task) in self.task_infos.iter_mut().enumerate() { let matches_exec = match &task.task_status { task_status::Status::Running(RunningTask { executor_id }) | task_status::Status::Successful(SuccessfulTask { @@ -1058,10 +1072,13 @@ impl RunningStage { failed_reason: Some(FailedReason::ResultLost(ResultLost {})), }); to_reschedule.extend(task.global_input_partition_ids.iter().copied()); + reset_task_ids.insert(task_id); reset += 1; } } self.pending.reschedule(to_reschedule); + self.runtime_stats_reports + .retain(|s| !reset_task_ids.contains(&s.producer_task_id)); reset } @@ -1724,4 +1741,83 @@ mod tests { let aggregated = operator_metrics.aggregate_by_name(); assert_eq!(aggregated.output_rows(), Some(2000)); } + + /// Build a bare-bones report tagged with a `partition_id` we can look + /// for in assertions — no sketches, no order-by; the purge cares only + /// about the wrapping `producer_task_id`. + fn make_report(marker_partition_id: u32) -> RuntimeStatsReport { + RuntimeStatsReport { + order_by: vec![], + partitions: vec![ + ballista_core::serde::protobuf::RuntimeStatsPartitionEntry { + partition_id: marker_partition_id, + row_count: 0, + sketch: None, + }, + ], + } + } + + /// When a task is reset for retry, its previously-appended runtime-stats + /// reports must be dropped so the retry's fresh sketches don't merge + /// with ghost data from the original attempt. Reports from *other* tasks + /// must survive. + #[test] + fn test_reset_task_info_purges_runtime_stats_reports() { + let mut stage = make_running_stage(2); + append_running_task(&mut stage, 0, "executor-1", vec![0]); + append_running_task(&mut stage, 1, "executor-1", vec![1]); + + stage.append_runtime_stats_reports(0, vec![make_report(100)]); + stage.append_runtime_stats_reports(1, vec![make_report(200)]); + assert_eq!(stage.runtime_stats_reports.len(), 2); + + stage.reset_task_info(0); + + assert_eq!(stage.runtime_stats_reports.len(), 1); + assert_eq!(stage.runtime_stats_reports[0].producer_task_id, 1); + assert_eq!( + stage.runtime_stats_reports[0].report.partitions[0].partition_id, + 200 + ); + } + + /// Executor loss resets every task the executor was hosting; the + /// runtime-stats reports those (previously-Successful) producers had + /// already contributed must be purged along with the task status. + /// Reports produced on surviving executors must be left alone. + #[test] + fn test_reset_tasks_purges_runtime_stats_reports_for_lost_executor() { + let mut stage = make_running_stage(3); + append_running_task(&mut stage, 0, "executor-1", vec![0]); + append_running_task(&mut stage, 1, "executor-1", vec![1]); + append_running_task(&mut stage, 2, "executor-2", vec![2]); + // Drain the pending queue so reschedules below are visible. + stage.pending.next_slice(3); + + // All three tasks made it to Successful and shipped reports. + for (task_id, executor) in + [(0, "executor-1"), (1, "executor-1"), (2, "executor-2")] + { + stage.task_infos[task_id].task_status = + task_status::Status::Successful(SuccessfulTask { + executor_id: executor.to_string(), + partitions: vec![], + runtime_stats: vec![], + }); + stage.append_runtime_stats_reports( + task_id, + vec![make_report(100 + task_id as u32)], + ); + } + assert_eq!(stage.runtime_stats_reports.len(), 3); + + // Simulate executor-1 heartbeat loss. + let reset_count = stage.reset_tasks("executor-1"); + assert_eq!(reset_count, 2); + + // Only executor-2's producer survives. + assert_eq!(stage.runtime_stats_reports.len(), 1); + assert_eq!(stage.runtime_stats_reports[0].producer_task_id, 2); + } }