viirya commented on code in PR #6071: URL: https://github.com/apache/datafusion-comet/pull/6071#discussion_r4127282625
########## native/core/src/execution/shared_pipeline.rs: ########## @@ -0,0 +1,2196 @@ +// 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 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<ScopedPlans> = LazyLock::new(ScopedPlans::default); + +#[derive(Default)] +struct ScopedPlans { + entries: Mutex<HashMap<Vec<u8>, Weak<SharedPipeline>>>, +} + +impl ScopedPlans { + fn get_or_build( + &self, + key: &[u8], + build: impl FnOnce() -> std::result::Result<Arc<SharedPipeline>, ExecutionError>, + ) -> std::result::Result<Arc<SharedPipeline>, 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()?; Review Comment: Addressed in `b02b2b9f3`. Failed builds now leave a tombstone, so subsequent lookups for that key fall back directly to private planning without repeating construction or logging. Tombstones share the registry’s 64-entry / 8-MiB key limits and can be evicted under capacity pressure. The warning is therefore once per **resident failed key**, rather than once for the executor’s lifetime. Keys that cannot be retained remain subject to those limits. Added regression coverage for failed-key suppression, capacity limits, and preservation of live entries. ########## native/core/src/execution/shared_pipeline.rs: ########## @@ -0,0 +1,2196 @@ +// 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 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<ScopedPlans> = LazyLock::new(ScopedPlans::default); + +#[derive(Default)] +struct ScopedPlans { + entries: Mutex<HashMap<Vec<u8>, Weak<SharedPipeline>>>, +} + +impl ScopedPlans { + fn get_or_build( + &self, + key: &[u8], + build: impl FnOnce() -> std::result::Result<Arc<SharedPipeline>, ExecutionError>, + ) -> std::result::Result<Arc<SharedPipeline>, 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::<usize>() + 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<u8> { + 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<String, String>, + batch_size: i32, + partition_count: i32, + task_cpus: i64, +) -> Vec<u8> { + fn append(key: &mut Vec<u8>, 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 +} + +/// 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<T>( + build: impl FnOnce() -> std::result::Result<T, ExecutionError>, +) -> Option<T> { + match build() { + Ok(plan) => Some(plan), + Err(error) => { + log::warn!("Cannot construct shared native plan; using a private plan: {error}"); + None + } + } +} + +// 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<SessionContext>, + partition_count: usize, +) -> std::result::Result<Arc<SharedPipeline>, 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<SparkPlan>, + scan_definitions: Vec<Operator>, + identity: Arc<()>, + partition_count: usize, + claimed_partitions: Mutex<HashSet<usize>>, +} + +type BoundAttempt = (Vec<ScanExec>, Arc<AttemptState>); + +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<SessionContext>, + ) -> std::result::Result<Arc<Self>, ExecutionError> { + Self::build_partitions(plan, session, 1) + } + + fn build_partitions( + plan: &Operator, + session: &Arc<SessionContext>, + partition_count: usize, + ) -> std::result::Result<Arc<Self>, 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) Review Comment: Addressed in `b02b2b9f3`. Setting `input_plans` now explicitly requests placeholder scans without importing task inputs. Shared construction no longer depends on `TEST_EXEC_CONTEXT_ID` to obtain that behavior. Added a regression test using a non-test execution ID and no task inputs. ########## native/core/src/execution/shared_pipeline.rs: ########## @@ -0,0 +1,2196 @@ +// 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 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<ScopedPlans> = LazyLock::new(ScopedPlans::default); + +#[derive(Default)] +struct ScopedPlans { + entries: Mutex<HashMap<Vec<u8>, Weak<SharedPipeline>>>, +} + +impl ScopedPlans { + fn get_or_build( + &self, + key: &[u8], + build: impl FnOnce() -> std::result::Result<Arc<SharedPipeline>, ExecutionError>, + ) -> std::result::Result<Arc<SharedPipeline>, 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::<usize>() + 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<u8> { + 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<String, String>, + batch_size: i32, + partition_count: i32, + task_cpus: i64, +) -> Vec<u8> { + fn append(key: &mut Vec<u8>, 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 +} + +/// 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<T>( + build: impl FnOnce() -> std::result::Result<T, ExecutionError>, +) -> Option<T> { + match build() { + Ok(plan) => Some(plan), + Err(error) => { + log::warn!("Cannot construct shared native plan; using a private plan: {error}"); + None + } + } +} + +// 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<SessionContext>, + partition_count: usize, +) -> std::result::Result<Arc<SharedPipeline>, 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<SparkPlan>, + scan_definitions: Vec<Operator>, + identity: Arc<()>, + partition_count: usize, + claimed_partitions: Mutex<HashSet<usize>>, +} + +type BoundAttempt = (Vec<ScanExec>, Arc<AttemptState>); + +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<SessionContext>, + ) -> std::result::Result<Arc<Self>, ExecutionError> { + Self::build_partitions(plan, session, 1) + } + + fn build_partitions( + plan: &Operator, + session: &Arc<SessionContext>, + partition_count: usize, + ) -> std::result::Result<Arc<Self>, 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; Review Comment: Addressed in `b02b2b9f3`. Removed the unreachable native-scan normalization and input-recording paths, made the deserialized `spark_plan` task-local without the unnecessary `Arc`, and removed the single-iteration sort-test loop and obsolete Top-K comments. The contributor guide now documents the admitted operators and expressions, together with the execution-state assumptions to audit when expanding admission or upgrading DataFusion. ########## native/core/src/execution/jni_api.rs: ########## @@ -1162,7 +1218,12 @@ pub extern "system" fn Java_org_apache_comet_Native_releasePlan( fn update_metrics(env: &mut Env, exec_context: &mut ExecutionContext) -> CometResult<()> { if let Some(native_query) = &exec_context.root_op { let metrics = exec_context.metrics.as_obj(); - update_comet_metric(env, metrics, native_query) + update_comet_metric( + env, + metrics, + native_query, + exec_context.shared_attempt.as_deref(), Review Comment: Addressed using the explicit-labeling option in `b02b2b9f3`. Shared-plan explain logs now say that their metrics are cumulative across the shared tree’s partitions and that task-local scan metrics are omitted. Private-plan logs are labeled task-local. This clarifies the existing display behavior; it does not implement a partition-specific explain renderer. SQL metric reporting continues to select the current attempt’s metrics. ########## 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<ScopedPlans> = LazyLock::new(ScopedPlans::default); + +#[derive(Default)] +struct ScopedPlans { + entries: Mutex<HashMap<Vec<u8>, Weak<SharedPipeline>>>, +} + +impl ScopedPlans { + fn get_or_build( + &self, + key: &[u8], + build: impl FnOnce() -> std::result::Result<Arc<SharedPipeline>, ExecutionError>, + ) -> std::result::Result<Arc<SharedPipeline>, 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::<usize>() + 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<u8> { + 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<String, String>, + batch_size: i32, + partition_count: i32, + task_cpus: i64, +) -> Vec<u8> { + fn append(key: &mut Vec<u8>, 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<u8> { + 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<SessionContext>, + partition_count: usize, +) -> std::result::Result<Arc<SharedPipeline>, 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<SparkPlan>, + scan_definitions: Vec<Operator>, + identity: Arc<()>, + partition_count: usize, + claimed_partitions: Mutex<HashSet<usize>>, +} + +type BoundAttempt = (Vec<ScanExec>, Arc<AttemptState>); + +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<SessionContext>, + ) -> std::result::Result<Arc<Self>, ExecutionError> { + Self::build_partitions(plan, session, 1) + } + + fn build_partitions( + plan: &Operator, + session: &Arc<SessionContext>, + partition_count: usize, + ) -> std::result::Result<Arc<Self>, 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<Self>, + planner: &PhysicalPlanner, + inputs: &mut Vec<Arc<Global<JObject<'static>>>>, + ) -> std::result::Result<(Vec<ScanExec>, Arc<AttemptState>), ExecutionError> { + self.bind_definitions( + planner, + inputs, + &self.scan_definitions.iter().collect::<Vec<_>>(), + 0, + ) + } + + pub fn try_bind_plan( + self: &Arc<Self>, + planner: &PhysicalPlanner, + inputs: &mut Vec<Arc<Global<JObject<'static>>>>, + task_plan: &Operator, + ) -> std::result::Result<Option<BoundAttempt>, 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<Self>, + planner: &PhysicalPlanner, + inputs: &mut Vec<Arc<Global<JObject<'static>>>>, + task_plan: &Operator, + ) -> std::result::Result<(Vec<ScanExec>, Arc<AttemptState>), 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<Self>, + planner: &PhysicalPlanner, + inputs: &mut Vec<Arc<Global<JObject<'static>>>>, + definitions: &[&Operator], + partition: usize, + ) -> std::result::Result<(Vec<ScanExec>, Arc<AttemptState>), 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<Arc<dyn ExecutionPlan>>, + _owner: Arc<SharedPipeline>, + identity: Arc<()>, + started: Vec<AtomicBool>, + partition: usize, +} + +impl AttemptState { + pub fn partition(&self) -> usize { + self.partition + } + + pub fn task_context(self: &Arc<Self>, session: &SessionContext) -> Arc<TaskContext> { + 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<dyn ExecutionPlan>) -> Option<MetricsSet> { + if let Some(input) = plan.downcast_ref::<SharedInputExec>() { + 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)) Review Comment: @andygrove Agreed. Partition-indexed retrieval addresses the snapshot/filtering cost, but it does not remove metric handles retained by the operators. I will include retained native bytes in the matched 8,192-partition rerun after integrating the upstream API, distinguishing bytes retained while the tree remains alive from allocation volume and peak RSS, and checking reclamation after the final owner drops. This remains outstanding. If retained memory is material, the partition cap remains an option to evaluate alongside its effect on reuse and construction cost. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
