perf: improve performance of update metrics

native/core/src/execution/jni_api.rs

@@ -99,8 +99,6 @@ struct ExecutionContext {
     pub debug_native: bool,
     /// Whether to write native plans with metrics to stdout
     pub explain_native: bool,
-    /// Map of metrics name -> jstring object to cache jni_NewStringUTF calls.
-    pub metrics_jstrings: HashMap<String, Arc<GlobalRef>>,
     /// Memory pool config
     pub memory_pool_config: MemoryPoolConfig,
 }
@@ -237,7 +235,6 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_createPlan(
             session_ctx: Arc::new(session),
             debug_native: debug_native == 1,
             explain_native: explain_native == 1,
-            metrics_jstrings: HashMap::new(),
             memory_pool_config,
         });
 
@@ -508,9 +505,6 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_executePlan(
             let next_item = exec_context.stream.as_mut().unwrap().next();
             let poll_output = exec_context.runtime.block_on(async { poll!(next_item) });
 
-            // Update metrics
-            update_metrics(&mut env, exec_context)?;
-
             match poll_output {
                 Poll::Ready(Some(output)) => {
                     // prepare output for FFI transfer
@@ -561,8 +555,12 @@ pub extern "system" fn Java_org_apache_comet_Native_releasePlan(
     _class: JClass,
     exec_context: jlong,
 ) {
-    try_unwrap_or_throw(&e, |_| unsafe {
+    try_unwrap_or_throw(&e, |mut env| unsafe {
         let execution_context = get_execution_context(exec_context);
+
+        // Update metrics
+        update_metrics(&mut env, execution_context)?;
+
         if execution_context.memory_pool_config.pool_type == MemoryPoolType::FairSpillTaskShared
             || execution_context.memory_pool_config.pool_type == MemoryPoolType::GreedyTaskShared
         {
@@ -586,10 +584,13 @@ pub extern "system" fn Java_org_apache_comet_Native_releasePlan(
 
 /// Updates the metrics of the query plan.
 fn update_metrics(env: &mut JNIEnv, exec_context: &mut ExecutionContext) -> CometResult<()> {
-    let native_query = exec_context.root_op.as_ref().unwrap();
-    let metrics = exec_context.metrics.as_obj();
-    let metrics_jstrings = &mut exec_context.metrics_jstrings;
-    update_comet_metric(env, metrics, native_query, metrics_jstrings)
+    if exec_context.root_op.is_some() {
+        let native_query = exec_context.root_op.as_ref().unwrap();
+        let metrics = exec_context.metrics.as_obj();
+        update_comet_metric(env, metrics, native_query)
+    } else {
+        Ok(())
+    }
 }
 
 fn convert_datatype_arrays(

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

@@ -16,14 +16,11 @@
 // under the License.
 
 use crate::execution::spark_plan::SparkPlan;
-use crate::jvm_bridge::jni_new_global_ref;
-use crate::{
-    errors::CometError,
-    jvm_bridge::{jni_call, jni_new_string},
-};
+use crate::{errors::CometError, jvm_bridge::jni_call};
 use datafusion::physical_plan::metrics::MetricValue;
-use jni::objects::{GlobalRef, JString};
+use datafusion_comet_proto::spark_metric::NativeMetricNode;
 use jni::{objects::JObject, JNIEnv};
+use prost::Message;
 use std::collections::HashMap;
 use std::sync::Arc;
 
@@ -34,10 +31,22 @@ pub fn update_comet_metric(
     env: &mut JNIEnv,
     metric_node: &JObject,
     spark_plan: &Arc<SparkPlan>,
-    metrics_jstrings: &mut HashMap<String, Arc<GlobalRef>>,
 ) -> Result<(), CometError> {
-    // combine all metrics from all native plans for this SparkPlan
-    let metrics = if spark_plan.additional_native_plans.is_empty() {
+    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) -> ())?;
+    }
+    Ok(())
+}
+
+pub fn to_native_metric_node(spark_plan: &Arc<SparkPlan>) -> Result<NativeMetricNode, CometError> {
+    let mut native_metric_node = NativeMetricNode {
+        metrics: HashMap::new(),
+        children: Vec::new(),
+    };
+
+    let node_metrics = if spark_plan.additional_native_plans.is_empty() {
         spark_plan.native_plan.metrics()
     } else {
         let mut metrics = spark_plan.native_plan.metrics().unwrap_or_default();
@@ -55,60 +64,21 @@ pub fn update_comet_metric(
         Some(metrics.aggregate_by_name())
     };
 
-    update_metrics(
-        env,
-        metric_node,
-        &metrics
-            .unwrap_or_default()
-            .iter()
-            .map(|m| m.value())
-            .map(|m| (m.name(), m.as_usize() as i64))
-            .collect::<Vec<_>>(),
-        metrics_jstrings,
-    )?;
+    // add metrics
+    node_metrics
+        .unwrap_or_default()
+        .iter()
+        .map(|m| m.value())
+        .map(|m| (m.name(), m.as_usize() as i64))
+        .for_each(|(name, value)| {
+            native_metric_node.metrics.insert(name.to_string(), value);
+        });
 
-    unsafe {
-        for (i, child_plan) in spark_plan.children().iter().enumerate() {
-            let child_metric_node: JObject = jni_call!(env,
-                comet_metric_node(metric_node).get_child_node(i as i32) -> JObject
-            )?;
-            if child_metric_node.is_null() {
-                continue;
-            }
-            update_comet_metric(env, &child_metric_node, child_plan, metrics_jstrings)?;
-        }
-    }
-    Ok(())
-}
+    // add children
+    spark_plan.children().iter().for_each(|child_plan| {
+        let child_node = to_native_metric_node(child_plan).unwrap();
+        native_metric_node.children.push(child_node);
+    });
 
-#[inline]
-fn update_metrics(
-    env: &mut JNIEnv,
-    metric_node: &JObject,
-    metric_values: &[(&str, i64)],
-    metrics_jstrings: &mut HashMap<String, Arc<GlobalRef>>,
-) -> Result<(), CometError> {
-    unsafe {
-        for &(name, value) in metric_values {
-            // Perform a lookup in the jstrings cache.
-            if let Some(map_global_ref) = metrics_jstrings.get(name) {
-                // Cache hit. Extract the jstring from the global ref.
-                let jobject = map_global_ref.as_obj();
-                let jstring = JString::from_raw(**jobject);
-                // Update the metrics using the jstring as a key.
-                jni_call!(env, comet_metric_node(metric_node).set(&jstring, value) -> ())?;
-            } else {
-                // Cache miss. Allocate a new string, promote to global ref, and insert into cache.
-                let local_jstring = jni_new_string!(env, &name)?;
-                let global_ref = jni_new_global_ref!(env, local_jstring)?;
-                let arc_global_ref = Arc::new(global_ref);
-                metrics_jstrings.insert(name.to_string(), Arc::clone(&arc_global_ref));
-                let jobject = arc_global_ref.as_obj();
-                let jstring = JString::from_raw(**jobject);
-                // Update the metrics using the jstring as a key.
-                jni_call!(env, comet_metric_node(metric_node).set(&jstring, value) -> ())?;
-            }
-        }
-    }
-    Ok(())
+    Ok(native_metric_node)
 }

native/core/src/jvm_bridge/comet_metric_node.rs

@@ -30,6 +30,8 @@ pub struct CometMetricNode<'a> {
     pub method_get_child_node_ret: ReturnType,
     pub method_set: JMethodID,
     pub method_set_ret: ReturnType,
+    pub method_set_all_from_bytes: JMethodID,
+    pub method_set_all_from_bytes_ret: ReturnType,
 }
 
 impl<'a> CometMetricNode<'a> {
@@ -47,6 +49,12 @@ impl<'a> CometMetricNode<'a> {
             method_get_child_node_ret: ReturnType::Object,
             method_set: env.get_method_id(Self::JVM_CLASS, "set", "(Ljava/lang/String;J)V")?,
             method_set_ret: ReturnType::Primitive(Primitive::Void),
+            method_set_all_from_bytes: env.get_method_id(
+                Self::JVM_CLASS,
+                "set_all_from_bytes",
+                "([B)V",
+            )?,
+            method_set_all_from_bytes_ret: ReturnType::Primitive(Primitive::Void),
             class,
         })
     }

native/core/src/jvm_bridge/mod.rs

@@ -46,13 +46,6 @@ macro_rules! jvalues {
     }}
 }
 
-/// Macro for create a new JNI string.
-macro_rules! jni_new_string {
-    ($env:expr, $value:expr) => {{
-        $crate::jvm_bridge::jni_map_error!($env, $env.new_string($value))
-    }};
-}
-
 /// Macro for calling a JNI method.
 /// The syntax is:
 /// jni_call!(env, comet_metric_node(metric_node).add(jname, value) -> ())?;
@@ -173,7 +166,6 @@ macro_rules! jni_new_global_ref {
 pub(crate) use jni_call;
 pub(crate) use jni_map_error;
 pub(crate) use jni_new_global_ref;
-pub(crate) use jni_new_string;
 pub(crate) use jni_static_call;
 pub(crate) use jvalues;
 

native/proto/build.rs

@@ -30,6 +30,7 @@ fn main() -> Result<()> {
     prost_build::Config::new().out_dir(out_dir).compile_protos(
         &[
             "src/proto/expr.proto",
+            "src/proto/metric.proto",
             "src/proto/partitioning.proto",
             "src/proto/operator.proto",
         ],

native/proto/src/lib.rs

@@ -36,3 +36,9 @@ pub mod spark_partitioning {
 pub mod spark_operator {
     include!(concat!("generated", "/spark.spark_operator.rs"));
 }
+
+// Include generated modules from .proto files.
+#[allow(missing_docs)]
+pub mod spark_metric {
+    include!(concat!("generated", "/spark.spark_metric.rs"));
+}

native/proto/src/proto/metric.proto

@@ -0,0 +1,29 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+
+
+syntax = "proto3";
+
+package spark.spark_metric;
+
+option java_package = "org.apache.comet.serde";
+
+message NativeMetricNode {
+  map<string, int64> metrics = 1;
+  repeated NativeMetricNode children = 2;
+}

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

@@ -19,11 +19,15 @@
 
 package org.apache.spark.sql.comet
 
+import scala.collection.JavaConverters._
+
 import org.apache.spark.SparkContext
 import org.apache.spark.internal.Logging
 import org.apache.spark.sql.execution.SparkPlan
 import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics}
 
+import org.apache.comet.serde.Metric
+
 /**
  * A node carrying SQL metrics from SparkPlan, and metrics of its children. Native code will call
  * [[getChildNode]] and [[set]] to update the metrics.
@@ -65,6 +69,20 @@ case class CometMetricNode(metrics: Map[String, SQLMetric], children: Seq[CometM
         logDebug(s"Non-existing metric: $metricName. Ignored")
     }
   }
+
+  private def set_all(metricNode: Metric.NativeMetricNode): Unit = {
+    metricNode.getMetricsMap.forEach((name, value) => {
+      set(name, value)
+    })
+    metricNode.getChildrenList.asScala.zip(children).foreach { case (child, childNode) =>
+      childNode.set_all(child)
+    }
+  }
+
+  def set_all_from_bytes(bytes: Array[Byte]): Unit = {
+    val metricNode = Metric.NativeMetricNode.parseFrom(bytes)
+    set_all(metricNode)
+  }
 }
 
 object CometMetricNode {