From 4059b205c46d4c553e37aa6a2424c6d50285350b Mon Sep 17 00:00:00 2001 From: Sergio Esteves Date: Fri, 25 Sep 2026 14:11:52 +0100 Subject: [PATCH 1/2] fix: finalize metrics for unexecuted tasks --- src/coordinator/query_coordinator.rs | 14 ++++++++++++++ src/metrics/task_metrics_rewriter.rs | 3 +++ tests/metrics_collection.rs | 1 - 3 files changed, 17 insertions(+), 1 deletion(-) diff --git a/src/coordinator/query_coordinator.rs b/src/coordinator/query_coordinator.rs index 4889f888f..475bae548 100644 --- a/src/coordinator/query_coordinator.rs +++ b/src/coordinator/query_coordinator.rs @@ -322,11 +322,13 @@ impl<'a> StageCoordinator<'a> { // metrics collection process might outlive the query's lifetime. #[allow(clippy::disallowed_methods)] tokio::spawn(async move { + let mut received_metrics = false; 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); + received_metrics = true; } } WorkerToCoordinatorMsg::LoadInfo(load_info) => { @@ -347,6 +349,18 @@ impl<'a> StageCoordinator<'a> { } } } + if !received_metrics { + if let Some(task_metrics) = task_metrics { + // An unexecuted task sends no metrics; still complete its wait. + task_metrics.insert( + task_key, + TaskMetrics { + pre_order_plan_metrics: vec![], + task_metrics: Default::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/tests/metrics_collection.rs b/tests/metrics_collection.rs index 911568eac..190c2347f 100644 --- a/tests/metrics_collection.rs +++ b/tests/metrics_collection.rs @@ -382,7 +382,6 @@ 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; From 5827fedb21411904ed1ca3bc31c72ec88e36dc86 Mon Sep 17 00:00:00 2001 From: Sergio Esteves Date: Fri, 25 Sep 2026 20:22:19 +0100 Subject: [PATCH 2/2] fix: finalize missing dynamic-filter reports --- src/coordinator/query_coordinator.rs | 29 +++++++++++----------------- src/protocol/worker_channel.rs | 2 +- tests/metrics_collection.rs | 9 ++++++++- 3 files changed, 20 insertions(+), 20 deletions(-) diff --git a/src/coordinator/query_coordinator.rs b/src/coordinator/query_coordinator.rs index 475bae548..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); @@ -322,13 +322,11 @@ impl<'a> StageCoordinator<'a> { // metrics collection process might outlive the query's lifetime. #[allow(clippy::disallowed_methods)] tokio::spawn(async move { - let mut received_metrics = false; 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); - received_metrics = true; + if let Some(store) = task_metrics.take() { + store.insert(task_key, v); } } WorkerToCoordinatorMsg::LoadInfo(load_info) => { @@ -340,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); } } @@ -349,17 +347,12 @@ impl<'a> StageCoordinator<'a> { } } } - if !received_metrics { - if let Some(task_metrics) = task_metrics { - // An unexecuted task sends no metrics; still complete its wait. - task_metrics.insert( - task_key, - TaskMetrics { - pre_order_plan_metrics: vec![], - task_metrics: Default::default(), - }, - ); - } + // 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/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 190c2347f..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; @@ -386,6 +386,7 @@ mod tests { { 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(); @@ -408,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),