You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/10/23 10:16:31 UTC
(datafusion-comet) 01/01: refactor: share native filter and projection metrics wrapper (#7056)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-7056-897f257149b9218be45493b868334a2f2803ea81
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
commit 6cc7b8625aa9f33c24de221182c59bc7fb58c97e
Author: KUAN-HAO HUANG <10...@users.noreply.github.com>
AuthorDate: Fri Oct 23 09:03:25 2026 +0000
refactor: share native filter and projection metrics wrapper (#7056)
---
.../contributor-guide/adding_a_new_operator.md | 7 +
native/operators/src/filter.rs | 212 ++--------
native/operators/src/lib.rs | 1 +
native/operators/src/metric_owning.rs | 437 +++++++++++++++++++++
native/operators/src/projection.rs | 196 +--------
5 files changed, 485 insertions(+), 368 deletions(-)
diff --git a/docs/source/contributor-guide/adding_a_new_operator.md b/docs/source/contributor-guide/adding_a_new_operator.md
index 0706e71f7d..e142aaa499 100644
--- a/docs/source/contributor-guide/adding_a_new_operator.md
+++ b/docs/source/contributor-guide/adding_a_new_operator.md
@@ -576,6 +576,13 @@ then the projection restores its order, duplicates and aliases. The predicate st
original input schema. Both native plans remain for metrics, including when they share a Spark
plan ID. Computed projections and projections that use every input column are unchanged.
+`CometFilterExec` and `CometProjectionExec` are typed aliases of `MetricOwningExec`, which owns
+the stable metrics for their Spark plan nodes. Execution-local child replacements share that
+metric set and register the executing DataFusion operator's fresh metric handles. Plan rewrites,
+including `reset_state`, create a new metric set. Both keep the operator's settings, such as the
+filter's runtime filter pushdown permission. The aliases retain their operator-specific accessors
+and distinct downcast types.
+
#### Implement the Operator
Create the operator implementation that the builder calls, either in an existing file or a new file in `native/core/src/execution/operators/`:
diff --git a/native/operators/src/filter.rs b/native/operators/src/filter.rs
index 6303824378..13d30a1675 100644
--- a/native/operators/src/filter.rs
+++ b/native/operators/src/filter.rs
@@ -15,213 +15,48 @@
// specific language governing permissions and limitations
// under the License.
-//! A DataFusion filter whose metrics remain owned by its Spark plan node.
-//!
-//! Comet normally keeps a one-to-one Spark/native plan tree so native metric
-//! handles map back to the corresponding Spark operator. Some execution-local
-//! rewrites need to replace a filter's child. DataFusion gives that replacement
-//! a new private metric set, so this adapter owns the stable metric set and
-//! registers the handles from the filter that actually executes.
+//! A DataFusion filter backed by the shared [MetricOwningExec] adapter.
-use std::fmt::Formatter;
use std::sync::Arc;
-use datafusion::common::tree_node::TreeNodeRecursion;
-use datafusion::common::{internal_err, Result, Statistics};
-use datafusion::execution::TaskContext;
use datafusion::physical_expr::PhysicalExpr;
-use datafusion::physical_plan::execution_plan::{
- CardinalityEffect, ChildrenPropertiesMode, ReplaceChildrenOptions,
-};
use datafusion::physical_plan::filter::FilterExec;
-use datafusion::physical_plan::metrics::{ExecutionPlanMetricsSet, MetricsSet};
-use datafusion::physical_plan::statistics::{ChildStats, StatisticsArgs};
-use datafusion::physical_plan::{
- DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, SendableRecordBatchStream,
-};
+use datafusion::physical_plan::ExecutionPlan;
-#[derive(Debug)]
-pub struct CometFilterExec {
- filter: FilterExec,
- metrics: ExecutionPlanMetricsSet,
+use crate::metric_owning::{MetricOwningExec, UnaryExec};
+
+pub type CometFilterExec = MetricOwningExec<FilterExec>;
+
+/// Comet settings that a filter keeps through every rebuild.
+#[derive(Clone, Debug, Default)]
+pub struct FilterSettings {
allow_runtime_filter_pushdown: bool,
}
-impl CometFilterExec {
- pub fn from_datafusion(filter: FilterExec) -> Self {
- Self {
- filter,
- metrics: ExecutionPlanMetricsSet::new(),
- allow_runtime_filter_pushdown: false,
- }
+impl UnaryExec for FilterExec {
+ type Settings = FilterSettings;
+
+ fn input(&self) -> &Arc<dyn ExecutionPlan> {
+ FilterExec::input(self)
}
+}
+impl CometFilterExec {
pub fn with_runtime_filter_pushdown(mut self, allowed: bool) -> Self {
- self.allow_runtime_filter_pushdown = allowed;
+ self.settings_mut().allow_runtime_filter_pushdown = allowed;
self
}
pub fn allows_runtime_filter_pushdown(&self) -> bool {
- self.allow_runtime_filter_pushdown
- }
-
- pub fn input(&self) -> &Arc<dyn ExecutionPlan> {
- self.filter.input()
+ self.settings().allow_runtime_filter_pushdown
}
pub fn predicate(&self) -> &Arc<dyn PhysicalExpr> {
- self.filter.predicate()
+ self.inner().predicate()
}
pub fn has_projection(&self) -> bool {
- self.filter.projection().is_some()
- }
-
- fn replace_input(
- &self,
- input: Arc<dyn ExecutionPlan>,
- options: ReplaceChildrenOptions,
- ) -> Result<FilterExec> {
- let replaced = Arc::new(self.filter.clone()).replace_children(vec![input], options)?;
- let Some(filter) = replaced.downcast_ref::<FilterExec>() else {
- return internal_err!("FilterExec child replacement changed its plan type");
- };
- Ok(filter.clone())
- }
-
- /// Replace the child for one execution while keeping the metric identity
- /// owned by the permanent Spark filter node.
- pub fn with_execution_input(
- &self,
- input: Arc<dyn ExecutionPlan>,
- ) -> Result<Arc<dyn ExecutionPlan>> {
- Ok(Arc::new(Self {
- filter: self.replace_input(
- input,
- ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute),
- )?,
- metrics: self.metrics.clone(),
- allow_runtime_filter_pushdown: self.allow_runtime_filter_pushdown,
- }))
- }
-}
-
-impl DisplayAs for CometFilterExec {
- fn fmt_as(&self, t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result {
- self.filter.fmt_as(t, f)
- }
-}
-
-impl ExecutionPlan for CometFilterExec {
- fn name(&self) -> &str {
- self.filter.name()
- }
-
- fn properties(&self) -> &Arc<PlanProperties> {
- self.filter.properties()
- }
-
- fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
- vec![self.filter.input()]
- }
-
- fn apply_expressions(
- &self,
- f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>,
- ) -> Result<TreeNodeRecursion> {
- self.filter.apply_expressions(f)
- }
-
- fn maintains_input_order(&self) -> Vec<bool> {
- self.filter.maintains_input_order()
- }
-
- fn replace_children(
- self: Arc<Self>,
- mut children: Vec<Arc<dyn ExecutionPlan>>,
- options: ReplaceChildrenOptions,
- ) -> Result<Arc<dyn ExecutionPlan>> {
- if children.len() != 1 {
- return internal_err!("CometFilterExec requires one child");
- }
- Ok(Arc::new(Self {
- filter: self.replace_input(children.remove(0), options)?,
- metrics: ExecutionPlanMetricsSet::new(),
- allow_runtime_filter_pushdown: self.allow_runtime_filter_pushdown,
- }))
- }
-
- fn with_new_children(
- self: Arc<Self>,
- children: Vec<Arc<dyn ExecutionPlan>>,
- ) -> Result<Arc<dyn ExecutionPlan>> {
- self.replace_children(
- children,
- ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute),
- )
- }
-
- fn execute(
- &self,
- partition: usize,
- context: Arc<TaskContext>,
- ) -> Result<SendableRecordBatchStream> {
- // Every execution gets fresh DataFusion filter metrics. Register their
- // handles on the stable set before returning the stream, including when
- // opening the child fails.
- let filter = self.replace_input(
- Arc::clone(self.filter.input()),
- ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute),
- )?;
- let result = filter.execute(partition, context);
- for metric in filter.metrics().unwrap_or_default().iter() {
- self.metrics.register(Arc::clone(metric));
- }
- result
- }
-
- fn metrics(&self) -> Option<MetricsSet> {
- Some(self.metrics.clone_inner())
- }
-
- fn child_stats_requests(&self, partition: Option<usize>) -> Vec<ChildStats> {
- self.filter.child_stats_requests(partition)
- }
-
- fn statistics_from_inputs(
- &self,
- input_stats: &[Arc<Statistics>],
- args: &StatisticsArgs,
- ) -> Result<Arc<Statistics>> {
- self.filter.statistics_from_inputs(input_stats, args)
- }
-
- fn cardinality_effect(&self) -> CardinalityEffect {
- self.filter.cardinality_effect()
- }
-
- fn fetch(&self) -> Option<usize> {
- self.filter.fetch()
- }
-
- fn with_fetch(&self, fetch: Option<usize>) -> Option<Arc<dyn ExecutionPlan>> {
- let filter = self.filter.with_fetch(fetch)?;
- let filter = filter.downcast_ref::<FilterExec>()?.clone();
- Some(Arc::new(Self {
- filter,
- metrics: ExecutionPlanMetricsSet::new(),
- allow_runtime_filter_pushdown: self.allow_runtime_filter_pushdown,
- }))
- }
-
- fn with_preserve_order(&self, preserve_order: bool) -> Option<Arc<dyn ExecutionPlan>> {
- let filter = self.filter.with_preserve_order(preserve_order)?;
- let filter = filter.downcast_ref::<FilterExec>()?.clone();
- Some(Arc::new(Self {
- filter,
- metrics: ExecutionPlanMetricsSet::new(),
- allow_runtime_filter_pushdown: self.allow_runtime_filter_pushdown,
- }))
+ self.inner().projection().is_some()
}
}
@@ -233,10 +68,15 @@ mod tests {
use arrow::datatypes::{DataType, Field, Schema};
use arrow::record_batch::RecordBatch;
use datafusion::common::stats::Precision;
+ use datafusion::common::tree_node::TreeNodeRecursion;
+ use datafusion::common::Result;
use datafusion::datasource::memory::MemorySourceConfig;
use datafusion::logical_expr::Operator;
use datafusion::physical_expr::expressions::{binary, col, lit};
- use datafusion::physical_plan::statistics::StatisticsContext;
+ use datafusion::physical_plan::execution_plan::{
+ ChildrenPropertiesMode, ReplaceChildrenOptions,
+ };
+ use datafusion::physical_plan::statistics::{StatisticsArgs, StatisticsContext};
fn filter() -> Result<FilterExec> {
let schema = Arc::new(Schema::new(vec![Field::new("key", DataType::Int32, false)]));
diff --git a/native/operators/src/lib.rs b/native/operators/src/lib.rs
index 97c80a9789..0a9fa33b94 100644
--- a/native/operators/src/lib.rs
+++ b/native/operators/src/lib.rs
@@ -28,6 +28,7 @@
mod expand;
mod explode;
mod filter;
+mod metric_owning;
mod projection;
mod range;
mod rank_limit;
diff --git a/native/operators/src/metric_owning.rs b/native/operators/src/metric_owning.rs
new file mode 100644
index 0000000000..32bb759697
--- /dev/null
+++ b/native/operators/src/metric_owning.rs
@@ -0,0 +1,437 @@
+// 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.
+
+//! Unary DataFusion operators whose metrics remain owned by their Spark plan node.
+//!
+//! Execution-local child replacements share the owner's metric set. Each execution
+//! registers the fresh DataFusion operator's handles on that set, even if opening
+//! its child fails. Plan rewrites, including `reset_state`, start a new metric set.
+//! Both keep the operator's Comet settings.
+
+use std::fmt::{Debug, Formatter};
+use std::sync::Arc;
+
+use datafusion::common::tree_node::TreeNodeRecursion;
+use datafusion::common::{internal_err, Result, Statistics};
+use datafusion::execution::TaskContext;
+use datafusion::physical_expr::PhysicalExpr;
+use datafusion::physical_plan::execution_plan::{
+ CardinalityEffect, ChildrenPropertiesMode, ReplaceChildrenOptions,
+};
+use datafusion::physical_plan::metrics::{ExecutionPlanMetricsSet, MetricsSet};
+use datafusion::physical_plan::statistics::{ChildStats, StatisticsArgs};
+use datafusion::physical_plan::{
+ DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, SendableRecordBatchStream,
+};
+
+/// The typed input accessor and Comet settings for the unary operators wrapped in this module.
+pub trait UnaryExec: ExecutionPlan + Clone {
+ /// Comet settings that child replacements and plan rewrites keep.
+ type Settings: Clone + Debug + Default + Send + Sync;
+
+ fn input(&self) -> &Arc<dyn ExecutionPlan>;
+}
+
+#[derive(Debug)]
+pub struct MetricOwningExec<T: UnaryExec> {
+ inner: T,
+ metrics: ExecutionPlanMetricsSet,
+ settings: T::Settings,
+}
+
+impl<T: UnaryExec> MetricOwningExec<T> {
+ pub fn from_datafusion(inner: T) -> Self {
+ Self {
+ inner,
+ metrics: ExecutionPlanMetricsSet::new(),
+ settings: T::Settings::default(),
+ }
+ }
+
+ pub(crate) fn inner(&self) -> &T {
+ &self.inner
+ }
+
+ pub(crate) fn settings(&self) -> &T::Settings {
+ &self.settings
+ }
+
+ pub(crate) fn settings_mut(&mut self) -> &mut T::Settings {
+ &mut self.settings
+ }
+
+ pub fn input(&self) -> &Arc<dyn ExecutionPlan> {
+ self.inner.input()
+ }
+
+ fn replace_input(
+ &self,
+ input: Arc<dyn ExecutionPlan>,
+ options: ReplaceChildrenOptions,
+ ) -> Result<T> {
+ let replaced = Arc::new(self.inner.clone()).replace_children(vec![input], options)?;
+ let Some(inner) = replaced.downcast_ref::<T>() else {
+ return internal_err!(
+ "{} child replacement changed its plan type",
+ self.inner.name()
+ );
+ };
+ Ok(inner.clone())
+ }
+
+ fn rebuild(&self, inner: T, metrics: ExecutionPlanMetricsSet) -> Arc<dyn ExecutionPlan> {
+ Arc::new(Self {
+ inner,
+ metrics,
+ settings: self.settings.clone(),
+ })
+ }
+
+ /// Replace the child for one execution while keeping the metric identity
+ /// owned by the permanent Spark plan node.
+ pub fn with_execution_input(
+ &self,
+ input: Arc<dyn ExecutionPlan>,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ let inner = self.replace_input(
+ input,
+ ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute),
+ )?;
+ Ok(self.rebuild(inner, self.metrics.clone()))
+ }
+}
+
+impl<T: UnaryExec> DisplayAs for MetricOwningExec<T> {
+ fn fmt_as(&self, t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result {
+ self.inner.fmt_as(t, f)
+ }
+}
+
+impl<T: UnaryExec> ExecutionPlan for MetricOwningExec<T> {
+ fn name(&self) -> &str {
+ self.inner.name()
+ }
+
+ fn properties(&self) -> &Arc<PlanProperties> {
+ self.inner.properties()
+ }
+
+ fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
+ vec![self.input()]
+ }
+
+ fn apply_expressions(
+ &self,
+ f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>,
+ ) -> Result<TreeNodeRecursion> {
+ self.inner.apply_expressions(f)
+ }
+
+ fn benefits_from_input_partitioning(&self) -> Vec<bool> {
+ self.inner.benefits_from_input_partitioning()
+ }
+
+ fn maintains_input_order(&self) -> Vec<bool> {
+ self.inner.maintains_input_order()
+ }
+
+ fn replace_children(
+ self: Arc<Self>,
+ mut children: Vec<Arc<dyn ExecutionPlan>>,
+ options: ReplaceChildrenOptions,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ if children.len() != 1 {
+ return internal_err!("{} requires one child", self.inner.name());
+ }
+ let inner = self.replace_input(children.remove(0), options)?;
+ Ok(self.rebuild(inner, ExecutionPlanMetricsSet::new()))
+ }
+
+ fn with_new_children(
+ self: Arc<Self>,
+ children: Vec<Arc<dyn ExecutionPlan>>,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ self.replace_children(
+ children,
+ ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute),
+ )
+ }
+
+ fn execute(
+ &self,
+ partition: usize,
+ context: Arc<TaskContext>,
+ ) -> Result<SendableRecordBatchStream> {
+ // Each execution gets fresh DataFusion metrics. Register their handles
+ // before returning the stream, including when opening the child fails.
+ let inner = self.replace_input(
+ Arc::clone(self.input()),
+ ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute),
+ )?;
+ let result = inner.execute(partition, context);
+ for metric in inner.metrics().unwrap_or_default().iter() {
+ self.metrics.register(Arc::clone(metric));
+ }
+ result
+ }
+
+ fn metrics(&self) -> Option<MetricsSet> {
+ Some(self.metrics.clone_inner())
+ }
+
+ fn child_stats_requests(&self, partition: Option<usize>) -> Vec<ChildStats> {
+ self.inner.child_stats_requests(partition)
+ }
+
+ fn statistics_from_inputs(
+ &self,
+ input_stats: &[Arc<Statistics>],
+ args: &StatisticsArgs,
+ ) -> Result<Arc<Statistics>> {
+ self.inner.statistics_from_inputs(input_stats, args)
+ }
+
+ fn cardinality_effect(&self) -> CardinalityEffect {
+ self.inner.cardinality_effect()
+ }
+
+ fn fetch(&self) -> Option<usize> {
+ self.inner.fetch()
+ }
+
+ fn with_fetch(&self, fetch: Option<usize>) -> Option<Arc<dyn ExecutionPlan>> {
+ let inner = self.inner.with_fetch(fetch)?;
+ let inner = inner.downcast_ref::<T>()?.clone();
+ Some(self.rebuild(inner, ExecutionPlanMetricsSet::new()))
+ }
+
+ fn with_preserve_order(&self, preserve_order: bool) -> Option<Arc<dyn ExecutionPlan>> {
+ let inner = self.inner.with_preserve_order(preserve_order)?;
+ let inner = inner.downcast_ref::<T>()?.clone();
+ Some(self.rebuild(inner, ExecutionPlanMetricsSet::new()))
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+
+ use arrow::array::Int32Array;
+ use arrow::datatypes::{DataType, Field, Schema};
+ use arrow::record_batch::RecordBatch;
+ use datafusion::datasource::memory::MemorySourceConfig;
+ use datafusion::logical_expr::Operator;
+ use datafusion::physical_expr::expressions::{binary, col, lit, IsNotNullExpr};
+ use datafusion::physical_plan::collect;
+ use datafusion::physical_plan::empty::EmptyExec;
+ use datafusion::physical_plan::filter::FilterExec;
+ use datafusion::physical_plan::projection::ProjectionExec;
+ use futures::TryStreamExt;
+
+ use crate::{CometFilterExec, CometProjectionExec};
+
+ fn input(replacement: bool) -> Result<Arc<dyn ExecutionPlan>> {
+ let schema = Arc::new(Schema::new(vec![Field::new("key", DataType::Int32, true)]));
+ let values = if replacement {
+ vec![vec![Some(1), None, Some(3)], vec![Some(4), None]]
+ } else {
+ vec![vec![Some(1), None], vec![None, Some(2)]]
+ };
+ let partitions = values
+ .into_iter()
+ .map(|values| {
+ RecordBatch::try_new(
+ Arc::clone(&schema),
+ vec![Arc::new(Int32Array::from(values))],
+ )
+ .map(|batch| vec![batch])
+ })
+ .collect::<std::result::Result<Vec<_>, _>>()?;
+ Ok(MemorySourceConfig::try_new_exec(&partitions, schema, None)?)
+ }
+
+ fn non_null(input: &Arc<dyn ExecutionPlan>) -> Arc<dyn PhysicalExpr> {
+ Arc::new(IsNotNullExpr::new(
+ col(input.schema().field(0).name(), &input.schema()).unwrap(),
+ ))
+ }
+
+ fn filter(input: Arc<dyn ExecutionPlan>) -> Result<FilterExec> {
+ FilterExec::try_new(non_null(&input), input)
+ }
+
+ fn projection(input: Arc<dyn ExecutionPlan>) -> Result<ProjectionExec> {
+ ProjectionExec::try_new(vec![(non_null(&input), "valid".into())], input)
+ }
+
+ fn rows(batches: &[RecordBatch]) -> usize {
+ batches.iter().map(RecordBatch::num_rows).sum()
+ }
+
+ async fn check_execution_metrics<T: UnaryExec>(
+ make_exec: fn(Arc<dyn ExecutionPlan>) -> Result<T>,
+ expected: [usize; 2],
+ ) -> Result<()> {
+ let original_input = input(false)?;
+ let owner = MetricOwningExec::from_datafusion(make_exec(Arc::clone(&original_input))?);
+ let replacement = input(true)?;
+ let execution = owner.with_execution_input(Arc::clone(&replacement))?;
+ assert!(Arc::ptr_eq(owner.input(), &original_input));
+ assert!(Arc::ptr_eq(
+ execution
+ .downcast_ref::<MetricOwningExec<T>>()
+ .unwrap()
+ .input(),
+ &replacement
+ ));
+
+ let context = Arc::new(TaskContext::default());
+ let stream = execution.execute(0, Arc::clone(&context))?;
+ // The original owner already holds the live handles before the stream is polled.
+ let registered = owner.metrics().unwrap();
+ assert_eq!(registered.output_rows(), Some(0));
+ assert_eq!(rows(&stream.try_collect::<Vec<_>>().await?), expected[0]);
+ assert_eq!(registered.output_rows(), Some(expected[0]));
+
+ for run in 1..=2 {
+ // collect executes both partitions, whose metrics must reach the same owner.
+ let batches = collect(Arc::clone(&execution), Arc::clone(&context)).await?;
+ assert_eq!(rows(&batches), expected.iter().sum::<usize>());
+ assert_eq!(
+ owner.metrics().unwrap().output_rows(),
+ Some(expected[0] + run * expected.iter().sum::<usize>())
+ );
+ }
+ // Later executions use fresh counters rather than updating the first execution's handles.
+ assert_eq!(registered.output_rows(), Some(expected[0]));
+ Ok(())
+ }
+
+ #[tokio::test]
+ async fn execution_rewrites_share_metrics_across_partitions_and_runs() -> Result<()> {
+ check_execution_metrics(filter, [2, 1]).await?;
+ check_execution_metrics(projection, [3, 2]).await
+ }
+
+ async fn check_plan_metrics<T: UnaryExec>(
+ make_exec: fn(Arc<dyn ExecutionPlan>) -> Result<T>,
+ expected: usize,
+ supports_fetch: bool,
+ ) -> Result<()> {
+ let owner = Arc::new(MetricOwningExec::from_datafusion(make_exec(input(true)?)?));
+ let context = Arc::new(TaskContext::default());
+ collect(
+ Arc::clone(&owner) as Arc<dyn ExecutionPlan>,
+ Arc::clone(&context),
+ )
+ .await?;
+ assert_eq!(owner.metrics().unwrap().output_rows(), Some(expected));
+
+ let mut rewrites = vec![
+ Arc::clone(&owner).replace_children(
+ vec![input(true)?],
+ ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute),
+ )?,
+ Arc::clone(&owner).replace_children(
+ vec![Arc::clone(owner.input())],
+ ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep),
+ )?,
+ Arc::clone(&owner).reset_state()?,
+ ];
+ let limited = owner.with_fetch(Some(1));
+ assert_eq!(limited.is_some(), supports_fetch);
+ if let Some(limited) = limited {
+ assert_eq!(limited.fetch(), Some(1));
+ rewrites.push(limited);
+ }
+ // The memory source does not support order-preserving rewrites for either operator.
+ assert!(owner.with_preserve_order(true).is_none());
+ for rewritten in rewrites {
+ assert!(rewritten.downcast_ref::<MetricOwningExec<T>>().is_some());
+ assert_eq!(rewritten.metrics().unwrap().iter().count(), 0);
+ let batches = collect(Arc::clone(&rewritten), Arc::clone(&context)).await?;
+ assert!(rows(&batches) > 0);
+ assert_eq!(
+ rewritten.metrics().unwrap().output_rows(),
+ Some(rows(&batches))
+ );
+ assert_eq!(owner.metrics().unwrap().output_rows(), Some(expected));
+ }
+ Ok(())
+ }
+
+ #[tokio::test]
+ async fn plan_rewrites_and_reset_start_new_metrics() -> Result<()> {
+ check_plan_metrics(filter, 3, true).await?;
+ check_plan_metrics(projection, 5, false).await
+ }
+
+ fn check_failed_execute<T: UnaryExec>(
+ make_exec: fn(Arc<dyn ExecutionPlan>) -> Result<T>,
+ ) -> Result<()> {
+ let child = Arc::new(EmptyExec::new(input(true)?.schema()));
+ let owner = MetricOwningExec::from_datafusion(make_exec(child)?);
+ let execution = owner.with_execution_input(Arc::clone(owner.input()))?;
+ assert_eq!(owner.metrics().unwrap().iter().count(), 0);
+ let error = execution
+ .execute(1, Arc::new(TaskContext::default()))
+ .err()
+ .unwrap();
+ assert!(error.to_string().contains("EmptyExec invalid partition"));
+ assert!(owner.metrics().unwrap().iter().count() > 0);
+ Ok(())
+ }
+
+ #[test]
+ fn failed_execute_registers_metrics_on_original_owner() -> Result<()> {
+ check_failed_execute(filter)?;
+ check_failed_execute(projection)
+ }
+
+ #[test]
+ fn partitioning_hints_match_each_inner_operator() -> Result<()> {
+ let filter = filter(input(true)?)?;
+ let expected = filter.benefits_from_input_partitioning();
+ assert_eq!(expected, vec![true]);
+ assert_eq!(
+ CometFilterExec::from_datafusion(filter).benefits_from_input_partitioning(),
+ expected
+ );
+
+ let source = input(true)?;
+ let schema = source.schema();
+ for (expr, expected) in [
+ (col("key", &schema)?, false),
+ (
+ binary(col("key", &schema)?, Operator::Plus, lit(5i32), &schema)?,
+ true,
+ ),
+ ] {
+ let projection =
+ ProjectionExec::try_new(vec![(expr, "key".into())], Arc::clone(&source))?;
+ assert_eq!(
+ projection.benefits_from_input_partitioning(),
+ vec![expected]
+ );
+ assert_eq!(
+ CometProjectionExec::from_datafusion(projection).benefits_from_input_partitioning(),
+ vec![expected]
+ );
+ }
+ Ok(())
+ }
+}
diff --git a/native/operators/src/projection.rs b/native/operators/src/projection.rs
index b35afcae31..a0344dfee7 100644
--- a/native/operators/src/projection.rs
+++ b/native/operators/src/projection.rs
@@ -15,198 +15,28 @@
// specific language governing permissions and limitations
// under the License.
-//! A DataFusion projection whose metrics remain owned by its Spark plan node.
-//!
-//! Comet normally keeps a one-to-one Spark/native plan tree so native metric
-//! handles map back to the corresponding Spark operator. Some execution-local
-//! rewrites need to replace a projection's child. DataFusion gives that replacement
-//! a new private metric set, so this adapter owns the stable metric set and
-//! registers the handles from the projection that actually executes.
+//! A DataFusion projection backed by the shared [MetricOwningExec] adapter.
-use std::fmt::Formatter;
use std::sync::Arc;
-use datafusion::common::tree_node::TreeNodeRecursion;
-use datafusion::common::{internal_err, Result, Statistics};
-use datafusion::execution::TaskContext;
-use datafusion::physical_expr::PhysicalExpr;
-use datafusion::physical_plan::execution_plan::{
- CardinalityEffect, ChildrenPropertiesMode, ReplaceChildrenOptions,
-};
-use datafusion::physical_plan::metrics::{ExecutionPlanMetricsSet, MetricsSet};
use datafusion::physical_plan::projection::ProjectionExec;
-use datafusion::physical_plan::statistics::{ChildStats, StatisticsArgs};
-use datafusion::physical_plan::{
- DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, SendableRecordBatchStream,
-};
+use datafusion::physical_plan::ExecutionPlan;
-#[derive(Debug)]
-pub struct CometProjectionExec {
- projection: ProjectionExec,
- metrics: ExecutionPlanMetricsSet,
-}
-
-impl CometProjectionExec {
- pub fn from_datafusion(projection: ProjectionExec) -> Self {
- Self {
- projection,
- metrics: ExecutionPlanMetricsSet::new(),
- }
- }
-
- pub fn input(&self) -> &Arc<dyn ExecutionPlan> {
- self.projection.input()
- }
-
- pub fn projection(&self) -> &ProjectionExec {
- &self.projection
- }
+use crate::metric_owning::{MetricOwningExec, UnaryExec};
- fn replace_input(
- &self,
- input: Arc<dyn ExecutionPlan>,
- options: ReplaceChildrenOptions,
- ) -> Result<ProjectionExec> {
- let replaced = Arc::new(self.projection.clone()).replace_children(vec![input], options)?;
- let Some(projection) = replaced.downcast_ref::<ProjectionExec>() else {
- return internal_err!("ProjectionExec child replacement changed its plan type");
- };
- Ok(projection.clone())
- }
+pub type CometProjectionExec = MetricOwningExec<ProjectionExec>;
- /// Replace the child for one execution while keeping the metric identity
- /// owned by the permanent Spark projection node.
- pub fn with_execution_input(
- &self,
- input: Arc<dyn ExecutionPlan>,
- ) -> Result<Arc<dyn ExecutionPlan>> {
- Ok(Arc::new(Self {
- projection: self.replace_input(
- input,
- ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute),
- )?,
- metrics: self.metrics.clone(),
- }))
- }
-}
+impl UnaryExec for ProjectionExec {
+ type Settings = ();
-impl DisplayAs for CometProjectionExec {
- fn fmt_as(&self, t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result {
- self.projection.fmt_as(t, f)
+ fn input(&self) -> &Arc<dyn ExecutionPlan> {
+ ProjectionExec::input(self)
}
}
-impl ExecutionPlan for CometProjectionExec {
- fn name(&self) -> &str {
- self.projection.name()
- }
-
- fn properties(&self) -> &Arc<PlanProperties> {
- self.projection.properties()
- }
-
- fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
- vec![self.projection.input()]
- }
-
- fn apply_expressions(
- &self,
- f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>,
- ) -> Result<TreeNodeRecursion> {
- self.projection.apply_expressions(f)
- }
-
- fn benefits_from_input_partitioning(&self) -> Vec<bool> {
- self.projection.benefits_from_input_partitioning()
- }
-
- fn maintains_input_order(&self) -> Vec<bool> {
- self.projection.maintains_input_order()
- }
-
- fn replace_children(
- self: Arc<Self>,
- mut children: Vec<Arc<dyn ExecutionPlan>>,
- options: ReplaceChildrenOptions,
- ) -> Result<Arc<dyn ExecutionPlan>> {
- if children.len() != 1 {
- return internal_err!("CometProjectionExec requires one child");
- }
- Ok(Arc::new(Self {
- projection: self.replace_input(children.remove(0), options)?,
- metrics: ExecutionPlanMetricsSet::new(),
- }))
- }
-
- fn with_new_children(
- self: Arc<Self>,
- children: Vec<Arc<dyn ExecutionPlan>>,
- ) -> Result<Arc<dyn ExecutionPlan>> {
- self.replace_children(
- children,
- ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute),
- )
- }
-
- fn execute(
- &self,
- partition: usize,
- context: Arc<TaskContext>,
- ) -> Result<SendableRecordBatchStream> {
- // Every execution gets fresh DataFusion projection metrics. Register their
- // handles on the stable set before returning the stream, including when
- // opening the child fails.
- let projection = self.replace_input(
- Arc::clone(self.projection.input()),
- ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute),
- )?;
- let result = projection.execute(partition, context);
- for metric in projection.metrics().unwrap_or_default().iter() {
- self.metrics.register(Arc::clone(metric));
- }
- result
- }
-
- fn metrics(&self) -> Option<MetricsSet> {
- Some(self.metrics.clone_inner())
- }
-
- fn child_stats_requests(&self, partition: Option<usize>) -> Vec<ChildStats> {
- self.projection.child_stats_requests(partition)
- }
-
- fn statistics_from_inputs(
- &self,
- input_stats: &[Arc<Statistics>],
- args: &StatisticsArgs,
- ) -> Result<Arc<Statistics>> {
- self.projection.statistics_from_inputs(input_stats, args)
- }
-
- fn cardinality_effect(&self) -> CardinalityEffect {
- self.projection.cardinality_effect()
- }
-
- fn fetch(&self) -> Option<usize> {
- self.projection.fetch()
- }
-
- fn with_fetch(&self, fetch: Option<usize>) -> Option<Arc<dyn ExecutionPlan>> {
- let projection = self.projection.with_fetch(fetch)?;
- let projection = projection.downcast_ref::<ProjectionExec>()?.clone();
- Some(Arc::new(Self {
- projection,
- metrics: ExecutionPlanMetricsSet::new(),
- }))
- }
-
- fn with_preserve_order(&self, preserve_order: bool) -> Option<Arc<dyn ExecutionPlan>> {
- let projection = self.projection.with_preserve_order(preserve_order)?;
- let projection = projection.downcast_ref::<ProjectionExec>()?.clone();
- Some(Arc::new(Self {
- projection,
- metrics: ExecutionPlanMetricsSet::new(),
- }))
+impl CometProjectionExec {
+ pub fn projection(&self) -> &ProjectionExec {
+ self.inner()
}
}
@@ -217,10 +47,12 @@ mod tests {
use arrow::datatypes::{DataType, Field, Schema};
use arrow::record_batch::RecordBatch;
use datafusion::common::stats::Precision;
+ use datafusion::common::tree_node::TreeNodeRecursion;
+ use datafusion::common::Result;
use datafusion::datasource::memory::MemorySourceConfig;
use datafusion::logical_expr::Operator;
use datafusion::physical_expr::expressions::{binary, col, lit};
- use datafusion::physical_plan::statistics::StatisticsContext;
+ use datafusion::physical_plan::statistics::{StatisticsArgs, StatisticsContext};
fn projection() -> Result<ProjectionExec> {
let schema = Arc::new(Schema::new(vec![Field::new("key", DataType::Int32, false)]));
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org