feat: Implement metrics for iceberg compat (apache#2615)

Author: EmilyMatt

common/src/main/java/org/apache/comet/parquet/Native.java

@@ -268,7 +268,8 @@ public static native long initRecordBatchReader(
       int batchSize,
       boolean caseSensitive,
       Map<String, String> objectStoreOptions,
-      CometFileKeyUnwrapper keyUnwrapper);
+      CometFileKeyUnwrapper keyUnwrapper,
+      Object metricsNode);
 
   // arrow native version of read batch
 

common/src/main/java/org/apache/comet/parquet/NativeBatchReader.java

@@ -114,6 +114,9 @@ public class NativeBatchReader extends RecordReader<Void, ColumnarBatch> impleme
   private InternalRow partitionValues;
   private PartitionedFile file;
   private final Map<String, SQLMetric> metrics;
+  // Unfortunately CometMetricNode is from the "spark" package and cannot be used directly here
+  // TODO: Move it to common package?
+  private Object metricsNode = null;
 
   private StructType sparkSchema;
   private StructType dataSchema;
@@ -214,7 +217,8 @@ private NativeBatchReader(AbstractColumnReader[] columnReaders) {
       boolean useLegacyDateTimestamp,
       StructType partitionSchema,
       InternalRow partitionValues,
-      Map<String, SQLMetric> metrics) {
+      Map<String, SQLMetric> metrics,
+      Object metricsNode) {
     this.conf = conf;
     this.capacity = capacity;
     this.sparkSchema = sparkSchema;
@@ -229,6 +233,7 @@ private NativeBatchReader(AbstractColumnReader[] columnReaders) {
     this.footer = footer;
     this.nativeFilter = nativeFilter;
     this.metrics = metrics;
+    this.metricsNode = metricsNode;
     this.taskContext = TaskContext$.MODULE$.get();
   }
 
@@ -436,7 +441,8 @@ public void init() throws Throwable {
               batchSize,
               caseSensitive,
               objectStoreOptions,
-              keyUnwrapper);
+              keyUnwrapper,
+              metricsNode);
     }
     isInitialized = true;
   }

native/core/src/execution/metrics/utils.rs

@@ -27,20 +27,24 @@ use std::sync::Arc;
 /// Updates the metrics of a CometMetricNode. This function is called recursively to
 /// update the metrics of all the children nodes. The metrics are pulled from the
 /// native execution plan and pushed to the Java side through JNI.
-pub fn update_comet_metric(
+pub(crate) fn update_comet_metric(
     env: &mut JNIEnv,
     metric_node: &JObject,
     spark_plan: &Arc<SparkPlan>,
 ) -> Result<(), CometError> {
-    unsafe {
-        let native_metric = to_native_metric_node(spark_plan);
-        let jbytes = env.byte_array_from_slice(&native_metric?.encode_to_vec())?;
-        jni_call!(env, comet_metric_node(metric_node).set_all_from_bytes(&jbytes) -> ())?;
+    if metric_node.is_null() {
+        return Ok(());
     }
-    Ok(())
+
+    let native_metric = to_native_metric_node(spark_plan);
+    let jbytes = env.byte_array_from_slice(&native_metric?.encode_to_vec())?;
+
+    unsafe { jni_call!(env, comet_metric_node(metric_node).set_all_from_bytes(&jbytes) -> ()) }
 }
 
-pub fn to_native_metric_node(spark_plan: &Arc<SparkPlan>) -> Result<NativeMetricNode, CometError> {
+pub(crate) fn to_native_metric_node(
+    spark_plan: &Arc<SparkPlan>,
+) -> Result<NativeMetricNode, CometError> {
     let mut native_metric_node = NativeMetricNode {
         metrics: HashMap::new(),
         children: Vec::new(),

native/core/src/execution/mod.rs

@@ -18,7 +18,7 @@
 //! PoC of vectorization execution through JNI to Rust.
 pub mod expressions;
 pub mod jni_api;
-mod metrics;
+pub(crate) mod metrics;
 pub mod operators;
 pub(crate) mod planner;
 pub mod serde;

native/core/src/parquet/mod.rs

@@ -46,9 +46,11 @@ use jni::{
 
 use self::util::jni::TypePromotionInfo;
 use crate::execution::jni_api::get_runtime;
+use crate::execution::metrics::utils::update_comet_metric;
 use crate::execution::operators::ExecutionError;
 use crate::execution::planner::PhysicalPlanner;
 use crate::execution::serde;
+use crate::execution::spark_plan::SparkPlan;
 use crate::execution::utils::SparkArrowConvert;
 use crate::jvm_bridge::{jni_new_global_ref, JVMClasses};
 use crate::parquet::data_type::AsBytes;
@@ -600,6 +602,8 @@ enum ParquetReaderState {
 }
 /// Parquet read context maintained across multiple JNI calls.
 struct BatchContext {
+    native_plan: Arc<SparkPlan>,
+    metrics_node: Arc<GlobalRef>,
     batch_stream: Option<SendableRecordBatchStream>,
     current_batch: Option<RecordBatch>,
     reader_state: ParquetReaderState,
@@ -697,6 +701,7 @@ pub unsafe extern "system" fn Java_org_apache_comet_parquet_Native_initRecordBat
     case_sensitive: jboolean,
     object_store_options: JObject,
     key_unwrapper_obj: JObject,
+    metrics_node: JObject,
 ) -> jlong {
     try_unwrap_or_throw(&e, |mut env| unsafe {
         JVMClasses::init(&mut env);
@@ -776,6 +781,8 @@ pub unsafe extern "system" fn Java_org_apache_comet_parquet_Native_initRecordBat
         let batch_stream = Some(scan.execute(partition_index, session_ctx.task_ctx())?);
 
         let ctx = BatchContext {
+            native_plan: Arc::new(SparkPlan::new(0, scan, vec![])),
+            metrics_node: Arc::new(jni_new_global_ref!(env, metrics_node)?),
             batch_stream,
             current_batch: None,
             reader_state: ParquetReaderState::Init,
@@ -791,7 +798,7 @@ pub extern "system" fn Java_org_apache_comet_parquet_Native_readNextRecordBatch(
     _jclass: JClass,
     handle: jlong,
 ) -> jint {
-    try_unwrap_or_throw(&e, |_env| {
+    try_unwrap_or_throw(&e, |mut env| {
         let context = get_batch_context(handle)?;
         let mut rows_read: i32 = 0;
         let batch_stream = context.batch_stream.as_mut().unwrap();
@@ -813,8 +820,11 @@ pub extern "system" fn Java_org_apache_comet_parquet_Native_readNextRecordBatch(
                 Poll::Ready(None) => {
                     // EOF
 
-                    // TODO: (ARROW NATIVE) We can update metrics here
-                    // crate::execution::jni_api::update_metrics(&mut env, exec_context)?;
+                    update_comet_metric(
+                        &mut env,
+                        context.metrics_node.as_obj(),
+                        &context.native_plan,
+                    )?;
 
                     context.current_batch = None;
                     context.reader_state = ParquetReaderState::Complete;

spark/src/main/scala/org/apache/comet/MetricsSupport.scala

@@ -19,38 +19,13 @@
 
 package org.apache.comet
 
-import org.apache.spark.SparkContext
-import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics}
+import org.apache.spark.sql.execution.metric.SQLMetric
 
 /**
  * A trait for Comet operators that support SQL metrics
  */
 trait MetricsSupport {
   protected var metrics: Map[String, SQLMetric] = Map.empty
 
-  def initMetrics(sparkContext: SparkContext): Map[String, SQLMetric] = {
-    metrics = Map(
-      "ParquetRowGroups" -> SQLMetrics.createMetric(
-        sparkContext,
-        "num of Parquet row groups read"),
-      "ParquetNativeDecodeTime" -> SQLMetrics.createNanoTimingMetric(
-        sparkContext,
-        "time spent in Parquet native decoding"),
-      "ParquetNativeLoadTime" -> SQLMetrics.createNanoTimingMetric(
-        sparkContext,
-        "time spent in loading Parquet native vectors"),
-      "ParquetLoadRowGroupTime" -> SQLMetrics.createNanoTimingMetric(
-        sparkContext,
-        "time spent in loading Parquet row groups"),
-      "ParquetInputFileReadTime" -> SQLMetrics.createNanoTimingMetric(
-        sparkContext,
-        "time spent in reading Parquet file from storage"),
-      "ParquetInputFileReadSize" -> SQLMetrics.createSizeMetric(
-        sparkContext,
-        "read size when reading Parquet file from storage (MB)"),
-      "ParquetInputFileReadThroughput" -> SQLMetrics.createAverageMetric(
-        sparkContext,
-        "read throughput when reading Parquet file from storage (MB/sec)"))
-    metrics
-  }
+  def getMetrics: Map[String, SQLMetric] = metrics
 }

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

@@ -30,6 +30,7 @@ import org.apache.spark.sql.SparkSession
 import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.util.CaseInsensitiveMap
 import org.apache.spark.sql.catalyst.util.RebaseDateTime.RebaseSpec
+import org.apache.spark.sql.comet.CometMetricNode
 import org.apache.spark.sql.execution.datasources.DataSourceUtils
 import org.apache.spark.sql.execution.datasources.PartitionedFile
 import org.apache.spark.sql.execution.datasources.RecordReaderIterator
@@ -56,10 +57,14 @@ import org.apache.comet.vector.CometVector
  *     in [[org.apache.comet.CometSparkSessionExtensions]]
  *   - `buildReaderWithPartitionValues`, so Spark calls Comet's Parquet reader to read values.
  */
-class CometParquetFileFormat(scanImpl: String)
+class CometParquetFileFormat(session: SparkSession, scanImpl: String)
     extends ParquetFileFormat
     with MetricsSupport
     with ShimSQLConf {
+  metrics =
+    CometMetricNode.nativeScanMetrics(session.sparkContext) ++ CometMetricNode.parquetScanMetrics(
+      session.sparkContext)
+
   override def shortName(): String = "parquet"
   override def toString: String = "CometParquet"
   override def hashCode(): Int = getClass.hashCode()
@@ -164,7 +169,8 @@ class CometParquetFileFormat(scanImpl: String)
             datetimeRebaseSpec.mode == CORRECTED,
             partitionSchema,
             file.partitionValues,
-            metrics.asJava)
+            metrics.asJava,
+            CometMetricNode(metrics))
           try {
             batchReader.init()
           } catch {

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

@@ -23,6 +23,7 @@ import scala.jdk.CollectionConverters._
 
 import org.apache.hadoop.conf.Configuration
 import org.apache.spark.sql.SparkSession
+import org.apache.spark.sql.comet.CometMetricNode
 import org.apache.spark.sql.connector.read.PartitionReaderFactory
 import org.apache.spark.sql.execution.datasources.parquet.ParquetOptions
 import org.apache.spark.sql.execution.datasources.v2.FileScan
@@ -34,6 +35,8 @@ import org.apache.spark.util.SerializableConfiguration
 
 import org.apache.comet.MetricsSupport
 
+// TODO: Consider creating a case class and patch SQL tests if needed, will make life easier.
+// currently hacking around this by setting the metrics within the object's apply method.
 trait CometParquetScan extends FileScan with MetricsSupport {
   def sparkSession: SparkSession
   def hadoopConf: Configuration
@@ -70,8 +73,8 @@ trait CometParquetScan extends FileScan with MetricsSupport {
 }
 
 object CometParquetScan {
-  def apply(scan: ParquetScan): CometParquetScan =
-    new ParquetScan(
+  def apply(session: SparkSession, scan: ParquetScan): CometParquetScan = {
+    val newScan = new ParquetScan(
       scan.sparkSession,
       scan.hadoopConf,
       scan.fileIndex,
@@ -82,4 +85,10 @@ object CometParquetScan {
       scan.options,
       partitionFilters = scan.partitionFilters,
       dataFilters = scan.dataFilters) with CometParquetScan
+
+    newScan.metrics = CometMetricNode.nativeScanMetrics(session.sparkContext) ++ CometMetricNode
+      .parquetScanMetrics(session.sparkContext)
+
+    newScan
+  }
 }

spark/src/main/scala/org/apache/comet/rules/CometScanRule.scala

@@ -203,7 +203,7 @@ case class CometScanRule(session: SparkSession) extends Rule[SparkPlan] with Com
           return withInfos(scanExec, fallbackReasons.toSet)
         }
 
-        if (scanImpl != CometConf.SCAN_NATIVE_COMET && encryptionEnabled(hadoopConf)) {
+        if (encryptionEnabled(hadoopConf) && scanImpl != CometConf.SCAN_NATIVE_COMET) {
           if (!isEncryptionConfigSupported(hadoopConf)) {
             return withInfos(scanExec, fallbackReasons.toSet)
           }
@@ -257,7 +257,7 @@ case class CometScanRule(session: SparkSession) extends Rule[SparkPlan] with Com
         }
 
         if (schemaSupported && partitionSchemaSupported && scan.pushedAggregate.isEmpty) {
-          val cometScan = CometParquetScan(scanExec.scan.asInstanceOf[ParquetScan])
+          val cometScan = CometParquetScan(session, scanExec.scan.asInstanceOf[ParquetScan])
           CometBatchScanExec(
             scanExec.copy(scan = cometScan),
             runtimeFilters = scanExec.runtimeFilters)

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

@@ -51,9 +51,13 @@ case class CometBatchScanExec(wrapped: BatchScanExec, runtimeFilters: Seq[Expres
   override lazy val inputRDD: RDD[InternalRow] = wrappedScan.inputRDD
 
   override def doExecuteColumnar(): RDD[ColumnarBatch] = {
+    val rdd = inputRDD.asInstanceOf[RDD[ColumnarBatch]]
+
+    // 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 = {
@@ -137,16 +141,12 @@ case class CometBatchScanExec(wrapped: BatchScanExec, runtimeFilters: Seq[Expres
     wrapped
   }
 
-  override lazy val metrics: Map[String, SQLMetric] = Map(
-    "numOutputRows" -> SQLMetrics.createMetric(sparkContext, "number of output rows"),
-    "scanTime" -> SQLMetrics.createNanoTimingMetric(
-      sparkContext,
-      "scan time")) ++ wrapped.customMetrics ++ {
-    wrapped.scan match {
-      case s: MetricsSupport => s.initMetrics(sparkContext)
+  override lazy val metrics: Map[String, SQLMetric] =
+    wrappedScan.customMetrics ++ CometMetricNode.baseScanMetrics(
+      session.sparkContext) ++ (scan match {
+      case s: MetricsSupport => s.getMetrics
       case _ => Map.empty
-    }
-  }
+    })
 
   @transient override lazy val partitions: Seq[Seq[InputPartition]] = wrappedScan.partitions