feat: Add additional metrics for shuffle write

common/src/main/scala/org/apache/comet/CometConf.scala

@@ -45,6 +45,9 @@ import org.apache.comet.shims.ShimCometConf
  */
 object CometConf extends ShimCometConf {
 
+  private val METRICS_GUIDE = "For more information, refer to the Comet Metrics " +
+    "Guide (https://datafusion.apache.org/comet/user-guide/metrics.html)"
+
   private val TUNING_GUIDE = "For more information, refer to the Comet Tuning " +
     "Guide (https://datafusion.apache.org/comet/user-guide/tuning.html)"
 
@@ -414,6 +417,12 @@ object CometConf extends ShimCometConf {
       .booleanConf
       .createWithDefault(false)
 
+  val COMET_ENABLE_DETAILED_METRICS: ConfigEntry[Boolean] =
+    conf("spark.comet.metrics.detailed")
+      .doc(s"Enable this option to see additional SQL metrics. $METRICS_GUIDE.")
+      .booleanConf
+      .createWithDefault(false)
+
   val COMET_EXPLAIN_FALLBACK_ENABLED: ConfigEntry[Boolean] =
     conf("spark.comet.explainFallback.enabled")
       .doc(

docs/source/index.rst

@@ -51,6 +51,7 @@ as a native runtime to achieve improvement in terms of query efficiency and quer
    Configuration Settings <user-guide/configs>
    Compatibility Guide <user-guide/compatibility>
    Tuning Guide <user-guide/tuning>
+   Metrics Guide <user-guide/metrics>
 
 .. _toc.contributor-guide-links:
 .. toctree::

docs/source/user-guide/configs.md

@@ -64,6 +64,7 @@ Comet provides the following configuration settings.
 | spark.comet.explainFallback.enabled | When this setting is enabled, Comet will provide logging explaining the reason(s) why a query stage cannot be executed natively. Set this to false to reduce the amount of logging. | false |
 | spark.comet.memory.overhead.factor | Fraction of executor memory to be allocated as additional non-heap memory per executor process for Comet. | 0.2 |
 | spark.comet.memory.overhead.min | Minimum amount of additional memory to be allocated per executor process for Comet, in MiB. | 402653184b |
+| spark.comet.metrics.detailed | Enable this option to see additional SQL metrics. For more information, refer to the Comet Metrics Guide (https://datafusion.apache.org/comet/user-guide/metrics.html). | false |
 | spark.comet.nativeLoadRequired | Whether to require Comet native library to load successfully when Comet is enabled. If not, Comet will silently fallback to Spark when it fails to load the native lib. Otherwise, an error will be thrown and the Spark job will be aborted. | false |
 | spark.comet.parquet.enable.directBuffer | Whether to use Java direct byte buffer when reading Parquet. | false |
 | spark.comet.parquet.read.io.adjust.readRange.skew | In the parallel reader, if the read ranges submitted are skewed in sizes, this option will cause the reader to break up larger read ranges into smaller ranges to reduce the skew. This will result in a slightly larger number of connections opened to the file system but may give improved performance. | false |

docs/source/user-guide/metrics.md

@@ -0,0 +1,63 @@
+<!---
+Licensed to the Apache Software Foundation (ASF) under one
+or more contributor license agreements.  See the NOTICE file
+distributed with this work for additional information
+regarding copyright ownership.  The ASF licenses this file
+to you under the Apache License, Version 2.0 (the
+"License"); you may not use this file except in compliance
+with the License.  You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing,
+software distributed under the License is distributed on an
+"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+KIND, either express or implied.  See the License for the
+specific language governing permissions and limitations
+under the License.
+-->
+
+# Comet Metrics
+
+## Spark SQL Metrics
+
+Set `spark.comet.metrics.detailed=true` to see all available Comet metrics.
+
+### CometScanExec
+
+`CometScanExec` uses nanoseconds for total scan time. Spark also measures scan time in nanoseconds but converts to
+milliseconds _per batch_ which can result in a large loss of precision, making it difficult to compare scan times
+between Spark and Comet.
+
+### Exchange
+
+Comet adds some additional metrics:
+
+| Metric                          | Description                                                                               |
+| ------------------------------- | ----------------------------------------------------------------------------------------- |
+| `native shuffle time`           | Total time spent in native shuffle writer, excluding the execution time of the input plan |
+| `native shuffle input time`     | Time spend executing the shuffle input plan and fetching batches.                         |
+| `shuffle wall time (inclusive)` | Total time executing the shuffle write, inclusive of executing the input plan.            |
+
+## Native Metrics
+
+Setting `spark.comet.explain.native.enabled=true` will cause native plans to be logged in each executor. Metrics are
+logged for each native plan (and there is one plan per task, so this is very verbose).
+
+Here is a guide to some of the native metrics.
+
+### ScanExec
+
+| Metric            | Description                                                                                         |
+| ----------------- | --------------------------------------------------------------------------------------------------- |
+| `elapsed_compute` | Total time spent in this operator, fetching batches from a JVM iterator.                            |
+| `jvm_fetch_time`  | Time spent in the JVM fetching input batches to be read by this `ScanExec` instance.                |
+| `arrow_ffi_time`  | Time spent using Arrow FFI to create Arrow batches from the memory addresses returned from the JVM. |
+
+### ShuffleWriterExec
+
+| Metric            | Description                                           |
+| ----------------- | ----------------------------------------------------- |
+| `elapsed_compute` | Total time excluding any child operators.             |
+| `input_time`      | Time spent executing input plan and fetching batches. |
+| `write_time`      | Time spent writing bytes to disk.                     |

docs/source/user-guide/tuning.md

@@ -102,28 +102,3 @@ native shuffle currently only supports `HashPartitioning` and `SinglePartitionin
 
 To enable native shuffle, set `spark.comet.exec.shuffle.mode` to `native`. If this mode is explicitly set,
 then any shuffle operations that cannot be supported in this mode will fall back to Spark.
-
-##  Metrics
-
-### Spark SQL Metrics
-
-Some Comet metrics are not directly comparable to Spark metrics in some cases:
-
-- `CometScanExec` uses nanoseconds for total scan time. Spark also measures scan time in nanoseconds but converts to
-  milliseconds _per batch_ which can result in a large loss of precision, making it difficult to compare scan times
-  between Spark and Comet.
-
-### Native Metrics
-
-Setting `spark.comet.explain.native.enabled=true` will cause native plans to be logged in each executor. Metrics are
-logged for each native plan (and there is one plan per task, so this is very verbose).
-
-Here is a guide to some of the native metrics.
-
-### ScanExec
-
-| Metric            | Description                                                                                         |
-| ----------------- | --------------------------------------------------------------------------------------------------- |
-| `elapsed_compute` | Total time spent in this operator, fetching batches from a JVM iterator.                            |
-| `jvm_fetch_time`  | Time spent in the JVM fetching input batches to be read by this `ScanExec` instance.                |
-| `arrow_ffi_time`  | Time spent using Arrow FFI to create Arrow batches from the memory addresses returned from the JVM. |

native/core/src/execution/operators/scan.rs

@@ -75,7 +75,7 @@ pub struct ScanExec {
     /// Metrics collector
     metrics: ExecutionPlanMetricsSet,
     /// Baseline metrics
-    baseline_metrics: BaselineMetrics,
+    pub(crate) baseline_metrics: BaselineMetrics,
     /// Time waiting for JVM input plan to execute and return batches
     jvm_fetch_time: Time,
     /// Time spent in FFI

native/core/src/execution/shuffle/shuffle_writer.rs

@@ -17,17 +17,6 @@
 
 //! Defines the External shuffle repartition plan.
 
-use std::{
-    any::Any,
-    fmt,
-    fmt::{Debug, Formatter},
-    fs::{File, OpenOptions},
-    io::{BufReader, BufWriter, Cursor, Read, Seek, SeekFrom, Write},
-    path::Path,
-    sync::Arc,
-    task::{Context, Poll},
-};
-
 use arrow::{datatypes::*, ipc::writer::StreamWriter};
 use async_trait::async_trait;
 use bytes::Buf;
@@ -59,7 +48,19 @@ use futures::executor::block_on;
 use futures::{lock::Mutex, Stream, StreamExt, TryFutureExt, TryStreamExt};
 use itertools::Itertools;
 use simd_adler32::Adler32;
+use std::time::Instant;
+use std::{
+    any::Any,
+    fmt,
+    fmt::{Debug, Formatter},
+    fs::{File, OpenOptions},
+    io::{BufReader, BufWriter, Cursor, Read, Seek, SeekFrom, Write},
+    path::Path,
+    sync::Arc,
+    task::{Context, Poll},
+};
 
+use crate::execution::operators::ScanExec;
 use crate::{
     common::bit::ceil,
     errors::{CometError, CometResult},
@@ -137,9 +138,21 @@ impl ExecutionPlan for ShuffleWriterExec {
         partition: usize,
         context: Arc<TaskContext>,
     ) -> Result<SendableRecordBatchStream> {
-        let input = self.input.execute(partition, Arc::clone(&context))?;
         let metrics = ShuffleRepartitionerMetrics::new(&self.metrics, 0);
 
+        if let Some(scan) = self.input.as_any().downcast_ref::<ScanExec>() {
+            // ScanExec starts executing and fetching data during query planning time so we
+            // need to capture that time here to ensure that we have accurate metrics
+            metrics
+                .input_time
+                .add(scan.baseline_metrics.elapsed_compute());
+        }
+
+        // execute the child plan
+        let start_time = Instant::now();
+        let input = self.input.execute(partition, Arc::clone(&context))?;
+        metrics.input_time.add_duration(start_time.elapsed());
+
         Ok(Box::pin(RecordBatchStreamAdapter::new(
             self.schema(),
             futures::stream::once(
@@ -271,7 +284,7 @@ impl PartitionBuffer {
         columns: &[ArrayRef],
         indices: &[usize],
         start_index: usize,
-        time_metric: &Time,
+        write_time: &Time,
     ) -> AppendRowStatus {
         let mut mem_diff = 0;
         let mut start = start_index;
@@ -293,7 +306,7 @@ impl PartitionBuffer {
                 });
             self.num_active_rows += end - start;
             if self.num_active_rows >= self.batch_size {
-                let mut timer = time_metric.timer();
+                let mut timer = write_time.timer();
                 let flush = self.flush();
                 if let Err(e) = flush {
                     return AppendRowStatus::MemDiff(Err(e));
@@ -628,6 +641,13 @@ struct ShuffleRepartitionerMetrics {
     /// metrics
     baseline: BaselineMetrics,
 
+    /// Time executing child plan and fetching batches
+    input_time: Time,
+
+    /// Time spent writing to disk. Maps to "shuffleWriteTime" in Spark SQL Metrics.
+    write_time: Time,
+
+    //other_time: Time,
     /// count of spills during the execution of the operator
     spill_count: Count,
 
@@ -640,8 +660,12 @@ struct ShuffleRepartitionerMetrics {
 
 impl ShuffleRepartitionerMetrics {
     fn new(metrics: &ExecutionPlanMetricsSet, partition: usize) -> Self {
+        let input_time = MetricBuilder::new(metrics).subset_time("input_time", partition);
+        let write_time = MetricBuilder::new(metrics).subset_time("write_time", partition);
         Self {
             baseline: BaselineMetrics::new(metrics, partition),
+            input_time,
+            write_time,
             spill_count: MetricBuilder::new(metrics).spill_count(partition),
             spilled_bytes: MetricBuilder::new(metrics).spilled_bytes(partition),
             data_size: MetricBuilder::new(metrics).counter("data_size", partition),
@@ -701,13 +725,18 @@ impl ShuffleRepartitioner {
     /// This function will slice input batch according to configured batch size and then
     /// shuffle rows into corresponding partition buffer.
     async fn insert_batch(&mut self, batch: RecordBatch) -> Result<()> {
+        let start_time = Instant::now();
         let mut start = 0;
         while start < batch.num_rows() {
             let end = (start + self.batch_size).min(batch.num_rows());
             let batch = batch.slice(start, end - start);
             self.partitioning_batch(batch).await?;
             start = end;
         }
+        self.metrics
+            .baseline
+            .elapsed_compute()
+            .add_duration(start_time.elapsed());
         Ok(())
     }
 
@@ -848,6 +877,7 @@ impl ShuffleRepartitioner {
 
     /// Writes buffered shuffled record batches into Arrow IPC bytes.
     async fn shuffle_write(&mut self) -> Result<SendableRecordBatchStream> {
+        let mut elapsed_compute = self.metrics.baseline.elapsed_compute().timer();
         let num_output_partitions = self.num_output_partitions;
         let buffered_partitions = &mut self.buffered_partitions;
         let mut output_batches: Vec<Vec<u8>> = vec![vec![]; num_output_partitions];
@@ -872,7 +902,7 @@ impl ShuffleRepartitioner {
             .map_err(|e| DataFusionError::Execution(format!("shuffle write error: {:?}", e)))?;
 
         for i in 0..num_output_partitions {
-            let mut timer = self.metrics.baseline.elapsed_compute().timer();
+            let mut timer = self.metrics.write_time.timer();
 
             offsets[i] = output_data.stream_position()?;
             output_data.write_all(&output_batches[i])?;
@@ -885,7 +915,7 @@ impl ShuffleRepartitioner {
             for spill in &output_spills {
                 let length = spill.offsets[i + 1] - spill.offsets[i];
                 if length > 0 {
-                    let mut timer = self.metrics.baseline.elapsed_compute().timer();
+                    let mut timer = self.metrics.write_time.timer();
 
                     let mut spill_file =
                         BufReader::new(File::open(spill.file.path()).map_err(|e| {
@@ -900,7 +930,7 @@ impl ShuffleRepartitioner {
                 }
             }
         }
-        let mut timer = self.metrics.baseline.elapsed_compute().timer();
+        let mut timer = self.metrics.write_time.timer();
         output_data.flush()?;
         timer.stop();
 
@@ -909,7 +939,7 @@ impl ShuffleRepartitioner {
             .stream_position()
             .map_err(|e| DataFusionError::Execution(format!("shuffle write error: {:?}", e)))?;
 
-        let mut timer = self.metrics.baseline.elapsed_compute().timer();
+        let mut timer = self.metrics.write_time.timer();
 
         let mut output_index =
             BufWriter::new(File::create(index_file).map_err(|e| {
@@ -927,6 +957,8 @@ impl ShuffleRepartitioner {
         let used = self.reservation.size();
         self.reservation.shrink(used);
 
+        elapsed_compute.stop();
+
         // shuffle writer always has empty output
         Ok(Box::pin(EmptyStream::try_new(Arc::clone(&self.schema))?))
     }
@@ -959,7 +991,7 @@ impl ShuffleRepartitioner {
             return Ok(0);
         }
 
-        let mut timer = self.metrics.baseline.elapsed_compute().timer();
+        let mut timer = self.metrics.write_time.timer();
 
         let spillfile = self
             .runtime
@@ -995,12 +1027,11 @@ impl ShuffleRepartitioner {
 
         let output = &mut self.buffered_partitions[partition_id];
 
-        let time_metric = self.metrics.baseline.elapsed_compute();
-
         // If the range of indices is not big enough, just appending the rows into
         // active array builders instead of directly adding them as a record batch.
         let mut start_index: usize = 0;
-        let mut output_ret = output.append_rows(columns, indices, start_index, time_metric);
+        let mut output_ret =
+            output.append_rows(columns, indices, start_index, &self.metrics.write_time);
 
         loop {
             match output_ret {
@@ -1017,10 +1048,9 @@ impl ShuffleRepartitioner {
                     let output = &mut self.buffered_partitions[partition_id];
                     output.reservation.free();
 
-                    let time_metric = self.metrics.baseline.elapsed_compute();
-
                     start_index = new_start;
-                    output_ret = output.append_rows(columns, indices, start_index, time_metric);
+                    output_ret =
+                        output.append_rows(columns, indices, start_index, &self.metrics.write_time);
 
                     if let AppendRowStatus::StartIndex(new_start) = output_ret {
                         if new_start == start_index {
@@ -1104,7 +1134,7 @@ async fn external_shuffle(
         context.session_config().batch_size(),
     );
 
-    while let Some(batch) = input.next().await {
+    while let Some(batch) = fetch_next_batch(&mut input, &repartitioner.metrics.input_time).await {
         // Block on the repartitioner to insert the batch and shuffle the rows
         // into the corresponding partition buffer.
         // Otherwise, pull the next batch from the input stream might overwrite the
@@ -1114,6 +1144,16 @@ async fn external_shuffle(
     repartitioner.shuffle_write().await
 }
 
+async fn fetch_next_batch(
+    input: &mut SendableRecordBatchStream,
+    input_time: &Time,
+) -> Option<Result<RecordBatch>> {
+    let mut input_time = input_time.timer();
+    let next_batch = input.next().await;
+    input_time.stop();
+    next_batch
+}
+
 fn new_array_builders(schema: &SchemaRef, batch_size: usize) -> Vec<Box<dyn ArrayBuilder>> {
     schema
         .fields()

spark/src/main/scala/org/apache/spark/sql/comet/CometCollectLimitExec.scala

@@ -57,7 +57,8 @@ case class CometCollectLimitExec(
     "dataSize" -> SQLMetrics.createSizeMetric(sparkContext, "data size"),
     "numPartitions" -> SQLMetrics.createMetric(
       sparkContext,
-      "number of partitions")) ++ readMetrics ++ writeMetrics
+      "number of partitions")) ++ readMetrics ++ writeMetrics ++ CometMetricNode.shuffleMetrics(
+    sparkContext)
 
   private lazy val serializer: Serializer =
     new UnsafeRowSerializer(child.output.size, longMetric("dataSize"))