diff --git a/src/coordinator/query_coordinator.rs b/src/coordinator/query_coordinator.rs index 4889f888f..c57fffceb 100644 --- a/src/coordinator/query_coordinator.rs +++ b/src/coordinator/query_coordinator.rs @@ -312,8 +312,8 @@ impl<'a> StageCoordinator<'a> { stage_id: self.stage_id, task_number: task_i, }; - let task_metrics = self.metrics_store.clone(); - let completed_dynamic_filter_store = self.completed_dynamic_filter_store.clone(); + let mut task_metrics = self.metrics_store.clone(); + let mut completed_dynamic_filter_store = self.completed_dynamic_filter_store.clone(); let dynamic_filter_registry = Arc::clone(self.dynamic_filter_registry); let (load_info_tx, load_info_rx) = tokio::sync::mpsc::unbounded_channel(); let mut load_info_tx_opt = Some(load_info_tx); @@ -325,8 +325,8 @@ impl<'a> StageCoordinator<'a> { while let Some(msg) = worker_to_coordinator_rx.recv().await { match msg { WorkerToCoordinatorMsg::TaskMetrics(v) => { - if let Some(task_metrics) = &task_metrics { - task_metrics.insert(task_key, v); + if let Some(store) = task_metrics.take() { + store.insert(task_key, v); } } WorkerToCoordinatorMsg::LoadInfo(load_info) => { @@ -338,7 +338,7 @@ impl<'a> StageCoordinator<'a> { let _ = load_info_tx_opt.take(); } WorkerToCoordinatorMsg::TaskCompletedDynamicFilters(filters) => { - if let Some(store) = &completed_dynamic_filter_store { + if let Some(store) = completed_dynamic_filter_store.take() { store.insert(task_key, filters); } } @@ -347,6 +347,13 @@ impl<'a> StageCoordinator<'a> { } } } + // An unexecuted task sends no final reports; still complete its waits. + if let Some(store) = task_metrics { + store.insert(task_key, TaskMetrics::default()); + } + if let Some(store) = completed_dynamic_filter_store { + store.insert(task_key, TaskCompletedDynamicFilters::default()); + } }); load_info_rx } diff --git a/src/metrics/task_metrics_rewriter.rs b/src/metrics/task_metrics_rewriter.rs index 4c45aae72..6e5b9f0b4 100644 --- a/src/metrics/task_metrics_rewriter.rs +++ b/src/metrics/task_metrics_rewriter.rs @@ -276,6 +276,9 @@ pub fn stage_metrics_rewriter( stage.num ); }; + if task_metrics.pre_order_plan_metrics.is_empty() { + continue; // The task was never executed, so there are no plan metrics to rewrite. + } let mut per_task_counter = 0usize; stage.plan.apply_with_dt_ctx(d_ctx, |node, _ctx| { diff --git a/src/protocol/worker_channel.rs b/src/protocol/worker_channel.rs index 4f5257cfc..78323b45c 100644 --- a/src/protocol/worker_channel.rs +++ b/src/protocol/worker_channel.rs @@ -164,7 +164,7 @@ pub struct TaskDynamicFilter { pub expression: MaybeEncoded>, } -#[derive(Clone, Debug)] +#[derive(Clone, Debug, Default)] pub struct TaskMetrics { /// Metrics for a single task's plan nodes in pre-order traversal order. /// The TaskKey is implicit — it is determined by the SetPlanRequest that diff --git a/tests/metrics_collection.rs b/tests/metrics_collection.rs index 911568eac..e1f9ad3aa 100644 --- a/tests/metrics_collection.rs +++ b/tests/metrics_collection.rs @@ -18,7 +18,7 @@ mod tests { use datafusion_distributed::{ DefaultSessionBuilder, DistributedExt, DistributedLeafExec, DistributedMetricsFormat, NetworkCoalesceExec, NetworkShuffleExec, WorkerQueryContext, display_plan_ascii, - rewrite_distributed_plan_with_metrics, + rewrite_distributed_plan_with_dynamic_filters, rewrite_distributed_plan_with_metrics, }; use futures::TryStreamExt; use std::sync::Arc; @@ -382,11 +382,11 @@ mod tests { /// Regression for #739: an empty build side can leave sampled probe tasks unexecuted. #[tokio::test] - #[ignore = "metrics rewrite hangs on planned but unexecuted tasks"] async fn metrics_rewrite_after_unexecuted_aqe_tasks() -> Result<(), Box> { let (mut ctx, _guard, _) = start_localhost_context(3, DefaultSessionBuilder).await; ctx.set_distributed_dynamic_task_count(true)?; + ctx = ctx.with_distributed_dynamic_filter_collection(true)?; register_parquet_tables(&ctx).await?; { let state = ctx.state_ref(); @@ -409,6 +409,12 @@ mod tests { .await?; assert_eq!(batches.iter().map(|b| b.num_rows()).sum::(), 0); + let task_ctx = ctx.task_ctx(); + let plan = tokio::time::timeout( + Duration::from_secs(10), + rewrite_distributed_plan_with_dynamic_filters(plan, &task_ctx), + ) + .await??; tokio::time::timeout( Duration::from_secs(10), rewrite_distributed_plan_with_metrics(plan, DistributedMetricsFormat::PerTask),