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..c1b53e7baf 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, 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, +}; 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..d654f93956 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,218 @@ 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()), + }) +} + +/// 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 +/// 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: &[TaskRuntimeStats], +) { + 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!( + "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 +1032,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/aqe/mod.rs b/ballista/scheduler/src/state/aqe/mod.rs index deec100639..1c92580e93 100644 --- a/ballista/scheduler/src/state/aqe/mod.rs +++ b/ballista/scheduler/src/state/aqe/mod.rs @@ -779,13 +779,24 @@ 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(task_id, 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 +826,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, diff --git a/ballista/scheduler/src/state/execution_graph.rs b/ballista/scheduler/src/state/execution_graph.rs index dd5704e769..8bf4fb884b 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,16 @@ 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(task_id, 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 +985,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..4643e3efdd 100644 --- a/ballista/scheduler/src/state/execution_stage.rs +++ b/ballista/scheduler/src/state/execution_stage.rs @@ -32,10 +32,13 @@ 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, 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 +221,13 @@ pub struct RunningStage { pub stage_metrics: Option>, /// [SessionConfig] used for this stage pub session_config: Arc, + /// 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 @@ -630,6 +640,7 @@ impl RunningStage { task_failure_numbers: vec![0; partitions], stage_metrics: None, session_config, + runtime_stats_reports: Vec::new(), } } @@ -825,6 +836,28 @@ impl RunningStage { true } + /// Accumulate the `RuntimeStatsReport`s a successful task shipped + /// 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.into_iter().map(|report| TaskRuntimeStats { + producer_task_id, + report, + })); + } + /// update and upsert the task metrics to the stage metrics pub fn update_task_metrics( &mut self, @@ -990,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(); @@ -1000,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 { @@ -1025,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 } @@ -1132,6 +1182,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 +1404,7 @@ mod tests { status: Some(task_status::Status::Successful(SuccessfulTask { executor_id: "executor-1".to_string(), partitions: vec![], + runtime_stats: vec![], })), metrics: vec![], } @@ -1687,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); + } } 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() + }, +) +```