sunchao commented on code in PR #6071:
URL: https://github.com/apache/datafusion-comet/pull/6071#discussion_r4064182764
##########
spark/src/main/scala/org/apache/comet/Native.scala:
##########
@@ -77,7 +77,8 @@ class Native extends NativeBase {
taskCPUs: Long,
keyUnwrapper: CometFileKeyUnwrapper,
taskContext: TaskContext,
- classLoader: ClassLoader): Long
+ classLoader: ClassLoader,
+ sharedPlanScope: String = ""): Long
Review Comment:
### Correctness
[P2] Could you reconcile the existing signature test with this API change?
`CometNativeShuffleSuite` still asserts that
`createPlan.getParameterTypes.last` is `ClassLoader`. Appending
`sharedPlanScope` makes it `String`, even with both new flags disabled. The
current Spark 4.1 shuffle job fails at `CometNativeShuffleSuite.scala:109` with
exactly that mismatch, leaving Required Checks failed. Please update the test
to the intended signature, or pass the scope through the existing configuration
argument if preserving the signature remains the contract.
[Failing current-merge
job](https://github.com/apache/datafusion-comet/actions/runs/35570336827/job/106243253739)
##########
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:
### Performance
[P2] Could each attempt retain or access only its own metric handles instead
of snapshotting the entire shared tree here? DataFusion 55.1.0's
`plan.metrics()` clones its `Vec<Arc<Metric>>` under the operator's mutex
before this filter runs. Every executed partition appends metrics, and
`releasePlan` calls this path for every task. When overlapping task waves or
one long task keep the tree alive, the nth task therefore clones and scans
metrics for all n partitions. Total completion-reporting work becomes quadratic
in the partitions processed by that tree, whereas private plans only inspect
their own metrics.
A focused benchmark of this branch and snapshot representation, using
dependency stand-ins and eight metrics per partition, measured median flush
times of 0.19 µs, 26.53 µs, 233.53 µs and 1.10 ms at 1, 1,024, 8,192 and 32,768
retained partitions. These are mechanism timings, not Spark query timings. The
64-entry and encoded-key limits do not bound the metrics in one live tree.
Please make the reporting cost independent of completed partitions and add a
matched large-stage benchmark with a long-lived task plus a cache-only control.
--
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]
