Add metrics to GpuUnion operator (#543)

* Add metrics to GpuUnion operator * minor fix * fix imports * addressed review comments Signed-off-by: Niranjan Artal <nartal@nvidia.com> * add line Signed-off-by: Niranjan Artal <nartal@nvidia.com>
NVIDIA · Aug 13, 2020 · ae4b00d · ae4b00d
1 parent e2a2d20
commit ae4b00d
Showing 1 changed file with 13 additions and 2 deletions.
diff --git a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/basicPhysicalOperators.scala b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/basicPhysicalOperators.scala
@@ -317,8 +317,19 @@ case class GpuUnionExec(children: Seq[SparkPlan]) extends SparkPlan with GpuExec
   override def doExecute(): RDD[InternalRow] =
     throw new IllegalStateException(s"Row-based execution should not occur for $this")
 
-  override def doExecuteColumnar(): RDD[ColumnarBatch] =
-    sparkContext.union(children.map(_.executeColumnar()))
+  override def doExecuteColumnar(): RDD[ColumnarBatch] = {
+    val numOutputRows = longMetric(NUM_OUTPUT_ROWS)
+    val numOutputBatches = longMetric(NUM_OUTPUT_BATCHES)
+    val totalTime = longMetric(TOTAL_TIME)
+
+    sparkContext.union(children.map(_.executeColumnar())).map { batch =>
+      withResource(new NvtxWithMetrics("Union", NvtxColor.CYAN, totalTime)) { _ =>
+        numOutputBatches += 1
+        numOutputRows += batch.numRows
+        batch
+      }
+    }
+  }
 }
 
 case class GpuCoalesceExec(numPartitions: Int, child: SparkPlan)