Fix metrics display on leaf nodes

src/execution_plans/metrics.rs

@@ -15,18 +15,11 @@ pub(crate) struct MetricsWrapperExec {
     inner: Arc<dyn ExecutionPlan>,
     /// metrics for this plan node.
     metrics: MetricsSet,
-    /// children is initially None. When used by the [TaskMetricsRewriter], the children will be updated
-    /// to point at other wrapped nodes.
-    children: Option<Vec<Arc<dyn ExecutionPlan>>>,
 }
 
 impl MetricsWrapperExec {
     pub(crate) fn new(inner: Arc<dyn ExecutionPlan>, metrics: MetricsSet) -> Self {
-        Self {
-            inner,
-            metrics,
-            children: None,
-        }
+        Self { inner, metrics }
     }
 }
 
@@ -53,22 +46,17 @@ impl ExecutionPlan for MetricsWrapperExec {
         }
     }
 
-    /// Retrusn
     fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
-        match &self.children {
-            Some(children) => children.iter().collect(),
-            None => self.inner.children(),
-        }
+        self.inner.children()
     }
 
     fn with_new_children(
         self: Arc<Self>,
         children: Vec<Arc<dyn ExecutionPlan>>,
     ) -> Result<Arc<dyn ExecutionPlan>> {
         Ok(Arc::new(MetricsWrapperExec {
-            inner: self.inner.clone(),
+            inner: Arc::clone(&self.inner).with_new_children(children.clone())?,
             metrics: self.metrics.clone(),
-            children: Some(children),
         }))
     }
 

src/execution_plans/network_broadcast.rs

@@ -9,6 +9,7 @@ use dashmap::DashMap;
 use datafusion::common::internal_datafusion_err;
 use datafusion::error::DataFusionError;
 use datafusion::execution::{SendableRecordBatchStream, TaskContext};
+use datafusion::physical_expr_common::metrics::MetricsSet;
 use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
 use datafusion::physical_plan::{
     DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, PlanProperties,
@@ -263,4 +264,8 @@ impl ExecutionPlan for NetworkBroadcastExec {
             futures::stream::select_all(streams),
         )))
     }
+
+    fn metrics(&self) -> Option<MetricsSet> {
+        Some(self.worker_connections.metrics.clone_inner())
+    }
 }

src/execution_plans/partition_isolator.rs

@@ -1,7 +1,10 @@
 use crate::DistributedTaskContext;
 use crate::common::require_one_child;
 use datafusion::execution::TaskContext;
+use datafusion::physical_expr_common::metrics::MetricsSet;
 use datafusion::physical_plan::ExecutionPlanProperties;
+use datafusion::physical_plan::metrics::{ExecutionPlanMetricsSet, MetricBuilder};
+use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
 use datafusion::{
     error::Result,
     execution::SendableRecordBatchStream,
@@ -10,6 +13,7 @@ use datafusion::{
         PlanProperties,
     },
 };
+use futures::TryStreamExt;
 use std::{fmt::Formatter, sync::Arc};
 
 /// This is a simple [ExecutionPlan] that isolates a set of N partitions from an input
@@ -52,6 +56,7 @@ pub struct PartitionIsolatorExec {
     pub(crate) input: Arc<dyn ExecutionPlan>,
     pub(crate) properties: PlanProperties,
     pub(crate) n_tasks: usize,
+    pub(crate) metrics: ExecutionPlanMetricsSet,
 }
 
 impl PartitionIsolatorExec {
@@ -69,6 +74,7 @@ impl PartitionIsolatorExec {
             input: input.clone(),
             properties,
             n_tasks,
+            metrics: ExecutionPlanMetricsSet::new(),
         }
     }
 
@@ -150,6 +156,7 @@ impl ExecutionPlan for PartitionIsolatorExec {
     ) -> Result<SendableRecordBatchStream> {
         let task_context = DistributedTaskContext::from_ctx(&context);
 
+        let metric = MetricBuilder::new(&self.metrics).output_rows(partition);
         let input_partitions = self.input.output_partitioning().partition_count();
 
         let partition_group = Self::partition_group(
@@ -169,13 +176,22 @@ impl ExecutionPlan for PartitionIsolatorExec {
                     Ok(Box::pin(EmptyRecordBatchStream::new(self.input.schema()))
                         as SendableRecordBatchStream)
                 } else {
-                    self.input.execute(*actual_partition_number, context)
+                    Ok(Box::pin(RecordBatchStreamAdapter::new(
+                        self.schema(),
+                        self.input
+                            .execute(*actual_partition_number, context)?
+                            .inspect_ok(move |v| metric.add(v.num_rows())),
+                    )))
                 }
             }
             None => Ok(Box::pin(EmptyRecordBatchStream::new(self.input.schema()))
                 as SendableRecordBatchStream),
         }
     }
+
+    fn metrics(&self) -> Option<MetricsSet> {
+        Some(self.metrics.clone_inner())
+    }
 }
 
 #[cfg(test)]

tests/metrics_collection.rs

@@ -1,7 +1,9 @@
 #[cfg(all(feature = "integration", test))]
 mod tests {
     use datafusion::catalog::memory::DataSourceExec;
+    use datafusion::common::assert_not_contains;
     use datafusion::common::tree_node::{Transformed, TreeNode, TreeNodeRecursion};
+    use datafusion::physical_plan::display::DisplayableExecutionPlan;
     use datafusion::physical_plan::{ExecutionPlan, execute_stream};
     use datafusion::prelude::SessionContext;
     use datafusion_distributed::test_utils::localhost::start_localhost_context;
@@ -163,6 +165,31 @@ mod tests {
         Ok(())
     }
 
+    #[tokio::test]
+    async fn test_metric_collection_display_all_have_metrics()
+    -> Result<(), Box<dyn std::error::Error>> {
+        let format = DistributedMetricsFormat::PerTask;
+        let (d_ctx, _guard, _) = start_localhost_context(3, DefaultSessionBuilder).await;
+
+        let query =
+            r#"SELECT count(*), "RainToday" FROM weather GROUP BY "RainToday" ORDER BY count(*)"#;
+
+        let s_ctx = SessionContext::default();
+        let (_, mut d_physical) = execute(&s_ctx, &d_ctx, query).await?;
+        d_physical = rewrite_distributed_plan_with_metrics(d_physical.clone(), format)?;
+
+        let display =
+            DisplayableExecutionPlan::with_metrics(d_physical.children().swap_remove(0).as_ref())
+                .indent(true)
+                .to_string();
+        assert_not_contains!(display, "metrics=[]");
+
+        let display = display_plan_ascii(d_physical.as_ref(), true);
+        assert_not_contains!(display, "metrics=[]");
+
+        Ok(())
+    }
+
     /// Looks for an [ExecutionPlan] that matches the provided type parameter `T` in
     /// both root nodes and compares its metrics.
     /// There might be more than one, so `index` determines which one is compared.