From 20aa89cad836450878d11aaeee37fbe248466548 Mon Sep 17 00:00:00 2001
From: Andy Grove <agrove@apache.org>
Date: Wed, 20 Nov 2024 17:57:34 -0700
Subject: [PATCH 01/10] Add custom metric for native shuffle read time

---
 .../src/execution/datafusion/shuffle_writer.rs | 18 ++++++++++++++++++
 native/core/src/execution/operators/scan.rs    |  5 +++--
 .../shuffle/CometShuffleExchangeExec.scala     |  8 +++++++-
 3 files changed, 28 insertions(+), 3 deletions(-)

diff --git a/native/core/src/execution/datafusion/shuffle_writer.rs b/native/core/src/execution/datafusion/shuffle_writer.rs
index 7587ff06dc..a080bf22b8 100644
--- a/native/core/src/execution/datafusion/shuffle_writer.rs
+++ b/native/core/src/execution/datafusion/shuffle_writer.rs
@@ -65,6 +65,7 @@ use crate::{
     errors::{CometError, CometResult},
 };
 use datafusion_comet_spark_expr::spark_hash::create_murmur3_hashes;
+use crate::execution::operators::ScanExec;
 
 /// The status of appending rows to a partition buffer.
 enum AppendRowStatus {
@@ -139,6 +140,13 @@ impl ExecutionPlan for ShuffleWriterExec {
     ) -> Result<SendableRecordBatchStream> {
         let input = self.input.execute(partition, Arc::clone(&context))?;
         let metrics = ShuffleRepartitionerMetrics::new(&self.metrics, 0);
+        let read_time = MetricBuilder::new(&self.metrics).subset_time("read_time", 0);
+
+        let scan_time = if let Some(scan) = self.input.as_any().downcast_ref::<ScanExec>() {
+            Some(scan.baseline_metrics.elapsed_compute().clone())
+        } else {
+            None
+        };
 
         Ok(Box::pin(RecordBatchStreamAdapter::new(
             self.schema(),
@@ -151,6 +159,8 @@ impl ExecutionPlan for ShuffleWriterExec {
                     self.partitioning.clone(),
                     metrics,
                     context,
+                    read_time,
+                    scan_time
                 )
                 .map_err(|e| ArrowError::ExternalError(Box::new(e))),
             )
@@ -1091,6 +1101,8 @@ async fn external_shuffle(
     partitioning: Partitioning,
     metrics: ShuffleRepartitionerMetrics,
     context: Arc<TaskContext>,
+    read_time: Time,
+    scan_time: Option<Time>
 ) -> Result<SendableRecordBatchStream> {
     let schema = input.schema();
     let mut repartitioner = ShuffleRepartitioner::new(
@@ -1111,6 +1123,12 @@ async fn external_shuffle(
         // current batch in the repartitioner.
         block_on(repartitioner.insert_batch(batch?))?;
     }
+
+    // copy ScanExec elapsed_compute into ShuffleWriterExec read_time
+    if let Some(t) = &scan_time {
+        read_time.add(t);
+    }
+
     repartitioner.shuffle_write().await
 }
 
diff --git a/native/core/src/execution/operators/scan.rs b/native/core/src/execution/operators/scan.rs
index a9bd954e99..ca1a0fb9f6 100644
--- a/native/core/src/execution/operators/scan.rs
+++ b/native/core/src/execution/operators/scan.rs
@@ -72,8 +72,9 @@ pub struct ScanExec {
     /// Cache of expensive-to-compute plan properties
     cache: PlanProperties,
     /// Metrics collector
-    metrics: ExecutionPlanMetricsSet,
-    baseline_metrics: BaselineMetrics,
+    pub(crate) metrics: ExecutionPlanMetricsSet,
+    /// Baseline metrics
+    pub(crate) baseline_metrics: BaselineMetrics,
 }
 
 impl ScanExec {
diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala
index 6430a7899f..a8ccc3b6e0 100644
--- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala
+++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala
@@ -77,6 +77,7 @@ case class CometShuffleExchangeExec(
     SQLShuffleReadMetricsReporter.createShuffleReadMetrics(sparkContext)
   override lazy val metrics: Map[String, SQLMetric] = Map(
     "dataSize" -> SQLMetrics.createSizeMetric(sparkContext, "data size"),
+    "read_time" -> SQLMetrics.createNanoTimingMetric(sparkContext, "read time"),
     "numPartitions" -> SQLMetrics.createMetric(
       sparkContext,
       "number of partitions")) ++ readMetrics ++ writeMetrics
@@ -480,7 +481,12 @@ class CometShuffleWriteProcessor(
       "output_rows" -> metrics(SQLShuffleWriteMetricsReporter.SHUFFLE_RECORDS_WRITTEN),
       "data_size" -> metrics("dataSize"),
       "elapsed_compute" -> metrics(SQLShuffleWriteMetricsReporter.SHUFFLE_WRITE_TIME))
-    val nativeMetrics = CometMetricNode(nativeSQLMetrics)
+
+    val nativeMetrics = if (metrics.contains("read_time")) {
+      CometMetricNode(nativeSQLMetrics ++ Map("read_time" -> metrics("read_time")))
+    } else {
+      CometMetricNode(nativeSQLMetrics)
+    }
 
     // Getting rid of the fake partitionId
     val newInputs = inputs.asInstanceOf[Iterator[_ <: Product2[Any, Any]]].map(_._2)

From 9f4aa8266cc80c9d5eddb0673a20f635e9d88139 Mon Sep 17 00:00:00 2001
From: Andy Grove <agrove@apache.org>
Date: Wed, 20 Nov 2024 18:24:39 -0700
Subject: [PATCH 02/10] format

---
 native/core/src/execution/datafusion/shuffle_writer.rs | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/native/core/src/execution/datafusion/shuffle_writer.rs b/native/core/src/execution/datafusion/shuffle_writer.rs
index a080bf22b8..289cf220ea 100644
--- a/native/core/src/execution/datafusion/shuffle_writer.rs
+++ b/native/core/src/execution/datafusion/shuffle_writer.rs
@@ -60,12 +60,12 @@ use futures::{lock::Mutex, Stream, StreamExt, TryFutureExt, TryStreamExt};
 use itertools::Itertools;
 use simd_adler32::Adler32;
 
+use crate::execution::operators::ScanExec;
 use crate::{
     common::bit::ceil,
     errors::{CometError, CometResult},
 };
 use datafusion_comet_spark_expr::spark_hash::create_murmur3_hashes;
-use crate::execution::operators::ScanExec;
 
 /// The status of appending rows to a partition buffer.
 enum AppendRowStatus {
@@ -160,7 +160,7 @@ impl ExecutionPlan for ShuffleWriterExec {
                     metrics,
                     context,
                     read_time,
-                    scan_time
+                    scan_time,
                 )
                 .map_err(|e| ArrowError::ExternalError(Box::new(e))),
             )
@@ -1102,7 +1102,7 @@ async fn external_shuffle(
     metrics: ShuffleRepartitionerMetrics,
     context: Arc<TaskContext>,
     read_time: Time,
-    scan_time: Option<Time>
+    scan_time: Option<Time>,
 ) -> Result<SendableRecordBatchStream> {
     let schema = input.schema();
     let mut repartitioner = ShuffleRepartitioner::new(

From 397c5c42b1b25989652d547496a8c586df27511c Mon Sep 17 00:00:00 2001
From: Andy Grove <agrove@apache.org>
Date: Wed, 20 Nov 2024 18:36:13 -0700
Subject: [PATCH 03/10] clippy

---
 .../core/src/execution/datafusion/shuffle_writer.rs  | 12 +++++++-----
 1 file changed, 7 insertions(+), 5 deletions(-)

diff --git a/native/core/src/execution/datafusion/shuffle_writer.rs b/native/core/src/execution/datafusion/shuffle_writer.rs
index 289cf220ea..9643dac2e5 100644
--- a/native/core/src/execution/datafusion/shuffle_writer.rs
+++ b/native/core/src/execution/datafusion/shuffle_writer.rs
@@ -142,11 +142,12 @@ impl ExecutionPlan for ShuffleWriterExec {
         let metrics = ShuffleRepartitionerMetrics::new(&self.metrics, 0);
         let read_time = MetricBuilder::new(&self.metrics).subset_time("read_time", 0);
 
-        let scan_time = if let Some(scan) = self.input.as_any().downcast_ref::<ScanExec>() {
-            Some(scan.baseline_metrics.elapsed_compute().clone())
-        } else {
-            None
-        };
+        // clone the ScanExec elapsed_compute metric
+        let scan_time = self
+            .input
+            .as_any()
+            .downcast_ref::<ScanExec>()
+            .map(|scan| scan.baseline_metrics.elapsed_compute().clone());
 
         Ok(Box::pin(RecordBatchStreamAdapter::new(
             self.schema(),
@@ -1093,6 +1094,7 @@ impl Debug for ShuffleRepartitioner {
     }
 }
 
+#[allow(clippy::too_many_arguments)]
 async fn external_shuffle(
     mut input: SendableRecordBatchStream,
     partition_id: usize,

From 4c29a992f4475b96c42013312626f4fcb0fd4a7a Mon Sep 17 00:00:00 2001
From: Andy Grove <agrove@apache.org>
Date: Thu, 21 Nov 2024 09:32:50 -0700
Subject: [PATCH 04/10] address feedback and add more docs

---
 docs/source/user-guide/tuning.md              | 15 +++++---
 .../execution/datafusion/shuffle_writer.rs    | 34 ++++++++-----------
 native/core/src/execution/operators/scan.rs   |  6 ++--
 3 files changed, 29 insertions(+), 26 deletions(-)

diff --git a/docs/source/user-guide/tuning.md b/docs/source/user-guide/tuning.md
index b1838ca801..3d30a82aa4 100644
--- a/docs/source/user-guide/tuning.md
+++ b/docs/source/user-guide/tuning.md
@@ -23,7 +23,7 @@ Comet provides some tuning options to help you get the best performance from you
 
 ## Memory Tuning
 
-Comet shares an off-heap memory pool between Spark and Comet. This requires setting `spark.memory.offHeap.enabled=true`. 
+Comet shares an off-heap memory pool between Spark and Comet. This requires setting `spark.memory.offHeap.enabled=true`.
 If this setting is not enabled, Comet will not accelerate queries and will fall back to Spark.
 
 Each executor will have a single memory pool which will be shared by all native plans being executed within that
@@ -105,8 +105,15 @@ then any shuffle operations that cannot be supported in this mode will fall back
 
 ## Metrics
 
-Comet metrics are not directly comparable to Spark metrics in some cases.
+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. In one case we saw total scan time
-of 41 seconds reported as 23 seconds for example.
+milliseconds _per batch_ which can result in a large loss of precision.
+
+Comet also adds some custom metrics:
+
+### ShuffleWriterExec
+
+| Metric      | Description                                                                                                                                                                             |
+| ----------- | --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
+| `read_time` | Measure the time it takes for ShuffleWriterExec to fetch an existing batch from the JVM. Note that this does not include the execution time of the query that produced the input batch. |
diff --git a/native/core/src/execution/datafusion/shuffle_writer.rs b/native/core/src/execution/datafusion/shuffle_writer.rs
index 9643dac2e5..e6d92d26ae 100644
--- a/native/core/src/execution/datafusion/shuffle_writer.rs
+++ b/native/core/src/execution/datafusion/shuffle_writer.rs
@@ -60,7 +60,6 @@ use futures::{lock::Mutex, Stream, StreamExt, TryFutureExt, TryStreamExt};
 use itertools::Itertools;
 use simd_adler32::Adler32;
 
-use crate::execution::operators::ScanExec;
 use crate::{
     common::bit::ceil,
     errors::{CometError, CometResult},
@@ -142,13 +141,6 @@ impl ExecutionPlan for ShuffleWriterExec {
         let metrics = ShuffleRepartitionerMetrics::new(&self.metrics, 0);
         let read_time = MetricBuilder::new(&self.metrics).subset_time("read_time", 0);
 
-        // clone the ScanExec elapsed_compute metric
-        let scan_time = self
-            .input
-            .as_any()
-            .downcast_ref::<ScanExec>()
-            .map(|scan| scan.baseline_metrics.elapsed_compute().clone());
-
         Ok(Box::pin(RecordBatchStreamAdapter::new(
             self.schema(),
             futures::stream::once(
@@ -161,7 +153,6 @@ impl ExecutionPlan for ShuffleWriterExec {
                     metrics,
                     context,
                     read_time,
-                    scan_time,
                 )
                 .map_err(|e| ArrowError::ExternalError(Box::new(e))),
             )
@@ -1104,7 +1095,6 @@ async fn external_shuffle(
     metrics: ShuffleRepartitionerMetrics,
     context: Arc<TaskContext>,
     read_time: Time,
-    scan_time: Option<Time>,
 ) -> Result<SendableRecordBatchStream> {
     let schema = input.schema();
     let mut repartitioner = ShuffleRepartitioner::new(
@@ -1118,17 +1108,21 @@ async fn external_shuffle(
         context.session_config().batch_size(),
     );
 
-    while let Some(batch) = input.next().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
-        // current batch in the repartitioner.
-        block_on(repartitioner.insert_batch(batch?))?;
-    }
+    loop {
+        let mut timer = read_time.timer();
+        let b = input.next().await;
+        timer.stop();
 
-    // copy ScanExec elapsed_compute into ShuffleWriterExec read_time
-    if let Some(t) = &scan_time {
-        read_time.add(t);
+        match b {
+            Some(batch_result) => {
+                // 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
+                // current batch in the repartitioner.
+                block_on(repartitioner.insert_batch(batch_result?))?;
+            }
+            _ => break,
+        }
     }
 
     repartitioner.shuffle_write().await
diff --git a/native/core/src/execution/operators/scan.rs b/native/core/src/execution/operators/scan.rs
index ca1a0fb9f6..badb865e6c 100644
--- a/native/core/src/execution/operators/scan.rs
+++ b/native/core/src/execution/operators/scan.rs
@@ -52,7 +52,9 @@ use std::{
 };
 
 /// ScanExec reads batches of data from Spark via JNI. The source of the scan could be a file
-/// scan or the result of reading a broadcast or shuffle exchange.
+/// scan or the result of reading a broadcast or shuffle exchange. ScanExec isn't invoked
+/// until the data is already available in the JVM. When CometExecIterator invokes
+/// Native.executePlan, it passes in the memory addresses of the input batches.
 #[derive(Debug, Clone)]
 pub struct ScanExec {
     /// The ID of the execution context that owns this subquery. We use this ID to retrieve the JVM
@@ -72,7 +74,7 @@ pub struct ScanExec {
     /// Cache of expensive-to-compute plan properties
     cache: PlanProperties,
     /// Metrics collector
-    pub(crate) metrics: ExecutionPlanMetricsSet,
+    metrics: ExecutionPlanMetricsSet,
     /// Baseline metrics
     pub(crate) baseline_metrics: BaselineMetrics,
 }

From f65e91febb5cbbe12e737c34f59559e9f9e9ae5e Mon Sep 17 00:00:00 2001
From: Andy Grove <agrove@apache.org>
Date: Thu, 21 Nov 2024 09:38:55 -0700
Subject: [PATCH 05/10] rename metric key

---
 docs/source/user-guide/tuning.md                       | 6 +++---
 native/core/src/execution/datafusion/shuffle_writer.rs | 8 ++++----
 2 files changed, 7 insertions(+), 7 deletions(-)

diff --git a/docs/source/user-guide/tuning.md b/docs/source/user-guide/tuning.md
index 3d30a82aa4..beefa75ba4 100644
--- a/docs/source/user-guide/tuning.md
+++ b/docs/source/user-guide/tuning.md
@@ -114,6 +114,6 @@ Comet also adds some custom metrics:
 
 ### ShuffleWriterExec
 
-| Metric      | Description                                                                                                                                                                             |
-| ----------- | --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
-| `read_time` | Measure the time it takes for ShuffleWriterExec to fetch an existing batch from the JVM. Note that this does not include the execution time of the query that produced the input batch. |
+| Metric           | Description                                                                                                                                                                             |
+|------------------| --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
+| `jvm_fetch_time` | Measure the time it takes for ShuffleWriterExec to fetch an existing batch from the JVM. Note that this does not include the execution time of the query that produced the input batch. |
diff --git a/native/core/src/execution/datafusion/shuffle_writer.rs b/native/core/src/execution/datafusion/shuffle_writer.rs
index e6d92d26ae..c79eeeb4a0 100644
--- a/native/core/src/execution/datafusion/shuffle_writer.rs
+++ b/native/core/src/execution/datafusion/shuffle_writer.rs
@@ -139,7 +139,7 @@ impl ExecutionPlan for ShuffleWriterExec {
     ) -> Result<SendableRecordBatchStream> {
         let input = self.input.execute(partition, Arc::clone(&context))?;
         let metrics = ShuffleRepartitionerMetrics::new(&self.metrics, 0);
-        let read_time = MetricBuilder::new(&self.metrics).subset_time("read_time", 0);
+        let jvm_fetch_time = MetricBuilder::new(&self.metrics).subset_time("jvm_fetch_time", 0);
 
         Ok(Box::pin(RecordBatchStreamAdapter::new(
             self.schema(),
@@ -152,7 +152,7 @@ impl ExecutionPlan for ShuffleWriterExec {
                     self.partitioning.clone(),
                     metrics,
                     context,
-                    read_time,
+                    jvm_fetch_time,
                 )
                 .map_err(|e| ArrowError::ExternalError(Box::new(e))),
             )
@@ -1094,7 +1094,7 @@ async fn external_shuffle(
     partitioning: Partitioning,
     metrics: ShuffleRepartitionerMetrics,
     context: Arc<TaskContext>,
-    read_time: Time,
+    jvm_fetch_time: Time,
 ) -> Result<SendableRecordBatchStream> {
     let schema = input.schema();
     let mut repartitioner = ShuffleRepartitioner::new(
@@ -1109,7 +1109,7 @@ async fn external_shuffle(
     );
 
     loop {
-        let mut timer = read_time.timer();
+        let mut timer = jvm_fetch_time.timer();
         let b = input.next().await;
         timer.stop();
 

From c686f1db28cbd42a5f94c127197e1faea992e11f Mon Sep 17 00:00:00 2001
From: Andy Grove <agrove@apache.org>
Date: Thu, 21 Nov 2024 09:40:36 -0700
Subject: [PATCH 06/10] fix

---
 docs/source/user-guide/tuning.md            | 6 +++---
 native/core/src/execution/operators/scan.rs | 2 +-
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/docs/source/user-guide/tuning.md b/docs/source/user-guide/tuning.md
index beefa75ba4..9bcc6cdfd0 100644
--- a/docs/source/user-guide/tuning.md
+++ b/docs/source/user-guide/tuning.md
@@ -114,6 +114,6 @@ Comet also adds some custom metrics:
 
 ### ShuffleWriterExec
 
-| Metric           | Description                                                                                                                                                                             |
-|------------------| --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
-| `jvm_fetch_time` | Measure the time it takes for ShuffleWriterExec to fetch an existing batch from the JVM. Note that this does not include the execution time of the query that produced the input batch. |
+| Metric           | Description                                                                                                                                                                               |
+| ---------------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
+| `jvm_fetch_time` | Measure the time it takes for `ShuffleWriterExec` to fetch an existing batch from the JVM. Note that this does not include the execution time of the query that produced the input batch. |
diff --git a/native/core/src/execution/operators/scan.rs b/native/core/src/execution/operators/scan.rs
index badb865e6c..2cb8a84d94 100644
--- a/native/core/src/execution/operators/scan.rs
+++ b/native/core/src/execution/operators/scan.rs
@@ -76,7 +76,7 @@ pub struct ScanExec {
     /// Metrics collector
     metrics: ExecutionPlanMetricsSet,
     /// Baseline metrics
-    pub(crate) baseline_metrics: BaselineMetrics,
+    baseline_metrics: BaselineMetrics,
 }
 
 impl ScanExec {

From b5b19b3fbd1be6cd9a2d2692229e040b05859abc Mon Sep 17 00:00:00 2001
From: Andy Grove <agrove@apache.org>
Date: Thu, 21 Nov 2024 09:42:21 -0700
Subject: [PATCH 07/10] fix

---
 .../comet/execution/shuffle/CometShuffleExchangeExec.scala  | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala
index a8ccc3b6e0..ac5e85a25e 100644
--- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala
+++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala
@@ -77,7 +77,7 @@ case class CometShuffleExchangeExec(
     SQLShuffleReadMetricsReporter.createShuffleReadMetrics(sparkContext)
   override lazy val metrics: Map[String, SQLMetric] = Map(
     "dataSize" -> SQLMetrics.createSizeMetric(sparkContext, "data size"),
-    "read_time" -> SQLMetrics.createNanoTimingMetric(sparkContext, "read time"),
+    "jvm_fetch_time" -> SQLMetrics.createNanoTimingMetric(sparkContext, "time fetching batches from JVM"),
     "numPartitions" -> SQLMetrics.createMetric(
       sparkContext,
       "number of partitions")) ++ readMetrics ++ writeMetrics
@@ -482,8 +482,8 @@ class CometShuffleWriteProcessor(
       "data_size" -> metrics("dataSize"),
       "elapsed_compute" -> metrics(SQLShuffleWriteMetricsReporter.SHUFFLE_WRITE_TIME))
 
-    val nativeMetrics = if (metrics.contains("read_time")) {
-      CometMetricNode(nativeSQLMetrics ++ Map("read_time" -> metrics("read_time")))
+    val nativeMetrics = if (metrics.contains("jvm_fetch_time")) {
+      CometMetricNode(nativeSQLMetrics ++ Map("jvm_fetch_time" -> metrics("time to fetch batches from JVM")))
     } else {
       CometMetricNode(nativeSQLMetrics)
     }

From 1fc34cd94deb2a98847b669e90a33cc66934f592 Mon Sep 17 00:00:00 2001
From: Andy Grove <agrove@apache.org>
Date: Thu, 21 Nov 2024 09:47:48 -0700
Subject: [PATCH 08/10] format

---
 .../execution/shuffle/CometShuffleExchangeExec.scala      | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala
index ac5e85a25e..1481a0400a 100644
--- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala
+++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala
@@ -77,7 +77,9 @@ case class CometShuffleExchangeExec(
     SQLShuffleReadMetricsReporter.createShuffleReadMetrics(sparkContext)
   override lazy val metrics: Map[String, SQLMetric] = Map(
     "dataSize" -> SQLMetrics.createSizeMetric(sparkContext, "data size"),
-    "jvm_fetch_time" -> SQLMetrics.createNanoTimingMetric(sparkContext, "time fetching batches from JVM"),
+    "jvm_fetch_time" -> SQLMetrics.createNanoTimingMetric(
+      sparkContext,
+      "time fetching batches from JVM"),
     "numPartitions" -> SQLMetrics.createMetric(
       sparkContext,
       "number of partitions")) ++ readMetrics ++ writeMetrics
@@ -483,7 +485,9 @@ class CometShuffleWriteProcessor(
       "elapsed_compute" -> metrics(SQLShuffleWriteMetricsReporter.SHUFFLE_WRITE_TIME))
 
     val nativeMetrics = if (metrics.contains("jvm_fetch_time")) {
-      CometMetricNode(nativeSQLMetrics ++ Map("jvm_fetch_time" -> metrics("time to fetch batches from JVM")))
+      CometMetricNode(
+        nativeSQLMetrics ++ Map("jvm_fetch_time" ->
+          metrics("time to fetch batches from JVM")))
     } else {
       CometMetricNode(nativeSQLMetrics)
     }

From 9adcad8ef1167eb15af0e815b456e02133c8077b Mon Sep 17 00:00:00 2001
From: Andy Grove <agrove@apache.org>
Date: Thu, 21 Nov 2024 09:54:40 -0700
Subject: [PATCH 09/10] fix

---
 .../sql/comet/execution/shuffle/CometShuffleExchangeExec.scala  | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala
index 1481a0400a..4c3f994f9e 100644
--- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala
+++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala
@@ -487,7 +487,7 @@ class CometShuffleWriteProcessor(
     val nativeMetrics = if (metrics.contains("jvm_fetch_time")) {
       CometMetricNode(
         nativeSQLMetrics ++ Map("jvm_fetch_time" ->
-          metrics("time to fetch batches from JVM")))
+          metrics("jvm_fetch_time")))
     } else {
       CometMetricNode(nativeSQLMetrics)
     }

From 5f86bc9d909bbf29f6c0db971b7b9fceefd1d4ed Mon Sep 17 00:00:00 2001
From: Andy Grove <agrove@apache.org>
Date: Thu, 21 Nov 2024 13:12:22 -0700
Subject: [PATCH 10/10] address feedback

---
 docs/source/user-guide/tuning.md | 13 +++++++------
 1 file changed, 7 insertions(+), 6 deletions(-)

diff --git a/docs/source/user-guide/tuning.md b/docs/source/user-guide/tuning.md
index 9bcc6cdfd0..e0d7c1bc93 100644
--- a/docs/source/user-guide/tuning.md
+++ b/docs/source/user-guide/tuning.md
@@ -105,15 +105,16 @@ then any shuffle operations that cannot be supported in this mode will fall back
 
 ## Metrics
 
-Some Comet metrics are not directly comparable to Spark metrics in some cases.
+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.
+- `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.
 
 Comet also adds some custom metrics:
 
 ### ShuffleWriterExec
 
-| Metric           | Description                                                                                                                                                                               |
-| ---------------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
-| `jvm_fetch_time` | Measure the time it takes for `ShuffleWriterExec` to fetch an existing batch from the JVM. Note that this does not include the execution time of the query that produced the input batch. |
+| Metric           | Description                                                                                                                                                                       |
+| ---------------- | --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
+| `jvm_fetch_time` | Measure the time it takes for `ShuffleWriterExec` to fetch batches from the JVM. Note that this does not include the execution time of the query that produced the input batches. |