feat: Share native DataFusion plans across tasks within a Spark stage attempt

native/core/src/execution/jni_api.rs

@@ -329,6 +329,7 @@ pub fn get_runtime() -> Handle {
 /// Must not be called from within the runtime's own worker threads, otherwise the shutdown
 /// would deadlock/panic.
 pub fn release_runtime() {
+    super::shared_pipeline::clear();
     let runtime = TOKIO_RUNTIME.lock().take();
     if let Some(runtime) = runtime {
         runtime.shutdown_timeout(Duration::from_secs(3));
@@ -394,8 +395,10 @@ fn collect_op_names<'a>(op: &'a Operator, names: &mut std::collections::BTreeSet
 struct ExecutionContext {
     /// The id of the execution context.
     pub id: i64,
-    /// The deserialized Spark plan
-    pub spark_plan: Operator,
+    /// Immutable plan definition; may be shared across task attempts. Execution state stays local.
+    pub spark_plan: Arc<Operator>,
+    shared_plan_key: Option<Vec<u8>>,
+    shared_attempt: Option<Arc<super::shared_pipeline::AttemptState>>,
     /// The number of partitions
     pub partition_count: usize,
     /// The DataFusion root operator converted from the `spark_plan`
@@ -479,6 +482,7 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_createPlan(
     key_unwrapper_obj: JObject,
     task_context_obj: JObject,
     class_loader_obj: JObject,
+    shared_plan_scope: JString,
 ) -> jlong {
     try_unwrap_or_throw(&e, |env| {
         // Deserialize Spark configs
@@ -507,7 +511,24 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_createPlan(
 
             // Deserialize query plan
             let bytes = env.convert_byte_array(serialized_query)?;
-            let spark_plan = serde::deserialize_op(bytes.as_slice())?;
+            let spark_plan = Arc::new(serde::deserialize_op(bytes.as_slice())?);
+
+            let shared_plan_scope = shared_plan_scope.try_to_string(env)?;
+            let shared_plan_key = (!shared_plan_scope.is_empty()
+                && spark_config.get_bool(super::spark_config::COMET_EXEC_SHARED_PLAN_ENABLED)
+                && super::shared_pipeline::supports(&spark_plan))
+            .then(|| {
+                super::shared_pipeline::scoped_key(
+                    shared_plan_scope.as_bytes(),
+                    &super::shared_pipeline::cache_key(
+                        &bytes,
+                        &spark_config,
+                        batch_size,
+                        partition_count,
+                        task_cpus,
+                    ),
+                )
+            });
 
             let metrics = Arc::new(jni_new_global_ref!(env, metrics_node)?);
 
@@ -615,6 +636,8 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_createPlan(
             let exec_context = Box::new(ExecutionContext {
                 id,
                 spark_plan,
+                shared_plan_key,
+                shared_attempt: None,
                 partition_count: partition_count as usize,
                 root_op: None,
                 scans: vec![],
@@ -967,11 +990,37 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_executePlan(
                         .with_shuffle_partition_pusher(
                             exec_context.shuffle_partition_pusher.clone(),
                         );
-                let (scans, shuffle_scans, root_op) = planner.create_plan(
-                    &exec_context.spark_plan,
-                    &mut exec_context.input_sources.clone(),
-                    exec_context.partition_count,
-                )?;
+                let shared = exec_context.shared_plan_key.as_ref().and_then(|key| {
+                    super::shared_pipeline::try_build(|| {
+                        super::shared_pipeline::get_or_build(
+                            key,
+                            &exec_context.spark_plan,
+                            &exec_context.session_ctx,
+                            exec_context.partition_count,
+                        )
+                    })
+                });
+                let binding = shared
+                    .as_ref()
+                    .map(|shared| {
+                        shared.try_bind_plan(
+                            &planner,
+                            &mut exec_context.input_sources.clone(),
+                            &exec_context.spark_plan,
+                        )
+                    })
+                    .transpose()?
+                    .flatten();
+                let (scans, shuffle_scans, root_op) = if let Some((scans, attempt)) = binding {
+                    exec_context.shared_attempt = Some(attempt);
+                    (scans, vec![], Arc::clone(&shared.as_ref().unwrap().root))
+                } else {
+                    planner.create_plan(
+                        &exec_context.spark_plan,
+                        &mut exec_context.input_sources.clone(),
+                        exec_context.partition_count,
+                    )?
+                };
                 let physical_plan_time = start.elapsed();
 
                 exec_context.plan_creation_time += physical_plan_time;
@@ -984,10 +1033,17 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_executePlan(
                     info!("Comet native query plan:\n{formatted_plan_str:}");
                 }
 
-                let task_ctx = exec_context.session_ctx.task_ctx();
-                // Each Comet native execution corresponds to a single Spark partition,
-                // so we should always execute partition 0.
-                let stream = root_op.native_plan.execute(0, task_ctx)?;
+                let task_ctx = match &exec_context.shared_attempt {
+                    Some(attempt) => attempt.task_context(&exec_context.session_ctx),
+                    None => exec_context.session_ctx.task_ctx(),
+                };
+                // Shared trees execute Spark data partitions. Private plans, including
+                // retries and speculation, retain their sole local partition 0.
+                let native_partition = exec_context
+                    .shared_attempt
+                    .as_ref()
+                    .map_or(0, |attempt| attempt.partition());
+                let stream = root_op.native_plan.execute(native_partition, task_ctx)?;
 
                 if exec_context.scans.is_empty() && exec_context.shuffle_scans.is_empty() {
                     // No JVM data sources — spawn onto tokio so the executor
@@ -1162,7 +1218,12 @@ pub extern "system" fn Java_org_apache_comet_Native_releasePlan(
 fn update_metrics(env: &mut Env, exec_context: &mut ExecutionContext) -> CometResult<()> {
     if let Some(native_query) = &exec_context.root_op {
         let metrics = exec_context.metrics.as_obj();
-        update_comet_metric(env, metrics, native_query)
+        update_comet_metric(
+            env,
+            metrics,
+            native_query,
+            exec_context.shared_attempt.as_deref(),
+        )
     } else {
         Ok(())
     }

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

@@ -16,8 +16,10 @@
 // under the License.
 
 use crate::errors::CometError;
+use crate::execution::shared_pipeline::AttemptState;
 use crate::execution::spark_plan::SparkPlan;
-use datafusion::physical_plan::metrics::MetricValue;
+use datafusion::physical_plan::metrics::{MetricValue, MetricsSet};
+use datafusion::physical_plan::ExecutionPlan;
 use datafusion_comet_proto::spark_metric::NativeMetricNode;
 use jni::{objects::JObject, Env};
 use prost::Message;
@@ -31,26 +33,50 @@ pub(crate) fn update_comet_metric(
     env: &mut Env,
     metric_node: &JObject,
     spark_plan: &Arc<SparkPlan>,
+    attempt: Option<&AttemptState>,
 ) -> Result<(), CometError> {
     if metric_node.is_null() {
         return Ok(());
     }
 
-    let native_metric = to_native_metric_node(spark_plan);
-    let jbytes = env.byte_array_from_slice(&native_metric?.encode_to_vec())?;
+    let native_metric = match attempt {
+        Some(attempt) => to_native_metric_node_with(spark_plan, &|plan| attempt.metrics_for(plan)),
+        None => to_native_metric_node(spark_plan),
+    };
+    let mut native_metric = native_metric?;
+    if attempt.is_some() {
+        mark_shared_plan_tasks(&mut native_metric);
+    }
+    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) -> ()) }
 }
 
+// SQLMetric.set replaces this task's previous report; periodic updates must not count as binds.
+// This counts tasks bound to a shared tree, not registry hits or simultaneous users of the tree.
+fn mark_shared_plan_tasks(node: &mut NativeMetricNode) {
+    node.metrics.insert("shared_plan_tasks".to_string(), 1);
+    for child in &mut node.children {
+        mark_shared_plan_tasks(child);
+    }
+}
+
 pub(crate) fn to_native_metric_node(
     spark_plan: &Arc<SparkPlan>,
+) -> Result<NativeMetricNode, CometError> {
+    to_native_metric_node_with(spark_plan, &|plan| plan.metrics())
+}
+
+pub(crate) fn to_native_metric_node_with(
+    spark_plan: &Arc<SparkPlan>,
+    metrics_for: &dyn Fn(&Arc<dyn ExecutionPlan>) -> Option<MetricsSet>,
 ) -> Result<NativeMetricNode, CometError> {
     let node_metrics = if spark_plan.additional_native_plans.is_empty() {
-        spark_plan.native_plan.metrics()
+        metrics_for(&spark_plan.native_plan)
     } else {
-        let mut metrics = spark_plan.native_plan.metrics().unwrap_or_default();
+        let mut metrics = metrics_for(&spark_plan.native_plan).unwrap_or_default();
         for plan in &spark_plan.additional_native_plans {
-            let additional_metrics = plan.metrics().unwrap_or_default();
+            let additional_metrics = metrics_for(plan).unwrap_or_default();
             for c in additional_metrics.iter() {
                 match c.value() {
                     MetricValue::OutputRows(_) => {
@@ -82,7 +108,7 @@ pub(crate) fn to_native_metric_node(
         .for_each(|m| insert_metric_value(&mut native_metric_node.metrics, m.value()));
 
     for child_plan in children {
-        let child_node = to_native_metric_node(child_plan)?;
+        let child_node = to_native_metric_node_with(child_plan, metrics_for)?;
         native_metric_node.children.push(child_node);
     }
 

native/core/src/execution/mod.rs

@@ -24,6 +24,7 @@ pub(crate) mod metrics;
 pub mod operators;
 pub(crate) mod planner;
 pub mod serde;
+mod shared_pipeline;
 pub use datafusion_comet_shuffle as shuffle;
 mod memory_pools;
 pub(crate) mod sort;

native/core/src/execution/planner.rs

@@ -285,8 +285,12 @@ pub struct BinaryExprOptions {
 
 pub const TEST_EXEC_CONTEXT_ID: i64 = -1;
 
+/// Input boundaries collected only while building a shared template.
+pub(super) type InputPlans = Arc<parking_lot::Mutex<Vec<Arc<dyn ExecutionPlan>>>>;
+
 /// The query planner for converting Spark query plans to DataFusion query plans.
 pub struct PhysicalPlanner {
+    input_plans: Option<InputPlans>,
     // The execution context id of this planner.
     exec_context_id: i64,
     partition: i32,
@@ -319,6 +323,7 @@ impl PhysicalPlanner {
     pub fn new(session_ctx: Arc<SessionContext>, partition: i32) -> Self {
         Self {
             exec_context_id: TEST_EXEC_CONTEXT_ID,
+            input_plans: None,
             session_ctx,
             partition,
             query_context_registry: datafusion_comet_spark_expr::create_query_context_map(),
@@ -329,6 +334,20 @@ impl PhysicalPlanner {
         }
     }
 
+    /// Record task-input boundaries while compiling a shared template. Kept only by the
+    /// builder; cached nodes copy immutable properties and never retain these input plans.
+    pub(super) fn with_input_plans(mut self, inputs: InputPlans) -> Self {
+        self.input_plans = Some(inputs);
+        self
+    }
+
+    fn input_plan(&self, plan: Arc<dyn ExecutionPlan>) -> Arc<dyn ExecutionPlan> {
+        if let Some(inputs) = &self.input_plans {
+            inputs.lock().push(Arc::clone(&plan));
+        }
+        plan
+    }
+
     /// Load the SQL text pool from the root operator of the plan about to be planned. Must be
     /// called with the *root* operator: the JVM only populates the pool there, and
     /// `QueryContext.sql_text_idx` values are indices into it.
@@ -1623,7 +1642,11 @@ impl PhysicalPlanner {
                     return Ok((
                         vec![],
                         vec![],
-                        Arc::new(SparkPlan::new(spark_plan.plan_id, empty_exec, vec![])),
+                        Arc::new(SparkPlan::new(
+                            spark_plan.plan_id,
+                            self.input_plan(empty_exec),
+                            vec![],
+                        )),
                     ));
                 }
 
@@ -1741,7 +1764,11 @@ impl PhysicalPlanner {
                 Ok((
                     vec![],
                     vec![],
-                    Arc::new(SparkPlan::new(spark_plan.plan_id, scan, vec![])),
+                    Arc::new(SparkPlan::new(
+                        spark_plan.plan_id,
+                        self.input_plan(scan),
+                        vec![],
+                    )),
                 ))
             }
             OpStruct::CsvScan(scan) => {
@@ -1823,7 +1850,11 @@ impl PhysicalPlanner {
                 Ok((
                     vec![scan.clone()],
                     vec![],
-                    Arc::new(SparkPlan::new(spark_plan.plan_id, Arc::new(scan), vec![])),
+                    Arc::new(SparkPlan::new(
+                        spark_plan.plan_id,
+                        self.input_plan(Arc::new(scan)),
+                        vec![],
+                    )),
                 ))
             }
             OpStruct::IcebergScan(scan) => {