fix: Various metrics bug fixes and improvements

native/core/src/execution/datafusion/mod.rs

@@ -21,4 +21,5 @@ pub mod expressions;
 mod operators;
 pub mod planner;
 pub mod shuffle_writer;
+pub(crate) mod spark_plan;
 mod util;

native/core/src/execution/datafusion/shuffle_writer.rs

@@ -139,7 +139,6 @@ impl ExecutionPlan for ShuffleWriterExec {
     ) -> Result<SendableRecordBatchStream> {
         let input = self.input.execute(partition, Arc::clone(&context))?;
         let metrics = ShuffleRepartitionerMetrics::new(&self.metrics, 0);
-        let jvm_fetch_time = MetricBuilder::new(&self.metrics).subset_time("jvm_fetch_time", 0);
 
         Ok(Box::pin(RecordBatchStreamAdapter::new(
             self.schema(),
@@ -152,7 +151,6 @@ impl ExecutionPlan for ShuffleWriterExec {
                     self.partitioning.clone(),
                     metrics,
                     context,
-                    jvm_fetch_time,
                 )
                 .map_err(|e| ArrowError::ExternalError(Box::new(e))),
             )
@@ -1094,7 +1092,6 @@ async fn external_shuffle(
     partitioning: Partitioning,
     metrics: ShuffleRepartitionerMetrics,
     context: Arc<TaskContext>,
-    jvm_fetch_time: Time,
 ) -> Result<SendableRecordBatchStream> {
     let schema = input.schema();
     let mut repartitioner = ShuffleRepartitioner::new(
@@ -1109,10 +1106,7 @@ async fn external_shuffle(
     );
 
     loop {
-        let mut timer = jvm_fetch_time.timer();
         let b = input.next().await;
-        timer.stop();
-
         match b {
             Some(batch_result) => {
                 // Block on the repartitioner to insert the batch and shuffle the rows

native/core/src/execution/datafusion/spark_plan.rs

@@ -0,0 +1,102 @@
+// 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.
+
+use crate::execution::operators::{CopyExec, ScanExec};
+use arrow_schema::SchemaRef;
+use datafusion::physical_plan::ExecutionPlan;
+use std::sync::Arc;
+
+/// Wrapper around a native plan that maps to a Spark plan and can optionally contain
+/// references to other native plans that should contribute to the Spark SQL metrics
+/// for the root plan (such as CopyExec and ScanExec nodes)
+#[derive(Debug, Clone)]
+pub(crate) struct SparkPlan {
+    /// Spark plan ID which is passed down in the protobuf
+    pub(crate) plan_id: u32,
+    /// The root native plan that was generated for this Spark plan
+    pub(crate) native_plan: Arc<dyn ExecutionPlan>,
+    /// Child Spark plans
+    pub(crate) children: Vec<Arc<SparkPlan>>,
+    /// Additional native plans that were generated for this Spark plan that we need
+    /// to collect metrics for (such as CopyExec and ScanExec)
+    pub(crate) additional_native_plans: Vec<Arc<dyn ExecutionPlan>>,
+}
+
+impl SparkPlan {
+    /// Create a SparkPlan that consists of a single native plan
+    pub(crate) fn new(
+        plan_id: u32,
+        native_plan: Arc<dyn ExecutionPlan>,
+        children: Vec<Arc<SparkPlan>>,
+    ) -> Self {
+        let mut additional_native_plans: Vec<Arc<dyn ExecutionPlan>> = vec![];
+        for child in &children {
+            collect_additional_plans(Arc::clone(&child.native_plan), &mut additional_native_plans);
+        }
+        Self {
+            plan_id,
+            native_plan,
+            children,
+            additional_native_plans,
+        }
+    }
+
+    /// Create a SparkPlan that consists of more than one native plan
+    pub(crate) fn new_with_additional(
+        plan_id: u32,
+        native_plan: Arc<dyn ExecutionPlan>,
+        children: Vec<Arc<SparkPlan>>,
+        additional_native_plans: Vec<Arc<dyn ExecutionPlan>>,
+    ) -> Self {
+        let mut accum: Vec<Arc<dyn ExecutionPlan>> = vec![];
+        for plan in &additional_native_plans {
+            accum.push(Arc::clone(plan));
+        }
+        for child in &children {
+            collect_additional_plans(Arc::clone(&child.native_plan), &mut accum);
+        }
+        Self {
+            plan_id,
+            native_plan,
+            children,
+            additional_native_plans: accum,
+        }
+    }
+
+    /// Get the schema of the native plan
+    pub(crate) fn schema(&self) -> SchemaRef {
+        self.native_plan.schema()
+    }
+
+    /// Get the child SparkPlan instances
+    pub(crate) fn children(&self) -> &Vec<Arc<SparkPlan>> {
+        &self.children
+    }
+}
+
+fn collect_additional_plans(
+    child: Arc<dyn ExecutionPlan>,
+    additional_native_plans: &mut Vec<Arc<dyn ExecutionPlan>>,
+) {
+    if child.as_any().is::<CopyExec>() {
+        additional_native_plans.push(Arc::clone(&child));
+        // CopyExec may be wrapping a ScanExec
+        collect_additional_plans(Arc::clone(child.children()[0]), additional_native_plans);
+    } else if child.as_any().is::<ScanExec>() {
+        additional_native_plans.push(Arc::clone(&child));
+    }
+}

native/core/src/execution/jni_api.rs

@@ -24,7 +24,7 @@ use datafusion::{
         disk_manager::DiskManagerConfig,
         runtime_env::{RuntimeConfig, RuntimeEnv},
     },
-    physical_plan::{display::DisplayableExecutionPlan, ExecutionPlan, SendableRecordBatchStream},
+    physical_plan::{display::DisplayableExecutionPlan, SendableRecordBatchStream},
     prelude::{SessionConfig, SessionContext},
 };
 use futures::poll;
@@ -59,6 +59,7 @@ use jni::{
 };
 use tokio::runtime::Runtime;
 
+use crate::execution::datafusion::spark_plan::SparkPlan;
 use crate::execution::operators::ScanExec;
 use log::info;
 
@@ -69,7 +70,7 @@ struct ExecutionContext {
     /// The deserialized Spark plan
     pub spark_plan: Operator,
     /// The DataFusion root operator converted from the `spark_plan`
-    pub root_op: Option<Arc<dyn ExecutionPlan>>,
+    pub root_op: Option<Arc<SparkPlan>>,
     /// The input sources for the DataFusion plan
     pub scans: Vec<ScanExec>,
     /// The global reference of input sources for the DataFusion plan
@@ -360,7 +361,7 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_executePlan(
 
             if exec_context.explain_native {
                 let formatted_plan_str =
-                    DisplayableExecutionPlan::new(root_op.as_ref()).indent(true);
+                    DisplayableExecutionPlan::new(root_op.native_plan.as_ref()).indent(true);
                 info!("Comet native query plan:\n{formatted_plan_str:}");
             }
 
@@ -369,6 +370,7 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_executePlan(
                 .root_op
                 .as_ref()
                 .unwrap()
+                .native_plan
                 .execute(0, task_ctx)?;
             exec_context.stream = Some(stream);
         } else {
@@ -400,7 +402,8 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_executePlan(
                     if exec_context.explain_native {
                         if let Some(plan) = &exec_context.root_op {
                             let formatted_plan_str =
-                                DisplayableExecutionPlan::with_metrics(plan.as_ref()).indent(true);
+                                DisplayableExecutionPlan::with_metrics(plan.native_plan.as_ref())
+                                    .indent(true);
                             info!(
                                 "Comet native query plan with metrics:\
                             \n[Stage {} Partition {}] plan creation (including CometScans fetching first batches) took {:?}:\

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

@@ -15,31 +15,50 @@
 // specific language governing permissions and limitations
 // under the License.
 
+use crate::execution::datafusion::spark_plan::SparkPlan;
 use crate::jvm_bridge::jni_new_global_ref;
 use crate::{
     errors::CometError,
     jvm_bridge::{jni_call, jni_new_string},
 };
-use datafusion::physical_plan::ExecutionPlan;
+use datafusion::physical_plan::metrics::MetricValue;
 use jni::objects::{GlobalRef, JString};
 use jni::{objects::JObject, JNIEnv};
 use std::collections::HashMap;
 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
-/// DataFusion execution plan and pushed to the Java side through JNI.
+/// native execution plan and pushed to the Java side through JNI.
 pub fn update_comet_metric(
     env: &mut JNIEnv,
     metric_node: &JObject,
-    execution_plan: &Arc<dyn ExecutionPlan>,
+    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() {
+        spark_plan.native_plan.metrics()
+    } else {
+        let mut metrics = spark_plan.native_plan.metrics().unwrap_or_default();
+        for plan in &spark_plan.additional_native_plans {
+            let additional_metrics = plan.metrics().unwrap_or_default();
+            for c in additional_metrics.iter() {
+                match c.value() {
+                    MetricValue::OutputRows(_) => {
+                        // we do not want to double count output rows
+                    }
+                    _ => metrics.push(c.to_owned()),
+                }
+            }
+        }
+        Some(metrics.aggregate_by_name())
+    };
+
     update_metrics(
         env,
         metric_node,
-        &execution_plan
-            .metrics()
+        &metrics
             .unwrap_or_default()
             .iter()
             .map(|m| m.value())
@@ -49,7 +68,7 @@ pub fn update_comet_metric(
     )?;
 
     unsafe {
-        for (i, child_plan) in execution_plan.children().iter().enumerate() {
+        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
             )?;

native/core/src/execution/operators/scan.rs

@@ -77,6 +77,8 @@ pub struct ScanExec {
     metrics: ExecutionPlanMetricsSet,
     /// Baseline metrics
     baseline_metrics: BaselineMetrics,
+    /// Timer
+    jvm_fetch_time: Time,
 }
 
 impl ScanExec {
@@ -88,6 +90,7 @@ impl ScanExec {
     ) -> Result<Self, CometError> {
         let metrics_set = ExecutionPlanMetricsSet::default();
         let baseline_metrics = BaselineMetrics::new(&metrics_set, 0);
+        let jvm_fetch_time = MetricBuilder::new(&metrics_set).subset_time("jvm_fetch_time", 0);
 
         // Scan's schema is determined by the input batch, so we need to set it before execution.
         // Note that we determine if arrays are dictionary-encoded based on the
@@ -97,8 +100,12 @@ impl ScanExec {
         // Dictionary-encoded primitive arrays are always unpacked.
         let first_batch = if let Some(input_source) = input_source.as_ref() {
             let mut timer = baseline_metrics.elapsed_compute().timer();
-            let batch =
-                ScanExec::get_next(exec_context_id, input_source.as_obj(), data_types.len())?;
+            let batch = ScanExec::get_next(
+                exec_context_id,
+                input_source.as_obj(),
+                data_types.len(),
+                &jvm_fetch_time,
+            )?;
             timer.stop();
             batch
         } else {
@@ -124,6 +131,7 @@ impl ScanExec {
             cache,
             metrics: metrics_set,
             baseline_metrics,
+            jvm_fetch_time,
             schema,
         })
     }
@@ -171,6 +179,7 @@ impl ScanExec {
                 self.exec_context_id,
                 self.input_source.as_ref().unwrap().as_obj(),
                 self.data_types.len(),
+                &self.jvm_fetch_time,
             )?;
             *current_batch = Some(next_batch);
         }
@@ -185,6 +194,7 @@ impl ScanExec {
         exec_context_id: i64,
         iter: &JObject,
         num_cols: usize,
+        jvm_fetch_time: &Time,
     ) -> Result<InputBatch, CometError> {
         if exec_context_id == TEST_EXEC_CONTEXT_ID {
             // This is a unit test. We don't need to call JNI.
@@ -197,6 +207,7 @@ impl ScanExec {
                 exec_context_id
             ))));
         }
+        let mut timer = jvm_fetch_time.timer();
 
         let mut env = JVMClasses::get_env()?;
 
@@ -255,6 +266,8 @@ impl ScanExec {
             }
         }
 
+        timer.stop();
+
         Ok(InputBatch::new(inputs, Some(num_rows as usize)))
     }
 }
@@ -324,7 +337,7 @@ impl ExecutionPlan for ScanExec {
             self.clone(),
             self.schema(),
             partition,
-            self.baseline_metrics.clone(),
+            self.jvm_fetch_time.clone(),
         )))
     }
 
@@ -365,28 +378,23 @@ struct ScanStream<'a> {
     scan: ScanExec,
     /// Schema representing the data
     schema: SchemaRef,
-    /// Metrics
-    baseline_metrics: BaselineMetrics,
     /// Cast options
     cast_options: CastOptions<'a>,
     /// elapsed time for casting columns to different data types during scan
     cast_time: Time,
+    /// elapsed time for fetching arrow arrays from JVM
+    jvm_fetch_time: Time,
 }
 
 impl<'a> ScanStream<'a> {
-    pub fn new(
-        scan: ScanExec,
-        schema: SchemaRef,
-        partition: usize,
-        baseline_metrics: BaselineMetrics,
-    ) -> Self {
+    pub fn new(scan: ScanExec, schema: SchemaRef, partition: usize, jvm_fetch_time: Time) -> Self {
         let cast_time = MetricBuilder::new(&scan.metrics).subset_time("cast_time", partition);
         Self {
             scan,
             schema,
-            baseline_metrics,
             cast_options: CastOptions::default(),
             cast_time,
+            jvm_fetch_time,
         }
     }
 
@@ -426,7 +434,7 @@ impl<'a> Stream for ScanStream<'a> {
     type Item = DataFusionResult<RecordBatch>;
 
     fn poll_next(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<Option<Self::Item>> {
-        let mut timer = self.baseline_metrics.elapsed_compute().timer();
+        let mut timer = self.jvm_fetch_time.timer();
         let mut scan_batch = self.scan.batch.try_lock().unwrap();
 
         let input_batch = &*scan_batch;
@@ -440,7 +448,6 @@ impl<'a> Stream for ScanStream<'a> {
         let result = match input_batch {
             InputBatch::EOF => Poll::Ready(None),
             InputBatch::Batch(columns, num_rows) => {
-                self.baseline_metrics.record_output(*num_rows);
                 let maybe_batch = self.build_record_batch(columns, *num_rows);
                 Poll::Ready(Some(maybe_batch))
             }