Add metrics to SchemaCastScanExec and VirtualExecutionPlan (#143)

Author: nuno-faria

datafusion-federation/src/schema_cast/mod.rs

@@ -3,6 +3,7 @@ use datafusion::arrow::datatypes::SchemaRef;
 use datafusion::common::Statistics;
 use datafusion::error::{DataFusionError, Result};
 use datafusion::execution::{SendableRecordBatchStream, TaskContext};
+use datafusion::physical_plan::metrics::{BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet};
 use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
 use datafusion::physical_plan::{
     DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, PlanProperties,
@@ -24,6 +25,7 @@ pub struct SchemaCastScanExec {
     input: Arc<dyn ExecutionPlan>,
     schema: SchemaRef,
     properties: PlanProperties,
+    metrics_set: ExecutionPlanMetricsSet,
 }
 
 impl SchemaCastScanExec {
@@ -41,6 +43,7 @@ impl SchemaCastScanExec {
             input,
             schema,
             properties,
+            metrics_set: ExecutionPlanMetricsSet::new(),
         }
     }
 }
@@ -101,14 +104,20 @@ impl ExecutionPlan for SchemaCastScanExec {
     ) -> Result<SendableRecordBatchStream> {
         let mut stream = self.input.execute(partition, context)?;
         let schema = Arc::clone(&self.schema);
+        let baseline_metrics = BaselineMetrics::new(&self.metrics_set, partition);
 
         Ok(Box::pin(RecordBatchStreamAdapter::new(
             Arc::clone(&schema),
             {
                 stream! {
                     while let Some(batch) = stream.next().await {
+                        let _timer = baseline_metrics.elapsed_compute().timer();
                         let batch = record_convert::try_cast_to(batch?, Arc::clone(&schema));
-                        yield batch.map_err(|e| { DataFusionError::External(Box::new(e)) });
+                        let batch = batch.map_err(|e| { DataFusionError::External(Box::new(e)) });
+                        if let Ok(ref b) = batch {
+                            baseline_metrics.output_rows().add(b.num_rows());
+                        }
+                        yield batch;
                     }
                 }
             },
@@ -118,4 +127,8 @@ impl ExecutionPlan for SchemaCastScanExec {
     fn partition_statistics(&self, partition: Option<usize>) -> Result<Statistics> {
         self.input.partition_statistics(partition)
     }
+
+    fn metrics(&self) -> Option<MetricsSet> {
+        Some(self.metrics_set.clone_inner())
+    }
 }

datafusion-federation/src/sql/executor.rs

@@ -5,7 +5,7 @@ use datafusion::{
     common::Statistics,
     error::Result,
     logical_expr::LogicalPlan,
-    physical_plan::SendableRecordBatchStream,
+    physical_plan::{metrics::MetricsSet, SendableRecordBatchStream},
     sql::{sqlparser::ast, unparser::dialect::Dialect},
 };
 use std::sync::Arc;
@@ -56,6 +56,11 @@ pub trait SQLExecutor: Sync + Send {
 
     /// Returns the schema of table_name within this [`SQLExecutor`]
     async fn get_table_schema(&self, table_name: &str) -> Result<SchemaRef>;
+
+    /// Returns the execution metrics, if available.
+    fn metrics(&self) -> Option<MetricsSet> {
+        None
+    }
 }
 
 impl fmt::Debug for dyn SQLExecutor {

datafusion-federation/src/sql/mod.rs

@@ -11,15 +11,18 @@ use analyzer::RewriteTableScanAnalyzer;
 use async_trait::async_trait;
 use datafusion::{
     arrow::datatypes::{Schema, SchemaRef},
-    common::tree_node::{Transformed, TreeNode},
-    common::Statistics,
+    common::{
+        tree_node::{Transformed, TreeNode},
+        Statistics,
+    },
     error::{DataFusionError, Result},
     execution::{context::SessionState, TaskContext},
     logical_expr::{Extension, LogicalPlan},
     optimizer::{optimizer::Optimizer, OptimizerConfig, OptimizerRule},
     physical_expr::EquivalenceProperties,
     physical_plan::{
         execution_plan::{Boundedness, EmissionType},
+        metrics::MetricsSet,
         DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, PlanProperties,
         SendableRecordBatchStream,
     },
@@ -365,6 +368,10 @@ impl ExecutionPlan for VirtualExecutionPlan {
     fn partition_statistics(&self, _partition: Option<usize>) -> Result<Statistics> {
         Ok(self.statistics.clone())
     }
+
+    fn metrics(&self) -> Option<MetricsSet> {
+        self.executor.metrics()
+    }
 }
 
 #[cfg(test)]