andygrove commented on code in PR #6071: URL: https://github.com/apache/datafusion-comet/pull/6071#discussion_r4115855158
########## 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: If construction fails deterministically, for example the `Unexpected operator in shared tree` case this fallback exists for, nothing records the failure. Every later task with that key repeats the failing build while holding the process-wide registry lock. All shared-mode tasks on the executor queue behind it, and each one logs a warning before planning privately. Could a failed key be remembered, for example with a tombstone entry that sends later lookups straight to private planning, with the warning logged once per key? ########## 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: `update_metrics` now reports only this attempt's metrics, but `log_plan_metrics` just below still renders `plan.native_plan` with `DisplayableExecutionPlan::with_metrics`. For a shared attempt that is the shared root, so the output aggregates every partition that has run on the tree so far. The log line carries this task's partition number, but the numbers are cumulative across the stage. The scan leaf also shows nothing because `SharedInputExec` has no metrics. This is the output people use when debugging spills. Could shared attempts log metrics selected through `AttemptState::metrics_for`, or at least say in the log line that the numbers are cumulative? ########## 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: `supports` rejects `NativeScan`, so this `file_partition` reset can never run. The same goes for the `NativeScan` arm in `input_definitions` and the two `self.input_plan(...)` hooks in the `NativeScan` branch of `planner.rs` (lines 1647 and 1769). The `spark_plan` doc comment at `jni_api.rs:398` still says the plan may be shared across task attempts, which stopped being true when the decoded cache went away, and the `Arc` is no longer needed. In the tests, `full_sort_and_cancellation_have_private_state` still loops over a single `(None, None)` and mentions Top-K's threshold. Could these leftovers from earlier revisions go? The admission code is the audit boundary for this feature, so it should only describe what is actually admitted. ########## 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: The builder relies on `PhysicalPlanner::new` defaulting to `TEST_EXEC_CONTEXT_ID`. That is what makes the `Scan` branch in `planner.rs` build a `ScanExec` with no input, and that branch is documented as unit-test only. A later change to the test path would then silently change production shared builds. Could the builder ask for placeholder scans explicitly, for example keyed off `input_plans` being set? ########## 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: One thing to check when you integrate apache/datafusion#25583. It removes the snapshot cost, but each executed partition still leaves its metric handles in every operator's `ExecutionPlanMetricsSet` until the tree drops. That's six baseline metrics per operator plus the operator-specific ones. In the long-task case one tree serves the whole stage, so this is native memory that grows with stage size and that no pool accounts for. Could the rerun of the 8,192-partition benchmark also report retained bytes? If it turns out to be material, the partition cap you prototyped bounds both the CPU and the memory side. -- 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]
