From bd2c151a02c823c92f315771a7a2153f7d03932d Mon Sep 17 00:00:00 2001 From: Liang-Chi Hsieh Date: Sun, 20 Sep 2026 23:33:16 -0700 Subject: [PATCH 1/3] feat: share native DataFusion plans within Spark stage attempts (#1204) --- native/core/src/execution/jni_api.rs | 89 +- native/core/src/execution/metrics/utils.rs | 25 +- native/core/src/execution/mod.rs | 2 + native/core/src/execution/plan_cache.rs | 455 ++++ native/core/src/execution/planner.rs | 37 +- native/core/src/execution/shared_pipeline.rs | 2154 +++++++++++++++++ native/core/src/execution/spark_config.rs | 2 + .../scala/org/apache/comet/CometConf.scala | 24 + .../org/apache/comet/CometExecIterator.scala | 23 +- .../main/scala/org/apache/comet/Native.scala | 3 +- .../apache/spark/sql/comet/CometExecRDD.scala | 6 +- .../apache/comet/exec/CometExecSuite.scala | 113 +- .../CometExecIteratorLifecycleSuite.scala | 60 +- 13 files changed, 2959 insertions(+), 34 deletions(-) create mode 100644 native/core/src/execution/plan_cache.rs create mode 100644 native/core/src/execution/shared_pipeline.rs diff --git a/native/core/src/execution/jni_api.rs b/native/core/src/execution/jni_api.rs index faf4d5dfda0..c5cc8c65fa2 100644 --- a/native/core/src/execution/jni_api.rs +++ b/native/core/src/execution/jni_api.rs @@ -329,6 +329,8 @@ 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::plan_cache::clear_plan_cache(); + super::shared_pipeline::clear(); let runtime = TOKIO_RUNTIME.lock().take(); if let Some(runtime) = runtime { runtime.shutdown_timeout(Duration::from_secs(3)); @@ -394,8 +396,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, + shared_plan_key: Option>, + shared_attempt: Option>, /// The number of partitions pub partition_count: usize, /// The DataFusion root operator converted from the `spark_plan` @@ -479,6 +483,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 +512,27 @@ 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 = super::plan_cache::decode_plan( + bytes.as_slice(), + spark_config.get_bool(super::spark_config::COMET_EXEC_PLAN_CACHE_ENABLED), + )?; + + 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( + &super::shared_pipeline::cache_bytes(&spark_plan, &bytes), + &spark_config, + batch_size, + partition_count, + task_cpus, + ), + ) + }); let metrics = Arc::new(jni_new_global_ref!(env, metrics_node)?); @@ -615,6 +640,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 +994,35 @@ 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 (scans, shuffle_scans, root_op) = + if let Some(key) = &exec_context.shared_plan_key { + let shared = super::shared_pipeline::get_or_build( + key, + &exec_context.spark_plan, + &exec_context.session_ctx, + exec_context.partition_count, + )?; + if let Some((scans, attempt)) = shared.try_bind_plan( + &planner, + &mut exec_context.input_sources.clone(), + &exec_context.spark_plan, + )? { + exec_context.shared_attempt = Some(attempt); + (scans, vec![], Arc::clone(&shared.root)) + } else { + planner.create_plan( + &exec_context.spark_plan, + &mut exec_context.input_sources.clone(), + exec_context.partition_count, + )? + } + } 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 +1035,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 +1220,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(()) } diff --git a/native/core/src/execution/metrics/utils.rs b/native/core/src/execution/metrics/utils.rs index 4e51dccb279..4634581ec79 100644 --- a/native/core/src/execution/metrics/utils.rs +++ b/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,12 +33,16 @@ pub(crate) fn update_comet_metric( env: &mut Env, metric_node: &JObject, spark_plan: &Arc, + attempt: Option<&AttemptState>, ) -> Result<(), CometError> { if metric_node.is_null() { return Ok(()); } - let native_metric = to_native_metric_node(spark_plan); + 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 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) -> ()) } @@ -44,13 +50,20 @@ pub(crate) fn update_comet_metric( pub(crate) fn to_native_metric_node( spark_plan: &Arc, +) -> Result { + to_native_metric_node_with(spark_plan, &|plan| plan.metrics()) +} + +pub(crate) fn to_native_metric_node_with( + spark_plan: &Arc, + metrics_for: &dyn Fn(&Arc) -> Option, ) -> Result { 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 +95,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); } diff --git a/native/core/src/execution/mod.rs b/native/core/src/execution/mod.rs index 55da2c733aa..327c7aeb5c3 100644 --- a/native/core/src/execution/mod.rs +++ b/native/core/src/execution/mod.rs @@ -22,8 +22,10 @@ pub mod jni_api; pub(crate) mod merge_as_partial; pub(crate) mod metrics; pub mod operators; +mod plan_cache; pub(crate) mod planner; pub mod serde; +mod shared_pipeline; pub use datafusion_comet_shuffle as shuffle; mod memory_pools; pub(crate) mod sort; diff --git a/native/core/src/execution/plan_cache.rs b/native/core/src/execution/plan_cache.rs new file mode 100644 index 00000000000..f9eccbdfe65 --- /dev/null +++ b/native/core/src/execution/plan_cache.rs @@ -0,0 +1,455 @@ +// 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. + +//! Executor-local reuse of immutable protobuf plans and bounded single-flight cache (#1204). +//! +//! The definition cache here owns protobuf data only. `shared_pipeline` separately uses the +//! generic cache for audited immutable physical trees. Neither cache may retain task resources. +//! Different partitions' scan payloads must remain different definition-cache keys. + +use std::collections::HashMap; +use std::sync::{Arc, LazyLock}; +use std::time::Instant; + +use datafusion_comet_proto::spark_operator::Operator; +use once_cell::sync::OnceCell; +use parking_lot::Mutex; + +use super::operators::ExecutionError; +use super::serde::deserialize_op; + +const MAX_ENTRIES: usize = 64; +const MAX_ENCODED_BYTES: usize = 8 * 1024 * 1024; + +// Process-local because tasks do not share a SessionContext. This cache owns only immutable +// protobuf data, never storage clients, JNI references or task resources. Exact bytes are +// compared, so neither hash collisions nor another query's configuration can change the decoded +// result. Retention is bounded by entry count and encoded bytes, and release_runtime clears it. +// The byte budget accounts for keys, not the decoded Rust heap (which can be larger). +static PLAN_CACHE: LazyLock> = + LazyLock::new(|| PlanCache::new(MAX_ENTRIES, MAX_ENCODED_BYTES)); + +pub(super) fn decode_plan( + bytes: &[u8], + cache_enabled: bool, +) -> Result, ExecutionError> { + if cache_enabled { + PLAN_CACHE.get_or_build(bytes, deserialize_op) + } else { + deserialize_op(bytes).map(Arc::new) + } +} + +pub(super) fn clear_plan_cache() { + PLAN_CACHE.clear(); +} + +type PlanSlot = Arc>>; + +struct CacheEntry { + plan: PlanSlot, + last_used: Instant, +} + +struct CacheState { + entries: HashMap, CacheEntry>, + encoded_bytes: usize, +} + +impl Default for CacheState { + fn default() -> Self { + Self { + entries: HashMap::new(), + encoded_bytes: 0, + } + } +} + +pub(super) struct PlanCache { + state: Mutex>, + max_entries: usize, + max_encoded_bytes: usize, +} + +impl PlanCache { + pub(super) fn new(max_entries: usize, max_encoded_bytes: usize) -> Self { + Self { + state: Mutex::new(CacheState::default()), + max_entries, + max_encoded_bytes, + } + } + + pub(super) fn get_or_build( + &self, + bytes: &[u8], + build: impl FnOnce(&[u8]) -> Result, + ) -> Result, ExecutionError> { + // A large plan must not evict the entire cache or circumvent the admission budget. + if self.max_entries == 0 || bytes.len() > self.max_encoded_bytes { + return build(bytes).map(Arc::new); + } + + let slot = { + let mut state = self.state.lock(); + if let Some(entry) = state.entries.get_mut(bytes) { + entry.last_used = Instant::now(); + Arc::clone(&entry.plan) + } else { + while state.entries.len() >= self.max_entries + || bytes.len() > self.max_encoded_bytes - state.encoded_bytes + { + let oldest = state + .entries + .iter() + .min_by_key(|(_, entry)| entry.last_used) + .map(|(key, _)| Arc::clone(key)) + .expect("an over-budget cache has an entry to evict"); + state.entries.remove(&oldest); + state.encoded_bytes -= oldest.len(); + } + let slot = Arc::new(OnceCell::new()); + state.entries.insert( + Arc::from(bytes), + CacheEntry { + plan: Arc::clone(&slot), + last_used: Instant::now(), + }, + ); + state.encoded_bytes += bytes.len(); + slot + } + }; + + // Only one successful build per resident entry, even on concurrent first use. The + // cache mutex is not held while building or waiting; unrelated plans can make progress. + let result = slot.get_or_try_init(|| build(bytes).map(Arc::new)).cloned(); + if result.is_err() { + // Do not retain malformed plans. An eviction/clear and a new insertion may have + // happened while building: this failed caller must not remove the replacement. + let mut state = self.state.lock(); + if state + .entries + .get(bytes) + .is_some_and(|entry| Arc::ptr_eq(&entry.plan, &slot)) + { + state.entries.remove(bytes); + state.encoded_bytes -= bytes.len(); + } + } + result + } + + pub(super) fn clear(&self) { + let old = std::mem::take(&mut *self.state.lock()); + // Active attempts retain their own Arc. Dropping the cache never invalidates a task. + drop(old); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use prost::Message; + use std::sync::atomic::{AtomicUsize, Ordering}; + use std::sync::{mpsc, Barrier}; + use std::time::Duration; + + fn encoded(id: u32) -> Vec { + Operator { + plan_id: id, + ..Default::default() + } + .encode_to_vec() + } + + #[test] + fn sequential_attempts_reuse_definition_after_previous_attempt_finishes() { + let cache = PlanCache::new(2, 1024); + let bytes = encoded(1); + let first = cache.get_or_build(&bytes, deserialize_op).unwrap(); + let weak = Arc::downgrade(&first); + drop(first); + let second = cache + .get_or_build(&bytes, |_| panic!("decoded again between task waves")) + .unwrap(); + assert!(Arc::ptr_eq(&weak.upgrade().unwrap(), &second)); + } + + #[test] + fn concurrent_first_use_decodes_once() { + let cache = PlanCache::new(2, 1024); + let decodes = AtomicUsize::new(0); + let start = Barrier::new(8); + let plans = std::thread::scope(|scope| { + let handles: Vec<_> = (0..8) + .map(|_| { + scope.spawn(|| { + start.wait(); + cache + .get_or_build(&encoded(7), |bytes| { + decodes.fetch_add(1, Ordering::SeqCst); + deserialize_op(bytes) + }) + .unwrap() + }) + }) + .collect(); + handles + .into_iter() + .map(|h| h.join().unwrap()) + .collect::>() + }); + assert_eq!(decodes.load(Ordering::SeqCst), 1); + assert!(plans.iter().all(|p| Arc::ptr_eq(p, &plans[0]))); + } + + #[test] + fn different_plan_bytes_do_not_alias() { + let cache = PlanCache::new(2, 1024); + let a = cache.get_or_build(&encoded(1), deserialize_op).unwrap(); + // Spark operator IDs alone are not a key: the payload can change across plans. + let different_payload = Operator { + plan_id: 1, + sql_text_pool: vec!["a different query".into()], + ..Default::default() + } + .encode_to_vec(); + let b = cache + .get_or_build(&different_payload, deserialize_op) + .unwrap(); + assert!(!Arc::ptr_eq(&a, &b)); + assert_eq!(a.plan_id, b.plan_id); + assert!(a.sql_text_pool.is_empty()); + assert_eq!(b.sql_text_pool, ["a different query"]); + } + + #[test] + fn least_recently_used_entry_is_evicted_without_invalidating_its_task() { + let cache = PlanCache::new(2, 1024); + let a = cache.get_or_build(&encoded(1), deserialize_op).unwrap(); + let b = cache.get_or_build(&encoded(2), deserialize_op).unwrap(); + cache.get_or_build(&encoded(1), deserialize_op).unwrap(); + cache.get_or_build(&encoded(3), deserialize_op).unwrap(); + let a_again = cache.get_or_build(&encoded(1), deserialize_op).unwrap(); + let b_again = cache.get_or_build(&encoded(2), deserialize_op).unwrap(); + assert!(Arc::ptr_eq(&a, &a_again)); + assert!(!Arc::ptr_eq(&b, &b_again)); + assert_eq!(b.plan_id, 2); + } + + #[test] + fn encoded_byte_budget_evicts_and_oversized_plans_bypass_cache() { + let bytes = encoded(1); + let cache = PlanCache::new(10, bytes.len()); + let a = cache.get_or_build(&bytes, deserialize_op).unwrap(); + cache.get_or_build(&encoded(2), deserialize_op).unwrap(); + let a_again = cache.get_or_build(&bytes, deserialize_op).unwrap(); + assert!(!Arc::ptr_eq(&a, &a_again)); + let large = Operator { + sql_text_pool: vec!["longer than the admission budget".to_owned()], + ..Default::default() + } + .encode_to_vec(); + let large_a = cache.get_or_build(&large, deserialize_op).unwrap(); + let large_b = cache.get_or_build(&large, deserialize_op).unwrap(); + assert!(!Arc::ptr_eq(&large_a, &large_b)); + assert!(Arc::ptr_eq( + &a_again, + &cache.get_or_build(&bytes, deserialize_op).unwrap() + )); + assert_eq!(cache.state.lock().encoded_bytes, bytes.len()); + } + + #[test] + fn malformed_plan_does_not_poison_or_occupy_cache() { + let cache = PlanCache::new(2, 1024); + for _ in 0..2 { + assert!(cache.get_or_build(&[0xff], deserialize_op).is_err()); + assert!(cache.state.lock().entries.is_empty()); + assert_eq!(cache.state.lock().encoded_bytes, 0); + } + assert_eq!( + cache + .get_or_build(&encoded(1), deserialize_op) + .unwrap() + .plan_id, + 1 + ); + } + + #[test] + fn clear_releases_idle_plans_but_keeps_active_attempts_valid() { + let cache = PlanCache::new(2, 1024); + let active = cache.get_or_build(&encoded(1), deserialize_op).unwrap(); + let idle = cache.get_or_build(&encoded(2), deserialize_op).unwrap(); + let idle_weak = Arc::downgrade(&idle); + drop(idle); + cache.clear(); + assert!(idle_weak.upgrade().is_none()); + assert_eq!(active.plan_id, 1); + assert_eq!(cache.state.lock().encoded_bytes, 0); + assert!(!Arc::ptr_eq( + &active, + &cache.get_or_build(&encoded(1), deserialize_op).unwrap() + )); + } + + #[test] + fn decoding_does_not_lock_out_unrelated_plans() { + let cache = PlanCache::new(2, 1024); + let (started_tx, started_rx) = mpsc::channel(); + let (release_tx, release_rx) = mpsc::channel(); + let (other_tx, other_rx) = mpsc::channel(); + std::thread::scope(|scope| { + let cache = &cache; + scope.spawn(move || { + cache + .get_or_build(&encoded(1), |bytes| { + started_tx.send(()).unwrap(); + release_rx.recv().unwrap(); + deserialize_op(bytes) + }) + .unwrap() + }); + started_rx.recv_timeout(Duration::from_secs(10)).unwrap(); + scope.spawn(|| { + cache.get_or_build(&encoded(2), deserialize_op).unwrap(); + other_tx.send(()).unwrap(); + }); + let progressed = other_rx.recv_timeout(Duration::from_secs(10)); + release_tx.send(()).unwrap(); + assert!( + progressed.is_ok(), + "unrelated decode blocked on the cache mutex" + ); + }); + } + + #[test] + fn failure_after_clear_does_not_remove_replacement() { + let cache = PlanCache::new(2, 1024); + let bytes = encoded(1); + let mut replacement = None; + let result = cache.get_or_build(&bytes, |_| { + cache.clear(); + replacement = Some(cache.get_or_build(&bytes, deserialize_op).unwrap()); + deserialize_op(&[0xff]) + }); + assert!(result.is_err()); + assert!(Arc::ptr_eq( + &replacement.unwrap(), + &cache + .get_or_build(&bytes, |_| panic!("replacement removed")) + .unwrap() + )); + } + + #[test] + fn shared_definition_keeps_partition_counters_inputs_and_metrics_attempt_local() { + use crate::execution::operators::InputBatch; + use crate::execution::planner::PhysicalPlanner; + use arrow::array::{Int32Array, Int64Array}; + use datafusion::prelude::SessionContext; + use datafusion_comet_proto::spark_expression::{ + expr::ExprStruct, DataType, EmptyExpr, Expr, + }; + use datafusion_comet_proto::spark_operator::{operator::OpStruct, Projection, Scan}; + use futures::StreamExt; + + let bytes = Operator { + children: vec![Operator { + op_struct: Some(OpStruct::Scan(Scan { + fields: vec![DataType { + type_id: 4, + type_info: None, + }], + source: "attempt-isolation".into(), + })), + ..Default::default() + }], + op_struct: Some(OpStruct::Projection(Projection { + project_list: vec![ + Expr { + expr_struct: Some(ExprStruct::SparkPartitionId(EmptyExpr {})), + ..Default::default() + }, + Expr { + expr_struct: Some(ExprStruct::MonotonicallyIncreasingId(EmptyExpr {})), + ..Default::default() + }, + ], + })), + ..Default::default() + } + .encode_to_vec(); + let cache = PlanCache::new(2, 4096); + let shared = cache.get_or_build(&bytes, deserialize_op).unwrap(); + let runtime = tokio::runtime::Runtime::new().unwrap(); + // The final attempt retries partition 7 while the earlier plans are still alive. + let mut roots = vec![]; + for (partition, rows) in [(7, 2), (9, 3), (7, 1)] { + let definition = cache.get_or_build(&bytes, deserialize_op).unwrap(); + assert!(Arc::ptr_eq(&shared, &definition)); + let session = Arc::new(SessionContext::new()); + let planner = PhysicalPlanner::new(Arc::clone(&session), partition); + let (mut scans, _, root) = planner.create_plan(&definition, &mut vec![], 10).unwrap(); + let mut stream = root.native_plan.execute(0, session.task_ctx()).unwrap(); + scans[0].set_input_batch(InputBatch::Batch( + vec![Arc::new(Int64Array::from(vec![42; rows]))], + rows, + )); + let batch = runtime.block_on(stream.next()).unwrap().unwrap(); + let partitions = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let ids = batch + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(partitions.values().as_ref(), vec![partition; rows]); + assert_eq!( + ids.values().as_ref(), + (0..rows) + .map(|i| ((partition as i64) << 33) + i as i64) + .collect::>() + ); + scans[0].set_input_batch(InputBatch::EOF); + assert!(runtime.block_on(stream.next()).is_none()); + assert_eq!( + root.native_plan.metrics().unwrap().output_rows(), + Some(rows) + ); + roots.push(root); + } + assert!(!Arc::ptr_eq(&roots[0].native_plan, &roots[2].native_plan)); + assert_eq!( + roots[0].native_plan.metrics().unwrap().output_rows(), + Some(2) + ); + } + + #[test] + fn disabled_cache_does_not_share_definitions() { + let a = decode_plan(&encoded(1), false).unwrap(); + let b = decode_plan(&encoded(1), false).unwrap(); + assert!(!Arc::ptr_eq(&a, &b)); + } +} diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index 8f030da455b..3ff7637fbb2 100644 --- a/native/core/src/execution/planner.rs +++ b/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>>>; + /// The query planner for converting Spark query plans to DataFusion query plans. pub struct PhysicalPlanner { + input_plans: Option, // 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, 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) -> Arc { + 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) => { diff --git a/native/core/src/execution/shared_pipeline.rs b/native/core/src/execution/shared_pipeline.rs new file mode 100644 index 00000000000..62603d4656c --- /dev/null +++ b/native/core/src/execution/shared_pipeline.rs @@ -0,0 +1,2154 @@ +// 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. + +//! Stage-attempt scoped DataFusion trees, using unmodified upstream operators. +//! Inputs are task-local; each shared partition executes at most once. + +use super::operators::{ExecutionError, ScanExec}; +use super::planner::PhysicalPlanner; +use super::spark_plan::SparkPlan; +use arrow::array::RecordBatch; +use arrow::datatypes::SchemaRef; +use datafusion::common::{internal_err, tree_node::TreeNodeRecursion, Result}; +use datafusion::execution::TaskContext; +use datafusion::physical_expr::PhysicalExpr; +use datafusion::physical_plan::aggregates::{AggregateExec, AggregateMode}; +use datafusion::physical_plan::filter::FilterExec; +use datafusion::physical_plan::joins::{HashJoinExec, PartitionMode}; +use datafusion::physical_plan::metrics::MetricsSet; +use datafusion::physical_plan::projection::ProjectionExec; +use datafusion::physical_plan::sorts::sort::SortExec; +use datafusion::physical_plan::{ + ChildrenPropertiesMode, DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, + PlanProperties, RecordBatchStream, ReplaceChildrenOptions, SendableRecordBatchStream, +}; +use datafusion::prelude::SessionContext; +use datafusion_comet_proto::spark_expression::{agg_expr, expr::ExprStruct, AggExpr, Expr}; +use datafusion_comet_proto::spark_operator::{operator::OpStruct, Operator}; +use futures::{Stream, StreamExt}; +use jni::objects::{Global, JObject}; +use parking_lot::Mutex; +use prost::Message; +use std::collections::{HashMap, HashSet}; +use std::fmt::Formatter; +use std::pin::Pin; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::{Arc, LazyLock, Weak}; +use std::task::{Context, Poll}; + +// The registry owns only weak references: the last task drops the physical tree and metrics. +// An executor has no reliable stage-completion callback, so idle gaps deliberately end reuse. +static PHYSICAL_PLANS: LazyLock = LazyLock::new(ScopedPlans::default); + +#[derive(Default)] +struct ScopedPlans { + entries: Mutex, Weak>>, +} + +impl ScopedPlans { + fn get_or_build( + &self, + key: &[u8], + build: impl FnOnce() -> std::result::Result, ExecutionError>, + ) -> std::result::Result, ExecutionError> { + let mut entries = self.entries.lock(); + entries.retain(|_, plan| plan.strong_count() != 0); + if let Some(plan) = entries.get(key).and_then(Weak::upgrade) { + return Ok(plan); + } + // First construction is serialized; failures never become resident entries. + let plan = build()?; + if entries.len() < 64 + && key.len() <= 8 * 1024 * 1024 + && entries.keys().map(Vec::len).sum::() + key.len() <= 8 * 1024 * 1024 + { + entries.insert(key.to_vec(), Arc::downgrade(&plan)); + } + Ok(plan) + } +} + +pub(super) fn clear() { + PHYSICAL_PLANS.entries.lock().clear(); +} + +/// JVM scope includes driver-generated block identity, stage ID and stage attempt. +pub(super) fn scoped_key(scope: &[u8], key: &[u8]) -> Vec { + let mut result = Vec::with_capacity(8 + scope.len() + key.len()); + result.extend_from_slice(&(scope.len() as u64).to_le_bytes()); + result.extend_from_slice(scope); + result.extend_from_slice(key); + result +} + +/// Length-prefix every field so distinct plans/configurations cannot alias. Configuration order +/// is immaterial. Partition index and attempt identity intentionally do not participate: admitted +/// expressions cannot depend on either. task_cpus also participates because it sets the session +/// target_partitions independently of serialized Spark config. Resources arrive through binding. +pub(super) fn cache_key( + bytes: &[u8], + config: &HashMap, + batch_size: i32, + partition_count: i32, + task_cpus: i64, +) -> Vec { + fn append(key: &mut Vec, bytes: &[u8]) { + key.extend_from_slice(&(bytes.len() as u64).to_le_bytes()); + key.extend_from_slice(bytes); + } + let mut key = Vec::new(); + append(&mut key, bytes); + key.extend_from_slice(&batch_size.to_le_bytes()); + key.extend_from_slice(&partition_count.to_le_bytes()); + key.extend_from_slice(&task_cpus.to_le_bytes()); + let mut entries: Vec<_> = config.iter().collect(); + entries.sort_unstable(); + for (name, value) in entries { + append(&mut key, name.as_bytes()); + append(&mut key, value.as_bytes()); + } + key +} + +pub(super) fn cache_bytes<'a>(plan: &Operator, original: &'a [u8]) -> std::borrow::Cow<'a, [u8]> { + fn has_files(plan: &Operator) -> bool { + matches!(plan.op_struct, Some(OpStruct::NativeScan(_))) + || plan.children.iter().any(has_files) + } + if has_files(plan) { + std::borrow::Cow::Owned(template_bytes(plan)) + } else { + std::borrow::Cow::Borrowed(original) + } +} + +/// Legacy file-list normalization. Native scans are currently rejected by admission, +/// so this does not expand the set of trees eligible for sharing. +pub(super) fn template_bytes(plan: &Operator) -> Vec { + fn normalize(plan: &mut Operator) { + if let Some(OpStruct::NativeScan(scan)) = plan.op_struct.as_mut() { + scan.file_partition = None; + } + for child in &mut plan.children { + normalize(child); + } + } + let mut template = plan.clone(); + normalize(&mut template); + template.encode_to_vec() +} + +// Preserve the planner's input_plan push order: children are planned left-to-right, including +// parse_join_parameters. HashJoin may swap physical children afterwards; convert_tree maps each +// original input Arc to this pre-swap slot, so binding must keep the protobuf child order. +fn input_definitions<'a>(plan: &'a Operator, result: &mut Vec<&'a Operator>) { + if matches!( + plan.op_struct, + Some(OpStruct::Scan(_) | OpStruct::NativeScan(_)) + ) { + result.push(plan); + } else { + for child in &plan.children { + input_definitions(child, result); + } + } +} + +pub(super) fn get_or_build( + key: &[u8], + plan: &Operator, + session: &Arc, + partition_count: usize, +) -> std::result::Result, ExecutionError> { + PHYSICAL_PLANS.get_or_build(key, || { + SharedPipeline::build_partitions(plan, session, partition_count) + }) +} + +pub(super) fn supports(plan: &Operator) -> bool { + match plan.op_struct.as_ref() { + Some(OpStruct::Scan(_)) => plan.children.is_empty(), + Some(OpStruct::NativeScan(_)) => false, + Some(OpStruct::Projection(project)) => { + plan.children.len() == 1 + && project.project_list.iter().all(supports_expr) + && supports(&plan.children[0]) + } + Some(OpStruct::Filter(filter)) => { + plan.children.len() == 1 + && filter.predicate.as_ref().is_some_and(supports_expr) + && supports(&plan.children[0]) + } + Some(OpStruct::HashJoin(join)) => { + plan.children.len() == 2 + && !join.dynamic_filter_enabled + && !join.null_aware_anti_join + && join.left_join_keys.iter().all(supports_expr) + && join.right_join_keys.iter().all(supports_expr) + && join.condition.as_ref().is_none_or(supports_expr) + && plan.children.iter().all(supports) + } + Some(OpStruct::Sort(sort)) => { + plan.children.len() == 1 + && !sort.sort_orders.is_empty() + && sort.fetch.is_none() + && sort.skip.is_none_or(|n| n == 0) + && sort.sort_orders.iter().all(supports_sort_order) + && supports(&plan.children[0]) + } + Some(OpStruct::HashAgg(agg)) => { + plan.children.len() == 1 + && agg.grouping_exprs.iter().all(supports_expr) + && agg.agg_exprs.iter().all(supports_aggregate) + && supports(&plan.children[0]) + } + _ => false, + } +} + +fn supports_sort_order(expr: &Expr) -> bool { + match expr.expr_struct.as_ref() { + Some(ExprStruct::SortOrder(order)) => order.child.as_deref().is_some_and(supports_expr), + _ => false, + } +} + +// DISTINCT is lowered by Spark to grouping/deduplication stages before serialization; AggExpr +// has no distinct flag. expr_modes changes how buffers are consumed, never which functions or +// child expressions are admitted here. The original planner supplies the merge definitions. +fn supports_aggregate(expr: &AggExpr) -> bool { + use agg_expr::ExprStruct::*; + expr.filter.as_ref().is_none_or(supports_expr) + && match expr.expr_struct.as_ref() { + Some(Count(e)) => !e.children.is_empty() && e.children.iter().all(supports_expr), + Some(Sum(e)) => e.child.as_ref().is_some_and(supports_expr), + Some(Avg(e)) => e.child.as_ref().is_some_and(supports_expr), + Some(Min(e)) => e.child.as_ref().is_some_and(supports_expr), + Some(Max(e)) => e.child.as_ref().is_some_and(supports_expr), + _ => false, + } +} + +fn supports_expr(expr: &Expr) -> bool { + match expr.expr_struct.as_ref() { + Some(ExprStruct::Bound(_) | ExprStruct::Literal(_)) => true, + Some(ExprStruct::Add(e) | ExprStruct::Subtract(e) | ExprStruct::Multiply(e)) => { + e.left.as_deref().is_some_and(supports_expr) + && e.right.as_deref().is_some_and(supports_expr) + } + Some( + ExprStruct::Eq(e) + | ExprStruct::Neq(e) + | ExprStruct::Gt(e) + | ExprStruct::GtEq(e) + | ExprStruct::Lt(e) + | ExprStruct::LtEq(e) + | ExprStruct::And(e) + | ExprStruct::Or(e), + ) => { + e.left.as_deref().is_some_and(supports_expr) + && e.right.as_deref().is_some_and(supports_expr) + } + Some(ExprStruct::IsNull(e) | ExprStruct::IsNotNull(e) | ExprStruct::Not(e)) => { + e.child.as_deref().is_some_and(supports_expr) + } + // In particular: RNGs, partition ID, subqueries, UDFs and unreviewed scalar functions. + _ => false, + } +} + +#[derive(Debug)] +pub(super) struct SharedPipeline { + pub root: Arc, + scan_definitions: Vec, + identity: Arc<()>, + partition_count: usize, + claimed_partitions: Mutex>, +} + +type BoundAttempt = (Vec, Arc); + +impl SharedPipeline { + /// Never release a claim: upstream metrics persist until the tree is dropped. + /// A repeated partition must execute on an ordinary private plan instead. + fn try_claim_partition(&self, partition: usize) -> bool { + partition < self.partition_count && self.claimed_partitions.lock().insert(partition) + } + + #[cfg(test)] + fn build( + plan: &Operator, + session: &Arc, + ) -> std::result::Result, ExecutionError> { + Self::build_partitions(plan, session, 1) + } + + fn build_partitions( + plan: &Operator, + session: &Arc, + partition_count: usize, + ) -> std::result::Result, ExecutionError> { + if partition_count == 0 || !supports(plan) { + return Err(ExecutionError::GeneralError( + "Unsupported shared native pipeline".into(), + )); + } + // TEST_EXEC_CONTEXT_ID is the planner's default. No task inputs/context are imported. + let input_plans = Arc::new(Mutex::new(Vec::new())); + let planner = PhysicalPlanner::new(Arc::clone(session), 0) + .with_sql_text_pool(plan) + .with_input_plans(Arc::clone(&input_plans)); + let (_, _, original) = planner.create_plan(plan, &mut vec![], 1)?; + let identity = Arc::new(()); + let mut mapping = Vec::new(); + convert_tree( + &original.native_plan, + &identity, + &input_plans.lock(), + &mut mapping, + partition_count, + )?; + let root = convert_spark_tree(&original, &mapping)?; + let mut definitions = Vec::new(); + input_definitions(plan, &mut definitions); + let scan_definitions = definitions + .into_iter() + .map(|p| { + let mut p = p.clone(); + if let Some(OpStruct::NativeScan(scan)) = p.op_struct.as_mut() { + scan.file_partition = None; + } + p + }) + .collect(); + Ok(Arc::new(Self { + root, + scan_definitions, + identity, + partition_count, + claimed_partitions: Mutex::new(HashSet::new()), + })) + } + + #[cfg(test)] + fn bind( + self: &Arc, + planner: &PhysicalPlanner, + inputs: &mut Vec>>>, + ) -> std::result::Result<(Vec, Arc), ExecutionError> { + self.bind_definitions( + planner, + inputs, + &self.scan_definitions.iter().collect::>(), + 0, + ) + } + + pub fn try_bind_plan( + self: &Arc, + planner: &PhysicalPlanner, + inputs: &mut Vec>>>, + task_plan: &Operator, + ) -> std::result::Result, ExecutionError> { + if !self.try_claim_partition(planner.partition() as usize) { + return Ok(None); + } + self.bind_plan(planner, inputs, task_plan).map(Some) + } + + fn bind_plan( + self: &Arc, + planner: &PhysicalPlanner, + inputs: &mut Vec>>>, + task_plan: &Operator, + ) -> std::result::Result<(Vec, Arc), ExecutionError> { + let mut definitions = Vec::new(); + input_definitions(task_plan, &mut definitions); + self.bind_definitions(planner, inputs, &definitions, planner.partition() as usize) + } + + fn bind_definitions( + self: &Arc, + planner: &PhysicalPlanner, + inputs: &mut Vec>>>, + definitions: &[&Operator], + partition: usize, + ) -> std::result::Result<(Vec, Arc), ExecutionError> { + if partition >= self.partition_count || definitions.len() != self.scan_definitions.len() { + return Err(ExecutionError::GeneralError( + "Shared input binding count mismatch".into(), + )); + } + let mut scans = Vec::new(); + let mut bound_inputs = Vec::new(); + for definition in definitions { + let (jvm_scans, _, input) = planner.create_plan(definition, inputs, 1)?; + scans.extend(jvm_scans); + bound_inputs.push(Arc::clone(&input.native_plan)); + } + let attempt = Arc::new(AttemptState { + inputs: bound_inputs, + _owner: Arc::clone(self), + identity: Arc::clone(&self.identity), + started: (0..definitions.len()) + .map(|_| AtomicBool::new(false)) + .collect(), + partition, + }); + Ok((scans, attempt)) + } +} + +/// Owned by one Spark task attempt. The shared tree contains no input readers, +/// memory pools or TaskContext belonging to an attempt. Metrics are partition-labelled. +#[derive(Debug)] +pub(crate) struct AttemptState { + inputs: Vec>, + _owner: Arc, + identity: Arc<()>, + started: Vec, + partition: usize, +} + +impl AttemptState { + pub fn partition(&self) -> usize { + self.partition + } + + pub fn task_context(self: &Arc, session: &SessionContext) -> Arc { + let context = TaskContext::from(session); + let config = context + .session_config() + .clone() + .with_extension(Arc::clone(self)); + Arc::new(context.with_session_config(config)) + } + + pub fn metrics_for(&self, plan: &Arc) -> Option { + if let Some(input) = plan.downcast_ref::() { + if !Arc::ptr_eq(&self.identity, &input.identity) { + return None; + } + self.inputs[input.index].metrics() + } else { + plan.metrics().map(|metrics| { + let mut selected = MetricsSet::new(); + for metric in metrics + .iter() + .filter(|m| m.partition() == Some(self.partition)) + { + selected.push(Arc::clone(metric)); + } + selected + }) + } + } +} + +/// The only replacement node: obtain this attempt's reader through TaskContext. +/// All operators above it are real, shared DataFusion ExecutionPlan instances. +#[derive(Debug)] +struct SharedInputExec { + index: usize, + identity: Arc<()>, + properties: Arc, +} + +impl DisplayAs for SharedInputExec { + fn fmt_as(&self, _: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result { + write!(f, "SharedInputExec: slot={}", self.index) + } +} + +impl ExecutionPlan for SharedInputExec { + fn name(&self) -> &str { + "SharedInputExec" + } + fn apply_expressions( + &self, + _: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + Ok(TreeNodeRecursion::Continue) + } + fn properties(&self) -> &Arc { + &self.properties + } + fn children(&self) -> Vec<&Arc> { + vec![] + } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + if children.is_empty() { + Ok(self) + } else { + internal_err!("Cannot add children to a shared input") + } + } + fn execute( + &self, + partition: usize, + context: Arc, + ) -> Result { + let attempt = context + .session_config() + .get_extension::() + .ok_or_else(|| { + datafusion::common::DataFusionError::Internal( + "Missing shared pipeline attempt binding".into(), + ) + })?; + if !Arc::ptr_eq(&self.identity, &attempt.identity) { + return internal_err!("Shared pipeline attempt belongs to a different plan"); + } + if partition != attempt.partition { + return internal_err!("Shared input partition does not match task binding"); + } + if attempt.started[self.index].swap(true, Ordering::AcqRel) { + return internal_err!("Shared pipeline attempt was already executed"); + } + let stream = attempt.inputs[self.index].execute(0, context)?; + // Keep the binding alive across asynchronous polling without retaining it + // in any shared operator. Cancellation drops the stream before its binding. + Ok(Box::pin(BoundInputStream { + stream, + _attempt: attempt, + })) + } +} + +struct BoundInputStream { + stream: SendableRecordBatchStream, + _attempt: Arc, +} +impl Stream for BoundInputStream { + type Item = Result; + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + self.stream.poll_next_unpin(cx) + } +} +impl RecordBatchStream for BoundInputStream { + fn schema(&self) -> SchemaRef { + self.stream.schema() + } +} + +type PlanMapping = Vec<(Arc, Arc)>; + +fn convert_tree( + plan: &Arc, + identity: &Arc<()>, + bound_inputs: &[Arc], + mapping: &mut PlanMapping, + partition_count: usize, +) -> Result> { + let shared: Arc = + if let Some(index) = bound_inputs.iter().position(|p| Arc::ptr_eq(p, plan)) { + Arc::new(SharedInputExec { + index, + identity: Arc::clone(identity), + properties: Arc::new( + plan.properties() + .as_ref() + .clone() + .with_partitioning(Partitioning::UnknownPartitioning(partition_count)), + ), + }) + } else { + let children = plan + .children() + .into_iter() + .map(|child| convert_tree(child, identity, bound_inputs, mapping, partition_count)) + .collect::>>()?; + // Rebuild only during cache construction to install shared input leaves. + // No constructor, reset_state or operator clone runs on task cache hits. + let rebuilt = Arc::clone(plan).replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + )?; + if let Some(sort) = rebuilt.downcast_ref::() { + if sort.fetch().is_some() { + return internal_err!("Top-K is not admitted"); + } + Arc::new(sort.clone().with_preserve_partitioning(true)) + } else if let Some(aggregate) = rebuilt.downcast_ref::() { + if *aggregate.mode() == AggregateMode::Final && partition_count > 1 { + // Each Spark task consumes its own already-shuffled partition. + Arc::new(AggregateExec::try_new( + AggregateMode::FinalPartitioned, + aggregate.group_expr().clone(), + aggregate.aggr_expr().to_vec(), + aggregate.filter_expr().to_vec(), + Arc::clone(aggregate.input()), + aggregate.input_schema(), + )?) + } else { + rebuilt + } + } else if let Some(join) = rebuilt.downcast_ref::() { + if *join.partition_mode() != PartitionMode::Partitioned { + return internal_err!("Only partitioned joins can share execution trees"); + } + rebuilt + } else if rebuilt.is::() || rebuilt.is::() { + rebuilt + } else { + return internal_err!("Unexpected operator in shared tree: {}", rebuilt.name()); + } + }; + mapping.push((Arc::clone(plan), Arc::clone(&shared))); + Ok(shared) +} + +fn convert_spark_tree(plan: &Arc, mapping: &PlanMapping) -> Result> { + let lookup = |old: &Arc| -> Result> { + mapping + .iter() + .find(|(from, _)| Arc::ptr_eq(from, old)) + .map(|(_, to)| Arc::clone(to)) + .ok_or_else(|| { + datafusion::common::DataFusionError::Internal( + "Missing shared pipeline metric node".into(), + ) + }) + }; + Ok(Arc::new(SparkPlan::new_with_additional( + plan.plan_id, + lookup(&plan.native_plan)?, + plan.children + .iter() + .map(|c| convert_spark_tree(c, mapping)) + .collect::>()?, + plan.additional_native_plans + .iter() + .map(lookup) + .collect::>()?, + ))) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::execution::metrics::utils::to_native_metric_node_with; + use crate::execution::operators::InputBatch; + use arrow::array::Int64Array; + use datafusion_comet_proto::spark_expression::{literal::Value, Literal}; + use datafusion_comet_proto::spark_expression::{ + BinaryExpr, BoundReference, DataType, EmptyExpr, MathExpr, + }; + use datafusion_comet_proto::spark_operator::{Filter, Projection, Scan}; + use futures::FutureExt; + use std::sync::{atomic::AtomicUsize, Barrier}; + + fn datatype() -> DataType { + DataType { + type_id: 4, + type_info: None, + } + } + fn expr(kind: ExprStruct) -> Expr { + Expr { + expr_struct: Some(kind), + ..Default::default() + } + } + fn column() -> Expr { + expr(ExprStruct::Bound(BoundReference { + index: 0, + datatype: Some(datatype()), + })) + } + fn literal(value: i64) -> Expr { + expr(ExprStruct::Literal(Literal { + value: Some(Value::LongVal(value)), + datatype: Some(datatype()), + is_null: false, + })) + } + fn pipeline() -> Operator { + Operator { + plan_id: 3, + op_struct: Some(OpStruct::Projection(Projection { + project_list: vec![expr(ExprStruct::Add(Box::new(MathExpr { + left: Some(Box::new(column())), + right: Some(Box::new(literal(10))), + return_type: Some(datatype()), + ..Default::default() + })))], + })), + children: vec![Operator { + plan_id: 2, + op_struct: Some(OpStruct::Filter(Filter { + predicate: Some(expr(ExprStruct::Gt(Box::new(BinaryExpr { + left: Some(Box::new(column())), + right: Some(Box::new(literal(0))), + })))), + })), + children: vec![Operator { + plan_id: 1, + op_struct: Some(OpStruct::Scan(Scan { + fields: vec![datatype()], + source: "shared-test".into(), + })), + ..Default::default() + }], + ..Default::default() + }], + ..Default::default() + } + } + fn feed(scan: &mut ScanExec, values: Vec>) { + let len = values.len(); + scan.set_input_batch(InputBatch::Batch( + vec![Arc::new(Int64Array::from(values))], + len, + )); + } + fn next(stream: &mut SendableRecordBatchStream) -> Vec { + let batch = stream + .next() + .now_or_never() + .expect("must not wait for another task") + .unwrap() + .unwrap(); + batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec() + } + + #[test] + fn concurrent_first_touch_builds_one_physical_tree_without_retaining_sessions() { + let cache = Arc::new(ScopedPlans::default()); + let builds = Arc::new(AtomicUsize::new(0)); + let barrier = Arc::new(Barrier::new(8)); + let handles: Vec<_> = (0..8) + .map(|_| { + let (cache, builds, barrier) = ( + Arc::clone(&cache), + Arc::clone(&builds), + Arc::clone(&barrier), + ); + std::thread::spawn(move || { + let session = Arc::new(SessionContext::new()); + let weak = Arc::downgrade(&session); + barrier.wait(); + let shared = cache + .get_or_build(b"same", || { + builds.fetch_add(1, Ordering::SeqCst); + SharedPipeline::build(&pipeline(), &session) + }) + .unwrap(); + drop(session); + assert!(weak.upgrade().is_none()); + shared + }) + }) + .collect(); + let plans: Vec<_> = handles.into_iter().map(|h| h.join().unwrap()).collect(); + assert_eq!(builds.load(Ordering::SeqCst), 1); + for plan in &plans { + assert!(Arc::ptr_eq( + &plans[0].root.native_plan, + &plan.root.native_plan + )); + } + } + + #[tokio::test] + async fn interleaved_attempts_cancellation_retry_and_metrics_are_isolated() { + let session = Arc::new(SessionContext::new()); + let shared = SharedPipeline::build(&pipeline(), &session).unwrap(); + let planner = PhysicalPlanner::new(Arc::clone(&session), 0); + let (mut scans_a, a) = shared.bind(&planner, &mut vec![]).unwrap(); + let retry_b = SharedPipeline::build(&pipeline(), &session).unwrap(); + let (mut scans_b, b) = retry_b.bind(&planner, &mut vec![]).unwrap(); + let mut stream_a = shared + .root + .native_plan + .execute(0, a.task_context(&session)) + .unwrap(); + let stream_b = retry_b + .root + .native_plan + .execute(0, b.task_context(&session)) + .unwrap(); + feed(&mut scans_a[0], vec![None, Some(-1), Some(1), Some(2)]); + feed(&mut scans_b[0], vec![Some(8)]); + assert!(stream_a.next().now_or_never().is_none()); + let weak = Arc::downgrade(&a); + drop(stream_a); + drop(scans_a); + drop(a); + assert!(weak.upgrade().is_none()); + let retry_c = SharedPipeline::build(&pipeline(), &session).unwrap(); + let (mut scans_c, c) = retry_c.bind(&planner, &mut vec![]).unwrap(); + let stream_c = retry_c + .root + .native_plan + .execute(0, c.task_context(&session)) + .unwrap(); + feed(&mut scans_c[0], vec![Some(30), Some(31)]); + let (out_b, out_c) = tokio::join!( + drain_inputs(scans_b, stream_b), + drain_inputs(scans_c, stream_c) + ); + assert_eq!(rows(&out_b), vec![vec![Some(18)]]); + assert_eq!(rows(&out_c), vec![vec![Some(40)], vec![Some(41)]]); + assert_eq!( + b.metrics_for(&retry_b.root.native_plan) + .unwrap() + .output_rows(), + Some(1) + ); + assert_eq!( + c.metrics_for(&retry_c.root.native_plan) + .unwrap() + .output_rows(), + Some(2) + ); + } + + #[test] + fn empty_and_null_batches_and_zero_column_projection_preserve_rows() { + let session = Arc::new(SessionContext::new()); + let mut definition = pipeline(); + let Some(OpStruct::Projection(project)) = definition.op_struct.as_mut() else { + unreachable!() + }; + project.project_list.clear(); + let shared = SharedPipeline::build(&definition, &session).unwrap(); + let planner = PhysicalPlanner::new(Arc::clone(&session), 99); + let (mut scans, attempt) = shared.bind(&planner, &mut vec![]).unwrap(); + let mut stream = shared + .root + .native_plan + .execute(0, attempt.task_context(&session)) + .unwrap(); + for batch in [vec![], vec![None, Some(-1)]] { + feed(&mut scans[0], batch); + assert!(stream.next().now_or_never().is_none()); + } + feed(&mut scans[0], vec![Some(2), Some(3)]); + assert!(stream.next().now_or_never().is_none()); + scans[0].set_input_batch(InputBatch::EOF); + let batch = stream.next().now_or_never().unwrap().unwrap().unwrap(); + assert_eq!((batch.num_columns(), batch.num_rows()), (0, 2)); + scans[0].set_input_batch(InputBatch::EOF); + assert!(stream.next().now_or_never().unwrap().is_none()); + } + + #[test] + fn bindings_reject_wrong_tree_missing_context_and_duplicate_execution() { + let session = Arc::new(SessionContext::new()); + let shared = SharedPipeline::build(&pipeline(), &session).unwrap(); + let other = SharedPipeline::build(&pipeline(), &session).unwrap(); + let planner = PhysicalPlanner::new(Arc::clone(&session), 0); + let (_, attempt) = shared.bind(&planner, &mut vec![]).unwrap(); + assert!(shared + .root + .native_plan + .execute(0, session.task_ctx()) + .is_err()); + assert!(shared + .root + .native_plan + .execute(1, attempt.task_context(&session)) + .is_err()); + assert!(other + .root + .native_plan + .execute(0, attempt.task_context(&session)) + .is_err()); + let _stream = shared + .root + .native_plan + .execute(0, attempt.task_context(&session)) + .unwrap(); + assert!(shared + .root + .native_plan + .execute(0, attempt.task_context(&session)) + .is_err()); + } + + #[test] + fn admission_is_recursive_and_configuration_is_part_of_identity() { + assert!(supports(&pipeline())); + for stateful in [ + ExprStruct::SparkPartitionId(EmptyExpr {}), + ExprStruct::MonotonicallyIncreasingId(EmptyExpr {}), + ] { + let mut plan = pipeline(); + let Some(OpStruct::Projection(project)) = plan.op_struct.as_mut() else { + unreachable!() + }; + let Some(ExprStruct::Add(add)) = project.project_list[0].expr_struct.as_mut() else { + unreachable!() + }; + add.left = Some(Box::new(expr(stateful))); + assert!(!supports(&plan)); + } + let mut plan = pipeline(); + plan.children.push(plan.children[0].clone()); + assert!(!supports(&plan)); + assert!(!supports(&Operator::default())); + let config: HashMap = HashMap::from([ + ("timezone".into(), "UTC".into()), + ("ansi".into(), "true".into()), + ]); + let reversed = config.iter().map(|(k, v)| (k.clone(), v.clone())).collect(); + let key = cache_key(b"plan", &config, 100, 8, 1); + assert_eq!(key, cache_key(b"plan", &reversed, 100, 8, 1)); + assert_ne!(key, cache_key(b"plan", &config, 101, 8, 1)); + assert_ne!(key, cache_key(b"plan", &config, 100, 9, 1)); + assert_ne!(key, cache_key(b"plan", &config, 100, 8, 2)); + assert_ne!(key, cache_key(b"other", &config, 100, 8, 1)); + assert_ne!(key, cache_key(b"plan", &HashMap::new(), 100, 8, 1)); + } + #[test] + fn stage_scope_and_partition_claims_isolate_retries() { + let cache = ScopedPlans::default(); + let session = Arc::new(SessionContext::new()); + let build = || SharedPipeline::build_partitions(&pipeline(), &session, 4); + let key = scoped_key(b"block-a:stage-1:attempt-0", b"plan"); + let first = cache.get_or_build(&key, build).unwrap(); + let same = cache + .get_or_build(&key, || panic!("same active scope")) + .unwrap(); + assert!(Arc::ptr_eq(&first, &same)); + assert!(first.try_claim_partition(0)); + assert!(same.try_claim_partition(2)); + assert!(!same.try_claim_partition(0)); + assert!(!same.try_claim_partition(4)); + for scope in [ + b"block-a:stage-1:attempt-1".as_slice(), + b"block-b:stage-1:attempt-0", + b"block-a:stage-2:attempt-0", + ] { + let other = cache + .get_or_build(&scoped_key(scope, b"plan"), build) + .unwrap(); + assert!(!Arc::ptr_eq(&first, &other)); + assert!(other.try_claim_partition(0)); + } + let planner = PhysicalPlanner::new(Arc::clone(&session), 1); + assert!(first + .try_bind_plan(&planner, &mut vec![], &pipeline()) + .unwrap() + .is_some()); + assert!(same + .try_bind_plan(&planner, &mut vec![], &pipeline()) + .unwrap() + .is_none()); + // Even an overlapping retry may execute privately without affecting the first tree. + let private = build().unwrap(); + assert!(!Arc::ptr_eq( + &first.root.native_plan, + &private.root.native_plan + )); + assert!(private.try_claim_partition(0)); + } + + #[test] + fn collapsed_spark_nodes_do_not_double_count_output_rows() { + let session = Arc::new(SessionContext::new()); + let mut definition = pipeline(); + definition.plan_id = definition.children[0].plan_id; + let shared = SharedPipeline::build(&definition, &session).unwrap(); + assert_eq!(shared.root.additional_native_plans.len(), 1); + assert_eq!(shared.root.children[0].plan_id, 1); + let planner = PhysicalPlanner::new(Arc::clone(&session), 0); + let (mut scans, attempt) = shared.bind(&planner, &mut vec![]).unwrap(); + let mut stream = shared + .root + .native_plan + .execute(0, attempt.task_context(&session)) + .unwrap(); + feed(&mut scans[0], vec![Some(-1), Some(2), Some(3)]); + assert!(stream.next().now_or_never().is_none()); + scans[0].set_input_batch(InputBatch::EOF); + assert_eq!(next(&mut stream), vec![12, 13]); + let metrics = + to_native_metric_node_with(&shared.root, &|p| attempt.metrics_for(p)).unwrap(); + assert_eq!(metrics.metrics["output_rows"], 2); + assert_eq!(metrics.metrics["selectivity_total"], 3); + assert_eq!(metrics.children[0].metrics["output_rows"], 3); + } + fn scan() -> Operator { + pipeline().children[0].children[0].clone() + } + fn sort_plan(fetch: Option, skip: Option) -> Operator { + use datafusion_comet_proto::spark_expression::SortOrder; + use datafusion_comet_proto::spark_operator::Sort; + Operator { + plan_id: 10, + children: vec![scan()], + op_struct: Some(OpStruct::Sort(Sort { + sort_orders: vec![expr(ExprStruct::SortOrder(Box::new(SortOrder { + child: Some(Box::new(column())), + direction: 0, + null_ordering: 1, + })))], + fetch, + skip, + })), + ..Default::default() + } + } + fn aggregate_plan() -> Operator { + use datafusion_comet_proto::spark_expression::{Count, Sum}; + use datafusion_comet_proto::spark_operator::HashAggregate; + Operator { + plan_id: 11, + children: vec![scan()], + op_struct: Some(OpStruct::HashAgg(HashAggregate { + grouping_exprs: vec![column()], + agg_exprs: vec![ + AggExpr { + expr_struct: Some(agg_expr::ExprStruct::Count(Count { + children: vec![column()], + })), + ..Default::default() + }, + AggExpr { + expr_struct: Some(agg_expr::ExprStruct::Sum(Sum { + child: Some(column()), + datatype: Some(datatype()), + eval_mode: 0, + })), + ..Default::default() + }, + ], + ..Default::default() + })), + ..Default::default() + } + } + fn join_plan(build_side: i32, join_type: i32) -> Operator { + use datafusion_comet_proto::spark_operator::HashJoin; + Operator { + plan_id: 12, + children: vec![ + scan(), + Operator { + plan_id: 2, + ..scan() + }, + ], + op_struct: Some(OpStruct::HashJoin(HashJoin { + left_join_keys: vec![column()], + right_join_keys: vec![column()], + join_type, + build_side, + ..Default::default() + })), + ..Default::default() + } + } + fn rows(batches: &[RecordBatch]) -> Vec>> { + use arrow::array::Array; + let mut rows = Vec::new(); + for batch in batches { + for row in 0..batch.num_rows() { + rows.push( + batch + .columns() + .iter() + .map(|c| { + let a = c.as_any().downcast_ref::().unwrap(); + (!a.is_null(row)).then(|| a.value(row)) + }) + .collect(), + ); + } + } + rows.sort(); + rows + } + async fn drain_inputs( + mut scans: Vec, + mut stream: SendableRecordBatchStream, + ) -> Vec { + // Mimic JNI's pull-on-Pending driver: mocked ScanStream does not wake a task when + // its slot is filled. Feed each consumed input independently, including build-right. + let mut ended = vec![false; scans.len()]; + let result = tokio::time::timeout(std::time::Duration::from_secs(10), async { + let mut batches = Vec::new(); + loop { + let next = futures::future::poll_fn(|cx| match stream.poll_next_unpin(cx) { + Poll::Pending => { + for (i, scan) in scans.iter_mut().enumerate() { + if !ended[i] && scan.batch.lock().unwrap().is_none() { + scan.set_input_batch(InputBatch::EOF); + ended[i] = true; + } + } + cx.waker().wake_by_ref(); + Poll::Pending + } + ready => ready, + }) + .await; + match next { + Some(batch) => batches.push(batch.unwrap()), + None => break, + } + } + batches + }) + .await; + result.expect("attempt must finish without other partitions") + } + + #[tokio::test] + async fn full_sort_and_cancellation_have_private_state() { + for (fetch, skip) in [(None, None)] { + let session = Arc::new(SessionContext::new()); + let shared = SharedPipeline::build(&sort_plan(fetch, skip), &session).unwrap(); + let planner = PhysicalPlanner::new(Arc::clone(&session), 7); + // Cancel with input buffered, before EOF. A later attempt must not inherit TopK's + // threshold, memory reservations, completion state or metrics. + let (mut scans, cancelled) = shared.bind(&planner, &mut vec![]).unwrap(); + feed(&mut scans[0], vec![Some(-100), Some(-200)]); + let mut stream = shared + .root + .native_plan + .execute(0, cancelled.task_context(&session)) + .unwrap(); + assert!(stream.next().now_or_never().is_none()); + let weak = Arc::downgrade(&cancelled); + drop(stream); + drop(cancelled); + drop(scans); + assert!(weak.upgrade().is_none()); + assert_eq!(session.runtime_env().memory_pool.reserved(), 0); + for values in [ + vec![Some(9), Some(3), Some(7), None], + vec![Some(100), Some(200)], + ] { + let shared = SharedPipeline::build(&sort_plan(fetch, skip), &session).unwrap(); + let (mut scans, attempt) = shared.bind(&planner, &mut vec![]).unwrap(); + feed(&mut scans[0], values.clone()); + let stream = shared + .root + .native_plan + .execute(0, attempt.task_context(&session)) + .unwrap(); + let batches = drain_inputs(scans, stream).await; + let mut expected = values; + expected.sort_by_key(|v| (v.is_none(), *v)); + if let Some(fetch) = fetch { + expected.truncate(fetch as usize); + } + let expected: Vec<_> = expected + .into_iter() + .skip(skip.unwrap_or(0) as usize) + .collect(); + let actual: Vec<_> = batches + .iter() + .flat_map(|b| { + b.column(0) + .as_any() + .downcast_ref::() + .unwrap() + .iter() + }) + .collect(); + assert_eq!(actual, expected); + assert_eq!( + attempt + .metrics_for(&shared.root.native_plan) + .unwrap() + .output_rows(), + Some(actual.len()) + ); + assert_eq!(session.runtime_env().memory_pool.reserved(), 0); + } + } + } + + #[test] + fn aggregate_modes_do_not_bypass_function_or_expression_admission() { + use datafusion_comet_proto::spark_expression::First; + // DISTINCT rewrites can emit uniform stages or mixed Partial/PartialMerge stages. + for (mode, expr_modes) in [(0, vec![]), (1, vec![]), (2, vec![2, 2]), (0, vec![0, 2])] { + let mut definition = aggregate_plan(); + let Some(OpStruct::HashAgg(agg)) = definition.op_struct.as_mut() else { + unreachable!() + }; + agg.mode = mode; + agg.expr_modes = expr_modes; + assert!(supports(&definition)); + let mut unsupported_function = definition.clone(); + let Some(OpStruct::HashAgg(agg)) = unsupported_function.op_struct.as_mut() else { + unreachable!() + }; + agg.agg_exprs[0].expr_struct = Some(agg_expr::ExprStruct::First(First { + child: Some(column()), + datatype: Some(datatype()), + ignore_nulls: false, + })); + assert!(!supports(&unsupported_function)); + let Some(OpStruct::HashAgg(agg)) = definition.op_struct.as_mut() else { + unreachable!() + }; + agg.agg_exprs[0].filter = Some(expr(ExprStruct::SparkPartitionId(EmptyExpr {}))); + assert!(!supports(&definition)); + } + } + + // Historical audit showing why repeated partition execution accumulates metrics. + // Production bindings prohibit this reuse; the audit deliberately bypasses that guard. + #[derive(Debug)] + struct AuditInput { + index: usize, + identity: Arc<()>, + properties: Arc, + } + + impl DisplayAs for AuditInput { + fn fmt_as(&self, _: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result { + write!(f, "AuditInput({})", self.index) + } + } + + impl ExecutionPlan for AuditInput { + fn apply_expressions( + &self, + _: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + Ok(TreeNodeRecursion::Continue) + } + fn name(&self) -> &str { + "AuditInput" + } + fn properties(&self) -> &Arc { + &self.properties + } + fn children(&self) -> Vec<&Arc> { + vec![] + } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + assert!(children.is_empty()); + Ok(self) + } + fn execute( + &self, + partition: usize, + context: Arc, + ) -> Result { + assert!(partition < 2); + let attempt = context + .session_config() + .get_extension::() + .unwrap(); + assert!(Arc::ptr_eq(&self.identity, &attempt.identity)); + // The DataFusion partition selects the execution lane, whereas the + // task's JVM input still has exactly one local partition. + attempt.inputs[self.index].execute(0, context) + } + } + + fn audit_tree( + plan: &Arc, + inputs: &[Arc], + identity: &Arc<()>, + ) -> Arc { + if let Some(index) = inputs.iter().position(|p| Arc::ptr_eq(p, plan)) { + return Arc::new(AuditInput { + index, + identity: Arc::clone(identity), + properties: Arc::new(plan.properties().as_ref().clone().with_partitioning( + datafusion::physical_plan::Partitioning::UnknownPartitioning(2), + )), + }); + } + let children = plan + .children() + .into_iter() + .map(|p| audit_tree(p, inputs, identity)) + .collect(); + let rewritten = Arc::clone(plan) + .replace_children( + children, + datafusion::physical_plan::ReplaceChildrenOptions::new( + datafusion::physical_plan::ChildrenPropertiesMode::Recompute, + ), + ) + .unwrap(); + if let Some(sort) = rewritten.downcast_ref::() { + assert!(sort.fetch().is_none()); + // Spark already defines each task's sort input; this is a local sort + // on each lane, not a new executor-wide global sort. + Arc::new(sort.clone().with_preserve_partitioning(true)) + } else { + rewritten + } + } + + fn audit_nodes(plan: &Arc) -> Vec> { + assert!(!plan.is::()); + let mut nodes = vec![Arc::clone(plan)]; + for child in plan.children() { + nodes.extend(audit_nodes(child)); + } + nodes + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn actual_datafusion_tree_is_reused_across_partitions_and_attempt_waves() { + let mut global = aggregate_plan(); + if let Some(OpStruct::HashAgg(agg)) = &mut global.op_struct { + agg.grouping_exprs.clear(); + // Wide aggregation: compiled expressions and AggregateExec itself + // must both survive all executions of the shared root. + agg.agg_exprs = vec![agg.agg_exprs[1].clone(); 64]; + } + let mut cases = vec![pipeline(), aggregate_plan(), global, sort_plan(None, None)]; + for side in 0..=1 { + for kind in 0..=5 { + cases.push(join_plan(side, kind)); + } + } + for definition in cases { + let session = Arc::new(SessionContext::new()); + let bindings = SharedPipeline::build(&definition, &session).unwrap(); + let input_plans = Arc::new(Mutex::new(Vec::new())); + let planner = PhysicalPlanner::new(Arc::clone(&session), 0); + let builder = PhysicalPlanner::new(Arc::clone(&session), 0) + .with_input_plans(Arc::clone(&input_plans)); + let (_, _, original) = builder.create_plan(&definition, &mut vec![], 1).unwrap(); + let root = audit_tree( + &original.native_plan, + &input_plans.lock(), + &bindings.identity, + ); + assert!(!root.is::()); + if matches!(definition.op_struct, Some(OpStruct::HashAgg(_))) { + assert!(root.is::()); + } + let original_nodes = audit_nodes(&root); + let mut cumulative_rows = [0; 2]; + let mut previous_metric_count = 0; + for wave in 0..2 { + let start = Arc::new(tokio::sync::Barrier::new(2)); + let mut executions = Vec::new(); + for partition in 0..2 { + let (mut scans, attempt) = bindings.bind(&planner, &mut vec![]).unwrap(); + let (mut reference_scans, _, reference) = + planner.create_plan(&definition, &mut vec![], 1).unwrap(); + for (i, scan) in scans.iter_mut().enumerate() { + let base = 100 * wave + 10 * partition as i64; + let values = vec![None, Some(base + i as i64), Some(base + 1)]; + feed(scan, values.clone()); + feed(&mut reference_scans[i], values); + } + let stream = root + .execute(partition, attempt.task_context(&session)) + .unwrap(); + let reference_stream = reference + .native_plan + .execute(0, session.task_ctx()) + .unwrap(); + let start = Arc::clone(&start); + executions.push(tokio::spawn(async move { + start.wait().await; + let (actual, expected) = tokio::join!( + drain_inputs(scans, stream), + drain_inputs(reference_scans, reference_stream) + ); + assert_eq!(rows(&actual), rows(&expected)); + let weak = Arc::downgrade(&attempt); + drop(attempt); + assert!(weak.upgrade().is_none()); + actual.iter().map(RecordBatch::num_rows).sum::() + })); + } + for (partition, execution) in executions.into_iter().enumerate() { + cumulative_rows[partition] += execution.await.unwrap(); + } + let current_nodes = audit_nodes(&root); + assert_eq!(current_nodes.len(), original_nodes.len()); + for (original, current) in original_nodes.iter().zip(¤t_nodes) { + assert!(Arc::ptr_eq(original, current)); + } + let metrics = root.metrics().unwrap(); + for (partition, expected) in cumulative_rows.iter().enumerate() { + let mut selected = MetricsSet::new(); + for metric in metrics.iter().filter(|m| m.partition() == Some(partition)) { + selected.push(Arc::clone(metric)); + } + assert_eq!(selected.output_rows(), Some(*expected)); + } + // Demonstrate the remaining integration gap: reusing a partition + // appends metrics, even after its stream and attempt are gone. + let metric_count = metrics.iter().count(); + assert!(metric_count > previous_metric_count); + previous_metric_count = metric_count; + assert_eq!(session.runtime_env().memory_pool.reserved(), 0); + } + } + } + + fn physical_nodes(plan: &Arc) -> Vec> { + let mut nodes = vec![Arc::clone(plan)]; + for child in plan.children() { + nodes.extend(physical_nodes(child)); + } + nodes + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn production_tree_shares_real_nodes_with_partition_metrics() { + let mut wide = aggregate_plan(); + if let Some(OpStruct::HashAgg(a)) = &mut wide.op_struct { + a.grouping_exprs.clear(); + a.agg_exprs = vec![a.agg_exprs[1].clone(); 64]; + } + let mut definitions = vec![pipeline(), aggregate_plan(), wide, sort_plan(None, None)]; + for side in 0..=1 { + for kind in 0..=5 { + definitions.push(join_plan(side, kind)); + } + } + for definition in definitions { + let session = Arc::new(SessionContext::new()); + let shared = SharedPipeline::build_partitions(&definition, &session, 4).unwrap(); + let original_nodes = physical_nodes(&shared.root.native_plan); + if matches!(definition.op_struct, Some(OpStruct::HashAgg(_))) { + assert!(shared.root.native_plan.is::()); + } + // Production claims each partition once. Retry/repeated-partition behavior + // is covered by stage_scope_and_partition_claims_isolate_retries. + let mut executions = Vec::new(); + let barrier = Arc::new(tokio::sync::Barrier::new(2)); + // Partitions 1 and 3 never execute here; operators must not wait for them. + for (attempt_no, partition) in [0, 2].into_iter().enumerate() { + let planner = PhysicalPlanner::new(Arc::clone(&session), partition); + let (mut scans, attempt) = shared + .try_bind_plan(&planner, &mut vec![], &definition) + .unwrap() + .unwrap(); + let (mut private_scans, _, private) = + planner.create_plan(&definition, &mut vec![], 1).unwrap(); + for (index, scan) in scans.iter_mut().enumerate() { + let base = attempt_no as i64 * 10; + let values = vec![None, Some(base + index as i64), Some(base + 1)]; + feed(scan, values.clone()); + feed(&mut private_scans[index], values); + } + let stream = shared + .root + .native_plan + .execute(partition as usize, attempt.task_context(&session)) + .unwrap(); + let private_stream = private.native_plan.execute(0, session.task_ctx()).unwrap(); + let root = Arc::clone(&shared.root.native_plan); + let barrier = Arc::clone(&barrier); + executions.push(tokio::spawn(async move { + barrier.wait().await; + let (actual, expected) = tokio::join!( + drain_inputs(scans, stream), + drain_inputs(private_scans, private_stream) + ); + assert_eq!(rows(&actual), rows(&expected)); + assert_eq!( + attempt.metrics_for(&root).unwrap().output_rows(), + Some(actual.iter().map(RecordBatch::num_rows).sum()) + ); + let weak = Arc::downgrade(&attempt); + drop(attempt); + assert!(weak.upgrade().is_none()); + })); + } + for execution in executions { + execution.await.unwrap(); + } + let current = physical_nodes(&shared.root.native_plan); + assert_eq!(original_nodes.len(), current.len()); + for (before, after) in original_nodes.iter().zip(current) { + assert!(Arc::ptr_eq(before, &after)); + } + assert_eq!(session.runtime_env().memory_pool.reserved(), 0); + } + } + + #[tokio::test] + async fn aggregate_stream_variants_resolve_attempt_metrics() { + for migration in [false, true] { + let mut config = datafusion::prelude::SessionConfig::new(); + config.options_mut().execution.enable_migration_aggregate = migration; + let session = Arc::new(SessionContext::new_with_config(config)); + for ordered in [false, true] { + let mut definition = aggregate_plan(); + if ordered { + definition.children[0] = sort_plan(None, None); + } + let shared = SharedPipeline::build(&definition, &session).unwrap(); + let planner = PhysicalPlanner::new(Arc::clone(&session), 0); + let (mut scans, attempt) = shared.bind(&planner, &mut vec![]).unwrap(); + let (mut private_scans, _, private) = + planner.create_plan(&definition, &mut vec![], 1).unwrap(); + let values = vec![Some(3), Some(1), Some(3), None]; + feed(&mut scans[0], values.clone()); + feed(&mut private_scans[0], values); + let stream = shared + .root + .native_plan + .execute(0, attempt.task_context(&session)) + .unwrap(); + let reference = private.native_plan.execute(0, session.task_ctx()).unwrap(); + let (actual, expected) = tokio::join!( + drain_inputs(scans, stream), + drain_inputs(private_scans, reference) + ); + assert_eq!(rows(&actual), rows(&expected)); + assert_eq!( + attempt + .metrics_for(&shared.root.native_plan) + .unwrap() + .output_rows(), + Some(3) + ); + + assert_eq!(session.runtime_env().memory_pool.reserved(), 0); + } + } + } + + fn min_max_plan(grouped: bool) -> Operator { + use datafusion_comet_proto::spark_expression::{Max, Min}; + let mut definition = aggregate_plan(); + let Some(OpStruct::HashAgg(a)) = &mut definition.op_struct else { + unreachable!() + }; + if !grouped { + a.grouping_exprs.clear(); + } + a.agg_exprs = vec![ + AggExpr { + expr_struct: Some(agg_expr::ExprStruct::Min(Min { + child: Some(column()), + datatype: Some(datatype()), + })), + ..Default::default() + }, + AggExpr { + expr_struct: Some(agg_expr::ExprStruct::Max(Max { + child: Some(column()), + datatype: Some(datatype()), + })), + ..Default::default() + }, + ]; + definition + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn min_max_share_real_operators_across_partitions() { + // Disjoint bounds, all-null and empty inputs detect accidental reuse of + // another partition's accumulator. Final/merge buffers are both i64. + for grouped in [false, true] { + for mode in 0..=2 { + for filtered in [false, true] { + // Spark applies aggregate filters before merging state buffers. + if filtered && mode != 0 { + continue; + } + let session = Arc::new(SessionContext::new()); + let mut definition = min_max_plan(grouped); + let Some(OpStruct::HashAgg(a)) = &mut definition.op_struct else { + unreachable!() + }; + a.mode = mode; + a.initial_input_buffer_offset = i32::from(grouped); + if filtered { + for aggregate in &mut a.agg_exprs { + aggregate.filter = Some(expr(ExprStruct::Gt(Box::new(BinaryExpr { + left: Some(Box::new(column())), + right: Some(Box::new(literal(-50))), + })))); + } + } + let width = if mode == 0 { + 1 + } else { + 2 + usize::from(grouped) + }; + let Some(OpStruct::Scan(scan)) = &mut definition.children[0].op_struct else { + unreachable!() + }; + scan.fields = vec![datatype(); width]; + assert!(supports(&definition)); + let shared = + SharedPipeline::build_partitions(&definition, &session, 5).unwrap(); + let original = physical_nodes(&shared.root.native_plan); + assert!(original[0].is::()); + let mut executions = Vec::new(); + let barrier = Arc::new(tokio::sync::Barrier::new(4)); + for (partition, values) in [ + (0, vec![Some(-100), Some(-10), None]), + (2, vec![Some(20), Some(200), None]), + (3, vec![None, None]), + (4, vec![]), + ] { + let planner = PhysicalPlanner::new(Arc::clone(&session), partition); + let (mut scans, attempt) = shared + .try_bind_plan(&planner, &mut vec![], &definition) + .unwrap() + .unwrap(); + let (mut private_scans, _, private) = + planner.create_plan(&definition, &mut vec![], 1).unwrap(); + let columns: Vec> = (0..width) + .map(|_| { + Arc::new(Int64Array::from(values.clone())) + as Arc + }) + .collect(); + scans[0].set_input_batch(InputBatch::Batch(columns.clone(), values.len())); + private_scans[0].set_input_batch(InputBatch::Batch(columns, values.len())); + let root = Arc::clone(&shared.root.native_plan); + let stream = root + .execute(partition as usize, attempt.task_context(&session)) + .unwrap(); + let reference = private.native_plan.execute(0, session.task_ctx()).unwrap(); + let barrier = Arc::clone(&barrier); + executions.push(tokio::spawn(async move { + barrier.wait().await; + let (actual, expected) = tokio::join!( + drain_inputs(scans, stream), + drain_inputs(private_scans, reference) + ); + assert_eq!(rows(&actual), rows(&expected)); + assert_eq!( + attempt.metrics_for(&root).unwrap().output_rows(), + Some(actual.iter().map(RecordBatch::num_rows).sum()) + ); + })); + } + for execution in executions { + execution.await.unwrap(); + } + for (before, after) in original + .iter() + .zip(physical_nodes(&shared.root.native_plan)) + { + assert!(Arc::ptr_eq(before, &after)); + } + assert_eq!(session.runtime_env().memory_pool.reserved(), 0); + } + } + } + } + + #[tokio::test] + async fn min_max_dynamic_bounds_do_not_replace_partition_accumulators() { + use datafusion::physical_expr::aggregate::AggregateExprBuilder; + use datafusion::physical_expr::expressions::Column; + let session = Arc::new(SessionContext::new()); + let definition = min_max_plan(false); + let mut shared = SharedPipeline::build_partitions(&definition, &session, 3).unwrap(); + let aggregate = shared + .root + .native_plan + .downcast_ref::() + .unwrap(); + // Comet normally wraps MIN/MAX arguments in CastExpr. Construct direct + // Column arguments so upstream really creates its plan-owned dynamic filter. + let expressions = aggregate + .aggr_expr() + .iter() + .map(|a| { + Arc::new( + AggregateExprBuilder::new( + Arc::new(a.fun().clone()), + vec![Arc::new(Column::new( + aggregate.input_schema().field(0).name(), + 0, + ))], + ) + .schema(aggregate.input_schema()) + .alias(a.name()) + .build() + .unwrap(), + ) + }) + .collect(); + let input = Arc::clone(aggregate.input()); + let raw: Arc = Arc::new( + AggregateExec::try_new( + AggregateMode::Partial, + aggregate.group_expr().clone(), + expressions, + vec![None, None], + Arc::clone(&input), + aggregate.input_schema(), + ) + .unwrap(), + ); + assert_eq!(raw.dynamic_expressions_produced().len(), 1); + let converted = convert_tree(&raw, &shared.identity, &[input], &mut vec![], 3).unwrap(); + assert_eq!(converted.dynamic_expressions_produced().len(), 1); + let owner = Arc::get_mut(&mut shared).unwrap(); + Arc::get_mut(&mut owner.root).unwrap().native_plan = converted; + let root = Arc::clone(&shared.root.native_plan); + let dynamic = root.dynamic_expressions_produced(); + let initial_filter = dynamic[0].to_string(); + // First establish tighter bounds, then execute a disjoint partition on + // exactly the same node. The second result must retain its own extrema. + for (partition, values, expected) in [ + (0, vec![Some(-100), Some(200)], vec![Some(-100), Some(200)]), + (2, vec![Some(20), Some(30)], vec![Some(20), Some(30)]), + ] { + let planner = PhysicalPlanner::new(Arc::clone(&session), partition); + let (mut scans, attempt) = shared + .try_bind_plan(&planner, &mut vec![], &definition) + .unwrap() + .unwrap(); + feed(&mut scans[0], values); + let stream = root + .execute(partition as usize, attempt.task_context(&session)) + .unwrap(); + assert_eq!(rows(&drain_inputs(scans, stream).await), vec![expected]); + assert_ne!(dynamic[0].to_string(), initial_filter); + assert_eq!(attempt.metrics_for(&root).unwrap().output_rows(), Some(1)); + assert!(Arc::ptr_eq(&root, &shared.root.native_plan)); + } + assert_eq!(session.runtime_env().memory_pool.reserved(), 0); + } + + #[tokio::test] + async fn average_stream_metrics_and_state_are_attempt_local() { + let session = Arc::new(SessionContext::new()); + let mut definition = aggregate_plan(); + if let Some(OpStruct::HashAgg(a)) = &mut definition.op_struct { + a.grouping_exprs.clear(); + a.agg_exprs = vec![AggExpr { + expr_struct: Some(agg_expr::ExprStruct::Avg( + datafusion_comet_proto::spark_expression::Avg { + child: Some(column()), + datatype: Some(DataType { + type_id: 6, + type_info: None, + }), + sum_datatype: Some(DataType { + type_id: 6, + type_info: None, + }), + eval_mode: 0, + }, + )), + ..Default::default() + }]; + } + let shared = SharedPipeline::build_partitions(&definition, &session, 2).unwrap(); + for partition in [0, 1] { + let planner = PhysicalPlanner::new(Arc::clone(&session), partition); + let (mut scans, attempt) = shared + .bind_plan(&planner, &mut vec![], &definition) + .unwrap(); + let (mut private_scans, _, private) = + planner.create_plan(&definition, &mut vec![], 1).unwrap(); + let values = vec![None, Some(2 + partition as i64), Some(4)]; + feed(&mut scans[0], values.clone()); + feed(&mut private_scans[0], values); + let stream = shared + .root + .native_plan + .execute(partition as usize, attempt.task_context(&session)) + .unwrap(); + let ordinary = private.native_plan.execute(0, session.task_ctx()).unwrap(); + let (actual, expected) = tokio::join!( + drain_inputs(scans, stream), + drain_inputs(private_scans, ordinary) + ); + assert_eq!(actual, expected); + assert_eq!( + attempt + .metrics_for(&shared.root.native_plan) + .unwrap() + .output_rows(), + Some(1) + ); + } + } + + #[tokio::test] + async fn final_and_partial_merge_aggregates_share_the_tree() { + use datafusion_comet_proto::spark_operator::AggregateMode as ProtoMode; + let session = Arc::new(SessionContext::new()); + let mut partial = aggregate_plan(); + if let Some(OpStruct::HashAgg(a)) = &mut partial.op_struct { + // COUNT has one i64 buffer, simplifying explicit merge inputs. + a.agg_exprs.truncate(1); + } + for mode in [ProtoMode::Final, ProtoMode::PartialMerge] { + let mut definition = partial.clone(); + if let Some(OpStruct::HashAgg(a)) = &mut definition.op_struct { + a.mode = mode as i32; + a.initial_input_buffer_offset = 1; + } + if let Some(OpStruct::Scan(scan)) = &mut definition.children[0].op_struct { + scan.fields = vec![datatype(), datatype()]; + } + let shared = SharedPipeline::build_partitions(&definition, &session, 3).unwrap(); + for partition in [0, 2] { + let planner = PhysicalPlanner::new(Arc::clone(&session), partition); + let (mut scans, attempt) = shared + .bind_plan(&planner, &mut vec![], &definition) + .unwrap(); + let (mut reference_scans, _, reference) = + planner.create_plan(&definition, &mut vec![], 1).unwrap(); + let columns: Vec> = vec![ + Arc::new(Int64Array::from(vec![1, 1, 2])), + Arc::new(Int64Array::from(vec![2, 3, 4])), + ]; + scans[0].set_input_batch(InputBatch::Batch(columns.clone(), 3)); + reference_scans[0].set_input_batch(InputBatch::Batch(columns, 3)); + let stream = shared + .root + .native_plan + .execute(partition as usize, attempt.task_context(&session)) + .unwrap(); + let ordinary = reference + .native_plan + .execute(0, session.task_ctx()) + .unwrap(); + let (actual, expected) = tokio::join!( + drain_inputs(scans, stream), + drain_inputs(reference_scans, ordinary) + ); + assert_eq!(rows(&actual), rows(&expected)); + assert_eq!( + rows(&actual), + vec![vec![Some(1), Some(5)], vec![Some(2), Some(4)]] + ); + assert_eq!( + attempt + .metrics_for(&shared.root.native_plan) + .unwrap() + .output_rows(), + Some(2) + ); + } + } + } + + #[tokio::test] + async fn completed_scopes_release_tree_and_metrics() { + let cache = ScopedPlans::default(); + let session = Arc::new(SessionContext::new()); + for value in 0..2000 { + let shared = cache + .get_or_build(b"scope", || { + SharedPipeline::build(&aggregate_plan(), &session) + }) + .unwrap(); + assert!(shared.try_claim_partition(0)); + let weak = Arc::downgrade(&shared); + let root = Arc::downgrade(&shared.root.native_plan); + let planner = PhysicalPlanner::new(Arc::clone(&session), 0); + let (mut scans, attempt) = shared.bind(&planner, &mut vec![]).unwrap(); + feed(&mut scans[0], vec![Some(value)]); + let stream = shared + .root + .native_plan + .execute(0, attempt.task_context(&session)) + .unwrap(); + drop(shared); + assert!(weak.upgrade().is_some()); + let output = drain_inputs(scans, stream).await; + assert_eq!(rows(&output).len(), 1); + assert_eq!( + attempt + .metrics_for(&root.upgrade().unwrap()) + .unwrap() + .output_rows(), + Some(1) + ); + drop(attempt); + assert!(weak.upgrade().is_none()); + assert!(root.upgrade().is_none()); + } + assert_eq!(session.runtime_env().memory_pool.reserved(), 0); + } + + #[tokio::test] + async fn grouped_aggregate_accumulators_and_metrics_are_attempt_local() { + let session = Arc::new(SessionContext::new()); + let shared = SharedPipeline::build_partitions(&aggregate_plan(), &session, 3).unwrap(); + for (partition, values) in [ + vec![Some(2), Some(2), None], + vec![Some(8), Some(8), Some(8)], + vec![], + ] + .into_iter() + .enumerate() + { + let planner = PhysicalPlanner::new(Arc::clone(&session), partition as i32); + let (mut scans, attempt) = shared + .try_bind_plan(&planner, &mut vec![], &aggregate_plan()) + .unwrap() + .unwrap(); + feed(&mut scans[0], values.clone()); + let stream = shared + .root + .native_plan + .execute(partition, attempt.task_context(&session)) + .unwrap(); + let batches = drain_inputs(scans, stream).await; + let (mut ordinary_scans, _, ordinary) = planner + .create_plan(&aggregate_plan(), &mut vec![], 1) + .unwrap(); + feed(&mut ordinary_scans[0], values); + let reference = drain_inputs( + ordinary_scans, + ordinary.native_plan.execute(0, session.task_ctx()).unwrap(), + ) + .await; + assert_eq!(rows(&batches), rows(&reference)); + assert_eq!( + attempt + .metrics_for(&shared.root.native_plan) + .unwrap() + .output_rows(), + Some(rows(&batches).len()) + ); + let weak = Arc::downgrade(&attempt); + drop(attempt); + assert!(weak.upgrade().is_none()); + assert_eq!(session.runtime_env().memory_pool.reserved(), 0); + } + } + + #[tokio::test] + async fn hash_join_build_sides_join_types_and_overlapping_attempts() { + for side in [0, 1] { + for kind in 0..=5 { + let session = Arc::new(SessionContext::new()); + let definition = join_plan(side, kind); + let shared = SharedPipeline::build_partitions(&definition, &session, 2).unwrap(); + let mut executions = Vec::new(); + for (partition, base) in [0, 100].into_iter().enumerate() { + let planner = PhysicalPlanner::new(Arc::clone(&session), partition as i32); + let left = vec![None, Some(base + 1), Some(base + 2), Some(base + 2)]; + let right = vec![Some(base + 2), Some(base + 3), None]; + let (mut scans, attempt) = shared + .try_bind_plan(&planner, &mut vec![], &definition) + .unwrap() + .unwrap(); + feed(&mut scans[0], left.clone()); + feed(&mut scans[1], right.clone()); + let stream = shared + .root + .native_plan + .execute(partition, attempt.task_context(&session)) + .unwrap(); + let (mut ordinary_scans, _, ordinary) = + planner.create_plan(&definition, &mut vec![], 1).unwrap(); + feed(&mut ordinary_scans[0], left); + feed(&mut ordinary_scans[1], right); + let ordinary_stream = + ordinary.native_plan.execute(0, session.task_ctx()).unwrap(); + executions.push(async move { + let (actual, expected) = tokio::join!( + drain_inputs(scans, stream), + drain_inputs(ordinary_scans, ordinary_stream) + ); + assert_eq!(rows(&actual), rows(&expected)); + (attempt, actual) + }); + } + let a = executions.remove(0); + let b = executions.remove(0); + let ((a, out_a), (b, out_b)) = tokio::join!(a, b); + assert_eq!( + a.metrics_for(&shared.root.native_plan) + .unwrap() + .output_rows(), + Some(rows(&out_a).len()) + ); + assert_eq!( + b.metrics_for(&shared.root.native_plan) + .unwrap() + .output_rows(), + Some(rows(&out_b).len()) + ); + let weak_a = Arc::downgrade(&a); + let weak_b = Arc::downgrade(&b); + drop(a); + drop(b); + assert!(weak_a.upgrade().is_none() && weak_b.upgrade().is_none()); + assert_eq!(session.runtime_env().memory_pool.reserved(), 0); + } + } + } + fn native_file_plan(path: &std::path::Path) -> Operator { + use datafusion_comet_proto::spark_operator::{ + NativeScan, NativeScanCommon, SparkFilePartition, SparkPartitionedFile, + SparkStructField, + }; + let field = SparkStructField { + name: "id".into(), + data_type: Some(datatype()), + nullable: true, + ..Default::default() + }; + let size = std::fs::metadata(path).unwrap().len() as i64; + let scan = Operator { + plan_id: 1, + op_struct: Some(OpStruct::NativeScan(NativeScan { + common: Some(NativeScanCommon { + required_schema: vec![field.clone()], + data_schema: vec![field], + projection_vector: vec![0], + session_timezone: "UTC".into(), + case_sensitive: true, + ..Default::default() + }), + file_partition: Some(SparkFilePartition { + partitioned_file: vec![SparkPartitionedFile { + file_path: format!("file://{}", path.display()), + start: 0, + length: size, + file_size: size, + ..Default::default() + }], + }), + ..Default::default() + })), + ..Default::default() + }; + let mut plan = pipeline(); + plan.children[0].children[0] = scan; + plan + } + + #[test] + fn unsupported_stateful_paths_fall_back_as_whole_blocks() { + assert!(!supports(&sort_plan(Some(2), None))); + assert!(!supports(&sort_plan(Some(3), Some(1)))); + let file = tempfile::NamedTempFile::new().unwrap(); + let native = native_file_plan(file.path()); + assert!(!supports(&native)); + assert!(SharedPipeline::build(&native, &Arc::new(SessionContext::new())).is_err()); + } + + #[tokio::test] + async fn sort_spills_with_attempt_memory_pool_and_releases_reservations() { + use datafusion::execution::{config::SessionConfig, runtime_env::RuntimeEnvBuilder}; + let config = SessionConfig::new().with_batch_size(128); + let reservation = config.options().execution.sort_spill_reservation_bytes; + let runtime = RuntimeEnvBuilder::new() + .with_memory_limit(reservation + 12288, 1.0) + .build_arc() + .unwrap(); + let session = Arc::new(SessionContext::new_with_config_rt(config, runtime)); + for _ in 0..2 { + let shared = SharedPipeline::build(&sort_plan(None, None), &session).unwrap(); + let planner = PhysicalPlanner::new(Arc::clone(&session), 7); + let (mut scans, attempt) = shared.bind(&planner, &mut vec![]).unwrap(); + let mut stream = shared + .root + .native_plan + .execute(0, attempt.task_context(&session)) + .unwrap(); + let mut remaining = 100; + let mut ended = false; + let mut actual = Vec::new(); + tokio::time::timeout(std::time::Duration::from_secs(20), async { + loop { + let batch = futures::future::poll_fn(|cx| match stream.poll_next_unpin(cx) { + Poll::Pending => { + if !ended && scans[0].batch.lock().unwrap().is_none() { + if remaining > 0 { + feed(&mut scans[0], (0..100).rev().map(Some).collect()); + remaining -= 1; + } else { + scans[0].set_input_batch(InputBatch::EOF); + ended = true; + } + } + cx.waker().wake_by_ref(); + Poll::Pending + } + ready => ready, + }) + .await; + match batch { + Some(batch) => actual.extend( + batch + .unwrap() + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .iter() + .copied(), + ), + None => break, + } + } + }) + .await + .unwrap(); + drop(stream); + assert_eq!(actual.len(), 10000); + assert!(actual.windows(2).all(|w| w[0] <= w[1])); + let metrics = attempt.metrics_for(&shared.root.native_plan).unwrap(); + assert_eq!(metrics.output_rows(), Some(10000)); + assert!(metrics.spill_count().unwrap_or(0) > 0); + assert_eq!(session.runtime_env().memory_pool.reserved(), 0); + } + } + #[tokio::test] + async fn cancelling_join_and_aggregate_drops_attempt_and_memory() { + for definition in [join_plan(0, 0), join_plan(1, 0), aggregate_plan()] { + let session = Arc::new(SessionContext::new()); + let shared = SharedPipeline::build(&definition, &session).unwrap(); + let planner = PhysicalPlanner::new(Arc::clone(&session), 7); + let (mut scans, attempt) = shared.bind(&planner, &mut vec![]).unwrap(); + for scan in &mut scans { + feed(scan, vec![Some(1), Some(2)]); + } + let mut stream = shared + .root + .native_plan + .execute(0, attempt.task_context(&session)) + .unwrap(); + assert!(stream.next().now_or_never().is_none()); + let weak = Arc::downgrade(&attempt); + drop(stream); + drop(scans); + drop(attempt); + assert!( + weak.upgrade().is_none(), + "execution kernel retained cancelled attempt" + ); + assert_eq!(session.runtime_env().memory_pool.reserved(), 0); + } + } +} diff --git a/native/core/src/execution/spark_config.rs b/native/core/src/execution/spark_config.rs index 4c2811cb5de..604c2a4ddde 100644 --- a/native/core/src/execution/spark_config.rs +++ b/native/core/src/execution/spark_config.rs @@ -18,6 +18,8 @@ use std::collections::HashMap; pub(crate) const COMET_TRACING_ENABLED: &str = "spark.comet.tracing.enabled"; +pub(crate) const COMET_EXEC_SHARED_PLAN_ENABLED: &str = "spark.comet.exec.sharedPlan.enabled"; +pub(crate) const COMET_EXEC_PLAN_CACHE_ENABLED: &str = "spark.comet.exec.planCache.enabled"; pub(crate) const COMET_DEBUG_ENABLED: &str = "spark.comet.debug.enabled"; pub(crate) const COMET_EXPLAIN_NATIVE_ENABLED: &str = "spark.comet.explain.native.enabled"; pub(crate) const COMET_MAX_TEMP_DIRECTORY_SIZE: &str = "spark.comet.maxTempDirectorySize"; diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index d2369a13c27..810e435df3c 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -215,6 +215,30 @@ object CometConf extends ShimCometConf { .booleanConf .createWithDefault(true) + val COMET_EXEC_SHARED_PLAN_ENABLED: ConfigEntry[Boolean] = + conf(s"$COMET_EXEC_CONFIG_PREFIX.sharedPlan.enabled") + .category(CATEGORY_EXEC) + .internal() + .doc( + "Share actual DataFusion JVM-input, filter/projection, full sort, COUNT/SUM/AVG/MIN/MAX " + + "aggregate and partitioned hash join operators within one executor, stage attempt " + + "and native block. Retries and repeated partitions use private plans. " + + "The registry has at most 64 weak entries and 8 MiB of encoded keys; trees and " + + "partition metrics are released with the last active task. Idle waves may rebuild.") + .booleanConf + .createWithDefault(false) + + val COMET_EXEC_PLAN_CACHE_ENABLED: ConfigEntry[Boolean] = + conf(s"$COMET_EXEC_CONFIG_PREFIX.planCache.enabled") + .category(CATEGORY_EXEC) + .internal() + .doc("Reuse immutable deserialized native plan definitions on an executor. Physical " + + "operators and execution state remain private to each task attempt. The cache holds " + + "at most 64 entries and 8 MiB of encoded plan keys; decoded heap usage is additional. " + + "Plans with different partition payloads are cached separately.") + .booleanConf + .createWithDefault(false) + val COMET_EXEC_PROJECT_ENABLED: ConfigEntry[Boolean] = createExecEnabledConfig("project", defaultValue = true) val COMET_EXEC_FILTER_ENABLED: ConfigEntry[Boolean] = diff --git a/spark/src/main/scala/org/apache/comet/CometExecIterator.scala b/spark/src/main/scala/org/apache/comet/CometExecIterator.scala index 4da95af18d0..40ad1ce2dc0 100644 --- a/spark/src/main/scala/org/apache/comet/CometExecIterator.scala +++ b/spark/src/main/scala/org/apache/comet/CometExecIterator.scala @@ -82,7 +82,8 @@ class CometExecIterator( shuffleBlockIterators: Map[Int, CometShuffleBlockIterator] = Map.empty, taskFilePaths: Seq[String] = Seq.empty, shufflePartitionPusher: Option[ShufflePartitionPusher] = None, - capturePartitionOffsets: Boolean = false) + capturePartitionOffsets: Boolean = false, + sharedPlanBlockId: Option[String] = None) extends Iterator[ColumnarBatch] with Logging { @@ -139,7 +140,8 @@ class CometExecIterator( // constructed on a Spark task thread (see `taskAttemptId` above); a JNI-attached Tokio // worker has neither. See CometUdfBridge.evaluate. TaskContext.get(), - Thread.currentThread().getContextClassLoader) + Thread.currentThread().getContextClassLoader, + CometExecIterator.sharedPlanScope(sharedPlanBlockId, TaskContext.get())) // Bind task-owned callbacks separately to preserve the existing createPlan JNI signature. try { @@ -362,6 +364,16 @@ class CometExecIterator( object CometExecIterator extends Logging { + // A missing block identity or a retry/speculative task must use a private native tree. + private[apache] def sharedPlanScope(blockId: Option[String], context: TaskContext): String = { + blockId + .filter(_ => context.attemptNumber() == 0) + .map { block => + s"$block:${context.stageId()}:${context.stageAttemptNumber()}" + } + .getOrElse("") + } + private def cometSqlConfs: Map[String, String] = SQLConf.get.getAllConfs.filter(_._1.startsWith(CometConf.COMET_PREFIX)) @@ -388,6 +400,13 @@ object CometExecIterator extends Logging { CometConf.COMET_PARQUET_ROW_FILTER_PUSHDOWN_ENABLED.key, CometConf.COMET_PARQUET_ROW_FILTER_PUSHDOWN_ENABLED.get(SQLConf.get).toString) + builder.putEntries( + CometConf.COMET_EXEC_PLAN_CACHE_ENABLED.key, + CometConf.COMET_EXEC_PLAN_CACHE_ENABLED.get(SQLConf.get).toString) + builder.putEntries( + CometConf.COMET_EXEC_SHARED_PLAN_ENABLED.key, + CometConf.COMET_EXEC_SHARED_PLAN_ENABLED.get(SQLConf.get).toString) + builder.build().toByteArray } diff --git a/spark/src/main/scala/org/apache/comet/Native.scala b/spark/src/main/scala/org/apache/comet/Native.scala index 93b396ce0f2..55de490db64 100644 --- a/spark/src/main/scala/org/apache/comet/Native.scala +++ b/spark/src/main/scala/org/apache/comet/Native.scala @@ -77,7 +77,8 @@ class Native extends NativeBase { taskCPUs: Long, keyUnwrapper: CometFileKeyUnwrapper, taskContext: TaskContext, - classLoader: ClassLoader): Long + classLoader: ClassLoader, + sharedPlanScope: String = ""): Long // scalastyle:on /** diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometExecRDD.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometExecRDD.scala index 1d876dfb83f..5756ea3b446 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometExecRDD.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometExecRDD.scala @@ -70,6 +70,9 @@ private[spark] class CometExecRDD( @transient perPartitionFilePaths: Array[Seq[String]] = Array.empty) extends RDD[ColumnarBatch](sc, inputRDDs.map(rdd => new OneToOneDependency(rdd))) { + // Generated on the driver and serialized unchanged to every task for this native block. + private val sharedPlanBlockId = java.util.UUID.randomUUID().toString + // Determine partition count: from inputs if available, otherwise from parameter private val numPartitions: Int = if (inputRDDs.nonEmpty) { inputRDDs.head.partitions.length @@ -136,7 +139,8 @@ private[spark] class CometExecRDD( broadcastedHadoopConfForEncryption, encryptedFilePaths, shuffleBlockIters, - taskFilePaths = partition.filePaths) + taskFilePaths = partition.filePaths, + sharedPlanBlockId = Some(sharedPlanBlockId)) // Register ScalarSubqueries so native code can look them up subqueries.foreach(sub => CometScalarSubquery.setSubquery(it.id, sub)) diff --git a/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala index 387785dfc5b..28cd8a0ad3c 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala @@ -63,7 +63,10 @@ class CometExecSuite extends CometTestBase { override protected def test(testName: String, testTags: Tag*)(testFun: => Any)(implicit pos: Position): Unit = { super.test(testName, testTags: _*) { - withSQLConf(CometConf.COMET_SHUFFLE_ENABLED.key -> "true") { + withSQLConf( + CometConf.COMET_SHUFFLE_ENABLED.key -> "true", + CometConf.COMET_EXEC_PLAN_CACHE_ENABLED.key -> "true", + CometConf.COMET_EXEC_SHARED_PLAN_ENABLED.key -> "true") { testFun } } @@ -90,6 +93,114 @@ class CometExecSuite extends CometTestBase { } } + test("native plan cache setting crosses JNI for both enabled and disabled execution") { + for (enabled <- Seq("true", "false")) { + withSQLConf(CometConf.COMET_EXEC_PLAN_CACHE_ENABLED.key -> enabled) { + val configs = ConfigMap.parseFrom(CometExecIterator.serializeCometSQLConfs()) + assert(configs.getEntriesMap.get(CometConf.COMET_EXEC_PLAN_CACHE_ENABLED.key) == enabled) + withParquetTable((0 until 32).map(i => (i, i + 1)), "plan_cache_input") { + checkSparkAnswerAndOperator( + sql("SELECT _1 + 1 FROM plan_cache_input WHERE _2 > 8"), + Seq(classOf[CometProjectExec])) + } + } + } + } + + test("shared native pipelines across task waves and AQE") { + for (enabled <- Seq("true", "false"); aqe <- Seq("true", "false")) { + withSQLConf( + CometConf.COMET_EXEC_SHARED_PLAN_ENABLED.key -> enabled, + CometConf.COMET_EXEC_PLAN_CACHE_ENABLED.key -> "false", + CometConf.COMET_EXPLAIN_NATIVE_ENABLED.key -> "true", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe, + SQLConf.ARROW_EXECUTION_MAX_RECORDS_PER_BATCH.key -> "17") { + val configs = ConfigMap.parseFrom(CometExecIterator.serializeCometSQLConfs()) + assert(configs.getEntriesMap.get(CometConf.COMET_EXEC_SHARED_PLAN_ENABLED.key) == enabled) + for (_ <- 0 until 2) { + val df = spark.range(0, 1000, 1, 16).where("id > 100").selectExpr("id + 10 AS value") + val (_, nativePlan) = checkSparkAnswerAndOperator(df, Seq(classOf[CometProjectExec])) + val projects = stripAQEPlan(nativePlan).collect { case p: CometProjectExec => p } + assert(projects.nonEmpty) + assert(projects.head.metrics("output_rows").value == 899L) + } + val empty = spark.range(0, 100, 1, 16).where("id < 0").selectExpr("id + 10 AS value") + checkSparkAnswerAndOperator(empty, Seq(classOf[CometProjectExec])) + // Stateful expressions in an otherwise eligible JVM-input block use private plans. + val stateful = spark + .range(0, 100, 1, 16) + .selectExpr("spark_partition_id() AS partition", "monotonically_increasing_id() AS id") + checkSparkAnswerAndOperator(stateful, Seq(classOf[CometProjectExec])) + } + } + } + + test("shared DataFusion stateful operators across Spark partitions") { + for (enabled <- Seq("true", "false"); aqe <- Seq("true", "false")) { + withSQLConf( + CometConf.COMET_EXEC_SHARED_PLAN_ENABLED.key -> enabled, + CometConf.COMET_EXPLAIN_NATIVE_ENABLED.key -> "true", + CometConf.COMET_SHUFFLE_MODE.key -> "jvm", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe, + SQLConf.ARROW_EXECUTION_MAX_RECORDS_PER_BATCH.key -> "17") { + val input = spark.range(0, 257, 1, 8).toDF() + checkSparkAnswerAndOperator( + input.sortWithinPartitions(desc("id")), + Seq(classOf[CometSortExec])) + checkSparkAnswerAndOperator(input.orderBy(desc("id")).limit(13)) + checkSparkAnswerAndOperator( + input.groupBy("id").agg(sum("id"), count("id"), min("id"), max("id"), avg("id"))) + checkSparkAnswerAndOperator(input.agg(sum("id"), count("id"), avg("id"))) + checkSparkAnswerAndOperator(input.agg(min("id"), max("id"))) + checkSparkAnswerAndOperator(spark.range(0).agg(min("id"), max("id"))) + checkSparkAnswerAndOperator(spark.range(0).agg(sum("id"), count("id"), avg("id"))) + val right = broadcast(spark.range(100, 300, 1, 4).withColumnRenamed("id", "key")) + checkSparkAnswerAndOperator(input.join(right, input("id") === right("key"))) + checkSparkAnswerAndOperator(input.join(right, input("id") === right("key"), "left")) + } + } + } + + test("shared DISTINCT and mixed PartialMerge aggregates") { + for (enabled <- Seq("true", "false"); aqe <- Seq("true", "false")) { + withSQLConf( + CometConf.COMET_EXEC_SHARED_PLAN_ENABLED.key -> enabled, + CometConf.COMET_EXPLAIN_NATIVE_ENABLED.key -> "true", + CometConf.COMET_SHUFFLE_MODE.key -> "jvm", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe, + SQLConf.ARROW_EXECUTION_MAX_RECORDS_PER_BATCH.key -> "17") { + // Repeated values and nulls span task boundaries. The exchange also separates the + // input's unreviewed CASE/modulo expressions from the candidate shared aggregates. + val input = spark + .range(0, 192, 1, 8) + .selectExpr( + "id % 2 AS g", + "CASE WHEN id % 7 = 0 THEN CAST(NULL AS BIGINT) ELSE id % 5 END AS v") + .repartition(4) + checkSparkAnswerAndOperator( + input.selectExpr("count(DISTINCT v)", "sum(DISTINCT v)", "avg(DISTINCT v)")) + checkSparkAnswerAndOperator( + input.groupBy("g").agg(countDistinct("v"), sum("v"), avg("v"), min("v"), max("v"))) + } + } + } + + test("native Parquet blocks use private plans with sharing enabled or disabled") { + for (enabled <- Seq("true", "false")) { + withSQLConf( + CometConf.COMET_EXEC_SHARED_PLAN_ENABLED.key -> enabled, + CometConf.COMET_EXPLAIN_NATIVE_ENABLED.key -> "true") { + withTempPath { path => + spark.range(0, 256, 1, 8).write.parquet(path.toString) + val df = spark.read.parquet(path.toString).selectExpr("id + 10 AS value") + checkSparkAnswerAndOperator( + df, + Seq(classOf[CometNativeScanExec], classOf[CometProjectExec])) + } + } + } + } + test("sample without replacement") { withParquetTable((0 until 1000).map(i => (i, i + 1)), "tbl") { val df = sql("SELECT * FROM tbl").sample(withReplacement = false, fraction = 0.3, seed = 42) diff --git a/spark/src/test/scala/org/apache/spark/CometExecIteratorLifecycleSuite.scala b/spark/src/test/scala/org/apache/spark/CometExecIteratorLifecycleSuite.scala index a5b98dd5988..735011ad6d5 100644 --- a/spark/src/test/scala/org/apache/spark/CometExecIteratorLifecycleSuite.scala +++ b/spark/src/test/scala/org/apache/spark/CometExecIteratorLifecycleSuite.scala @@ -24,6 +24,9 @@ import java.lang.ref.WeakReference import java.util.Properties import java.util.concurrent.atomic.AtomicBoolean +import org.scalactic.source.Position +import org.scalatest.Tag + import org.apache.spark.executor.TaskMetrics import org.apache.spark.memory.{TaskMemoryManager, TestMemoryManager} import org.apache.spark.sql.CometTestBase @@ -43,16 +46,33 @@ import org.apache.comet.serde.OperatorOuterClass */ class CometExecIteratorLifecycleSuite extends CometTestBase { - private def withTaskContext[T](taskAttemptId: Long)(f: => T): T = { + // Retaining decoded definitions or shared trees must not retain task memory managers or JNI refs, + // including when session setup or iterator teardown fails. + override protected def test(testName: String, testTags: Tag*)(testFun: => Any)(implicit + pos: Position): Unit = { + super.test(testName, testTags: _*) { + withSQLConf( + CometConf.COMET_EXEC_PLAN_CACHE_ENABLED.key -> "true", + CometConf.COMET_EXEC_SHARED_PLAN_ENABLED.key -> "true") { + testFun + } + } + } + + private def withTaskContext[T]( + taskAttemptId: Long, + stageId: Int = 0, + stageAttemptNumber: Int = 0, + attemptNumber: Int = 0)(f: => T): T = { val memoryManager = new TestMemoryManager(new SparkConf()) val taskMemoryManager = new TaskMemoryManager(memoryManager, taskAttemptId) val taskContext = new TaskContextImpl( - stageId = 0, - stageAttemptNumber = 0, + stageId = stageId, + stageAttemptNumber = stageAttemptNumber, partitionId = 0, numPartitions = 1, taskAttemptId = taskAttemptId, - attemptNumber = 0, + attemptNumber = attemptNumber, taskMemoryManager = taskMemoryManager, localProperties = new Properties, metricsSystem = null, @@ -68,6 +88,27 @@ class CometExecIteratorLifecycleSuite extends CometTestBase { } } + test("shared plan scope isolates stages, stage attempts, blocks and task retries") { + def scope( + block: Option[String], + task: Long, + stage: Int, + stageAttempt: Int, + attempt: Int): String = { + withTaskContext(task, stage, stageAttempt, attempt) { + CometExecIterator.sharedPlanScope(block, TaskContext.get()) + } + } + val first = scope(Some("block-a"), 10L, 3, 0, 0) + assert(first.nonEmpty) + assert(first == scope(Some("block-a"), 11L, 3, 0, 0)) + assert(first != scope(Some("block-a"), 12L, 4, 0, 0)) + assert(first != scope(Some("block-a"), 13L, 3, 1, 0)) + assert(first != scope(Some("block-b"), 14L, 3, 0, 0)) + assert(scope(Some("block-a"), 15L, 3, 0, 1).isEmpty) + assert(scope(None, 16L, 3, 0, 0).isEmpty) + } + /** Retries GC until every weak reference clears or the deadline passes; returns survivors. */ private def survivorsAfterGc(refs: Seq[WeakReference[_]]): Int = { val deadline = System.nanoTime() + 30L * 1000 * 1000 * 1000 @@ -86,6 +127,8 @@ class CometExecIteratorLifecycleSuite extends CometTestBase { val badConfigs = ConfigMap .newBuilder() .putEntries("spark.comet.datafusion.no_such_namespace.option", "1") + .putEntries(CometConf.COMET_EXEC_PLAN_CACHE_ENABLED.key, "true") + .putEntries(CometConf.COMET_EXEC_SHARED_PLAN_ENABLED.key, "true") .build() .toByteArray @@ -192,12 +235,15 @@ class CometExecIteratorLifecycleSuite extends CometTestBase { schema, CometArrowStream.NATIVE_TIMEZONE, "lifecycle-test") - val limitOp = - CometExecUtils.getLimitNativePlan(Seq(PrettyAttribute("test", LongType)), 100).get + val scanOp = + CometExecUtils + .getLimitNativePlan(Seq(PrettyAttribute("test", LongType)), 100) + .get + .getChildren(0) val iter = CometExec.getCometIterator( Array(stream.asInstanceOf[Object]), 1, - limitOp, + scanOp, new ThrowingMetricNode, 1, 0, From c2e48d2a4225e38436b5669e1af881b468d09838 Mon Sep 17 00:00:00 2001 From: Liang-Chi Hsieh Date: Mon, 21 Sep 2026 10:54:11 -0700 Subject: [PATCH 2/3] test: update createPlan JNI signature assertion --- .../org/apache/comet/exec/CometNativeShuffleSuite.scala | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala index 353ee66d4d1..f8c3b30c012 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala @@ -34,7 +34,7 @@ import org.apache.arrow.memory.ArrowBuf import org.apache.arrow.vector.ipc.ArrowReader import org.apache.arrow.vector.types.pojo.{Field, Schema} import org.apache.hadoop.fs.Path -import org.apache.spark.SparkEnv +import org.apache.spark.{SparkEnv, TaskContext} import org.apache.spark.sql.{CometTestBase, DataFrame, Dataset, Row} import org.apache.spark.sql.catalyst.expressions.AttributeReference import org.apache.spark.sql.catalyst.plans.logical.LocalRelation @@ -104,9 +104,11 @@ class CometNativeShuffleSuite extends CometTestBase with AdaptiveSparkPlanHelper .toByteArray } - test("native shuffle callback registration preserves the existing createPlan JNI signature") { + test("native plan and shuffle callback JNI signatures include task context and sharing scope") { val createPlan = classOf[Native].getDeclaredMethods.find(_.getName == "createPlan").get - assert(createPlan.getParameterTypes.last == classOf[ClassLoader]) + assert( + createPlan.getParameterTypes.takeRight(3).toSeq == + Seq(classOf[TaskContext], classOf[ClassLoader], classOf[String])) val registration = classOf[Native] .getDeclaredMethod( From ffdbbd2df833c99902c44dbc1f436520342badf3 Mon Sep 17 00:00:00 2001 From: Liang-Chi Hsieh Date: Tue, 22 Sep 2026 11:57:22 -0700 Subject: [PATCH 3/3] fix: remove decoded caching and verify shared-plan execution paths Remove the decoded-plan cache and its config after the standalone measurements showed little benefit. Shared construction failures now log and use private planning before any task input streams are imported. Report shared_plan_tasks through SQL metrics and assert sharing and fallback paths in JVM tests. Restore default-mode suite coverage and the original limit lifecycle case; exercise scan cleanup separately with real shared binding. Synthetic TaskContexts now receive SQLConf through task local properties. --- native/core/src/execution/jni_api.rs | 54 +-- native/core/src/execution/metrics/utils.rs | 15 +- native/core/src/execution/mod.rs | 1 - native/core/src/execution/plan_cache.rs | 455 ------------------ native/core/src/execution/shared_pipeline.rs | 92 +++- native/core/src/execution/spark_config.rs | 1 - .../scala/org/apache/comet/CometConf.scala | 11 - .../org/apache/comet/CometExecIterator.scala | 3 - .../spark/sql/comet/CometMetricNode.scala | 3 + .../apache/comet/exec/CometExecSuite.scala | 113 ++++- .../CometExecIteratorLifecycleSuite.scala | 303 ++++++------ 11 files changed, 354 insertions(+), 697 deletions(-) delete mode 100644 native/core/src/execution/plan_cache.rs diff --git a/native/core/src/execution/jni_api.rs b/native/core/src/execution/jni_api.rs index c5cc8c65fa2..6049ca88536 100644 --- a/native/core/src/execution/jni_api.rs +++ b/native/core/src/execution/jni_api.rs @@ -329,7 +329,6 @@ 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::plan_cache::clear_plan_cache(); super::shared_pipeline::clear(); let runtime = TOKIO_RUNTIME.lock().take(); if let Some(runtime) = runtime { @@ -512,10 +511,7 @@ 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 = super::plan_cache::decode_plan( - bytes.as_slice(), - spark_config.get_bool(super::spark_config::COMET_EXEC_PLAN_CACHE_ENABLED), - )?; + 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() @@ -525,7 +521,7 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_createPlan( super::shared_pipeline::scoped_key( shared_plan_scope.as_bytes(), &super::shared_pipeline::cache_key( - &super::shared_pipeline::cache_bytes(&spark_plan, &bytes), + &bytes, &spark_config, batch_size, partition_count, @@ -994,35 +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) = - if let Some(key) = &exec_context.shared_plan_key { - let shared = super::shared_pipeline::get_or_build( + 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, - )?; - if let Some((scans, attempt)) = shared.try_bind_plan( + ) + }) + }); + let binding = shared + .as_ref() + .map(|shared| { + shared.try_bind_plan( &planner, &mut exec_context.input_sources.clone(), &exec_context.spark_plan, - )? { - exec_context.shared_attempt = Some(attempt); - (scans, vec![], Arc::clone(&shared.root)) - } else { - planner.create_plan( - &exec_context.spark_plan, - &mut exec_context.input_sources.clone(), - exec_context.partition_count, - )? - } - } else { - planner.create_plan( - &exec_context.spark_plan, - &mut exec_context.input_sources.clone(), - exec_context.partition_count, - )? - }; + ) + }) + .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; diff --git a/native/core/src/execution/metrics/utils.rs b/native/core/src/execution/metrics/utils.rs index 4634581ec79..148e3ffdf54 100644 --- a/native/core/src/execution/metrics/utils.rs +++ b/native/core/src/execution/metrics/utils.rs @@ -43,11 +43,24 @@ pub(crate) fn update_comet_metric( Some(attempt) => to_native_metric_node_with(spark_plan, &|plan| attempt.metrics_for(plan)), None => to_native_metric_node(spark_plan), }; - let jbytes = env.byte_array_from_slice(&native_metric?.encode_to_vec())?; + 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, ) -> Result { diff --git a/native/core/src/execution/mod.rs b/native/core/src/execution/mod.rs index 327c7aeb5c3..d174e6dbbc5 100644 --- a/native/core/src/execution/mod.rs +++ b/native/core/src/execution/mod.rs @@ -22,7 +22,6 @@ pub mod jni_api; pub(crate) mod merge_as_partial; pub(crate) mod metrics; pub mod operators; -mod plan_cache; pub(crate) mod planner; pub mod serde; mod shared_pipeline; diff --git a/native/core/src/execution/plan_cache.rs b/native/core/src/execution/plan_cache.rs deleted file mode 100644 index f9eccbdfe65..00000000000 --- a/native/core/src/execution/plan_cache.rs +++ /dev/null @@ -1,455 +0,0 @@ -// 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. - -//! Executor-local reuse of immutable protobuf plans and bounded single-flight cache (#1204). -//! -//! The definition cache here owns protobuf data only. `shared_pipeline` separately uses the -//! generic cache for audited immutable physical trees. Neither cache may retain task resources. -//! Different partitions' scan payloads must remain different definition-cache keys. - -use std::collections::HashMap; -use std::sync::{Arc, LazyLock}; -use std::time::Instant; - -use datafusion_comet_proto::spark_operator::Operator; -use once_cell::sync::OnceCell; -use parking_lot::Mutex; - -use super::operators::ExecutionError; -use super::serde::deserialize_op; - -const MAX_ENTRIES: usize = 64; -const MAX_ENCODED_BYTES: usize = 8 * 1024 * 1024; - -// Process-local because tasks do not share a SessionContext. This cache owns only immutable -// protobuf data, never storage clients, JNI references or task resources. Exact bytes are -// compared, so neither hash collisions nor another query's configuration can change the decoded -// result. Retention is bounded by entry count and encoded bytes, and release_runtime clears it. -// The byte budget accounts for keys, not the decoded Rust heap (which can be larger). -static PLAN_CACHE: LazyLock> = - LazyLock::new(|| PlanCache::new(MAX_ENTRIES, MAX_ENCODED_BYTES)); - -pub(super) fn decode_plan( - bytes: &[u8], - cache_enabled: bool, -) -> Result, ExecutionError> { - if cache_enabled { - PLAN_CACHE.get_or_build(bytes, deserialize_op) - } else { - deserialize_op(bytes).map(Arc::new) - } -} - -pub(super) fn clear_plan_cache() { - PLAN_CACHE.clear(); -} - -type PlanSlot = Arc>>; - -struct CacheEntry { - plan: PlanSlot, - last_used: Instant, -} - -struct CacheState { - entries: HashMap, CacheEntry>, - encoded_bytes: usize, -} - -impl Default for CacheState { - fn default() -> Self { - Self { - entries: HashMap::new(), - encoded_bytes: 0, - } - } -} - -pub(super) struct PlanCache { - state: Mutex>, - max_entries: usize, - max_encoded_bytes: usize, -} - -impl PlanCache { - pub(super) fn new(max_entries: usize, max_encoded_bytes: usize) -> Self { - Self { - state: Mutex::new(CacheState::default()), - max_entries, - max_encoded_bytes, - } - } - - pub(super) fn get_or_build( - &self, - bytes: &[u8], - build: impl FnOnce(&[u8]) -> Result, - ) -> Result, ExecutionError> { - // A large plan must not evict the entire cache or circumvent the admission budget. - if self.max_entries == 0 || bytes.len() > self.max_encoded_bytes { - return build(bytes).map(Arc::new); - } - - let slot = { - let mut state = self.state.lock(); - if let Some(entry) = state.entries.get_mut(bytes) { - entry.last_used = Instant::now(); - Arc::clone(&entry.plan) - } else { - while state.entries.len() >= self.max_entries - || bytes.len() > self.max_encoded_bytes - state.encoded_bytes - { - let oldest = state - .entries - .iter() - .min_by_key(|(_, entry)| entry.last_used) - .map(|(key, _)| Arc::clone(key)) - .expect("an over-budget cache has an entry to evict"); - state.entries.remove(&oldest); - state.encoded_bytes -= oldest.len(); - } - let slot = Arc::new(OnceCell::new()); - state.entries.insert( - Arc::from(bytes), - CacheEntry { - plan: Arc::clone(&slot), - last_used: Instant::now(), - }, - ); - state.encoded_bytes += bytes.len(); - slot - } - }; - - // Only one successful build per resident entry, even on concurrent first use. The - // cache mutex is not held while building or waiting; unrelated plans can make progress. - let result = slot.get_or_try_init(|| build(bytes).map(Arc::new)).cloned(); - if result.is_err() { - // Do not retain malformed plans. An eviction/clear and a new insertion may have - // happened while building: this failed caller must not remove the replacement. - let mut state = self.state.lock(); - if state - .entries - .get(bytes) - .is_some_and(|entry| Arc::ptr_eq(&entry.plan, &slot)) - { - state.entries.remove(bytes); - state.encoded_bytes -= bytes.len(); - } - } - result - } - - pub(super) fn clear(&self) { - let old = std::mem::take(&mut *self.state.lock()); - // Active attempts retain their own Arc. Dropping the cache never invalidates a task. - drop(old); - } -} - -#[cfg(test)] -mod tests { - use super::*; - use prost::Message; - use std::sync::atomic::{AtomicUsize, Ordering}; - use std::sync::{mpsc, Barrier}; - use std::time::Duration; - - fn encoded(id: u32) -> Vec { - Operator { - plan_id: id, - ..Default::default() - } - .encode_to_vec() - } - - #[test] - fn sequential_attempts_reuse_definition_after_previous_attempt_finishes() { - let cache = PlanCache::new(2, 1024); - let bytes = encoded(1); - let first = cache.get_or_build(&bytes, deserialize_op).unwrap(); - let weak = Arc::downgrade(&first); - drop(first); - let second = cache - .get_or_build(&bytes, |_| panic!("decoded again between task waves")) - .unwrap(); - assert!(Arc::ptr_eq(&weak.upgrade().unwrap(), &second)); - } - - #[test] - fn concurrent_first_use_decodes_once() { - let cache = PlanCache::new(2, 1024); - let decodes = AtomicUsize::new(0); - let start = Barrier::new(8); - let plans = std::thread::scope(|scope| { - let handles: Vec<_> = (0..8) - .map(|_| { - scope.spawn(|| { - start.wait(); - cache - .get_or_build(&encoded(7), |bytes| { - decodes.fetch_add(1, Ordering::SeqCst); - deserialize_op(bytes) - }) - .unwrap() - }) - }) - .collect(); - handles - .into_iter() - .map(|h| h.join().unwrap()) - .collect::>() - }); - assert_eq!(decodes.load(Ordering::SeqCst), 1); - assert!(plans.iter().all(|p| Arc::ptr_eq(p, &plans[0]))); - } - - #[test] - fn different_plan_bytes_do_not_alias() { - let cache = PlanCache::new(2, 1024); - let a = cache.get_or_build(&encoded(1), deserialize_op).unwrap(); - // Spark operator IDs alone are not a key: the payload can change across plans. - let different_payload = Operator { - plan_id: 1, - sql_text_pool: vec!["a different query".into()], - ..Default::default() - } - .encode_to_vec(); - let b = cache - .get_or_build(&different_payload, deserialize_op) - .unwrap(); - assert!(!Arc::ptr_eq(&a, &b)); - assert_eq!(a.plan_id, b.plan_id); - assert!(a.sql_text_pool.is_empty()); - assert_eq!(b.sql_text_pool, ["a different query"]); - } - - #[test] - fn least_recently_used_entry_is_evicted_without_invalidating_its_task() { - let cache = PlanCache::new(2, 1024); - let a = cache.get_or_build(&encoded(1), deserialize_op).unwrap(); - let b = cache.get_or_build(&encoded(2), deserialize_op).unwrap(); - cache.get_or_build(&encoded(1), deserialize_op).unwrap(); - cache.get_or_build(&encoded(3), deserialize_op).unwrap(); - let a_again = cache.get_or_build(&encoded(1), deserialize_op).unwrap(); - let b_again = cache.get_or_build(&encoded(2), deserialize_op).unwrap(); - assert!(Arc::ptr_eq(&a, &a_again)); - assert!(!Arc::ptr_eq(&b, &b_again)); - assert_eq!(b.plan_id, 2); - } - - #[test] - fn encoded_byte_budget_evicts_and_oversized_plans_bypass_cache() { - let bytes = encoded(1); - let cache = PlanCache::new(10, bytes.len()); - let a = cache.get_or_build(&bytes, deserialize_op).unwrap(); - cache.get_or_build(&encoded(2), deserialize_op).unwrap(); - let a_again = cache.get_or_build(&bytes, deserialize_op).unwrap(); - assert!(!Arc::ptr_eq(&a, &a_again)); - let large = Operator { - sql_text_pool: vec!["longer than the admission budget".to_owned()], - ..Default::default() - } - .encode_to_vec(); - let large_a = cache.get_or_build(&large, deserialize_op).unwrap(); - let large_b = cache.get_or_build(&large, deserialize_op).unwrap(); - assert!(!Arc::ptr_eq(&large_a, &large_b)); - assert!(Arc::ptr_eq( - &a_again, - &cache.get_or_build(&bytes, deserialize_op).unwrap() - )); - assert_eq!(cache.state.lock().encoded_bytes, bytes.len()); - } - - #[test] - fn malformed_plan_does_not_poison_or_occupy_cache() { - let cache = PlanCache::new(2, 1024); - for _ in 0..2 { - assert!(cache.get_or_build(&[0xff], deserialize_op).is_err()); - assert!(cache.state.lock().entries.is_empty()); - assert_eq!(cache.state.lock().encoded_bytes, 0); - } - assert_eq!( - cache - .get_or_build(&encoded(1), deserialize_op) - .unwrap() - .plan_id, - 1 - ); - } - - #[test] - fn clear_releases_idle_plans_but_keeps_active_attempts_valid() { - let cache = PlanCache::new(2, 1024); - let active = cache.get_or_build(&encoded(1), deserialize_op).unwrap(); - let idle = cache.get_or_build(&encoded(2), deserialize_op).unwrap(); - let idle_weak = Arc::downgrade(&idle); - drop(idle); - cache.clear(); - assert!(idle_weak.upgrade().is_none()); - assert_eq!(active.plan_id, 1); - assert_eq!(cache.state.lock().encoded_bytes, 0); - assert!(!Arc::ptr_eq( - &active, - &cache.get_or_build(&encoded(1), deserialize_op).unwrap() - )); - } - - #[test] - fn decoding_does_not_lock_out_unrelated_plans() { - let cache = PlanCache::new(2, 1024); - let (started_tx, started_rx) = mpsc::channel(); - let (release_tx, release_rx) = mpsc::channel(); - let (other_tx, other_rx) = mpsc::channel(); - std::thread::scope(|scope| { - let cache = &cache; - scope.spawn(move || { - cache - .get_or_build(&encoded(1), |bytes| { - started_tx.send(()).unwrap(); - release_rx.recv().unwrap(); - deserialize_op(bytes) - }) - .unwrap() - }); - started_rx.recv_timeout(Duration::from_secs(10)).unwrap(); - scope.spawn(|| { - cache.get_or_build(&encoded(2), deserialize_op).unwrap(); - other_tx.send(()).unwrap(); - }); - let progressed = other_rx.recv_timeout(Duration::from_secs(10)); - release_tx.send(()).unwrap(); - assert!( - progressed.is_ok(), - "unrelated decode blocked on the cache mutex" - ); - }); - } - - #[test] - fn failure_after_clear_does_not_remove_replacement() { - let cache = PlanCache::new(2, 1024); - let bytes = encoded(1); - let mut replacement = None; - let result = cache.get_or_build(&bytes, |_| { - cache.clear(); - replacement = Some(cache.get_or_build(&bytes, deserialize_op).unwrap()); - deserialize_op(&[0xff]) - }); - assert!(result.is_err()); - assert!(Arc::ptr_eq( - &replacement.unwrap(), - &cache - .get_or_build(&bytes, |_| panic!("replacement removed")) - .unwrap() - )); - } - - #[test] - fn shared_definition_keeps_partition_counters_inputs_and_metrics_attempt_local() { - use crate::execution::operators::InputBatch; - use crate::execution::planner::PhysicalPlanner; - use arrow::array::{Int32Array, Int64Array}; - use datafusion::prelude::SessionContext; - use datafusion_comet_proto::spark_expression::{ - expr::ExprStruct, DataType, EmptyExpr, Expr, - }; - use datafusion_comet_proto::spark_operator::{operator::OpStruct, Projection, Scan}; - use futures::StreamExt; - - let bytes = Operator { - children: vec![Operator { - op_struct: Some(OpStruct::Scan(Scan { - fields: vec![DataType { - type_id: 4, - type_info: None, - }], - source: "attempt-isolation".into(), - })), - ..Default::default() - }], - op_struct: Some(OpStruct::Projection(Projection { - project_list: vec![ - Expr { - expr_struct: Some(ExprStruct::SparkPartitionId(EmptyExpr {})), - ..Default::default() - }, - Expr { - expr_struct: Some(ExprStruct::MonotonicallyIncreasingId(EmptyExpr {})), - ..Default::default() - }, - ], - })), - ..Default::default() - } - .encode_to_vec(); - let cache = PlanCache::new(2, 4096); - let shared = cache.get_or_build(&bytes, deserialize_op).unwrap(); - let runtime = tokio::runtime::Runtime::new().unwrap(); - // The final attempt retries partition 7 while the earlier plans are still alive. - let mut roots = vec![]; - for (partition, rows) in [(7, 2), (9, 3), (7, 1)] { - let definition = cache.get_or_build(&bytes, deserialize_op).unwrap(); - assert!(Arc::ptr_eq(&shared, &definition)); - let session = Arc::new(SessionContext::new()); - let planner = PhysicalPlanner::new(Arc::clone(&session), partition); - let (mut scans, _, root) = planner.create_plan(&definition, &mut vec![], 10).unwrap(); - let mut stream = root.native_plan.execute(0, session.task_ctx()).unwrap(); - scans[0].set_input_batch(InputBatch::Batch( - vec![Arc::new(Int64Array::from(vec![42; rows]))], - rows, - )); - let batch = runtime.block_on(stream.next()).unwrap().unwrap(); - let partitions = batch - .column(0) - .as_any() - .downcast_ref::() - .unwrap(); - let ids = batch - .column(1) - .as_any() - .downcast_ref::() - .unwrap(); - assert_eq!(partitions.values().as_ref(), vec![partition; rows]); - assert_eq!( - ids.values().as_ref(), - (0..rows) - .map(|i| ((partition as i64) << 33) + i as i64) - .collect::>() - ); - scans[0].set_input_batch(InputBatch::EOF); - assert!(runtime.block_on(stream.next()).is_none()); - assert_eq!( - root.native_plan.metrics().unwrap().output_rows(), - Some(rows) - ); - roots.push(root); - } - assert!(!Arc::ptr_eq(&roots[0].native_plan, &roots[2].native_plan)); - assert_eq!( - roots[0].native_plan.metrics().unwrap().output_rows(), - Some(2) - ); - } - - #[test] - fn disabled_cache_does_not_share_definitions() { - let a = decode_plan(&encoded(1), false).unwrap(); - let b = decode_plan(&encoded(1), false).unwrap(); - assert!(!Arc::ptr_eq(&a, &b)); - } -} diff --git a/native/core/src/execution/shared_pipeline.rs b/native/core/src/execution/shared_pipeline.rs index 62603d4656c..9cb4dc208bc 100644 --- a/native/core/src/execution/shared_pipeline.rs +++ b/native/core/src/execution/shared_pipeline.rs @@ -42,7 +42,6 @@ use datafusion_comet_proto::spark_operator::{operator::OpStruct, Operator}; use futures::{Stream, StreamExt}; use jni::objects::{Global, JObject}; use parking_lot::Mutex; -use prost::Message; use std::collections::{HashMap, HashSet}; use std::fmt::Formatter; use std::pin::Pin; @@ -124,32 +123,18 @@ pub(super) fn cache_key( key } -pub(super) fn cache_bytes<'a>(plan: &Operator, original: &'a [u8]) -> std::borrow::Cow<'a, [u8]> { - fn has_files(plan: &Operator) -> bool { - matches!(plan.op_struct, Some(OpStruct::NativeScan(_))) - || plan.children.iter().any(has_files) - } - if has_files(plan) { - std::borrow::Cow::Owned(template_bytes(plan)) - } else { - std::borrow::Cow::Borrowed(original) - } -} - -/// Legacy file-list normalization. Native scans are currently rejected by admission, -/// so this does not expand the set of trees eligible for sharing. -pub(super) fn template_bytes(plan: &Operator) -> Vec { - fn normalize(plan: &mut Operator) { - if let Some(OpStruct::NativeScan(scan)) = plan.op_struct.as_mut() { - scan.file_partition = None; - } - for child in &mut plan.children { - normalize(child); +/// Only shared construction is recoverable: it has not imported task-owned input streams. +/// Binding or execution errors must propagate rather than retrying consumed resources. +pub(super) fn try_build( + build: impl FnOnce() -> std::result::Result, +) -> Option { + match build() { + Ok(plan) => Some(plan), + Err(error) => { + log::warn!("Cannot construct shared native plan; using a private plan: {error}"); + None } } - let mut template = plan.clone(); - normalize(&mut template); - template.encode_to_vec() } // Preserve the planner's input_plan push order: children are planned left-to-right, including @@ -735,6 +720,63 @@ mod tests { .to_vec() } + #[tokio::test] + async fn failed_shared_conversion_uses_private_plan() { + use datafusion::physical_plan::limit::GlobalLimitExec; + + let session = Arc::new(SessionContext::new()); + let definition = pipeline(); + let planner = PhysicalPlanner::new(Arc::clone(&session), 0); + let inputs = Arc::new(Mutex::new(Vec::new())); + let builder = + PhysicalPlanner::new(Arc::clone(&session), 0).with_input_plans(Arc::clone(&inputs)); + let (_, _, original) = builder.create_plan(&definition, &mut vec![], 1).unwrap(); + // Simulate a planner adding a wrapper not handled by shared conversion. Admission of + // the protobuf still succeeds and the ordinary planner can execute it correctly. + assert!(supports(&definition)); + let wrapped: Arc = Arc::new(GlobalLimitExec::new( + Arc::clone(&original.native_plan), + 0, + None, + )); + let registry = ScopedPlans::default(); + let shared = try_build(|| { + registry.get_or_build(b"failed", || { + let error = convert_tree(&wrapped, &Arc::new(()), &inputs.lock(), &mut vec![], 2) + .unwrap_err(); + assert!(error + .to_string() + .contains("Unexpected operator in shared tree")); + Err(error.into()) + }) + }); + assert!(shared.is_none()); + assert!(registry.entries.lock().is_empty()); + let (mut scans, _, private) = planner.create_plan(&definition, &mut vec![], 1).unwrap(); + feed(&mut scans[0], vec![Some(-1), Some(1), Some(2)]); + let stream = private.native_plan.execute(0, session.task_ctx()).unwrap(); + let batches = drain_inputs(scans, stream).await; + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 2); + let values: Vec<_> = batches + .iter() + .flat_map(|batch| { + batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec() + }) + .collect(); + assert_eq!(values, vec![11, 12]); + // The error was not cached; another construction of this key can succeed. + assert!(try_build(|| registry.get_or_build(b"failed", || { + SharedPipeline::build(&definition, &session) + })) + .is_some()); + } + #[test] fn concurrent_first_touch_builds_one_physical_tree_without_retaining_sessions() { let cache = Arc::new(ScopedPlans::default()); diff --git a/native/core/src/execution/spark_config.rs b/native/core/src/execution/spark_config.rs index 604c2a4ddde..7eb2e4b3189 100644 --- a/native/core/src/execution/spark_config.rs +++ b/native/core/src/execution/spark_config.rs @@ -19,7 +19,6 @@ use std::collections::HashMap; pub(crate) const COMET_TRACING_ENABLED: &str = "spark.comet.tracing.enabled"; pub(crate) const COMET_EXEC_SHARED_PLAN_ENABLED: &str = "spark.comet.exec.sharedPlan.enabled"; -pub(crate) const COMET_EXEC_PLAN_CACHE_ENABLED: &str = "spark.comet.exec.planCache.enabled"; pub(crate) const COMET_DEBUG_ENABLED: &str = "spark.comet.debug.enabled"; pub(crate) const COMET_EXPLAIN_NATIVE_ENABLED: &str = "spark.comet.explain.native.enabled"; pub(crate) const COMET_MAX_TEMP_DIRECTORY_SIZE: &str = "spark.comet.maxTempDirectorySize"; diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index 810e435df3c..3757e9ef3ed 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -228,17 +228,6 @@ object CometConf extends ShimCometConf { .booleanConf .createWithDefault(false) - val COMET_EXEC_PLAN_CACHE_ENABLED: ConfigEntry[Boolean] = - conf(s"$COMET_EXEC_CONFIG_PREFIX.planCache.enabled") - .category(CATEGORY_EXEC) - .internal() - .doc("Reuse immutable deserialized native plan definitions on an executor. Physical " + - "operators and execution state remain private to each task attempt. The cache holds " + - "at most 64 entries and 8 MiB of encoded plan keys; decoded heap usage is additional. " + - "Plans with different partition payloads are cached separately.") - .booleanConf - .createWithDefault(false) - val COMET_EXEC_PROJECT_ENABLED: ConfigEntry[Boolean] = createExecEnabledConfig("project", defaultValue = true) val COMET_EXEC_FILTER_ENABLED: ConfigEntry[Boolean] = diff --git a/spark/src/main/scala/org/apache/comet/CometExecIterator.scala b/spark/src/main/scala/org/apache/comet/CometExecIterator.scala index 40ad1ce2dc0..e25467715cf 100644 --- a/spark/src/main/scala/org/apache/comet/CometExecIterator.scala +++ b/spark/src/main/scala/org/apache/comet/CometExecIterator.scala @@ -400,9 +400,6 @@ object CometExecIterator extends Logging { CometConf.COMET_PARQUET_ROW_FILTER_PUSHDOWN_ENABLED.key, CometConf.COMET_PARQUET_ROW_FILTER_PUSHDOWN_ENABLED.get(SQLConf.get).toString) - builder.putEntries( - CometConf.COMET_EXEC_PLAN_CACHE_ENABLED.key, - CometConf.COMET_EXEC_PLAN_CACHE_ENABLED.get(SQLConf.get).toString) builder.putEntries( CometConf.COMET_EXEC_SHARED_PLAN_ENABLED.key, CometConf.COMET_EXEC_SHARED_PLAN_ENABLED.get(SQLConf.get).toString) diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometMetricNode.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometMetricNode.scala index 806c9d00eda..5ecd2a44beb 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometMetricNode.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometMetricNode.scala @@ -254,6 +254,7 @@ object CometMetricNode { */ def baselineMetrics(sc: SparkContext): Map[String, SQLMetric] = { Map( + "shared_plan_tasks" -> SQLMetrics.createMetric(sc, "tasks using a shared native plan"), "output_rows" -> SQLMetrics.createMetric(sc, "number of output rows"), "elapsed_compute" -> SQLMetrics.createNanoTimingMetric( sc, @@ -315,6 +316,7 @@ object CometMetricNode { */ def nativeScanMetrics(sc: SparkContext): Map[String, SQLMetric] = { Map( + "shared_plan_tasks" -> SQLMetrics.createMetric(sc, "tasks using a shared native plan"), "output_rows" -> SQLMetrics.createMetric(sc, "number of output rows"), "time_elapsed_opening" -> SQLMetrics.createNanoTimingMetric(sc, "Wall clock time elapsed for file opening"), @@ -438,6 +440,7 @@ object CometMetricNode { */ def joinMetrics(sc: SparkContext): Map[String, SQLMetric] = { Map( + "shared_plan_tasks" -> SQLMetrics.createMetric(sc, "tasks using a shared native plan"), "build_time" -> SQLMetrics.createNanoTimingMetric(sc, "Total time for collecting build-side of join"), "build_input_batches" -> diff --git a/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala index 28cd8a0ad3c..6350ccc221b 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala @@ -63,10 +63,7 @@ class CometExecSuite extends CometTestBase { override protected def test(testName: String, testTags: Tag*)(testFun: => Any)(implicit pos: Position): Unit = { super.test(testName, testTags: _*) { - withSQLConf( - CometConf.COMET_SHUFFLE_ENABLED.key -> "true", - CometConf.COMET_EXEC_PLAN_CACHE_ENABLED.key -> "true", - CometConf.COMET_EXEC_SHARED_PLAN_ENABLED.key -> "true") { + withSQLConf(CometConf.COMET_SHUFFLE_ENABLED.key -> "true") { testFun } } @@ -93,25 +90,10 @@ class CometExecSuite extends CometTestBase { } } - test("native plan cache setting crosses JNI for both enabled and disabled execution") { - for (enabled <- Seq("true", "false")) { - withSQLConf(CometConf.COMET_EXEC_PLAN_CACHE_ENABLED.key -> enabled) { - val configs = ConfigMap.parseFrom(CometExecIterator.serializeCometSQLConfs()) - assert(configs.getEntriesMap.get(CometConf.COMET_EXEC_PLAN_CACHE_ENABLED.key) == enabled) - withParquetTable((0 until 32).map(i => (i, i + 1)), "plan_cache_input") { - checkSparkAnswerAndOperator( - sql("SELECT _1 + 1 FROM plan_cache_input WHERE _2 > 8"), - Seq(classOf[CometProjectExec])) - } - } - } - } - test("shared native pipelines across task waves and AQE") { for (enabled <- Seq("true", "false"); aqe <- Seq("true", "false")) { withSQLConf( CometConf.COMET_EXEC_SHARED_PLAN_ENABLED.key -> enabled, - CometConf.COMET_EXEC_PLAN_CACHE_ENABLED.key -> "false", CometConf.COMET_EXPLAIN_NATIVE_ENABLED.key -> "true", SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe, SQLConf.ARROW_EXECUTION_MAX_RECORDS_PER_BATCH.key -> "17") { @@ -123,6 +105,9 @@ class CometExecSuite extends CometTestBase { val projects = stripAQEPlan(nativePlan).collect { case p: CometProjectExec => p } assert(projects.nonEmpty) assert(projects.head.metrics("output_rows").value == 899L) + assert( + projects.head.metrics("shared_plan_tasks").value == + (if (enabled == "true") 16L else 0L)) } val empty = spark.range(0, 100, 1, 16).where("id < 0").selectExpr("id + 10 AS value") checkSparkAnswerAndOperator(empty, Seq(classOf[CometProjectExec])) @@ -130,12 +115,91 @@ class CometExecSuite extends CometTestBase { val stateful = spark .range(0, 100, 1, 16) .selectExpr("spark_partition_id() AS partition", "monotonically_increasing_id() AS id") - checkSparkAnswerAndOperator(stateful, Seq(classOf[CometProjectExec])) + val (_, statefulPlan) = + checkSparkAnswerAndOperator(stateful, Seq(classOf[CometProjectExec])) + assert( + stripAQEPlan(statefulPlan) + .collect { case p: CometProjectExec => + p.metrics("shared_plan_tasks").value + } + .forall(_ == 0L)) } } } - test("shared DataFusion stateful operators across Spark partitions") { + test("shared binds identify eligible stateful operators") { + for (enabled <- Seq("false", "true"); aqe <- Seq("false", "true")) { + withSQLConf( + CometConf.COMET_EXEC_SHARED_PLAN_ENABLED.key -> enabled, + // Deliberately arrange JVM inputs for these positive admission tests. Default + // ShuffleScan fallback is checked separately below. + CometConf.COMET_SHUFFLE_DIRECT_READ_ENABLED.key -> "false", + CometConf.COMET_SHUFFLE_MODE.key -> "jvm", + CometConf.COMET_EXEC_JOIN_DYNAMIC_FILTER_ENABLED.key -> "false", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe, + SQLConf.COALESCE_PARTITIONS_ENABLED.key -> "false", + SQLConf.SHUFFLE_PARTITIONS.key -> "4") { + val input = spark.range(0, 64, 1, 4).toDF() + val (_, sorted) = checkSparkAnswerAndOperator( + input.sortWithinPartitions(desc("id")), + Seq(classOf[CometSortExec])) + val sorts = collect(sorted) { case p: CometSortExec => p } + assert(sorts.nonEmpty) + assert( + sorts.forall(p => (p.metrics("shared_plan_tasks").value > 0) == (enabled == "true"))) + + val (_, aggregated) = + checkSparkAnswerAndOperator(input.repartition(4).groupBy("id").count()) + val finals = collect(aggregated) { + case p: CometHashAggregateExec + if p.modes.contains(org.apache.spark.sql.catalyst.expressions.aggregate.Final) => + p + } + assert(finals.nonEmpty) + assert( + finals.forall(p => (p.metrics("shared_plan_tasks").value > 0) == (enabled == "true"))) + + val left = input.repartition(4) + val right = broadcast(spark.range(0, 32, 1, 2).withColumnRenamed("id", "key")) + val (_, joined) = + checkSparkAnswerAndOperator(left.join(right, left("id") === right("key"))) + val joins = collect(joined) { case p: CometBroadcastHashJoinExec => p } + assert(joins.nonEmpty) + assert( + joins.forall(p => (p.metrics("shared_plan_tasks").value > 0) == (enabled == "true"))) + } + } + } + + test("ShuffleScan blocks use private plans with sharing enabled or disabled") { + for (enabled <- Seq("false", "true")) { + withSQLConf( + CometConf.COMET_EXEC_SHARED_PLAN_ENABLED.key -> enabled, + CometConf.COMET_SHUFFLE_DIRECT_READ_ENABLED.key -> "true", + CometConf.COMET_SHUFFLE_MODE.key -> "native", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true") { + withParquetTable((0 until 64).map(i => (i, i.toLong)), "shared_shuffle_input") { + val (_, plan) = checkSparkAnswerAndOperator( + sql("SELECT * FROM shared_shuffle_input") + .repartition(1, $"_1") + .sortWithinPartitions($"_1".desc)) + val sorts = collect(plan) { case p: CometSortExec => p } + assert(sorts.nonEmpty) + assert( + sorts.exists(_.serializedPlanOpt.plan.exists { bytes => + org.apache.comet.serde.OperatorOuterClass.Operator + .parseFrom(bytes) + .toString + .contains("shuffle_scan") + }), + s"Expected a serialized ShuffleScan input:\n$plan") + assert(sorts.forall(_.metrics("shared_plan_tasks").value == 0L)) + } + } + } + } + + test("stateful operator results with sharing enabled or disabled") { for (enabled <- Seq("true", "false"); aqe <- Seq("true", "false")) { withSQLConf( CometConf.COMET_EXEC_SHARED_PLAN_ENABLED.key -> enabled, @@ -161,7 +225,7 @@ class CometExecSuite extends CometTestBase { } } - test("shared DISTINCT and mixed PartialMerge aggregates") { + test("DISTINCT and mixed PartialMerge results with sharing enabled or disabled") { for (enabled <- Seq("true", "false"); aqe <- Seq("true", "false")) { withSQLConf( CometConf.COMET_EXEC_SHARED_PLAN_ENABLED.key -> enabled, @@ -193,9 +257,12 @@ class CometExecSuite extends CometTestBase { withTempPath { path => spark.range(0, 256, 1, 8).write.parquet(path.toString) val df = spark.read.parquet(path.toString).selectExpr("id + 10 AS value") - checkSparkAnswerAndOperator( + val (_, plan) = checkSparkAnswerAndOperator( df, Seq(classOf[CometNativeScanExec], classOf[CometProjectExec])) + val projects = stripAQEPlan(plan).collect { case p: CometProjectExec => p } + assert(projects.nonEmpty) + assert(projects.forall(_.metrics("shared_plan_tasks").value == 0L)) } } } diff --git a/spark/src/test/scala/org/apache/spark/CometExecIteratorLifecycleSuite.scala b/spark/src/test/scala/org/apache/spark/CometExecIteratorLifecycleSuite.scala index 735011ad6d5..b0a6d0d6581 100644 --- a/spark/src/test/scala/org/apache/spark/CometExecIteratorLifecycleSuite.scala +++ b/spark/src/test/scala/org/apache/spark/CometExecIteratorLifecycleSuite.scala @@ -24,15 +24,13 @@ import java.lang.ref.WeakReference import java.util.Properties import java.util.concurrent.atomic.AtomicBoolean -import org.scalactic.source.Position -import org.scalatest.Tag - import org.apache.spark.executor.TaskMetrics import org.apache.spark.memory.{TaskMemoryManager, TestMemoryManager} import org.apache.spark.sql.CometTestBase import org.apache.spark.sql.catalyst.expressions.PrettyAttribute -import org.apache.spark.sql.comet.{CometExec, CometExecUtils, CometMetricNode} +import org.apache.spark.sql.comet.{CometExecUtils, CometMetricNode} import org.apache.spark.sql.comet.execution.arrow.CometArrowStream +import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.{LongType, StructField, StructType} import org.apache.comet.{CometConf, CometExecIterator, CometShuffleBlockIterator, Native} @@ -46,19 +44,6 @@ import org.apache.comet.serde.OperatorOuterClass */ class CometExecIteratorLifecycleSuite extends CometTestBase { - // Retaining decoded definitions or shared trees must not retain task memory managers or JNI refs, - // including when session setup or iterator teardown fails. - override protected def test(testName: String, testTags: Tag*)(testFun: => Any)(implicit - pos: Position): Unit = { - super.test(testName, testTags: _*) { - withSQLConf( - CometConf.COMET_EXEC_PLAN_CACHE_ENABLED.key -> "true", - CometConf.COMET_EXEC_SHARED_PLAN_ENABLED.key -> "true") { - testFun - } - } - } - private def withTaskContext[T]( taskAttemptId: Long, stageId: Int = 0, @@ -66,6 +51,12 @@ class CometExecIteratorLifecycleSuite extends CometTestBase { attemptNumber: Int = 0)(f: => T): T = { val memoryManager = new TestMemoryManager(new SparkConf()) val taskMemoryManager = new TaskMemoryManager(memoryManager, taskAttemptId) + // Spark propagates SQLConf through task local properties. SQLConf.get reads those once + // a TaskContext is installed, so an empty Properties would silently disable sharing. + val taskProperties = new Properties + SQLConf.get.getAllConfs.foreach { case (key, value) => + taskProperties.setProperty(key, value) + } val taskContext = new TaskContextImpl( stageId = stageId, stageAttemptNumber = stageAttemptNumber, @@ -74,7 +65,7 @@ class CometExecIteratorLifecycleSuite extends CometTestBase { taskAttemptId = taskAttemptId, attemptNumber = attemptNumber, taskMemoryManager = taskMemoryManager, - localProperties = new Properties, + localProperties = taskProperties, metricsSystem = null, taskMetrics = TaskMetrics.empty, cpus = 1, @@ -119,150 +110,164 @@ class CometExecIteratorLifecycleSuite extends CometTestBase { refs.count(_.get() != null) } - test("createPlan failure releases the task-shared memory pool reference") { - val nativeLib = new Native() - val emptyPlan = OperatorOuterClass.Operator.newBuilder().build().toByteArray - // An unknown DataFusion config makes createPlan fail while building the session context, - // which happens after the task-shared memory pool has been registered for the task. - val badConfigs = ConfigMap - .newBuilder() - .putEntries("spark.comet.datafusion.no_such_namespace.option", "1") - .putEntries(CometConf.COMET_EXEC_PLAN_CACHE_ENABLED.key, "true") - .putEntries(CometConf.COMET_EXEC_SHARED_PLAN_ENABLED.key, "true") - .build() - .toByteArray + for (enabled <- Seq("false", "true")) { + test( + s"createPlan failure releases the task-shared memory pool reference (sharing=$enabled)") { + val nativeLib = new Native() + val emptyPlan = OperatorOuterClass.Operator.newBuilder().build().toByteArray + // An unknown DataFusion config makes createPlan fail while building the session context, + // which happens after the task-shared memory pool has been registered for the task. + val badConfigs = ConfigMap + .newBuilder() + .putEntries("spark.comet.datafusion.no_such_namespace.option", "1") + .putEntries(CometConf.COMET_EXEC_SHARED_PLAN_ENABLED.key, enabled) + .build() + .toByteArray - val managerRefs = (0 until 10).map { i => - // Unique synthetic task attempt ids keep each iteration's pool entry independent. - val taskAttemptId = 4200000L + i - withTaskContext(taskAttemptId) { - val manager = new CometTaskMemoryManager(i, taskAttemptId) - val thrown = intercept[Throwable] { - nativeLib.createPlan( - i, - Array.empty[Object], - emptyPlan, - badConfigs, - 1, - CometMetricNode(Map.empty), - 0L, - manager, - Array(System.getProperty("java.io.tmpdir")), - 8192, - true, - "fair_unified", - 64L << 20, - 64L << 20, - taskAttemptId, - 1L, - null, - null, - null) + val managerRefs = (0 until 10).map { i => + // Unique synthetic task attempt ids keep each iteration's pool entry independent. + val taskAttemptId = 4200000L + i + withTaskContext(taskAttemptId) { + val manager = new CometTaskMemoryManager(i, taskAttemptId) + val thrown = intercept[Throwable] { + nativeLib.createPlan( + i, + Array.empty[Object], + emptyPlan, + badConfigs, + 1, + CometMetricNode(Map.empty), + 0L, + manager, + Array(System.getProperty("java.io.tmpdir")), + 8192, + true, + "fair_unified", + 64L << 20, + 64L << 20, + taskAttemptId, + 1L, + null, + null, + null) + } + // Guard against a vacuous pass: the failure must be the injected config error thrown + // inside createPlan, not e.g. an UnsatisfiedLinkError from a missing native library. + assert( + thrown.getMessage != null && thrown.getMessage.contains("no_such_namespace"), + s"expected the injected DataFusion config failure, got: $thrown") + new WeakReference(manager) } - // Guard against a vacuous pass: the failure must be the injected config error thrown - // inside createPlan, not e.g. an UnsatisfiedLinkError from a missing native library. - assert( - thrown.getMessage != null && thrown.getMessage.contains("no_such_namespace"), - s"expected the injected DataFusion config failure, got: $thrown") - new WeakReference(manager) } + + // A stranded TASK_SHARED_MEMORY_POOLS entry holds a JNI global ref to the + // CometTaskMemoryManager, so the manager staying reachable means the pool leaked. + val survivors = survivorsAfterGc(managerRefs) + assert( + survivors == 0, + s"$survivors of ${managerRefs.size} CometTaskMemoryManagers stayed reachable: " + + "createPlan failure leaked their task-shared memory pool references") } - // A stranded TASK_SHARED_MEMORY_POOLS entry holds a JNI global ref to the - // CometTaskMemoryManager, so the manager staying reachable means the pool leaked. - val survivors = survivorsAfterGc(managerRefs) - assert( - survivors == 0, - s"$survivors of ${managerRefs.size} CometTaskMemoryManagers stayed reachable: " + - "createPlan failure leaked their task-shared memory pool references") - } + test(s"close() releases the plan when teardown throws (sharing=$enabled)") { + withSQLConf(CometConf.COMET_EXEC_SHARED_PLAN_ENABLED.key -> enabled) { + withTaskContext(4300000L) { + val boom = new java.io.IOException("injected shuffle block close failure") + val throwingBlockIter = + new CometShuffleBlockIterator(new ByteArrayInputStream(Array.emptyByteArray)) { + override def close(): Unit = throw boom + } + @volatile var laterInputClosed = false + val trackingBlockIter = + new CometShuffleBlockIterator(new ByteArrayInputStream(Array.emptyByteArray)) { + override def close(): Unit = { + laterInputClosed = true + super.close() + } + } + val limitOp = + CometExecUtils.getLimitNativePlan(Seq(PrettyAttribute("test", LongType)), 100).get + val iter = new CometExecIterator( + id = 1L, + inputObjects = Array.empty[Object], + numOutputCols = 1, + protobufQueryPlan = limitOp.toByteArray, + nativeMetrics = CometMetricNode(Map.empty), + numParts = 1, + partitionIndex = 0, + shuffleBlockIterators = Map(0 -> throwingBlockIter, 1 -> trackingBlockIter)) - test("close() is idempotent and still releases the plan when teardown throws") { - withTaskContext(4300000L) { - val boom = new java.io.IOException("injected shuffle block close failure") - val throwingBlockIter = - new CometShuffleBlockIterator(new ByteArrayInputStream(Array.emptyByteArray)) { - override def close(): Unit = throw boom + val thrown = intercept[java.io.IOException](iter.close()) + assert(thrown eq boom) + // One input's close failure must not skip the remaining resources: this close() is the only + // chance to release them, since the task-completion retry is a no-op once `closed` is set. + assert( + laterInputClosed, + "a later shuffle input was not closed after an earlier one threw") + // The first close() must have marked the iterator closed and released the plan despite the + // teardown failure: a second close() re-running releasePlan would free the native + // execution context twice, and skipping the release would strand it. + iter.close() } - @volatile var laterInputClosed = false - val trackingBlockIter = - new CometShuffleBlockIterator(new ByteArrayInputStream(Array.emptyByteArray)) { - override def close(): Unit = { - laterInputClosed = true - super.close() - } - } - val limitOp = - CometExecUtils.getLimitNativePlan(Seq(PrettyAttribute("test", LongType)), 100).get - val iter = new CometExecIterator( - id = 1L, - inputObjects = Array.empty[Object], - numOutputCols = 1, - protobufQueryPlan = limitOp.toByteArray, - nativeMetrics = CometMetricNode(Map.empty), - numParts = 1, - partitionIndex = 0, - shuffleBlockIterators = Map(0 -> throwingBlockIter, 1 -> trackingBlockIter)) - - val thrown = intercept[java.io.IOException](iter.close()) - assert(thrown eq boom) - // One input's close failure must not skip the remaining resources: this close() is the only - // chance to release them, since the task-completion retry is a no-op once `closed` is set. - assert(laterInputClosed, "a later shuffle input was not closed after an earlier one threw") - // The first close() must have marked the iterator closed and released the plan despite the - // teardown failure: a second close() re-running releasePlan would free the native - // execution context twice, and skipping the release would strand it. - iter.close() + } } } - test("releasePlan frees the native context even when the final metrics update fails") { - // Disable the periodic metrics updates inside executePlan, so the only metrics update -- and - // therefore the only place the injected failure can fire -- is the one in releasePlan. - withSQLConf(CometConf.COMET_METRICS_UPDATE_INTERVAL.key -> "0") { - withTaskContext(4400000L) { - val failMetrics = new AtomicBoolean(false) - class ThrowingMetricNode extends CometMetricNode(Map.empty, Nil) { - override def set_all_from_bytes(bytes: Array[Byte]): Unit = { - if (failMetrics.get()) { - throw new IllegalStateException("injected metrics update failure") + for (enabled <- Seq("false", "true"); useScan <- Seq(false, true)) { + test(s"releasePlan frees context after metrics failure (sharing=$enabled, scan=$useScan)") { + // Disable the periodic metrics updates inside executePlan, so the only metrics update -- and + // therefore the only place the injected failure can fire -- is the one in releasePlan. + withSQLConf( + CometConf.COMET_METRICS_UPDATE_INTERVAL.key -> "0", + CometConf.COMET_EXEC_SHARED_PLAN_ENABLED.key -> enabled) { + withTaskContext(4400000L) { + val failMetrics = new AtomicBoolean(false) + var usedSharedPlan = false + class ThrowingMetricNode extends CometMetricNode(Map.empty, Nil) { + override def set_all_from_bytes(bytes: Array[Byte]): Unit = { + usedSharedPlan = org.apache.comet.serde.Metric.NativeMetricNode + .parseFrom(bytes) + .getMetricsMap + .containsKey("shared_plan_tasks") + if (failMetrics.get()) { + throw new IllegalStateException("injected metrics update failure") + } } } - } - val schema = StructType(Seq(StructField("test", LongType, nullable = false))) - val stream = CometArrowStream.fromColumnarBatchIter( - Iterator.empty, - schema, - CometArrowStream.NATIVE_TIMEZONE, - "lifecycle-test") - val scanOp = - CometExecUtils - .getLimitNativePlan(Seq(PrettyAttribute("test", LongType)), 100) - .get - .getChildren(0) - val iter = CometExec.getCometIterator( - Array(stream.asInstanceOf[Object]), - 1, - scanOp, - new ThrowingMetricNode, - 1, - 0, - None, - Seq.empty) + val schema = StructType(Seq(StructField("test", LongType, nullable = false))) + val stream = CometArrowStream.fromColumnarBatchIter( + Iterator.empty, + schema, + CometArrowStream.NATIVE_TIMEZONE, + "lifecycle-test") + val limitOp = + CometExecUtils.getLimitNativePlan(Seq(PrettyAttribute("test", LongType)), 100).get + val operator = if (useScan) limitOp.getChildren(0) else limitOp + val iter = new CometExecIterator( + id = 2L, + inputObjects = Array(stream.asInstanceOf[Object]), + numOutputCols = 1, + protobufQueryPlan = operator.toByteArray, + nativeMetrics = new ThrowingMetricNode, + numParts = 1, + partitionIndex = 0, + sharedPlanBlockId = Some("metrics-failure-test")) - failMetrics.set(true) - // Exhausting the iterator closes it, and the close propagates the metrics failure thrown - // by the native releasePlan call. - val thrown = intercept[Throwable](iter.hasNext) - // Guard against a vacuous pass: the failure must be the injected one, thrown from the - // releasePlan metrics update (the only metrics update left with the interval disabled). - assert( - thrown.getMessage != null && thrown.getMessage.contains( - "injected metrics update failure"), - s"expected the injected metrics update failure, got: $thrown") - // The metrics failure must not have left the iterator open or the native context alive: a - // second close() must be a no-op instead of calling releasePlan again. - iter.close() + failMetrics.set(true) + // Exhausting the iterator closes it, and the close propagates the metrics failure thrown + // by the native releasePlan call. + val thrown = intercept[Throwable](iter.hasNext) + // Guard against a vacuous pass: the failure must be the injected one, thrown from the + // releasePlan metrics update (the only metrics update left with the interval disabled). + assert( + thrown.getMessage != null && thrown.getMessage.contains( + "injected metrics update failure"), + s"expected the injected metrics update failure, got: $thrown") + assert(usedSharedPlan == (useScan && enabled == "true")) + // The metrics failure must not have left the iterator open or the native context alive: a + // second close() must be a no-op instead of calling releasePlan again. + iter.close() + } } } }