datafusion_distributed/metrics/
task_metrics_rewriter.rs

1use crate::common::TreeNodeExt;
2use crate::coordinator::{DistributedExec, MetricsStore};
3use crate::distributed_planner::NetworkBoundaryExt;
4use crate::execution_plans::MetricsWrapperExec;
5use crate::metrics::DISTRIBUTED_DATAFUSION_TASK_ID_LABEL;
6use crate::metrics::collect_plan_metrics;
7use crate::stage::{LocalStage, Stage};
8use crate::{DistributedTaskContext, TaskKey};
9use datafusion::common::HashMap;
10use datafusion::common::plan_err;
11use datafusion::common::tree_node::Transformed;
12use datafusion::common::tree_node::TreeNode;
13use datafusion::common::tree_node::TreeNodeRecursion;
14use datafusion::error::Result;
15use datafusion::physical_plan::internal_err;
16use datafusion::physical_plan::metrics::{Label, Metric, MetricsSet};
17use datafusion::physical_plan::{ChildrenPropertiesMode, ExecutionPlan, ReplaceChildrenOptions};
18use std::sync::Arc;
19
20/// Format to use when displaying metrics for a distributed plan.
21#[derive(Debug, Clone, Copy, PartialEq, Eq)]
22pub enum DistributedMetricsFormat {
23    /// Metrics are aggregated across all tasks. ex. a `output_rows=X` represents the output rows for all tasks.
24    Aggregated,
25
26    /// Metrics are labeled with their task id and displayed per task. ex. `output_rows` is shown as
27    /// `output_rows={0:.., 1:..}`, one entry per task.
28    PerTask,
29}
30
31impl DistributedMetricsFormat {
32    pub(crate) fn to_rewrite_ctx(self, task_id: u64) -> RewriteCtx {
33        match self {
34            DistributedMetricsFormat::Aggregated => RewriteCtx::default(),
35            DistributedMetricsFormat::PerTask => RewriteCtx::from_task_id(task_id),
36        }
37    }
38}
39
40/// Rewrites a distributed plan with metrics. Does nothing if the root node is not a [DistributedExec].
41/// Returns an error if the distributed plan was not executed.
42///
43/// Waits for all worker task metrics to arrive before rewriting, so the result is always complete.
44pub async fn rewrite_distributed_plan_with_metrics(
45    plan: Arc<dyn ExecutionPlan>,
46    format: DistributedMetricsFormat,
47) -> Result<Arc<dyn ExecutionPlan>> {
48    let Some(distributed_exec) = plan.downcast_ref::<DistributedExec>() else {
49        return Ok(plan);
50    };
51
52    distributed_exec.wait_for_metrics().await;
53
54    let Some(metrics_collection) = distributed_exec.metrics_store.clone() else {
55        return Ok(plan);
56    };
57
58    let head_stage = distributed_exec.head_stage()?;
59    let task_metrics = collect_plan_metrics(&head_stage)?;
60
61    // Rewrite the DistributedExec's child plan with metrics.
62    let dist_exec_plan_with_metrics = rewrite_local_plan_with_metrics(
63        format.to_rewrite_ctx(0), // Task id is 0 for the DistributedExec plan
64        distributed_exec.plan_for_viz()?,
65        task_metrics,
66    )?;
67
68    let transformed = dist_exec_plan_with_metrics.transform_down(|plan| {
69        // Transform all stages using NetworkShuffleExec and NetworkCoalesceExec as barriers.
70        if let Some(network_boundary) = plan.as_network_boundary() {
71            let Stage::Local(stage) = network_boundary.input_stage() else {
72                return plan_err!("Stage was not in Local state");
73            };
74            // This transform is a bit inefficient because we traverse the plan nodes twice
75            // For now, we are okay with trading off performance for simplicity.
76            let plan_with_metrics =
77                stage_metrics_rewriter(stage, Arc::clone(&metrics_collection), format)?;
78            let network_boundary = network_boundary.with_input_stage(Stage::Local(LocalStage {
79                query_id: stage.query_id,
80                num: stage.num,
81                plan: plan_with_metrics,
82                tasks: stage.tasks,
83                metrics_set: stage.metrics_set.clone(),
84            }))?;
85            let network_boundary =
86                MetricsWrapperExec::new(network_boundary, plan.metrics().unwrap_or_default());
87            return Ok(Transformed::yes(Arc::new(network_boundary)));
88        }
89
90        Ok(Transformed::no(plan))
91    })?;
92    plan.replace_children(
93        vec![transformed.data],
94        ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute),
95    )
96}
97
98/// Extra information for rewriting local plans.
99#[derive(Default)]
100pub struct RewriteCtx {
101    /// Used to rename metrics for the current task.
102    pub task_id: Option<u64>,
103}
104
105impl RewriteCtx {
106    pub(crate) fn from_task_id(task_id: u64) -> RewriteCtx {
107        RewriteCtx {
108            task_id: Some(task_id),
109        }
110    }
111
112    /// Rewrites the [MetricsSet] depending on the context.
113    pub(crate) fn maybe_rewrite_node_metics(&self, node_metrics: MetricsSet) -> MetricsSet {
114        if let Some(task_id) = self.task_id {
115            return annotate_metrics_set_with_task_id(node_metrics, task_id);
116        }
117        node_metrics
118    }
119}
120
121/// Adds task id labels to all metrics in the provided [MetricsSet].
122///
123/// TODO: This re-allocates the vec of metrics by creating a new [MetricsSet]. It also
124/// reallocates the labels vec for each metric. Can we avoid this?
125/// See https://github.com/apache/datafusion/issues/19959
126pub fn annotate_metrics_set_with_task_id(metrics_set: MetricsSet, task_id: u64) -> MetricsSet {
127    let mut result = MetricsSet::new();
128
129    for metric in metrics_set.iter() {
130        let mut labels = metric.labels().to_vec();
131        labels.push(Label::new(
132            DISTRIBUTED_DATAFUSION_TASK_ID_LABEL,
133            task_id.to_string(),
134        ));
135        result.push(Arc::new(Metric::new_with_labels(
136            metric.value().clone(),
137            metric.partition(),
138            labels,
139        )));
140    }
141
142    result
143}
144
145/// Rewrites a local plan with metrics, stopping at network boundaries.
146///
147/// Example:
148///
149/// AggregateExec [output_rows = 1, elapsed_compute = 100]
150///  └── ProjectionExec [output_rows = 2, elapsed_compute = 200]
151///      └── NetworkShuffleExec [bytes_transferred = 100, max_mem_used = 100]
152///
153/// The result will be:
154///
155/// MetricsWrapperExec (wrapped: AggregateExec) [output_rows = 1, elapsed_compute = 100]
156///  └── MetricsWrapperExec (wrapped: ProjectionExec) [output_rows = 2, elapsed_compute = 200]
157///      └── MetricsWrapperExec (wrapped: NetworkShuffleExec) [bytes_transferred = 100, max_mem_used = 100]
158pub fn rewrite_local_plan_with_metrics(
159    ctx: RewriteCtx,
160    plan: Arc<dyn ExecutionPlan>,
161    metrics: Vec<MetricsSet>,
162) -> Result<Arc<dyn ExecutionPlan>> {
163    let mut idx = 0;
164    Ok(plan
165        .transform_down(|node| {
166            if idx >= metrics.len() {
167                return internal_err!("not enough metrics provided to rewrite plan");
168            }
169            let mut node_metrics = metrics[idx].clone();
170
171            node_metrics = ctx.maybe_rewrite_node_metics(node_metrics);
172
173            idx += 1;
174            Ok(Transformed::new(
175                Arc::new(MetricsWrapperExec::new(node.clone(), node_metrics)),
176                true,
177                if node.is_network_boundary() {
178                    TreeNodeRecursion::Jump
179                } else {
180                    TreeNodeRecursion::Continue
181                },
182            ))
183        })?
184        .data)
185}
186
187/// Enriches a stage with metrics from each task by re-writing the plan using
188/// [MetricsWrapperExec] nodes.
189///
190/// Example:
191///
192/// For a stage with 2 tasks:
193///
194/// Task 1:
195/// AggregateExec [output_rows = 1, elapsed_compute = 100]
196///  └── ProjectionExec [output_rows = 2, elapsed_compute = 200]
197///      └── NetworkShuffleExec [bytes_transferred = 100, max_mem_used = 100]
198///
199/// Task 2:
200/// AggregateExec [output_rows = 3, elapsed_compute = 300]
201///  └── ProjectionExec [output_rows = 4, elapsed_compute = 400]
202///      └── NetworkShuffleExec [bytes_transferred = 200, max_mem_used = 200]
203///
204/// The result will be:
205///
206/// MetricsWrapperExec (wrapped: AggregateExec) [output_rows = 1, output_rows = 3, elapsed_compute = 100, elapsed_compute = 300]
207///  └── MetricsWrapperExec (wrapped: ProjectionExec) [output_rows = 2, output_rows = 4, elapsed_compute = 200, elapsed_compute = 400]
208///      └── MetricsWrapperExec (wrapped: NetworkShuffleExec) [bytes_transferred = 100, bytes_transferred = 200, max_mem_used = 100, max_mem_used = 200]
209///
210/// Note: Metrics may be aggregated by name (ex. output_rows) automatically by various datafusion utils.
211pub fn stage_metrics_rewriter(
212    stage: &LocalStage,
213    metrics_collection: Arc<MetricsStore>,
214    format: DistributedMetricsFormat,
215) -> Result<Arc<dyn ExecutionPlan>> {
216    // Phase 1 — accumulate per-task metrics into a map keyed by node identity.
217    //
218    // For each task, the plan is traversed with `apply_with_dt_ctx`, which visits nodes in pre-order
219    // traversal, ignoring branches that do not belong to the recursed DistributedTaskContext
220    // (e.g., because of the presence of ChildrenIsolatorUnionExec).
221    //
222    // The raw allocation address of each `Arc<dyn ExecutionPlan>` as the node key.
223    // The planning plan is not modified between traversals, so these addresses are stable.
224    let mut node_metrics_map: HashMap<usize, MetricsSet> = HashMap::new();
225
226    for task_id in 0..stage.tasks {
227        let d_ctx = DistributedTaskContext {
228            task_index: task_id,
229            task_count: stage.tasks,
230        };
231        let task_key = TaskKey {
232            query_id: stage.query_id,
233            stage_id: stage.num,
234            task_number: task_id,
235        };
236        let Some(task_metrics) = metrics_collection.get(&task_key) else {
237            return internal_err!(
238                "not enough metrics provided to rewrite task: missing metrics for task {} in stage {}",
239                task_id,
240                stage.num
241            );
242        };
243
244        let mut per_task_counter = 0usize;
245        stage.plan.apply_with_dt_ctx(d_ctx, |node, _ctx| {
246            if per_task_counter >= task_metrics.pre_order_plan_metrics.len() {
247                return internal_err!(
248                    "not enough metrics provided to rewrite task: {} metrics provided",
249                    task_metrics.pre_order_plan_metrics.len()
250                );
251            }
252
253            let mut node_metrics = task_metrics.pre_order_plan_metrics[per_task_counter].clone();
254            let rewrite_ctx = format.to_rewrite_ctx(task_id as u64);
255            node_metrics = rewrite_ctx.maybe_rewrite_node_metics(node_metrics);
256
257            let id = Arc::as_ptr(node) as *const () as usize;
258            let entry = node_metrics_map.entry(id).or_default();
259            for metric in node_metrics.iter().map(Arc::clone) {
260                entry.push(metric);
261            }
262
263            per_task_counter += 1;
264            Ok(TreeNodeRecursion::Continue)
265        })?;
266    }
267
268    // Phase 2 — rewrite: wrap every node with its accumulated metrics.
269    // Nodes that were inactive for all tasks (never visited in phase 1) get empty metrics.
270    Arc::clone(&stage.plan)
271        .transform_down(|plan| {
272            let id = Arc::as_ptr(&plan) as *const () as usize;
273            let metrics = node_metrics_map.remove(&id).unwrap_or_default();
274            Ok(Transformed::new(
275                Arc::new(MetricsWrapperExec::new(plan.clone(), metrics)),
276                true,
277                match plan.is_network_boundary() {
278                    true => TreeNodeRecursion::Jump,
279                    false => TreeNodeRecursion::Continue,
280                },
281            ))
282        })
283        .map(|v| v.data)
284}
285
286#[cfg(test)]
287mod tests {
288    use crate::DistributedExt;
289    use crate::coordinator::MetricsStore;
290    use crate::metrics::DISTRIBUTED_DATAFUSION_TASK_ID_LABEL;
291    use crate::metrics::task_metrics_rewriter::MetricsWrapperExec;
292    use crate::metrics::task_metrics_rewriter::{
293        annotate_metrics_set_with_task_id, stage_metrics_rewriter,
294    };
295    use crate::metrics::{DistributedMetricsFormat, rewrite_distributed_plan_with_metrics};
296    use crate::stage::LocalStage;
297    use crate::test_utils::in_memory_channel_resolver::{
298        InMemoryChannelResolver, InMemoryWorkerResolver,
299    };
300    use crate::test_utils::metrics::make_test_metrics_set_from_seed;
301    use crate::test_utils::plans::count_plan_nodes_up_to_network_boundary;
302    use crate::test_utils::session_context::register_temp_parquet_table;
303    use crate::{DistributedExec, SessionStateBuilderExt, TaskKey, TaskMetrics};
304    use datafusion::arrow::array::{Int32Array, StringArray};
305    use datafusion::arrow::datatypes::{DataType, Field, Schema};
306    use datafusion::arrow::record_batch::RecordBatch;
307    use datafusion::execution::SessionStateBuilder;
308    use datafusion::physical_plan::empty::EmptyExec;
309    use datafusion::physical_plan::metrics::{Count, Label, Metric, MetricValue, MetricsSet};
310    use datafusion::physical_plan::{ExecutionPlan, collect};
311    use datafusion::prelude::SessionConfig;
312    use datafusion::prelude::SessionContext;
313    use itertools::Itertools;
314    use std::sync::Arc;
315    use test_case::test_case;
316    use uuid::Uuid;
317
318    async fn make_test_ctx() -> SessionContext {
319        make_test_ctx_inner(false).await
320    }
321
322    async fn make_test_distributed_ctx() -> SessionContext {
323        make_test_ctx_inner(true).await
324    }
325
326    /// Creates a non-distributed session context and registers two tables:
327    /// - table1 (id: int, name: string)
328    /// - table2 (id: int, name: string, phone: string, balance: float64)
329    async fn make_test_ctx_inner(distributed: bool) -> SessionContext {
330        let config = SessionConfig::new().with_target_partitions(4);
331        let mut builder = SessionStateBuilder::new()
332            .with_default_features()
333            .with_config(config);
334
335        if distributed {
336            builder = builder
337                .with_distributed_worker_resolver(InMemoryWorkerResolver::new(10))
338                .with_distributed_channel_resolver(InMemoryChannelResolver::default())
339                .with_distributed_metrics_collection(true)
340                .unwrap()
341                .with_distributed_planner()
342                .with_distributed_desired_task_count_handler(2)
343        }
344
345        let state = builder.build();
346        let ctx = SessionContext::from(state);
347
348        // Create test data for table1
349        let schema1 = Arc::new(Schema::new(vec![
350            Field::new("id", DataType::Int32, false),
351            Field::new("name", DataType::Utf8, false),
352        ]));
353
354        let batches1 = vec![
355            RecordBatch::try_new(
356                schema1.clone(),
357                vec![
358                    Arc::new(Int32Array::from(vec![1, 2, 3])),
359                    Arc::new(StringArray::from(vec!["a", "b", "c"])),
360                ],
361            )
362            .unwrap(),
363        ];
364
365        // Create test data for table2 with extended schema
366        let schema2 = Arc::new(Schema::new(vec![
367            Field::new("id", DataType::Int32, false),
368            Field::new("name", DataType::Utf8, false),
369            Field::new("phone", DataType::Utf8, false),
370            Field::new("balance", DataType::Float64, false),
371        ]));
372
373        let batches2 = vec![
374            RecordBatch::try_new(
375                schema2.clone(),
376                vec![
377                    Arc::new(Int32Array::from(vec![1, 2, 3])),
378                    Arc::new(StringArray::from(vec![
379                        "customer1",
380                        "customer2",
381                        "customer3",
382                    ])),
383                    Arc::new(StringArray::from(vec![
384                        "13-123-4567",
385                        "31-456-7890",
386                        "23-789-0123",
387                    ])),
388                    Arc::new(datafusion::arrow::array::Float64Array::from(vec![
389                        100.5, 250.0, 50.25,
390                    ])),
391                ],
392            )
393            .unwrap(),
394        ];
395
396        // Register the test data as parquet tables
397        let _ = register_temp_parquet_table("table1", schema1, batches1, &ctx)
398            .await
399            .unwrap();
400
401        let _ = register_temp_parquet_table("table2", schema2, batches2, &ctx)
402            .await
403            .unwrap();
404
405        ctx
406    }
407
408    fn make_test_stage(plan: Arc<dyn ExecutionPlan>) -> LocalStage {
409        LocalStage {
410            query_id: Uuid::new_v4(),
411            num: 2,
412            plan,
413            tasks: 4,
414            metrics_set: Default::default(),
415        }
416    }
417
418    fn collect_metrics_from_plan(plan: &Arc<dyn ExecutionPlan>, metrics: &mut Vec<MetricsSet>) {
419        metrics.extend(plan.metrics());
420        for child in plan.children() {
421            collect_metrics_from_plan(child, metrics);
422        }
423    }
424
425    fn metrics_set_eq(a: &MetricsSet, b: &MetricsSet) -> bool {
426        println!("a: {a:?}");
427        println!("b: {b:?}");
428        a.iter().count() == b.iter().count()
429            && a.iter().zip(b.iter()).all(|(a, b)| {
430                a.value() == b.value() && a.partition() == b.partition() && a.labels() == b.labels()
431            })
432    }
433
434    /// Asserts that we successfully re-write the metrics of a plan generated from the provided SQL query.
435    /// Also asserts that the order which metrics are collected from a plan matches the order which
436    /// they are re-written (ie. ensures we don't assign metrics to the wrong nodes)
437    ///
438    /// Only tests single node plans since the [TaskMetricsRewriter] stops on [NetworkBoundary].
439    async fn run_stage_metrics_rewriter_test(sql: &str, format: DistributedMetricsFormat) {
440        // Generate the plan
441        let ctx = make_test_ctx().await;
442        let plan = ctx
443            .sql(sql)
444            .await
445            .unwrap()
446            .create_physical_plan()
447            .await
448            .unwrap();
449
450        let stage = make_test_stage(plan.clone());
451
452        let num_metrics_per_task_per_node = 4;
453
454        // Generate metrics for each task and store them in the map.
455        let metrics_collection = MetricsStore::from_entries((0..stage.tasks).map(|task_id| {
456            let task_key = TaskKey {
457                query_id: stage.query_id,
458                stage_id: stage.num,
459                task_number: task_id,
460            };
461            let metrics = (0..count_plan_nodes_up_to_network_boundary(&plan))
462                .map(|node_id| {
463                    make_test_metrics_set_from_seed(
464                        (node_id * task_id) as u64,
465                        num_metrics_per_task_per_node,
466                    )
467                })
468                .collect::<Vec<MetricsSet>>();
469            let task_metrics = TaskMetrics {
470                task_metrics: MetricsSet::new(),
471                pre_order_plan_metrics: metrics,
472            };
473            (task_key, task_metrics)
474        }));
475        let metrics_collection = Arc::new(metrics_collection);
476
477        // Rewrite the plan.
478        let rewritten_plan =
479            stage_metrics_rewriter(&stage, metrics_collection.clone(), format).unwrap();
480
481        // Collect metrics from the plan.
482        let mut actual_metrics = vec![];
483        collect_metrics_from_plan(&rewritten_plan, &mut actual_metrics);
484        assert_eq!(
485            actual_metrics.len(),
486            count_plan_nodes_up_to_network_boundary(&plan)
487        );
488
489        // Assert that metrics from all tasks are present.
490        // actual_stage_node_metrics_set contains metrics for all task ex. [output_rows=1, elapsed_compute=1, output_rows=2, elapsed_compute=2...]
491        for (node_id, actual_stage_node_metrics_set) in actual_metrics.iter().enumerate() {
492            // actual_task_node_metrics_set contains metrics for one task ex. [output_rows=1, elapsed_compute=1]
493            for (task_id, actual_task_node_metrics_set) in actual_stage_node_metrics_set
494                .iter()
495                .chunks(num_metrics_per_task_per_node)
496                .into_iter()
497                .enumerate()
498            {
499                let expected_task_node_metrics = metrics_collection
500                    .get(&TaskKey {
501                        query_id: stage.query_id,
502                        stage_id: stage.num,
503                        task_number: task_id,
504                    })
505                    .unwrap()
506                    .pre_order_plan_metrics[node_id]
507                    .clone();
508
509                let mut actual_metrics_set = MetricsSet::new();
510                actual_task_node_metrics_set
511                    .for_each(|metric| actual_metrics_set.push(metric.clone()));
512
513                let mut expected_metrics_set = expected_task_node_metrics;
514
515                if format == DistributedMetricsFormat::PerTask {
516                    // Add task ids labels. We expect the actual metrics to be annotated by the
517                    // rewriter when using DistributedMetricsFormat::PerTask
518                    expected_metrics_set =
519                        annotate_metrics_set_with_task_id(expected_metrics_set, task_id as u64);
520                }
521                assert!(metrics_set_eq(&actual_metrics_set, &expected_metrics_set));
522            }
523        }
524    }
525
526    #[test_case(DistributedMetricsFormat::Aggregated ; "aggregated_metrics")]
527    #[test_case(DistributedMetricsFormat::PerTask ; "per_task_metrics")]
528    #[tokio::test]
529    async fn test_stage_metrics_rewriter_1(format: DistributedMetricsFormat) {
530        run_stage_metrics_rewriter_test(
531            "SELECT sum(balance) / 7.0 as avg_yearly from table2 group by name",
532            format,
533        )
534        .await;
535    }
536
537    #[test_case(DistributedMetricsFormat::Aggregated ; "aggregated_metrics")]
538    #[test_case(DistributedMetricsFormat::PerTask ; "per_task_metrics")]
539    #[tokio::test]
540    async fn test_stage_metrics_rewriter_2(format: DistributedMetricsFormat) {
541        run_stage_metrics_rewriter_test("SELECT id, COUNT(*) as count FROM table1 WHERE id > 1 GROUP BY id ORDER BY id LIMIT 10", format).await;
542    }
543
544    #[test_case(DistributedMetricsFormat::Aggregated ; "aggregated_metrics")]
545    #[test_case(DistributedMetricsFormat::PerTask ; "per_task_metrics")]
546    #[tokio::test]
547    async fn test_stage_metrics_rewriter_3(format: DistributedMetricsFormat) {
548        run_stage_metrics_rewriter_test(
549            "SELECT sum(balance) / 7.0 as avg_yearly
550            FROM table2
551            WHERE name LIKE 'customer%'
552              AND balance < (
553                SELECT 0.2 * avg(balance)
554                FROM table2 t2_inner
555                WHERE t2_inner.id = table2.id
556              )",
557            format,
558        )
559        .await;
560    }
561
562    #[tokio::test]
563    async fn test_rewrite_unexecuted_distributed_plan_with_metrics_err() {
564        let ctx = make_test_distributed_ctx().await;
565        let plan = ctx
566            .sql("SELECT id, COUNT(*) as count FROM table1 WHERE id > 1 GROUP BY id ORDER BY id LIMIT 10")
567            .await
568            .unwrap()
569            .create_physical_plan()
570            .await
571            .unwrap();
572        assert!(plan.is::<DistributedExec>());
573        assert!(
574            rewrite_distributed_plan_with_metrics(plan, DistributedMetricsFormat::Aggregated)
575                .await
576                .is_err()
577        );
578    }
579
580    // Assert every plan node has at least one metric except partition isolators, network boundary nodes, and the root DistributedExec node.
581    fn assert_metrics_present_in_plan(plan: &Arc<dyn ExecutionPlan>) {
582        if let Some(metrics) = plan.metrics() {
583            assert!(metrics.iter().count() > 0);
584        } else {
585            assert!(plan.is::<DistributedExec>());
586        }
587        for child in plan.children() {
588            assert_metrics_present_in_plan(child);
589        }
590    }
591
592    #[tokio::test]
593    async fn test_executed_distributed_plan_has_metrics() {
594        let ctx = make_test_distributed_ctx().await;
595        let plan = ctx
596            .sql("SELECT id, COUNT(*) as count FROM table1 WHERE id > 1 GROUP BY id ORDER BY id LIMIT 10")
597            .await
598            .unwrap()
599            .create_physical_plan()
600            .await
601            .unwrap();
602        collect(plan.clone(), ctx.task_ctx()).await.unwrap();
603        assert!(plan.is::<DistributedExec>());
604        let rewritten_plan =
605            rewrite_distributed_plan_with_metrics(plan, DistributedMetricsFormat::Aggregated)
606                .await
607                .unwrap();
608        assert_metrics_present_in_plan(&rewritten_plan);
609    }
610
611    #[test]
612    // An important feature of DF execution plans which we want to preserve is the ability
613    // to traverse a plan and collect metrics from specific nodes. To do this, the wrapper must
614    // allow access to the inner node. This test asserts that we support this.
615    fn test_wrapped_node_is_accessible() {
616        let example_node = Arc::new(EmptyExec::new(Arc::new(Schema::new(vec![Field::new(
617            "id",
618            DataType::Int32,
619            false,
620        )]))));
621
622        let wrapped = MetricsWrapperExec::new(example_node, MetricsSet::new());
623        assert_eq!(wrapped.name(), "EmptyExec");
624        assert!(wrapped.inner().is::<EmptyExec>());
625    }
626
627    #[test]
628    fn test_annotate_metrics_set_with_task_id_output_rows() {
629        // Create a MetricsSet with an OutputRows metric
630        let mut metrics_set = MetricsSet::new();
631        let count = Count::new();
632        count.add(1234);
633        let labels = vec![Label::new("operator", "scan")];
634        metrics_set.push(Arc::new(Metric::new_with_labels(
635            MetricValue::OutputRows(count),
636            Some(0),
637            labels,
638        )));
639
640        let task_id = 42;
641        let annotated = annotate_metrics_set_with_task_id(metrics_set, task_id);
642
643        // Verify we have one metric
644        assert_eq!(annotated.iter().count(), 1);
645
646        let metric = annotated.iter().next().unwrap();
647
648        // Verify metric type is preserved (OutputRows)
649        match metric.value() {
650            MetricValue::OutputRows(count) => {
651                assert_eq!(count.value(), 1234);
652            }
653            other => panic!("Expected OutputRows, got {:?}", other.name()),
654        }
655
656        // Verify partition is preserved
657        assert_eq!(metric.partition(), Some(0));
658
659        // Verify original labels are preserved and task_id label is added
660        let labels: Vec<_> = metric.labels().iter().collect();
661        assert_eq!(labels.len(), 2);
662        assert_eq!(labels[0].name(), "operator");
663        assert_eq!(labels[0].value(), "scan");
664        assert_eq!(labels[1].name(), DISTRIBUTED_DATAFUSION_TASK_ID_LABEL);
665        assert_eq!(labels[1].value(), "42");
666    }
667}
datafusion_distributed/metrics/task_metrics_rewriter.rs

datafusion_distributed/metrics/
task_metrics_rewriter.rs
