feat: Added metric

Author: Weijun-H

datafusion/physical-plan/src/repartition/mod.rs

@@ -57,7 +57,7 @@ use datafusion_physical_expr_common::sort_expr::LexOrdering;
 use futures::stream::Stream;
 use futures::{ready, FutureExt, StreamExt, TryStreamExt};
 use log::trace;
-use on_demand_repartition::OnDemandRepartitionExec;
+use on_demand_repartition::{OnDemandRepartitionExec, OnDemandRepartitionMetrics};
 use parking_lot::Mutex;
 
 mod distributor_channels;
@@ -220,8 +220,6 @@ impl RepartitionExecState {
                 })
                 .collect();
 
-            let r_metrics = RepartitionMetrics::new(i, num_output_partitions, &metrics);
-
             let input_task = if enable_pull_based {
                 let partition_rx = if preserve_order {
                     partition_receivers.clone().expect(
@@ -233,6 +231,9 @@ impl RepartitionExecState {
                         "partition_receivers must be provided when preserve_order is disabled",
                     )[0].clone()
                 };
+                let r_metrics =
+                    OnDemandRepartitionMetrics::new(i, num_output_partitions, &metrics);
+
                 SpawnedTask::spawn(OnDemandRepartitionExec::pull_from_input(
                     Arc::clone(&input),
                     i,
@@ -243,6 +244,9 @@ impl RepartitionExecState {
                     Arc::clone(&context),
                 ))
             } else {
+                let r_metrics =
+                    RepartitionMetrics::new(i, num_output_partitions, &metrics);
+
                 SpawnedTask::spawn(RepartitionExec::pull_from_input(
                     Arc::clone(&input),
                     i,
@@ -612,7 +616,7 @@ pub struct RepartitionExec {
 }
 
 #[derive(Debug, Clone)]
-pub(crate) struct RepartitionMetrics {
+struct RepartitionMetrics {
     /// Time in nanos to execute child operator and fetch batches
     fetch_time: metrics::Time,
     /// Repartitioning elapsed time in nanos

datafusion/physical-plan/src/repartition/on_demand_repartition.rs

@@ -30,11 +30,11 @@ use std::{any::Any, vec};
 use super::metrics::{ExecutionPlanMetricsSet, MetricsSet};
 use super::{
     DisplayAs, ExecutionPlanProperties, MaybeBatch, RecordBatchStream,
-    RepartitionExecBase, RepartitionMetrics, SendableRecordBatchStream,
+    RepartitionExecBase, SendableRecordBatchStream,
 };
 use crate::common::SharedMemoryReservation;
 use crate::execution_plan::CardinalityEffect;
-use crate::metrics::BaselineMetrics;
+use crate::metrics::{self, BaselineMetrics, MetricBuilder};
 use crate::repartition::distributor_channels::{
     DistributionReceiver, DistributionSender,
 };
@@ -398,22 +398,38 @@ impl OnDemandRepartitionExec {
         partition: usize,
         buffer_tx: tokio::sync::mpsc::Sender<RecordBatch>,
         context: Arc<TaskContext>,
+        fetch_time: metrics::Time,
+        send_buffer_time: metrics::Time,
     ) -> Result<()> {
+        let timer = fetch_time.timer();
         let mut stream = input.execute(partition, context).map_err(|e| {
             internal_datafusion_err!(
                 "Error executing input partition {} for on demand repartitioning: {}",
                 partition,
                 e
             )
         })?;
-        while let Some(batch) = stream.next().await {
-            buffer_tx.send(batch?).await.map_err(|e| {
-                internal_datafusion_err!(
-                    "Error sending batch to buffer channel for partition {}: {}",
-                    partition,
-                    e
-                )
-            })?;
+        timer.done();
+
+        loop {
+            let timer = fetch_time.timer();
+            let batch = stream.next().await;
+            timer.done();
+
+            // send the batch to the buffer channel
+            if let Some(batch) = batch {
+                let timer = send_buffer_time.timer();
+                buffer_tx.send(batch?).await.map_err(|e| {
+                    internal_datafusion_err!(
+                        "Error sending batch to buffer channel for partition {}: {}",
+                        partition,
+                        e
+                    )
+                })?;
+                timer.done();
+            } else {
+                break;
+            }
         }
 
         Ok(())
@@ -432,7 +448,7 @@ impl OnDemandRepartitionExec {
         >,
         partitioning: Partitioning,
         output_partition_rx: Receiver<usize>,
-        metrics: RepartitionMetrics,
+        metrics: OnDemandRepartitionMetrics,
         context: Arc<TaskContext>,
     ) -> Result<()> {
         // execute the child operator in a separate task
@@ -442,6 +458,8 @@ impl OnDemandRepartitionExec {
             partition,
             buffer_tx,
             Arc::clone(&context),
+            metrics.fetch_time,
+            metrics.send_buffer_time,
         ));
 
         // While there are still outputs to send to, keep pulling inputs
@@ -501,6 +519,50 @@ impl OnDemandRepartitionExec {
     }
 }
 
+#[derive(Debug, Clone)]
+pub(crate) struct OnDemandRepartitionMetrics {
+    /// Time in nanos to execute child operator and fetch batches
+    fetch_time: metrics::Time,
+    /// Time in nanos for sending resulting batches to buffer channels.
+    send_buffer_time: metrics::Time,
+    /// Time in nanos for sending resulting batches to channels.
+    ///
+    /// One metric per output partition.
+    send_time: Vec<metrics::Time>,
+}
+
+impl OnDemandRepartitionMetrics {
+    pub fn new(
+        input_partition: usize,
+        num_output_partitions: usize,
+        metrics: &ExecutionPlanMetricsSet,
+    ) -> Self {
+        // Time in nanos to execute child operator and fetch batches
+        let fetch_time =
+            MetricBuilder::new(metrics).subset_time("fetch_time", input_partition);
+
+        // Time in nanos for sending resulting batches to channels
+        let send_time = (0..num_output_partitions)
+            .map(|output_partition| {
+                let label =
+                    metrics::Label::new("outputPartition", output_partition.to_string());
+                MetricBuilder::new(metrics)
+                    .with_label(label)
+                    .subset_time("send_time", input_partition)
+            })
+            .collect();
+
+        // Time in nanos for sending resulting batches to buffer channels
+        let send_buffer_time =
+            MetricBuilder::new(metrics).subset_time("send_buffer_time", input_partition);
+        Self {
+            fetch_time,
+            send_time,
+            send_buffer_time,
+        }
+    }
+}
+
 /// This struct converts a receiver to a stream.
 /// Receiver receives data on an SPSC channel.
 struct OnDemandPerPartitionStream {