viirya commented on code in PR #6071: URL: https://github.com/apache/datafusion-comet/pull/6071#discussion_r4067192287
########## 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: Thanks for identifying this. After investigating the alternatives, we plan to propose an upstream DataFusion PR to add partition-specific metrics access. The intent is to support indexed retrieval in `ExecutionPlanMetricsSet` and expose it through `ExecutionPlan`, so Comet can retrieve a task's metrics without cloning and scanning metrics from all previously executed partitions. A convenience API that still filters a full snapshot would not address the underlying cost. We also experimented with bounding the number of partitions sharing each tree. That removed the large-stage regression in our local benchmark, but it sacrifices plan reuse and introduces scheduling-dependent tradeoffs. We would prefer to address the metrics API in DataFusion rather than make that workaround part of Comet's design. Both new configs are currently disabled by default, so this does not affect users running with the defaults. Users explicitly enabling physical-plan sharing can still encounter this issue; it remains unresolved until the upstream API improvement is available and integrated. We will follow up with the DataFusion PR and rerun the matched large-stage/long-lived-task benchmark against the cache-only control after integration. -- 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]
