viirya commented on code in PR #6071: URL: https://github.com/apache/datafusion-comet/pull/6071#discussion_r4064985175
########## 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: Agreed. I reproduced this with actual DataFusion operators in a release build, with partition 0's stream kept alive while the remaining partitions execute and complete. For a 16-expression projection/filter tree, the median cumulative snapshot/filter time over three runs was 51.3 ms at 1,024 partitions and 3.90 s at 8,192 partitions. The matched decoded-cache-only control took 0.142 ms and 1.11 ms, respectively. These are native mechanism measurements, not Spark query timings. The current DataFusion metrics API returns a full snapshot, so merely caching each task's handles after an initial full scan would move the quadratic work rather than remove it. I am evaluating options without modifying DataFusion and am adding the matched Spark cache-only versus sharing comparison with a long-lived task. This issue remains unresolved; the existing benchmark results do not justify the current reporting path. -- 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]
