fixes for streaming

Author: EmilyMatt

spark/src/main/scala/org/apache/comet/parquet/CometParquetFileFormat.scala

@@ -61,11 +61,9 @@ class CometParquetFileFormat(session: SparkSession, scanImpl: String)
     extends ParquetFileFormat
     with MetricsSupport
     with ShimSQLConf {
-  if (scanImpl == CometConf.SCAN_NATIVE_COMET) {
-    metrics = CometMetricNode.parquetScanMetrics(session.sparkContext)
-  } else {
-    metrics = CometMetricNode.nativeScanMetrics(session.sparkContext)
-  }
+  metrics =
+    CometMetricNode.nativeScanMetrics(session.sparkContext) ++ CometMetricNode.parquetScanMetrics(
+      session.sparkContext)
 
   override def shortName(): String = "parquet"
   override def toString: String = "CometParquet"

spark/src/main/scala/org/apache/comet/parquet/CometParquetScan.scala

@@ -86,8 +86,8 @@ object CometParquetScan {
       partitionFilters = scan.partitionFilters,
       dataFilters = scan.dataFilters) with CometParquetScan
 
-    // TODO: Change this when we support DataFusion reader for Parquet in V2
-    newScan.metrics = CometMetricNode.parquetScanMetrics(session.sparkContext)
+    newScan.metrics = CometMetricNode.nativeScanMetrics(session.sparkContext) ++ CometMetricNode
+      .parquetScanMetrics(session.sparkContext)
 
     newScan
   }

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

@@ -53,15 +53,11 @@ case class CometBatchScanExec(wrapped: BatchScanExec, runtimeFilters: Seq[Expres
   override def doExecuteColumnar(): RDD[ColumnarBatch] = {
     val rdd = inputRDD.asInstanceOf[RDD[ColumnarBatch]]
 
-    // Can skip the following logic if we're using different metrics to calculate this,
-    // e.g., Datafusion reader metrics.
-    if (Seq("numOutputRows", "scanTime").exists(metric => !metrics.contains(metric))) {
-      return rdd
-    }
-
+    // These metrics are important for streaming solutions.
+    // despite there being similar metrics published by the native reader.
     val numOutputRows = longMetric("numOutputRows")
     val scanTime = longMetric("scanTime")
-    inputRDD.asInstanceOf[RDD[ColumnarBatch]].mapPartitionsInternal { batches =>
+    rdd.mapPartitionsInternal { batches =>
       new Iterator[ColumnarBatch] {
 
         override def hasNext: Boolean = {

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

@@ -98,19 +98,8 @@ object CometMetricNode {
         "total time (in ms) spent in this operator"))
   }
 
-  /**
-   * SQL Metrics for Comet native ScanExec
-   */
-  def scanMetrics(sc: SparkContext): Map[String, SQLMetric] = {
-    Map(
-      "cast_time" ->
-        SQLMetrics.createNanoTimingMetric(sc, "Total time for casting columns"))
-  }
-
   def parquetScanMetrics(sc: SparkContext): Map[String, SQLMetric] = {
     Map(
-      "numOutputRows" -> SQLMetrics.createMetric(sc, "number of output rows"),
-      "scanTime" -> SQLMetrics.createNanoTimingMetric(sc, "scan time"),
       "ParquetRowGroups" -> SQLMetrics.createMetric(sc, "num of Parquet row groups read"),
       "ParquetNativeDecodeTime" -> SQLMetrics.createNanoTimingMetric(
         sc,
@@ -134,6 +123,10 @@ object CometMetricNode {
 
   def nativeScanMetrics(sc: SparkContext): Map[String, SQLMetric] = {
     Map(
+      // Spark metrics for streaming pipelines
+      "numOutputRows" -> SQLMetrics.createMetric(sc, "number of output rows"),
+      "scanTime" -> SQLMetrics.createNanoTimingMetric(sc, "scan time"),
+      // Datafusion reader metrics
       "output_rows" -> SQLMetrics.createMetric(sc, "number of output rows"),
       "time_elapsed_opening" ->
         SQLMetrics.createNanoTimingMetric(sc, "Wall clock time elapsed for file opening"),

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

@@ -259,15 +259,11 @@ case class CometScanExec(
   protected override def doExecuteColumnar(): RDD[ColumnarBatch] = {
     val rdd = inputRDD.asInstanceOf[RDD[ColumnarBatch]]
 
-    // Can skip the following logic if we're using different metrics to calculate this,
-    // e.g., Datafusion reader metrics.
-    if (Seq("numOutputRows", "scanTime").exists(metric => !metrics.contains(metric))) {
-      return rdd
-    }
-
+    // These metrics are important for streaming solutions.
+    // despite there being similar metrics published by the native reader.
     val numOutputRows = longMetric("numOutputRows")
     val scanTime = longMetric("scanTime")
-    inputRDD.asInstanceOf[RDD[ColumnarBatch]].mapPartitionsInternal { batches =>
+    rdd.mapPartitionsInternal { batches =>
       new Iterator[ColumnarBatch] {
 
         override def hasNext: Boolean = {