improvement ShuffleReader metrics (apache#12)

Co-authored-by: zhaokuo03 <[email protected]>
kecookier · Sep 25, 2023 · 6387e56 · 6387e56
1 parent 96c2c70
commit 6387e56
Show file tree

Hide file tree

Showing 5 changed files with 40 additions and 8 deletions.
diff --git a/backends-velox/src/main/scala/io/glutenproject/backendsapi/velox/MetricsHandler.scala b/backends-velox/src/main/scala/io/glutenproject/backendsapi/velox/MetricsHandler.scala
@@ -246,6 +246,8 @@ class MetricsHandler extends MetricsApi with Logging {
       "compressTime" -> SQLMetrics.createNanoTimingMetric(sparkContext, "totaltime to compress"),
       "prepareTime" -> SQLMetrics.createNanoTimingMetric(sparkContext, "totaltime to prepare"),
       "decompressTime" -> SQLMetrics.createNanoTimingMetric(sparkContext, "totaltime_decompress"),
+      "ipcTime" -> SQLMetrics.createNanoTimingMetric(sparkContext, "totaltime_ipc"),
+      "deserializeTime" -> SQLMetrics.createNanoTimingMetric(sparkContext, "totaltime_deserialize"),
       "avgReadBatchNumRows" -> SQLMetrics
         .createAverageMetric(sparkContext, "avg read batch num rows"),
       "numInputRows" -> SQLMetrics.createMetric(sparkContext, "number of input rows"),

diff --git a/backends-velox/src/main/scala/io/glutenproject/backendsapi/velox/SparkPlanExecHandler.scala b/backends-velox/src/main/scala/io/glutenproject/backendsapi/velox/SparkPlanExecHandler.scala
@@ -253,13 +253,15 @@ class SparkPlanExecHandler extends SparkPlanExecApi {
     val readBatchNumRows = metrics("avgReadBatchNumRows")
     val numOutputRows = metrics("numOutputRows")
     val decompressTime = metrics("decompressTime")
+    val ipcTime = metrics("ipcTime")
+    val deserializeTime = metrics("deserializeTime")
     if (GlutenConfig.getConf.isUseCelebornShuffleManager) {
       val clazz = ClassUtils.getClass("org.apache.spark.shuffle.CelebornColumnarBatchSerializer")
       val constructor =
         clazz.getConstructor(classOf[StructType], classOf[SQLMetric], classOf[SQLMetric])
       constructor.newInstance(schema, readBatchNumRows, numOutputRows).asInstanceOf[Serializer]
     } else {
-      new ColumnarBatchSerializer(schema, readBatchNumRows, numOutputRows, decompressTime)
+      new ColumnarBatchSerializer(schema, readBatchNumRows, numOutputRows, decompressTime, ipcTime, deserializeTime)
     }
   }
 

diff --git a/cpp/core/jni/JniWrapper.cc b/cpp/core/jni/JniWrapper.cc
@@ -76,6 +76,8 @@ static jmethodID veloxColumnarBatchScannerNext;
 
 static jclass shuffleReaderMetricsClass;
 static jmethodID shuffleReaderMetricsSetDecompressTime;
+static jmethodID shuffleReaderMetricsSetIpcTime;
+static jmethodID shuffleReaderMetricsSetDeserializeTime;
 
 class JavaInputStreamAdaptor final : public arrow::io::InputStream {
  public:
@@ -293,6 +295,9 @@ jint JNI_OnLoad(JavaVM* vm, void* reserved) {
       createGlobalClassReferenceOrError(env, "Lio/glutenproject/vectorized/ShuffleReaderMetrics;");
   shuffleReaderMetricsSetDecompressTime =
       getMethodIdOrError(env, shuffleReaderMetricsClass, "setDecompressTime", "(J)V");
+  shuffleReaderMetricsSetIpcTime = getMethodIdOrError(env, shuffleReaderMetricsClass, "setIpcTime", "(J)V");
+  shuffleReaderMetricsSetDeserializeTime =
+      getMethodIdOrError(env, shuffleReaderMetricsClass, "SetDeserializeTime", "(J)V");
 
   return jniVersion;
 }
@@ -1056,10 +1061,8 @@ JNIEXPORT void JNICALL Java_io_glutenproject_vectorized_ShuffleReaderJniWrapper_
 
   auto reader = executionCtx->getShuffleReader(shuffleReaderHandle);
   env->CallVoidMethod(metrics, shuffleReaderMetricsSetDecompressTime, reader->getDecompressTime());
-
-  // TODO zhaokuo
-  // env->CallVoidMethod(metrics, shuffleReaderMetricsSetIpcTime, reader->getIpcTime());
-  // env->CallVoidMethod(metrics, shuffleReaderMetricsSetDeserializeTime, reader->getDeserializeTime());
+  env->CallVoidMethod(metrics, shuffleReaderMetricsSetIpcTime, reader->getIpcTime());
+  env->CallVoidMethod(metrics, shuffleReaderMetricsSetDeserializeTime, reader->getDeserializeTime());
 
   checkException(env);
   JNI_METHOD_END()

diff --git a/gluten-data/src/main/java/io/glutenproject/vectorized/ShuffleReaderMetrics.java b/gluten-data/src/main/java/io/glutenproject/vectorized/ShuffleReaderMetrics.java
@@ -18,6 +18,8 @@
 
 public class ShuffleReaderMetrics {
   private long decompressTime;
+  private long ipcTime;
+  private long deserializeTime;
 
   public void setDecompressTime(long decompressTime) {
     this.decompressTime = decompressTime;
@@ -26,4 +28,20 @@ public void setDecompressTime(long decompressTime) {
   public long getDecompressTime() {
     return decompressTime;
   }
+
+  public void setIpcTime(long ipcTime) {
+    this.ipcTime = ipcTime;
+  }
+
+  public long getIpcTime() {
+    return ipcTime;
+  }
+
+  public void SetDeserializeTime(long ipcTime) {
+    this.deserializeTime = deserializeTime;
+  }
+
+  public long getDeserializeTime() {
+    return deserializeTime;
+  }
 }
diff --git a/gluten-data/src/main/scala/io/glutenproject/vectorized/ColumnarBatchSerializer.scala b/gluten-data/src/main/scala/io/glutenproject/vectorized/ColumnarBatchSerializer.scala
@@ -45,7 +45,9 @@ class ColumnarBatchSerializer(
     schema: StructType,
     readBatchNumRows: SQLMetric,
     numOutputRows: SQLMetric,
-    decompressTime: SQLMetric)
+    decompressTime: SQLMetric,
+    ipcTime: SQLMetric,
+    deserializeTime: SQLMetric)
   extends Serializer
   with Serializable {
 
@@ -54,7 +56,8 @@ class ColumnarBatchSerializer(
 
   /** Creates a new [[SerializerInstance]]. */
   override def newInstance(): SerializerInstance = {
-    new ColumnarBatchSerializerInstance(schema, readBatchNumRows, numOutputRows, decompressTime)
+    new ColumnarBatchSerializerInstance(schema, readBatchNumRows, numOutputRows, decompressTime,
+      ipcTime, deserializeTime)
   }
 
   override def supportsRelocationOfSerializedObjects: Boolean = supportsRelocation
@@ -64,7 +67,9 @@ private class ColumnarBatchSerializerInstance(
     schema: StructType,
     readBatchNumRows: SQLMetric,
     numOutputRows: SQLMetric,
-    decompressTime: SQLMetric)
+    decompressTime: SQLMetric,
+    ipcTime: SQLMetric,
+    deserializeTime: SQLMetric)
   extends SerializerInstance
   with Logging {
 
@@ -106,6 +111,8 @@ private class ColumnarBatchSerializerInstance(
         shuffleReaderHandle,
         readerMetrics)
       decompressTime += readerMetrics.getDecompressTime
+      ipcTime += readerMetrics.getIpcTime
+      deserializeTime += readerMetrics.getDeserializeTime
 
       cSchema.close()
       ShuffleReaderJniWrapper.INSTANCE.close(executionCtxHandle, shuffleReaderHandle)