[GLUTEN-6695][CH] Introduce shuffleWallTime in CHMetricsApi to calcul…

…ate the overall shuffle write time
apache · Aug 2, 2024 · d97b06b · d97b06b
1 parent bedb2ee
commit d97b06b
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 0 deletions.
diff --git a/...nds-clickhouse/src/main/scala/org/apache/gluten/backendsapi/clickhouse/CHMetricsApi.scala b/...nds-clickhouse/src/main/scala/org/apache/gluten/backendsapi/clickhouse/CHMetricsApi.scala
@@ -222,6 +222,7 @@ class CHMetricsApi extends MetricsApi with Logging with LogLevelUtil {
       "spillTime" -> SQLMetrics.createNanoTimingMetric(sparkContext, "time to spill"),
       "compressTime" -> SQLMetrics.createNanoTimingMetric(sparkContext, "time to compress"),
       "prepareTime" -> SQLMetrics.createNanoTimingMetric(sparkContext, "time to prepare"),
+      "shuffleWallTime" -> SQLMetrics.createNanoTimingMetric(sparkContext, "shuffle wall time"),
       "avgReadBatchNumRows" -> SQLMetrics
         .createAverageMetric(sparkContext, "avg read batch num rows"),
       "numInputRows" -> SQLMetrics.createMetric(sparkContext, "number of input rows"),

diff --git a/.../clickhouse/src/main/scala/org/apache/spark/shuffle/CHCelebornColumnarShuffleWriter.scala b/.../clickhouse/src/main/scala/org/apache/spark/shuffle/CHCelebornColumnarShuffleWriter.scala
@@ -63,7 +63,9 @@ class CHCelebornColumnarShuffleWriter[K, V](
       } else {
         initShuffleWriter(cb)
         val col = cb.column(0).asInstanceOf[CHColumnVector]
+        val startTime = System.nanoTime()
         jniWrapper.split(nativeShuffleWriter, col.getBlockAddress)
+        dep.metrics("shuffleWallTime").add(System.nanoTime() - startTime)
         dep.metrics("numInputRows").add(cb.numRows)
         dep.metrics("inputBatches").add(1)
         // This metric is important, AQE use it to decide if EliminateLimit
@@ -77,8 +79,10 @@ class CHCelebornColumnarShuffleWriter[K, V](
       return
     }
 
+    val startTime = System.nanoTime()
     splitResult = jniWrapper.stop(nativeShuffleWriter)
 
+    dep.metrics("shuffleWallTime").add(System.nanoTime() - startTime)
     dep.metrics("splitTime").add(splitResult.getSplitTime)
     dep.metrics("IOTime").add(splitResult.getDiskWriteTime)
     dep.metrics("serializeTime").add(splitResult.getSerializationTime)