From ad862d5e4a9975d86d6c80e2b9093431a79d5589 Mon Sep 17 00:00:00 2001 From: Justin O'Dwyer Date: Sat, 25 Jul 2026 23:48:23 -0400 Subject: [PATCH 1/9] Create replace_children ExecutionPlan method. --- .../physical-plan/src/execution_plan.rs | 59 +++++++++++++++---- 1 file changed, 47 insertions(+), 12 deletions(-) diff --git a/datafusion/physical-plan/src/execution_plan.rs b/datafusion/physical-plan/src/execution_plan.rs index 3b9d5d258a838..63be0f1d385bd 100644 --- a/datafusion/physical-plan/src/execution_plan.rs +++ b/datafusion/physical-plan/src/execution_plan.rs @@ -247,6 +247,21 @@ pub trait ExecutionPlan: Any + Debug + DisplayAs + Send + Sync { /// joins). fn children(&self) -> Vec<&Arc>; + /// Returns a clone of the existing plan with the children replaced, skipping + /// recomputation of plan properties if possible as indicated by the hint. + fn replace_children( + self: Arc, + children: Vec>, + hint: ChildrenPropertiesHint, + ) -> Result> { + match hint { + ChildrenPropertiesHint::SameProperties => { + self.with_new_children_and_same_properties(children) + } + ChildrenPropertiesHint::Recompute => self.with_new_children(children), + } + } + /// Returns a new `ExecutionPlan` where all existing children were replaced /// by the `children`, in order fn with_new_children( @@ -306,7 +321,7 @@ pub trait ExecutionPlan: Any + Debug + DisplayAs + Send + Sync { /// If the `ExecutionPlan` does not support changing its partitioning, /// returns `Ok(None)` (the default). /// - /// It is the `ExecutionPlan` can increase its partitioning, but not to the + /// If the `ExecutionPlan` can increase its partitioning, but not to /// `target_partitions`, it may return an ExecutionPlan with fewer /// partitions. This might happen, for example, if each new partition would /// be too small to be efficiently processed individually. @@ -427,7 +442,7 @@ pub trait ExecutionPlan: Any + Debug + DisplayAs + Send + Sync { /// partition: usize, /// context: Arc, /// ) -> Result { - /// // use functions from futures crate convert the batch into a stream + /// // use functions from futures crate to convert the batch into a stream /// let fut = futures::future::ready(Ok(self.batch.clone())); /// let stream = futures::stream::once(fut); /// Ok(Box::pin(RecordBatchStreamAdapter::new( @@ -659,7 +674,7 @@ pub trait ExecutionPlan: Any + Debug + DisplayAs + Send + Sync { /// up the plan that `DataSourceExec` can actually bind the filters. /// /// The default implementation bars all parent filters from being pushed down and adds no new filters. - /// This is the safest option, making filter pushdown opt-in on a per-node pasis. + /// This is the safest option, making filter pushdown opt-in on a per-node basis. /// /// There are two different phases in filter pushdown, which some operators may handle the same and some differently. /// Depending on the phase the operator may or may not be allowed to modify the plan. @@ -846,6 +861,18 @@ pub trait ExecutionPlan: Any + Debug + DisplayAs + Send + Sync { } } +// A hint from `with_new_children_if_necessary` to `replace_children` indicating +// whether the properties of the new children must be recomputed. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ChildrenPropertiesHint { + // The properties of the new children are identical to the properties of the + // existing children, so we can skip recomuptation of plan properties. + SameProperties, + // The properties of the new children are different from the properties of the + // existing children, so we must recompute the properties from scratch. + Recompute, +} + impl dyn ExecutionPlan { /// Returns `true` if the plan is of type `T`. /// @@ -1176,12 +1203,10 @@ pub(crate) fn emission_type_from_children<'a>( } } -/// Stores certain, often expensive to compute, plan properties used in query -/// optimization. +/// Stores plan properties used in query optimization. /// -/// These properties are stored a single structure to permit this information to -/// be computed once and then those cached results used multiple times without -/// recomputation (aka a cache) +/// Serves as a cache for these properties, which are often +/// expensive to compute. #[derive(Debug, Clone)] pub struct PlanProperties { /// See [ExecutionPlanProperties::equivalence_properties] @@ -1396,11 +1421,12 @@ pub fn with_new_children_if_necessary( } // Layer 2: same child properties → reuse `PlanProperties` cache. if has_same_children_properties(plan.as_ref(), &children)? { - return plan.with_new_children_and_same_properties(children); + return plan + .replace_children(children, ChildrenPropertiesHint::SameProperties); } } // Layer 3: full recompute. - plan.with_new_children(children) + plan.replace_children(children, ChildrenPropertiesHint::Recompute) } /// Return a [`DisplayableExecutionPlan`] wrapper around an @@ -2209,9 +2235,18 @@ mod tests { let parent = Arc::new(WithChildrenTestParentDefault::new(Arc::clone(&leaf_a))); let parent_dyn: Arc = Arc::clone(&parent) as _; - // Distinct child Arc but same `PlanProperties` Arc — the helper + // Using the same child means we return the original plan Arc verbatim, so even when + // no override exists for `with_new_children_and_same_properties` we do not recompute. + let out = with_new_children_if_necessary( + Arc::clone(&parent_dyn), + vec![Arc::clone(&leaf_a)], + )?; + assert!(Arc::ptr_eq(&out, &parent_dyn)); + assert_eq!(parent.recompute_calls.load(Ordering::SeqCst), 0); + + // Using a distinct child but the same `PlanProperties` Arc means the helper // enters the "same properties" branch and calls the trait method, - // whose default forwards to `with_new_children`. + // whose default forwards to `with_new_children`, causing recomputation. let out = with_new_children_if_necessary( Arc::clone(&parent_dyn), vec![Arc::clone(&leaf_b)], From 3b75d74fb10c7db9be670879345122d0934673f9 Mon Sep 17 00:00:00 2001 From: Justin O'Dwyer Date: Sun, 26 Jul 2026 18:06:29 -0400 Subject: [PATCH 2/9] Migrate call sites. --- .../src/hash_join_buffering.rs | 24 ++++++++++++------- .../src/output_requirements.rs | 3 ++- .../physical-plan/src/execution_plan.rs | 12 +++++----- 3 files changed, 23 insertions(+), 16 deletions(-) diff --git a/datafusion/physical-optimizer/src/hash_join_buffering.rs b/datafusion/physical-optimizer/src/hash_join_buffering.rs index 7a198cac13fc9..44ebddefffd51 100644 --- a/datafusion/physical-optimizer/src/hash_join_buffering.rs +++ b/datafusion/physical-optimizer/src/hash_join_buffering.rs @@ -19,9 +19,9 @@ use crate::PhysicalOptimizerRule; use datafusion_common::JoinSide; use datafusion_common::config::ConfigOptions; use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode}; -use datafusion_physical_plan::ExecutionPlan; use datafusion_physical_plan::buffer::BufferExec; use datafusion_physical_plan::joins::HashJoinExec; +use datafusion_physical_plan::{ExecutionPlan, with_new_children_if_necessary}; use std::sync::Arc; /// Looks for all the [HashJoinExec]s in the plan and places a [BufferExec] node with the @@ -74,19 +74,25 @@ impl PhysicalOptimizerRule for HashJoinBuffering { if node.left.is::() { return Ok(Transformed::no(plan)); } - plan.with_new_children(vec![ - Arc::new(BufferExec::new(Arc::clone(&node.left), capacity)), - Arc::clone(&node.right), - ])? + with_new_children_if_necessary( + plan, + vec![ + Arc::new(BufferExec::new(Arc::clone(&node.left), capacity)), + Arc::clone(&node.right), + ], + )? } else { // Do not stack BufferExec nodes together. if node.right.is::() { return Ok(Transformed::no(plan)); } - plan.with_new_children(vec![ - Arc::clone(&node.left), - Arc::new(BufferExec::new(Arc::clone(&node.right), capacity)), - ])? + with_new_children_if_necessary( + plan, + vec![ + Arc::clone(&node.left), + Arc::new(BufferExec::new(Arc::clone(&node.right), capacity)), + ], + )? }, )) }) diff --git a/datafusion/physical-optimizer/src/output_requirements.rs b/datafusion/physical-optimizer/src/output_requirements.rs index b9d0d06da1dda..0583e5ee885a5 100644 --- a/datafusion/physical-optimizer/src/output_requirements.rs +++ b/datafusion/physical-optimizer/src/output_requirements.rs @@ -42,6 +42,7 @@ use datafusion_physical_plan::sorts::sort_preserving_merge::SortPreservingMergeE use datafusion_physical_plan::{ ChildStats, DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, PlanProperties, SendableRecordBatchStream, StatisticsArgs, + with_new_children_if_necessary, }; /// This rule either adds or removes [`OutputRequirements`]s to/from the physical @@ -460,7 +461,7 @@ fn require_top_ordering_helper( require_top_ordering_helper(Arc::clone(&children[idx]))?; if is_changed { children[idx] = new_child; - return Ok((plan.with_new_children(children)?, true)); + return Ok((with_new_children_if_necessary(plan, children)?, true)); } } Ok((plan, false)) diff --git a/datafusion/physical-plan/src/execution_plan.rs b/datafusion/physical-plan/src/execution_plan.rs index 63be0f1d385bd..ba27399f9ae41 100644 --- a/datafusion/physical-plan/src/execution_plan.rs +++ b/datafusion/physical-plan/src/execution_plan.rs @@ -861,15 +861,15 @@ pub trait ExecutionPlan: Any + Debug + DisplayAs + Send + Sync { } } -// A hint from `with_new_children_if_necessary` to `replace_children` indicating -// whether the properties of the new children must be recomputed. +/// A hint from `with_new_children_if_necessary` to `replace_children` indicating +/// whether the properties of the new children must be recomputed. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum ChildrenPropertiesHint { - // The properties of the new children are identical to the properties of the - // existing children, so we can skip recomuptation of plan properties. + /// The properties of the new children are identical to the properties of the + /// existing children, so we can skip recomputation of plan properties. SameProperties, - // The properties of the new children are different from the properties of the - // existing children, so we must recompute the properties from scratch. + /// The properties of the new children are different from the properties of the + /// existing children, so we must recompute the properties from scratch. Recompute, } From 1fdc7abf448e1c8ca198dd756393466c26bb7f69 Mon Sep 17 00:00:00 2001 From: Justin O'Dwyer Date: Sun, 26 Jul 2026 18:36:04 -0400 Subject: [PATCH 3/9] Start migrating to replace_children. --- datafusion/physical-plan/src/lib.rs | 6 +-- .../src/sorts/sort_preserving_merge.rs | 40 ++++++++++------- datafusion/physical-plan/src/unnest.rs | 44 ++++++++++++------- 3 files changed, 55 insertions(+), 35 deletions(-) diff --git a/datafusion/physical-plan/src/lib.rs b/datafusion/physical-plan/src/lib.rs index 8cba650b79770..a8dd87332bb71 100644 --- a/datafusion/physical-plan/src/lib.rs +++ b/datafusion/physical-plan/src/lib.rs @@ -45,9 +45,9 @@ pub use crate::distribution_requirements::{ ChildSatisfactionOptions, InputDistributionRequirements, }; pub use crate::execution_plan::{ - ExecutionPlan, ExecutionPlanProperties, PlanProperties, collect, collect_partitioned, - displayable, execute_input_stream, execute_stream, execute_stream_partitioned, - get_plan_string, with_new_children_if_necessary, + ChildrenPropertiesHint, ExecutionPlan, ExecutionPlanProperties, PlanProperties, + collect, collect_partitioned, displayable, execute_input_stream, execute_stream, + execute_stream_partitioned, get_plan_string, with_new_children_if_necessary, }; pub use crate::metrics::Metric; pub use crate::ordering::InputOrderMode; diff --git a/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs b/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs index 2add3e1eb82f0..0a6571ffaf74d 100644 --- a/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs +++ b/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs @@ -26,9 +26,9 @@ use crate::projection::{ProjectionExec, make_with_child, update_ordering}; use crate::sorts::streaming_merge::StreamingMergeBuilder; use crate::statistics::{ChildStats, StatisticsArgs}; use crate::{ - DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, ExecutionPlanProperties, - Partitioning, PlanProperties, SendableRecordBatchStream, Statistics, - check_if_same_properties, + ChildrenPropertiesHint, DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, + ExecutionPlanProperties, Partitioning, PlanProperties, SendableRecordBatchStream, + Statistics, }; use datafusion_common::{Result, assert_eq_or_internal_err, internal_err}; @@ -281,26 +281,36 @@ impl ExecutionPlan for SortPreservingMergeExec { vec![&self.input] } - fn with_new_children( + fn replace_children( self: Arc, mut children: Vec>, + hint: ChildrenPropertiesHint, ) -> Result> { - check_if_same_properties!(self, children); - Ok(Arc::new( - SortPreservingMergeExec::new(self.expr.clone(), children.swap_remove(0)) - .with_fetch(self.fetch), - )) + match hint { + ChildrenPropertiesHint::SameProperties => Ok(Arc::new(Self { + input: children.swap_remove(0), + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(&*self) + })), + ChildrenPropertiesHint::Recompute => Ok(Arc::new( + SortPreservingMergeExec::new(self.expr.clone(), children.swap_remove(0)) + .with_fetch(self.fetch), + )), + } + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children(children, ChildrenPropertiesHint::Recompute) } fn with_new_children_and_same_properties( self: Arc, - mut children: Vec>, + children: Vec>, ) -> Result> { - Ok(Arc::new(Self { - input: children.swap_remove(0), - metrics: ExecutionPlanMetricsSet::new(), - ..Self::clone(&*self) - })) + self.replace_children(children, ChildrenPropertiesHint::SameProperties) } fn execute( diff --git a/datafusion/physical-plan/src/unnest.rs b/datafusion/physical-plan/src/unnest.rs index fbe849229941a..6086515ff1834 100644 --- a/datafusion/physical-plan/src/unnest.rs +++ b/datafusion/physical-plan/src/unnest.rs @@ -28,8 +28,8 @@ use super::metrics::{ use super::{DisplayAs, ExecutionPlanProperties, PlanProperties}; use crate::stream::EmptyRecordBatchStream; use crate::{ - DisplayFormatType, Distribution, ExecutionPlan, RecordBatchStream, - SendableRecordBatchStream, check_if_same_properties, + ChildrenPropertiesHint, DisplayFormatType, Distribution, ExecutionPlan, + RecordBatchStream, SendableRecordBatchStream, }; use arrow::array::{ @@ -227,29 +227,39 @@ impl ExecutionPlan for UnnestExec { vec![&self.input] } - fn with_new_children( + fn replace_children( self: Arc, mut children: Vec>, + hint: ChildrenPropertiesHint, ) -> Result> { - check_if_same_properties!(self, children); - Ok(Arc::new(UnnestExec::new( - children.swap_remove(0), - self.list_column_indices.clone(), - self.struct_column_indices.clone(), - Arc::clone(&self.schema), - self.options.clone(), - )?)) + match hint { + ChildrenPropertiesHint::SameProperties => Ok(Arc::new(Self { + input: children.swap_remove(0), + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(&*self) + })), + ChildrenPropertiesHint::Recompute => Ok(Arc::new(UnnestExec::new( + children.swap_remove(0), + self.list_column_indices.clone(), + self.struct_column_indices.clone(), + Arc::clone(&self.schema), + self.options.clone(), + )?)), + } + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children(children, ChildrenPropertiesHint::Recompute) } fn with_new_children_and_same_properties( self: Arc, - mut children: Vec>, + children: Vec>, ) -> Result> { - Ok(Arc::new(Self { - input: children.swap_remove(0), - metrics: ExecutionPlanMetricsSet::new(), - ..Self::clone(&*self) - })) + self.replace_children(children, ChildrenPropertiesHint::SameProperties) } fn required_input_distribution(&self) -> Vec { From 9b872efedd4d5ad3bf5634696f5d99e76ab3a5d7 Mon Sep 17 00:00:00 2001 From: Justin O'Dwyer Date: Mon, 27 Jul 2026 14:16:38 -0400 Subject: [PATCH 4/9] Migrate implementations to replace_children. --- .../physical-plan/src/aggregates/mod.rs | 79 +++++++++------- datafusion/physical-plan/src/async_func.rs | 44 +++++---- datafusion/physical-plan/src/buffer.rs | 36 ++++--- .../physical-plan/src/coalesce_batches.rs | 39 +++++--- .../physical-plan/src/coalesce_partitions.rs | 40 +++++--- datafusion/physical-plan/src/coop.rs | 39 +++++--- .../physical-plan/src/execution_plan.rs | 17 ++++ datafusion/physical-plan/src/filter.rs | 41 +++++--- .../physical-plan/src/joins/cross_join.rs | 53 +++++++---- .../src/joins/nested_loop_join.rs | 74 ++++++++------- .../src/joins/piecewise_merge_join/exec.rs | 93 +++++++++++-------- .../src/joins/sort_merge_join/exec.rs | 63 ++++++++----- .../src/joins/symmetric_hash_join.rs | 59 +++++++----- datafusion/physical-plan/src/limit.rs | 73 +++++++++------ datafusion/physical-plan/src/projection.rs | 42 ++++++--- .../physical-plan/src/repartition/mod.rs | 51 ++++++---- .../physical-plan/src/sorts/partial_sort.rs | 51 ++++++---- .../src/sorts/sort_preserving_merge.rs | 3 +- datafusion/physical-plan/src/union.rs | 64 +++++++++---- datafusion/physical-plan/src/unnest.rs | 3 +- .../src/windows/bounded_window_agg_exec.rs | 47 ++++++---- .../src/windows/window_agg_exec.rs | 45 +++++---- 22 files changed, 667 insertions(+), 389 deletions(-) diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index 787cd4f03ff6c..3197d822bbfdd 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -146,16 +146,6 @@ use std::borrow::Cow; use std::sync::Arc; use super::{DisplayAs, ExecutionPlanProperties, PlanProperties}; -use crate::aggregates::{ - aggregate_stream::AggregateStream, - grouped_hash_stream::GroupedHashAggregateStream, - grouped_topk_stream::GroupedTopKAggregateStream, - hash_stream::{FinalHashAggregateStream, PartialHashAggregateStream}, - ordered_final_stream::OrderedFinalAggregateStream, - ordered_partial_stream::OrderedPartialAggregateStream, - partial_reduce_stream::PartialReduceHashAggregateStream, - single_stream::SingleHashAggregateStream, -}; use crate::execution_plan::{CardinalityEffect, EmissionType}; use crate::filter_pushdown::{ ChildFilterDescription, ChildPushdownResult, FilterDescription, FilterPushdownPhase, @@ -163,9 +153,23 @@ use crate::filter_pushdown::{ }; use crate::metrics::{ExecutionPlanMetricsSet, MetricsSet}; use crate::statistics::{ChildStats, StatisticsArgs}; +use crate::{ + ChildrenPropertiesHint, + aggregates::{ + aggregate_stream::AggregateStream, + grouped_hash_stream::GroupedHashAggregateStream, + grouped_topk_stream::GroupedTopKAggregateStream, + hash_stream::{FinalHashAggregateStream, PartialHashAggregateStream}, + ordered_final_stream::OrderedFinalAggregateStream, + ordered_partial_stream::OrderedPartialAggregateStream, + partial_reduce_stream::PartialReduceHashAggregateStream, + single_stream::SingleHashAggregateStream, + }, + validate_child_count, +}; use crate::{ DisplayFormatType, Distribution, ExecutionPlan, InputDistributionRequirements, - InputOrderMode, SendableRecordBatchStream, Statistics, check_if_same_properties, + InputOrderMode, SendableRecordBatchStream, Statistics, }; use datafusion_common::config::ConfigOptions; use datafusion_physical_expr::utils::collect_columns; @@ -1961,36 +1965,47 @@ impl ExecutionPlan for AggregateExec { vec![&self.input] } + fn replace_children( + self: Arc, + mut children: Vec>, + hint: ChildrenPropertiesHint, + ) -> Result> { + validate_child_count!(self, children); + match hint { + ChildrenPropertiesHint::SameProperties => Ok(Arc::new(Self { + input: children.swap_remove(0), + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(&*self) + })), + ChildrenPropertiesHint::Recompute => { + let mut me = AggregateExec::try_new_with_schema( + self.mode, + Arc::clone(&self.group_by), + self.aggr_expr.to_vec(), + Arc::clone(&self.filter_expr), + Arc::clone(&children[0]), + Arc::clone(&self.input_schema), + Arc::clone(&self.schema), + )?; + me.limit_options = self.limit_options; + me.dynamic_filter.clone_from(&self.dynamic_filter); + Ok(Arc::new(me)) + } + } + } + fn with_new_children( self: Arc, children: Vec>, ) -> Result> { - check_if_same_properties!(self, children); - - let mut me = AggregateExec::try_new_with_schema( - self.mode, - Arc::clone(&self.group_by), - self.aggr_expr.to_vec(), - Arc::clone(&self.filter_expr), - Arc::clone(&children[0]), - Arc::clone(&self.input_schema), - Arc::clone(&self.schema), - )?; - me.limit_options = self.limit_options; - me.dynamic_filter.clone_from(&self.dynamic_filter); - - Ok(Arc::new(me)) + self.replace_children(children, ChildrenPropertiesHint::Recompute) } fn with_new_children_and_same_properties( self: Arc, - mut children: Vec>, + children: Vec>, ) -> Result> { - Ok(Arc::new(Self { - input: children.swap_remove(0), - metrics: ExecutionPlanMetricsSet::new(), - ..Self::clone(&*self) - })) + self.replace_children(children, ChildrenPropertiesHint::SameProperties) } fn execute( diff --git a/datafusion/physical-plan/src/async_func.rs b/datafusion/physical-plan/src/async_func.rs index e13a5b986aa2c..214b7703450e0 100644 --- a/datafusion/physical-plan/src/async_func.rs +++ b/datafusion/physical-plan/src/async_func.rs @@ -19,8 +19,8 @@ use crate::coalesce::LimitedBatchCoalescer; use crate::metrics::{ExecutionPlanMetricsSet, MetricsSet}; use crate::stream::{EmptyRecordBatchStream, RecordBatchStreamAdapter}; use crate::{ - DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, PlanProperties, - check_if_same_properties, + ChildrenPropertiesHint, DisplayAs, DisplayFormatType, ExecutionPlan, + ExecutionPlanProperties, PlanProperties, validate_child_count, }; use arrow::array::RecordBatch; use arrow_schema::{FieldRef, Fields, Schema, SchemaRef}; @@ -153,31 +153,37 @@ impl ExecutionPlan for AsyncFuncExec { vec![&self.input] } - fn with_new_children( + fn replace_children( self: Arc, mut children: Vec>, + hint: ChildrenPropertiesHint, ) -> Result> { - assert_eq_or_internal_err!( - children.len(), - 1, - "AsyncFuncExec wrong number of children" - ); - check_if_same_properties!(self, children); - Ok(Arc::new(AsyncFuncExec::try_new( - self.async_exprs.clone(), - children.swap_remove(0), - )?)) + validate_child_count!(self, children); + match hint { + ChildrenPropertiesHint::SameProperties => Ok(Arc::new(Self { + input: children.swap_remove(0), + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(&*self) + })), + ChildrenPropertiesHint::Recompute => Ok(Arc::new(AsyncFuncExec::try_new( + self.async_exprs.clone(), + children.swap_remove(0), + )?)), + } + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children(children, ChildrenPropertiesHint::Recompute) } fn with_new_children_and_same_properties( self: Arc, - mut children: Vec>, + children: Vec>, ) -> Result> { - Ok(Arc::new(Self { - input: children.swap_remove(0), - metrics: ExecutionPlanMetricsSet::new(), - ..Self::clone(&*self) - })) + self.replace_children(children, ChildrenPropertiesHint::SameProperties) } fn execute( diff --git a/datafusion/physical-plan/src/buffer.rs b/datafusion/physical-plan/src/buffer.rs index 3be331a1ee1ba..1e82d1f2dddc4 100644 --- a/datafusion/physical-plan/src/buffer.rs +++ b/datafusion/physical-plan/src/buffer.rs @@ -27,8 +27,8 @@ use crate::projection::ProjectionExec; use crate::statistics::{ChildStats, StatisticsArgs}; use crate::stream::RecordBatchStreamAdapter; use crate::{ - DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, SortOrderPushdownResult, - check_if_same_properties, + ChildrenPropertiesHint, DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, + SortOrderPushdownResult, validate_child_count, }; use arrow::array::RecordBatch; use datafusion_common::config::ConfigOptions; @@ -159,26 +159,36 @@ impl ExecutionPlan for BufferExec { vec![&self.input] } - fn with_new_children( + fn replace_children( self: Arc, mut children: Vec>, + hint: ChildrenPropertiesHint, ) -> Result> { - check_if_same_properties!(self, children); - if children.len() != 1 { - return plan_err!("BufferExec can only have one child"); + validate_child_count!(self, children); + match hint { + ChildrenPropertiesHint::SameProperties => Ok(Arc::new(Self { + input: children.swap_remove(0), + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(&*self) + })), + ChildrenPropertiesHint::Recompute => { + Ok(Arc::new(Self::new(children.swap_remove(0), self.capacity))) + } } - Ok(Arc::new(Self::new(children.swap_remove(0), self.capacity))) + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children(children, ChildrenPropertiesHint::Recompute) } fn with_new_children_and_same_properties( self: Arc, - mut children: Vec>, + children: Vec>, ) -> Result> { - Ok(Arc::new(Self { - input: children.swap_remove(0), - metrics: ExecutionPlanMetricsSet::new(), - ..Self::clone(&*self) - })) + self.replace_children(children, ChildrenPropertiesHint::SameProperties) } fn execute( diff --git a/datafusion/physical-plan/src/coalesce_batches.rs b/datafusion/physical-plan/src/coalesce_batches.rs index c5b91767777f2..a4dd239e7906e 100644 --- a/datafusion/physical-plan/src/coalesce_batches.rs +++ b/datafusion/physical-plan/src/coalesce_batches.rs @@ -27,8 +27,8 @@ use crate::projection::ProjectionExec; use crate::statistics::{ChildStats, StatisticsArgs}; use crate::stream::EmptyRecordBatchStream; use crate::{ - DisplayFormatType, ExecutionPlan, RecordBatchStream, SendableRecordBatchStream, - check_if_same_properties, + ChildrenPropertiesHint, DisplayFormatType, ExecutionPlan, RecordBatchStream, + SendableRecordBatchStream, validate_child_count, }; use arrow::datatypes::SchemaRef; @@ -173,26 +173,37 @@ impl ExecutionPlan for CoalesceBatchesExec { vec![false] } - fn with_new_children( + fn replace_children( self: Arc, mut children: Vec>, + hint: ChildrenPropertiesHint, ) -> Result> { - check_if_same_properties!(self, children); - Ok(Arc::new( - CoalesceBatchesExec::new(children.swap_remove(0), self.target_batch_size) - .with_fetch(self.fetch), - )) + validate_child_count!(self, children); + match hint { + ChildrenPropertiesHint::SameProperties => Ok(Arc::new(Self { + input: children.swap_remove(0), + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(&*self) + })), + ChildrenPropertiesHint::Recompute => Ok(Arc::new( + CoalesceBatchesExec::new(children.swap_remove(0), self.target_batch_size) + .with_fetch(self.fetch), + )), + } + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children(children, ChildrenPropertiesHint::Recompute) } fn with_new_children_and_same_properties( self: Arc, - mut children: Vec>, + children: Vec>, ) -> Result> { - Ok(Arc::new(Self { - input: children.swap_remove(0), - metrics: ExecutionPlanMetricsSet::new(), - ..Self::clone(&*self) - })) + self.replace_children(children, ChildrenPropertiesHint::SameProperties) } fn execute( diff --git a/datafusion/physical-plan/src/coalesce_partitions.rs b/datafusion/physical-plan/src/coalesce_partitions.rs index f9694e0d16817..788a00c7ecb58 100644 --- a/datafusion/physical-plan/src/coalesce_partitions.rs +++ b/datafusion/physical-plan/src/coalesce_partitions.rs @@ -31,7 +31,10 @@ use crate::filter_pushdown::{FilterDescription, FilterPushdownPhase}; use crate::projection::{ProjectionExec, make_with_child}; use crate::sort_pushdown::SortOrderPushdownResult; use crate::statistics::{ChildStats, StatisticsArgs}; -use crate::{DisplayFormatType, ExecutionPlan, Partitioning, check_if_same_properties}; +use crate::{ + ChildrenPropertiesHint, DisplayFormatType, ExecutionPlan, Partitioning, + validate_child_count, +}; use datafusion_physical_expr_common::sort_expr::PhysicalSortExpr; use datafusion_common::config::ConfigOptions; @@ -143,25 +146,38 @@ impl ExecutionPlan for CoalescePartitionsExec { vec![false] } - fn with_new_children( + fn replace_children( self: Arc, mut children: Vec>, + hint: ChildrenPropertiesHint, ) -> Result> { - check_if_same_properties!(self, children); - let mut plan = CoalescePartitionsExec::new(children.swap_remove(0)); - plan.fetch = self.fetch; - Ok(Arc::new(plan)) + validate_child_count!(self, children); + match hint { + ChildrenPropertiesHint::SameProperties => Ok(Arc::new(Self { + input: children.swap_remove(0), + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(&*self) + })), + ChildrenPropertiesHint::Recompute => { + let mut plan = CoalescePartitionsExec::new(children.swap_remove(0)); + plan.fetch = self.fetch; + Ok(Arc::new(plan)) + } + } + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children(children, ChildrenPropertiesHint::Recompute) } fn with_new_children_and_same_properties( self: Arc, - mut children: Vec>, + children: Vec>, ) -> Result> { - Ok(Arc::new(Self { - input: children.swap_remove(0), - metrics: ExecutionPlanMetricsSet::new(), - ..Self::clone(&*self) - })) + self.replace_children(children, ChildrenPropertiesHint::SameProperties) } fn execute( diff --git a/datafusion/physical-plan/src/coop.rs b/datafusion/physical-plan/src/coop.rs index a5b57f546bbfa..28eeb61b5cc5f 100644 --- a/datafusion/physical-plan/src/coop.rs +++ b/datafusion/physical-plan/src/coop.rs @@ -86,8 +86,9 @@ use crate::filter_pushdown::{ use crate::projection::ProjectionExec; use crate::statistics::{ChildStats, StatisticsArgs}; use crate::{ - DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, RecordBatchStream, - SendableRecordBatchStream, SortOrderPushdownResult, check_if_same_properties, + ChildrenPropertiesHint, DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, + RecordBatchStream, SendableRecordBatchStream, SortOrderPushdownResult, + validate_child_count, }; use arrow::record_batch::RecordBatch; use arrow_schema::Schema; @@ -267,27 +268,35 @@ impl ExecutionPlan for CooperativeExec { vec![&self.input] } - fn with_new_children( + fn replace_children( self: Arc, mut children: Vec>, + hint: ChildrenPropertiesHint, ) -> Result> { - assert_eq_or_internal_err!( - children.len(), - 1, - "CooperativeExec requires exactly one child" - ); - check_if_same_properties!(self, children); - Ok(Arc::new(CooperativeExec::new(children.swap_remove(0)))) + validate_child_count!(self, children); + match hint { + ChildrenPropertiesHint::SameProperties => Ok(Arc::new(Self { + input: children.swap_remove(0), + ..Self::clone(&*self) + })), + ChildrenPropertiesHint::Recompute => { + Ok(Arc::new(CooperativeExec::new(children.swap_remove(0)))) + } + } + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children(children, ChildrenPropertiesHint::Recompute) } fn with_new_children_and_same_properties( self: Arc, - mut children: Vec>, + children: Vec>, ) -> Result> { - Ok(Arc::new(Self { - input: children.swap_remove(0), - ..Self::clone(&*self) - })) + self.replace_children(children, ChildrenPropertiesHint::SameProperties) } fn execute( diff --git a/datafusion/physical-plan/src/execution_plan.rs b/datafusion/physical-plan/src/execution_plan.rs index ba27399f9ae41..372625f3c23da 100644 --- a/datafusion/physical-plan/src/execution_plan.rs +++ b/datafusion/physical-plan/src/execution_plan.rs @@ -1707,6 +1707,23 @@ macro_rules! check_if_same_properties { }; } +/// Helper macro to validate that replacement children match a plan's existing +/// child count. +/// +/// This is useful for [`ExecutionPlan::replace_children`] implementations that +/// no longer call [`check_if_same_properties`] but still need to preserve the +/// same child-count validation behavior. +#[macro_export] +macro_rules! validate_child_count { + ($plan: expr, $children: expr) => { + datafusion_common::assert_eq_or_internal_err!( + $children.len(), + $plan.children().len(), + "Wrong number of children" + ); + }; +} + /// Utility function yielding a string representation of the given [`ExecutionPlan`]. pub fn get_plan_string(plan: &Arc) -> Vec { let formatted = displayable(plan.as_ref()).indent(true).to_string(); diff --git a/datafusion/physical-plan/src/filter.rs b/datafusion/physical-plan/src/filter.rs index d367be16eb6ed..77fac49321c6a 100644 --- a/datafusion/physical-plan/src/filter.rs +++ b/datafusion/physical-plan/src/filter.rs @@ -28,7 +28,6 @@ use super::{ ColumnStatistics, DisplayAs, ExecutionPlanProperties, PlanProperties, RecordBatchStream, SendableRecordBatchStream, Statistics, }; -use crate::check_if_same_properties; use crate::coalesce::{LimitedBatchCoalescer, PushBatchStatus}; use crate::common::can_project; use crate::execution_plan::CardinalityEffect; @@ -44,6 +43,7 @@ use crate::projection::{ }; use crate::statistics::{ChildStats, StatisticsArgs, StatisticsContext}; use crate::stream::EmptyRecordBatchStream; +use crate::{ChildrenPropertiesHint, validate_child_count}; use crate::{ DisplayFormatType, ExecutionPlan, metrics::{BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet, RatioMetrics}, @@ -539,27 +539,40 @@ impl ExecutionPlan for FilterExec { vec![true] } - fn with_new_children( + fn replace_children( self: Arc, mut children: Vec>, + hint: ChildrenPropertiesHint, ) -> Result> { - check_if_same_properties!(self, children); - let new_input = children.swap_remove(0); - FilterExecBuilder::from(&*self) - .with_input(new_input) - .build() - .map(|e| Arc::new(e) as _) + validate_child_count!(self, children); + match hint { + ChildrenPropertiesHint::SameProperties => Ok(Arc::new(Self { + input: children.swap_remove(0), + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(&*self) + })), + ChildrenPropertiesHint::Recompute => { + let new_input = children.swap_remove(0); + FilterExecBuilder::from(&*self) + .with_input(new_input) + .build() + .map(|e| Arc::new(e) as _) + } + } + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children(children, ChildrenPropertiesHint::Recompute) } fn with_new_children_and_same_properties( self: Arc, - mut children: Vec>, + children: Vec>, ) -> Result> { - Ok(Arc::new(Self { - input: children.swap_remove(0), - metrics: ExecutionPlanMetricsSet::new(), - ..Self::clone(&*self) - })) + self.replace_children(children, ChildrenPropertiesHint::SameProperties) } fn execute( diff --git a/datafusion/physical-plan/src/joins/cross_join.rs b/datafusion/physical-plan/src/joins/cross_join.rs index 1a631aac980ab..b21d86a860405 100644 --- a/datafusion/physical-plan/src/joins/cross_join.rs +++ b/datafusion/physical-plan/src/joins/cross_join.rs @@ -34,9 +34,9 @@ use crate::projection::{ use crate::statistics::{ChildStats, StatisticsArgs}; use crate::stream::EmptyRecordBatchStream; use crate::{ - ColumnStatistics, DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, - ExecutionPlanProperties, PlanProperties, RecordBatchStream, - SendableRecordBatchStream, Statistics, check_if_same_properties, handle_state, + ChildrenPropertiesHint, ColumnStatistics, DisplayAs, DisplayFormatType, Distribution, + ExecutionPlan, ExecutionPlanProperties, PlanProperties, RecordBatchStream, + SendableRecordBatchStream, Statistics, handle_state, validate_child_count, }; use arrow::array::{RecordBatch, RecordBatchOptions}; @@ -290,32 +290,45 @@ impl ExecutionPlan for CrossJoinExec { Some(self.metrics.clone_inner()) } + fn replace_children( + self: Arc, + mut children: Vec>, + hint: ChildrenPropertiesHint, + ) -> Result> { + validate_child_count!(self, children); + match hint { + ChildrenPropertiesHint::SameProperties => { + let left = children.swap_remove(0); + let right = children.swap_remove(0); + + Ok(Arc::new(Self { + left, + right, + metrics: ExecutionPlanMetricsSet::new(), + left_fut: Default::default(), + cache: Arc::clone(&self.cache), + schema: Arc::clone(&self.schema), + })) + } + ChildrenPropertiesHint::Recompute => Ok(Arc::new(CrossJoinExec::new( + Arc::clone(&children[0]), + Arc::clone(&children[1]), + ))), + } + } + fn with_new_children( self: Arc, children: Vec>, ) -> Result> { - check_if_same_properties!(self, children); - Ok(Arc::new(CrossJoinExec::new( - Arc::clone(&children[0]), - Arc::clone(&children[1]), - ))) + self.replace_children(children, ChildrenPropertiesHint::Recompute) } fn with_new_children_and_same_properties( self: Arc, - mut children: Vec>, + children: Vec>, ) -> Result> { - let left = children.swap_remove(0); - let right = children.swap_remove(0); - - Ok(Arc::new(Self { - left, - right, - metrics: ExecutionPlanMetricsSet::new(), - left_fut: Default::default(), - cache: Arc::clone(&self.cache), - schema: Arc::clone(&self.schema), - })) + self.replace_children(children, ChildrenPropertiesHint::SameProperties) } fn reset_state(self: Arc) -> Result> { diff --git a/datafusion/physical-plan/src/joins/nested_loop_join.rs b/datafusion/physical-plan/src/joins/nested_loop_join.rs index 515dcc2931c05..68a52664f2185 100644 --- a/datafusion/physical-plan/src/joins/nested_loop_join.rs +++ b/datafusion/physical-plan/src/joins/nested_loop_join.rs @@ -44,9 +44,9 @@ use crate::projection::{ }; use crate::statistics::{ChildStats, StatisticsArgs}; use crate::{ - DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, ExecutionPlanProperties, - PlanProperties, RecordBatchStream, SendableRecordBatchStream, - check_if_same_properties, + ChildrenPropertiesHint, DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, + ExecutionPlanProperties, PlanProperties, RecordBatchStream, + SendableRecordBatchStream, validate_child_count, }; use arrow::array::{ @@ -562,43 +562,55 @@ impl ExecutionPlan for NestedLoopJoinExec { vec![&self.left, &self.right] } + fn replace_children( + self: Arc, + mut children: Vec>, + hint: ChildrenPropertiesHint, + ) -> Result> { + validate_child_count!(self, children); + match hint { + ChildrenPropertiesHint::SameProperties => { + let left = children.swap_remove(0); + let right = children.swap_remove(0); + Ok(Arc::new(Self { + left, + right, + metrics: ExecutionPlanMetricsSet::new(), + build_side_data: Default::default(), + left_spill_data: Arc::new(OnceAsync::default()), + cache: Arc::clone(&self.cache), + filter: self.filter.clone(), + join_type: self.join_type, + join_schema: Arc::clone(&self.join_schema), + column_indices: self.column_indices.clone(), + projection: self.projection.clone(), + })) + } + ChildrenPropertiesHint::Recompute => Ok(Arc::new( + NestedLoopJoinExecBuilder::new( + Arc::clone(&children[0]), + Arc::clone(&children[1]), + self.join_type, + ) + .with_filter(self.filter.clone()) + .with_projection_ref(self.projection.clone()) + .build()?, + )), + } + } + fn with_new_children( self: Arc, children: Vec>, ) -> Result> { - check_if_same_properties!(self, children); - Ok(Arc::new( - NestedLoopJoinExecBuilder::new( - Arc::clone(&children[0]), - Arc::clone(&children[1]), - self.join_type, - ) - .with_filter(self.filter.clone()) - .with_projection_ref(self.projection.clone()) - .build()?, - )) + self.replace_children(children, ChildrenPropertiesHint::Recompute) } fn with_new_children_and_same_properties( self: Arc, - mut children: Vec>, + children: Vec>, ) -> Result> { - let left = children.swap_remove(0); - let right = children.swap_remove(0); - - Ok(Arc::new(Self { - left, - right, - metrics: ExecutionPlanMetricsSet::new(), - build_side_data: Default::default(), - left_spill_data: Arc::new(OnceAsync::default()), - cache: Arc::clone(&self.cache), - filter: self.filter.clone(), - join_type: self.join_type, - join_schema: Arc::clone(&self.join_schema), - column_indices: self.column_indices.clone(), - projection: self.projection.clone(), - })) + self.replace_children(children, ChildrenPropertiesHint::SameProperties) } fn execute( diff --git a/datafusion/physical-plan/src/joins/piecewise_merge_join/exec.rs b/datafusion/physical-plan/src/joins/piecewise_merge_join/exec.rs index 5ec564295ece1..70a9ed095797a 100644 --- a/datafusion/physical-plan/src/joins/piecewise_merge_join/exec.rs +++ b/datafusion/physical-plan/src/joins/piecewise_merge_join/exec.rs @@ -52,7 +52,8 @@ use crate::joins::piecewise_merge_join::utils::{ use crate::joins::utils::asymmetric_join_output_partitioning; use crate::metrics::MetricsSet; use crate::{ - DisplayAs, DisplayFormatType, ExecutionPlanProperties, check_if_same_properties, + ChildrenPropertiesHint, DisplayAs, DisplayFormatType, ExecutionPlanProperties, + validate_child_count, }; use crate::{ ExecutionPlan, PlanProperties, @@ -508,56 +509,74 @@ impl ExecutionPlan for PiecewiseMergeJoinExec { } } + fn replace_children( + self: Arc, + mut children: Vec>, + hint: ChildrenPropertiesHint, + ) -> Result> { + validate_child_count!(self, children); + match hint { + ChildrenPropertiesHint::SameProperties => { + let buffered = children.swap_remove(0); + let streamed = children.swap_remove(0); + Ok(Arc::new(Self { + buffered, + streamed, + on: self.on.clone(), + operator: self.operator, + join_type: self.join_type, + schema: Arc::clone(&self.schema), + left_child_plan_required_order: self + .left_child_plan_required_order + .clone(), + right_batch_required_orders: self.right_batch_required_orders.clone(), + sort_options: self.sort_options, + cache: Arc::clone(&self.cache), + num_partitions: self.num_partitions, + + // Re-set state. + metrics: ExecutionPlanMetricsSet::new(), + buffered_fut: Default::default(), + })) + } + ChildrenPropertiesHint::Recompute => match &children[..] { + [left, right] => Ok(Arc::new(PiecewiseMergeJoinExec::try_new( + Arc::clone(left), + Arc::clone(right), + self.on.clone(), + self.operator, + self.join_type, + self.num_partitions, + )?)), + _ => internal_err!( + "PiecewiseMergeJoin should have 2 children, found {}", + children.len() + ), + }, + } + } + fn with_new_children( self: Arc, children: Vec>, ) -> Result> { - check_if_same_properties!(self, children); - match &children[..] { - [left, right] => Ok(Arc::new(PiecewiseMergeJoinExec::try_new( - Arc::clone(left), - Arc::clone(right), - self.on.clone(), - self.operator, - self.join_type, - self.num_partitions, - )?)), - _ => internal_err!( - "PiecewiseMergeJoin should have 2 children, found {}", - children.len() - ), - } + self.replace_children(children, ChildrenPropertiesHint::Recompute) } fn with_new_children_and_same_properties( self: Arc, - mut children: Vec>, + children: Vec>, ) -> Result> { - let buffered = children.swap_remove(0); - let streamed = children.swap_remove(0); - Ok(Arc::new(Self { - buffered, - streamed, - on: self.on.clone(), - operator: self.operator, - join_type: self.join_type, - schema: Arc::clone(&self.schema), - left_child_plan_required_order: self.left_child_plan_required_order.clone(), - right_batch_required_orders: self.right_batch_required_orders.clone(), - sort_options: self.sort_options, - cache: Arc::clone(&self.cache), - num_partitions: self.num_partitions, - - // Re-set state. - metrics: ExecutionPlanMetricsSet::new(), - buffered_fut: Default::default(), - })) + self.replace_children(children, ChildrenPropertiesHint::SameProperties) } fn reset_state(self: Arc) -> Result> { let buffered = Arc::clone(&self.buffered); let streamed = Arc::clone(&self.streamed); - self.with_new_children_and_same_properties(vec![buffered, streamed]) + self.replace_children( + vec![buffered, streamed], + ChildrenPropertiesHint::SameProperties, + ) } fn execute( diff --git a/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs b/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs index 00cac069aae5d..0ee6e86acf615 100644 --- a/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs +++ b/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs @@ -40,9 +40,9 @@ use crate::projection::{ use crate::spill::spill_manager::SpillManager; use crate::statistics::{ChildStats, StatisticsArgs}; use crate::{ - DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, ExecutionPlanProperties, - InputDistributionRequirements, PlanProperties, SendableRecordBatchStream, Statistics, - check_if_same_properties, + ChildrenPropertiesHint, DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, + ExecutionPlanProperties, InputDistributionRequirements, PlanProperties, + SendableRecordBatchStream, Statistics, validate_child_count, }; use arrow::compute::SortOptions; @@ -441,37 +441,50 @@ impl ExecutionPlan for SortMergeJoinExec { vec![&self.left, &self.right] } + fn replace_children( + self: Arc, + mut children: Vec>, + hint: ChildrenPropertiesHint, + ) -> Result> { + validate_child_count!(self, children); + match hint { + ChildrenPropertiesHint::SameProperties => { + let left = children.swap_remove(0); + let right = children.swap_remove(0); + Ok(Arc::new(Self { + left, + right, + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(&*self) + })) + } + ChildrenPropertiesHint::Recompute => match &children[..] { + [left, right] => Ok(Arc::new(SortMergeJoinExec::try_new( + Arc::clone(left), + Arc::clone(right), + self.on.clone(), + self.filter.clone(), + self.join_type, + self.sort_options.clone(), + self.null_equality, + )?)), + _ => internal_err!("SortMergeJoin wrong number of children"), + }, + } + } + fn with_new_children( self: Arc, children: Vec>, ) -> Result> { - check_if_same_properties!(self, children); - match &children[..] { - [left, right] => Ok(Arc::new(SortMergeJoinExec::try_new( - Arc::clone(left), - Arc::clone(right), - self.on.clone(), - self.filter.clone(), - self.join_type, - self.sort_options.clone(), - self.null_equality, - )?)), - _ => internal_err!("SortMergeJoin wrong number of children"), - } + self.replace_children(children, ChildrenPropertiesHint::Recompute) } fn with_new_children_and_same_properties( self: Arc, - mut children: Vec>, + children: Vec>, ) -> Result> { - let left = children.swap_remove(0); - let right = children.swap_remove(0); - Ok(Arc::new(Self { - left, - right, - metrics: ExecutionPlanMetricsSet::new(), - ..Self::clone(&*self) - })) + self.replace_children(children, ChildrenPropertiesHint::SameProperties) } fn execute( diff --git a/datafusion/physical-plan/src/joins/symmetric_hash_join.rs b/datafusion/physical-plan/src/joins/symmetric_hash_join.rs index 95f4c35871431..f0ba8343e3e66 100644 --- a/datafusion/physical-plan/src/joins/symmetric_hash_join.rs +++ b/datafusion/physical-plan/src/joins/symmetric_hash_join.rs @@ -31,7 +31,6 @@ use std::sync::Arc; use std::task::{Context, Poll}; use std::vec; -use crate::check_if_same_properties; use crate::common::SharedMemoryReservation; use crate::execution_plan::{boundedness_from_children, emission_type_from_children}; use crate::joins::stream_join_utils::{ @@ -50,6 +49,7 @@ use crate::projection::{ JoinData, ProjectionExec, try_pushdown_through_join_with_column_indices, }; use crate::stream::EmptyRecordBatchStream; +use crate::{ChildrenPropertiesHint, validate_child_count}; use crate::{ DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, ExecutionPlanProperties, InputDistributionRequirements, PlanProperties, RecordBatchStream, @@ -452,36 +452,51 @@ impl ExecutionPlan for SymmetricHashJoinExec { vec![&self.left, &self.right] } + fn replace_children( + self: Arc, + mut children: Vec>, + hint: ChildrenPropertiesHint, + ) -> Result> { + validate_child_count!(self, children); + match hint { + ChildrenPropertiesHint::SameProperties => { + let left = children.swap_remove(0); + let right = children.swap_remove(0); + Ok(Arc::new(Self { + left, + right, + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(&*self) + })) + } + ChildrenPropertiesHint::Recompute => { + Ok(Arc::new(SymmetricHashJoinExec::try_new( + Arc::clone(&children[0]), + Arc::clone(&children[1]), + self.on.clone(), + self.filter.clone(), + &self.join_type, + self.null_equality, + self.left_sort_exprs.clone(), + self.right_sort_exprs.clone(), + self.mode, + )?)) + } + } + } + fn with_new_children( self: Arc, children: Vec>, ) -> Result> { - check_if_same_properties!(self, children); - Ok(Arc::new(SymmetricHashJoinExec::try_new( - Arc::clone(&children[0]), - Arc::clone(&children[1]), - self.on.clone(), - self.filter.clone(), - &self.join_type, - self.null_equality, - self.left_sort_exprs.clone(), - self.right_sort_exprs.clone(), - self.mode, - )?)) + self.replace_children(children, ChildrenPropertiesHint::Recompute) } fn with_new_children_and_same_properties( self: Arc, - mut children: Vec>, + children: Vec>, ) -> Result> { - let left = children.swap_remove(0); - let right = children.swap_remove(0); - Ok(Arc::new(Self { - left, - right, - metrics: ExecutionPlanMetricsSet::new(), - ..Self::clone(&*self) - })) + self.replace_children(children, ChildrenPropertiesHint::SameProperties) } fn metrics(&self) -> Option { diff --git a/datafusion/physical-plan/src/limit.rs b/datafusion/physical-plan/src/limit.rs index ddce680fc18ad..431e0803dfe4d 100644 --- a/datafusion/physical-plan/src/limit.rs +++ b/datafusion/physical-plan/src/limit.rs @@ -29,8 +29,8 @@ use super::{ use crate::execution_plan::{Boundedness, CardinalityEffect}; use crate::statistics::{ChildStats, StatisticsArgs}; use crate::{ - DisplayFormatType, Distribution, ExecutionPlan, Partitioning, - check_if_same_properties, + ChildrenPropertiesHint, DisplayFormatType, Distribution, ExecutionPlan, Partitioning, + validate_child_count, }; use arrow::datatypes::SchemaRef; @@ -167,27 +167,38 @@ impl ExecutionPlan for GlobalLimitExec { vec![false] } - fn with_new_children( + fn replace_children( self: Arc, mut children: Vec>, + hint: ChildrenPropertiesHint, ) -> Result> { - check_if_same_properties!(self, children); - Ok(Arc::new(GlobalLimitExec::new( - children.swap_remove(0), - self.skip, - self.fetch, - ))) + validate_child_count!(self, children); + match hint { + ChildrenPropertiesHint::SameProperties => Ok(Arc::new(Self { + input: children.swap_remove(0), + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(&*self) + })), + ChildrenPropertiesHint::Recompute => Ok(Arc::new(GlobalLimitExec::new( + children.swap_remove(0), + self.skip, + self.fetch, + ))), + } + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children(children, ChildrenPropertiesHint::Recompute) } fn with_new_children_and_same_properties( self: Arc, - mut children: Vec>, + children: Vec>, ) -> Result> { - Ok(Arc::new(Self { - input: children.swap_remove(0), - metrics: ExecutionPlanMetricsSet::new(), - ..Self::clone(&*self) - })) + self.replace_children(children, ChildrenPropertiesHint::SameProperties) } fn execute( @@ -398,29 +409,37 @@ impl ExecutionPlan for LocalLimitExec { vec![true] } - fn with_new_children( + fn replace_children( self: Arc, - children: Vec>, + mut children: Vec>, + hint: ChildrenPropertiesHint, ) -> Result> { - check_if_same_properties!(self, children); - match children.len() { - 1 => Ok(Arc::new(LocalLimitExec::new( + validate_child_count!(self, children); + match hint { + ChildrenPropertiesHint::SameProperties => Ok(Arc::new(Self { + input: children.swap_remove(0), + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(&*self) + })), + ChildrenPropertiesHint::Recompute => Ok(Arc::new(LocalLimitExec::new( Arc::clone(&children[0]), self.fetch, ))), - _ => internal_err!("LocalLimitExec wrong number of children"), } } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children(children, ChildrenPropertiesHint::Recompute) + } + fn with_new_children_and_same_properties( self: Arc, - mut children: Vec>, + children: Vec>, ) -> Result> { - Ok(Arc::new(Self { - input: children.swap_remove(0), - metrics: ExecutionPlanMetricsSet::new(), - ..Self::clone(&*self) - })) + self.replace_children(children, ChildrenPropertiesHint::SameProperties) } fn execute( diff --git a/datafusion/physical-plan/src/projection.rs b/datafusion/physical-plan/src/projection.rs index 42501f22395b4..f830044619466 100644 --- a/datafusion/physical-plan/src/projection.rs +++ b/datafusion/physical-plan/src/projection.rs @@ -34,7 +34,10 @@ use crate::filter_pushdown::{ }; use crate::joins::utils::{ColumnIndex, JoinFilter, JoinOn, JoinOnRef}; use crate::statistics::{ChildStats, StatisticsArgs}; -use crate::{DisplayFormatType, ExecutionPlan, PhysicalExpr, check_if_same_properties}; +use crate::{ + ChildrenPropertiesHint, DisplayFormatType, ExecutionPlan, PhysicalExpr, + validate_child_count, +}; use std::collections::HashMap; use std::pin::Pin; use std::sync::Arc; @@ -302,27 +305,38 @@ impl ExecutionPlan for ProjectionExec { vec![&self.input] } - fn with_new_children( + fn replace_children( self: Arc, mut children: Vec>, + hint: ChildrenPropertiesHint, ) -> Result> { - check_if_same_properties!(self, children); - ProjectionExec::try_from_projector( - self.projector.clone(), - children.swap_remove(0), - ) - .map(|p| Arc::new(p) as _) + validate_child_count!(self, children); + match hint { + ChildrenPropertiesHint::SameProperties => Ok(Arc::new(Self { + input: children.swap_remove(0), + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(&*self) + })), + ChildrenPropertiesHint::Recompute => ProjectionExec::try_from_projector( + self.projector.clone(), + children.swap_remove(0), + ) + .map(|p| Arc::new(p) as _), + } + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children(children, ChildrenPropertiesHint::Recompute) } fn with_new_children_and_same_properties( self: Arc, - mut children: Vec>, + children: Vec>, ) -> Result> { - Ok(Arc::new(Self { - input: children.swap_remove(0), - metrics: ExecutionPlanMetricsSet::new(), - ..Self::clone(&*self) - })) + self.replace_children(children, ChildrenPropertiesHint::SameProperties) } fn execute( diff --git a/datafusion/physical-plan/src/repartition/mod.rs b/datafusion/physical-plan/src/repartition/mod.rs index 3473aad9b3fc0..b7ce64e9ca9a2 100644 --- a/datafusion/physical-plan/src/repartition/mod.rs +++ b/datafusion/physical-plan/src/repartition/mod.rs @@ -15,7 +15,7 @@ // specific language governing permissions and limitations // under the License. -//! This file implements the [`RepartitionExec`] operator, which maps N input +//! This file implements the [`RepartitionExec`] operator, which maps N input //! partitions to M output partitions based on a partitioning scheme, optionally //! maintaining the order of the input rows in the output. @@ -43,8 +43,8 @@ use crate::spill::spill_pool::{self, SpillPoolSink, SpillPoolWriter}; use crate::statistics::{ChildStats, StatisticsArgs}; use crate::stream::{EmptyRecordBatchStream, RecordBatchStreamAdapter}; use crate::{ - DisplayFormatType, ExecutionPlan, Partitioning, PlanProperties, Statistics, - check_if_same_properties, + ChildrenPropertiesHint, DisplayFormatType, ExecutionPlan, Partitioning, + PlanProperties, Statistics, validate_child_count, }; use arrow::array::{Array, PrimitiveArray, RecordBatch, RecordBatchOptions}; @@ -1337,31 +1337,44 @@ impl ExecutionPlan for RepartitionExec { vec![&self.input] } - fn with_new_children( + fn replace_children( self: Arc, mut children: Vec>, + hint: ChildrenPropertiesHint, ) -> Result> { - check_if_same_properties!(self, children); - let mut repartition = RepartitionExec::try_new( - children.swap_remove(0), - self.partitioning().clone(), - )?; - if self.preserve_order { - repartition = repartition.with_preserve_order(); + validate_child_count!(self, children); + match hint { + ChildrenPropertiesHint::SameProperties => Ok(Arc::new(Self { + input: children.swap_remove(0), + metrics: ExecutionPlanMetricsSet::new(), + state: Default::default(), + ..Self::clone(&*self) + })), + ChildrenPropertiesHint::Recompute => { + let mut repartition = RepartitionExec::try_new( + children.swap_remove(0), + self.partitioning().clone(), + )?; + if self.preserve_order { + repartition = repartition.with_preserve_order(); + } + Ok(Arc::new(repartition)) + } } - Ok(Arc::new(repartition)) + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children(children, ChildrenPropertiesHint::Recompute) } fn with_new_children_and_same_properties( self: Arc, - mut children: Vec>, + children: Vec>, ) -> Result> { - Ok(Arc::new(Self { - input: children.swap_remove(0), - metrics: ExecutionPlanMetricsSet::new(), - state: Default::default(), - ..Self::clone(&*self) - })) + self.replace_children(children, ChildrenPropertiesHint::SameProperties) } fn benefits_from_input_partitioning(&self) -> Vec { diff --git a/datafusion/physical-plan/src/sorts/partial_sort.rs b/datafusion/physical-plan/src/sorts/partial_sort.rs index 3a4ddd2fbeaf7..b125824f628d7 100644 --- a/datafusion/physical-plan/src/sorts/partial_sort.rs +++ b/datafusion/physical-plan/src/sorts/partial_sort.rs @@ -61,9 +61,9 @@ use crate::sorts::sort::sort_batch; use crate::statistics::{ChildStats, StatisticsArgs}; use crate::stream::EmptyRecordBatchStream; use crate::{ - DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, ExecutionPlanProperties, - Partitioning, PlanProperties, SendableRecordBatchStream, Statistics, - check_if_same_properties, + ChildrenPropertiesHint, DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, + ExecutionPlanProperties, Partitioning, PlanProperties, SendableRecordBatchStream, + Statistics, validate_child_count, }; use arrow::compute::concat_batches; @@ -307,31 +307,44 @@ impl ExecutionPlan for PartialSortExec { vec![&self.input] } + fn replace_children( + self: Arc, + mut children: Vec>, + hint: ChildrenPropertiesHint, + ) -> Result> { + validate_child_count!(self, children); + match hint { + ChildrenPropertiesHint::SameProperties => Ok(Arc::new(Self { + input: children.swap_remove(0), + metrics_set: ExecutionPlanMetricsSet::new(), + ..Self::clone(&*self) + })), + ChildrenPropertiesHint::Recompute => { + let new_partial_sort = PartialSortExec::new( + self.expr.clone(), + Arc::clone(&children[0]), + self.common_prefix_length, + ) + .with_fetch(self.fetch) + .with_preserve_partitioning(self.preserve_partitioning); + + Ok(Arc::new(new_partial_sort)) + } + } + } + fn with_new_children( self: Arc, children: Vec>, ) -> Result> { - check_if_same_properties!(self, children); - let new_partial_sort = PartialSortExec::new( - self.expr.clone(), - Arc::clone(&children[0]), - self.common_prefix_length, - ) - .with_fetch(self.fetch) - .with_preserve_partitioning(self.preserve_partitioning); - - Ok(Arc::new(new_partial_sort)) + self.replace_children(children, ChildrenPropertiesHint::Recompute) } fn with_new_children_and_same_properties( self: Arc, - mut children: Vec>, + children: Vec>, ) -> Result> { - Ok(Arc::new(Self { - input: children.swap_remove(0), - metrics_set: ExecutionPlanMetricsSet::new(), - ..Self::clone(&*self) - })) + self.replace_children(children, ChildrenPropertiesHint::SameProperties) } fn execute( diff --git a/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs b/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs index 0a6571ffaf74d..eb9730ddf1f19 100644 --- a/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs +++ b/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs @@ -28,7 +28,7 @@ use crate::statistics::{ChildStats, StatisticsArgs}; use crate::{ ChildrenPropertiesHint, DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, ExecutionPlanProperties, Partitioning, PlanProperties, SendableRecordBatchStream, - Statistics, + Statistics, validate_child_count, }; use datafusion_common::{Result, assert_eq_or_internal_err, internal_err}; @@ -286,6 +286,7 @@ impl ExecutionPlan for SortPreservingMergeExec { mut children: Vec>, hint: ChildrenPropertiesHint, ) -> Result> { + validate_child_count!(self, children); match hint { ChildrenPropertiesHint::SameProperties => Ok(Arc::new(Self { input: children.swap_remove(0), diff --git a/datafusion/physical-plan/src/union.rs b/datafusion/physical-plan/src/union.rs index 4722329ea55a4..830e24a944a9c 100644 --- a/datafusion/physical-plan/src/union.rs +++ b/datafusion/physical-plan/src/union.rs @@ -31,7 +31,6 @@ use super::{ PlanProperties, RecordBatchStream, SendableRecordBatchStream, Statistics, metrics::{ExecutionPlanMetricsSet, MetricsSet}, }; -use crate::check_if_same_properties; use crate::execution_plan::{ CardinalityEffect, InvariantLevel, boundedness_from_children, check_default_invariants, emission_type_from_children, @@ -45,6 +44,7 @@ use crate::metrics::BaselineMetrics; use crate::projection::{ProjectionExec, make_with_child}; use crate::statistics::{ChildStats, StatisticsArgs}; use crate::stream::ObservedStream; +use crate::{ChildrenPropertiesHint, validate_child_count}; use arrow::datatypes::{Field, Schema, SchemaRef}; use arrow::record_batch::RecordBatch; @@ -252,23 +252,34 @@ impl ExecutionPlan for UnionExec { self.inputs.iter().collect() } + fn replace_children( + self: Arc, + children: Vec>, + hint: ChildrenPropertiesHint, + ) -> Result> { + validate_child_count!(self, children); + match hint { + ChildrenPropertiesHint::SameProperties => Ok(Arc::new(Self { + inputs: children, + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(&*self) + })), + ChildrenPropertiesHint::Recompute => UnionExec::try_new(children), + } + } + fn with_new_children( self: Arc, children: Vec>, ) -> Result> { - check_if_same_properties!(self, children); - UnionExec::try_new(children) + self.replace_children(children, ChildrenPropertiesHint::Recompute) } fn with_new_children_and_same_properties( self: Arc, children: Vec>, ) -> Result> { - Ok(Arc::new(Self { - inputs: children, - metrics: ExecutionPlanMetricsSet::new(), - ..Self::clone(&*self) - })) + self.replace_children(children, ChildrenPropertiesHint::SameProperties) } fn execute( @@ -624,28 +635,41 @@ impl ExecutionPlan for InterleaveExec { vec![false; self.inputs().len()] } + fn replace_children( + self: Arc, + children: Vec>, + hint: ChildrenPropertiesHint, + ) -> Result> { + validate_child_count!(self, children); + match hint { + ChildrenPropertiesHint::SameProperties => Ok(Arc::new(Self { + inputs: children, + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(&*self) + })), + ChildrenPropertiesHint::Recompute => { + // New children are no longer interleavable, which might be a bug of optimization rewrite. + assert_or_internal_err!( + can_interleave(children.iter()), + "Can not create InterleaveExec: new children can not be interleaved" + ); + Ok(Arc::new(InterleaveExec::try_new(children)?)) + } + } + } + fn with_new_children( self: Arc, children: Vec>, ) -> Result> { - // New children are no longer interleavable, which might be a bug of optimization rewrite. - assert_or_internal_err!( - can_interleave(children.iter()), - "Can not create InterleaveExec: new children can not be interleaved" - ); - check_if_same_properties!(self, children); - Ok(Arc::new(InterleaveExec::try_new(children)?)) + self.replace_children(children, ChildrenPropertiesHint::Recompute) } fn with_new_children_and_same_properties( self: Arc, children: Vec>, ) -> Result> { - Ok(Arc::new(Self { - inputs: children, - metrics: ExecutionPlanMetricsSet::new(), - ..Self::clone(&*self) - })) + self.replace_children(children, ChildrenPropertiesHint::SameProperties) } fn execute( diff --git a/datafusion/physical-plan/src/unnest.rs b/datafusion/physical-plan/src/unnest.rs index 6086515ff1834..b5da0ae3bc621 100644 --- a/datafusion/physical-plan/src/unnest.rs +++ b/datafusion/physical-plan/src/unnest.rs @@ -29,7 +29,7 @@ use super::{DisplayAs, ExecutionPlanProperties, PlanProperties}; use crate::stream::EmptyRecordBatchStream; use crate::{ ChildrenPropertiesHint, DisplayFormatType, Distribution, ExecutionPlan, - RecordBatchStream, SendableRecordBatchStream, + RecordBatchStream, SendableRecordBatchStream, validate_child_count, }; use arrow::array::{ @@ -232,6 +232,7 @@ impl ExecutionPlan for UnnestExec { mut children: Vec>, hint: ChildrenPropertiesHint, ) -> Result> { + validate_child_count!(self, children); match hint { ChildrenPropertiesHint::SameProperties => Ok(Arc::new(Self { input: children.swap_remove(0), diff --git a/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs b/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs index d5863080895f6..6919870687ed4 100644 --- a/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs +++ b/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs @@ -35,10 +35,10 @@ use crate::windows::{ window_equivalence_properties, }; use crate::{ - ColumnStatistics, DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, - ExecutionPlanProperties, InputDistributionRequirements, InputOrderMode, - PlanProperties, RecordBatchStream, SendableRecordBatchStream, Statistics, WindowExpr, - check_if_same_properties, + ChildrenPropertiesHint, ColumnStatistics, DisplayAs, DisplayFormatType, Distribution, + ExecutionPlan, ExecutionPlanProperties, InputDistributionRequirements, + InputOrderMode, PlanProperties, RecordBatchStream, SendableRecordBatchStream, + Statistics, WindowExpr, validate_child_count, }; use arrow::compute::take_record_batch; @@ -340,28 +340,41 @@ impl ExecutionPlan for BoundedWindowAggExec { vec![true] } + fn replace_children( + self: Arc, + mut children: Vec>, + hint: ChildrenPropertiesHint, + ) -> Result> { + validate_child_count!(self, children); + match hint { + ChildrenPropertiesHint::SameProperties => Ok(Arc::new(Self { + input: children.swap_remove(0), + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(&*self) + })), + ChildrenPropertiesHint::Recompute => { + Ok(Arc::new(BoundedWindowAggExec::try_new( + self.window_expr.clone(), + Arc::clone(&children[0]), + self.input_order_mode.clone(), + self.can_repartition, + )?)) + } + } + } + fn with_new_children( self: Arc, children: Vec>, ) -> Result> { - check_if_same_properties!(self, children); - Ok(Arc::new(BoundedWindowAggExec::try_new( - self.window_expr.clone(), - Arc::clone(&children[0]), - self.input_order_mode.clone(), - self.can_repartition, - )?)) + self.replace_children(children, ChildrenPropertiesHint::Recompute) } fn with_new_children_and_same_properties( self: Arc, - mut children: Vec>, + children: Vec>, ) -> Result> { - Ok(Arc::new(Self { - input: children.swap_remove(0), - metrics: ExecutionPlanMetricsSet::new(), - ..Self::clone(&*self) - })) + self.replace_children(children, ChildrenPropertiesHint::SameProperties) } fn execute( diff --git a/datafusion/physical-plan/src/windows/window_agg_exec.rs b/datafusion/physical-plan/src/windows/window_agg_exec.rs index 81838300cf5c7..3f7ac1b11f351 100644 --- a/datafusion/physical-plan/src/windows/window_agg_exec.rs +++ b/datafusion/physical-plan/src/windows/window_agg_exec.rs @@ -33,10 +33,10 @@ use crate::windows::{ window_equivalence_properties, }; use crate::{ - ColumnStatistics, DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, - ExecutionPlanProperties, InputDistributionRequirements, PhysicalExpr, PlanProperties, - RecordBatchStream, SendableRecordBatchStream, Statistics, WindowExpr, - check_if_same_properties, + ChildrenPropertiesHint, ColumnStatistics, DisplayAs, DisplayFormatType, Distribution, + ExecutionPlan, ExecutionPlanProperties, InputDistributionRequirements, PhysicalExpr, + PlanProperties, RecordBatchStream, SendableRecordBatchStream, Statistics, WindowExpr, + validate_child_count, }; use arrow::array::ArrayRef; @@ -246,27 +246,38 @@ impl ExecutionPlan for WindowAggExec { } } - fn with_new_children( + fn replace_children( self: Arc, mut children: Vec>, + hint: ChildrenPropertiesHint, ) -> Result> { - check_if_same_properties!(self, children); - Ok(Arc::new(WindowAggExec::try_new( - self.window_expr.clone(), - children.swap_remove(0), - true, - )?)) + validate_child_count!(self, children); + match hint { + ChildrenPropertiesHint::SameProperties => Ok(Arc::new(Self { + input: children.swap_remove(0), + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(&*self) + })), + ChildrenPropertiesHint::Recompute => Ok(Arc::new(WindowAggExec::try_new( + self.window_expr.clone(), + children.swap_remove(0), + true, + )?)), + } + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children(children, ChildrenPropertiesHint::Recompute) } fn with_new_children_and_same_properties( self: Arc, - mut children: Vec>, + children: Vec>, ) -> Result> { - Ok(Arc::new(Self { - input: children.swap_remove(0), - metrics: ExecutionPlanMetricsSet::new(), - ..Self::clone(&*self) - })) + self.replace_children(children, ChildrenPropertiesHint::SameProperties) } fn execute( From 20df17de2e96a39915a2237d9ff59f841fb27182 Mon Sep 17 00:00:00 2001 From: Justin O'Dwyer Date: Tue, 28 Jul 2026 12:40:51 -0400 Subject: [PATCH 5/9] Migrate more implementations to replace_children. --- .../custom_data_source/custom_datasource.rs | 14 +++++-- .../memory_pool_execution_plan.rs | 12 +++++- .../proto/composed_extension_codec.rs | 25 +++++++++-- .../examples/relation_planner/table_sample.rs | 15 ++++++- datafusion/catalog/src/memory/table.rs | 16 +++++-- datafusion/core/src/physical_planner.rs | 42 ++++++++++++++++--- datafusion/core/tests/fuzz_cases/once_exec.rs | 16 +++++-- .../tests/user_defined/insert_operation.rs | 12 +++++- .../tests/user_defined/user_defined_plan.rs | 11 ++++- datafusion/ffi/src/execution_plan.rs | 24 +++++++++-- datafusion/ffi/src/tests/async_provider.rs | 14 +++++-- .../physical-optimizer/src/ensure_coop.rs | 13 +++++- .../src/output_requirements.rs | 14 +++++-- datafusion/physical-plan/src/async_func.rs | 2 +- datafusion/physical-plan/src/buffer.rs | 2 +- datafusion/physical-plan/src/coop.rs | 2 +- datafusion/physical-plan/src/limit.rs | 2 +- .../physical-plan/src/placeholder_row.rs | 14 +++++-- .../physical-plan/src/scalar_subquery.rs | 24 +++++++++-- .../src/sorts/partitioned_topk.rs | 11 ++++- datafusion/physical-plan/src/sorts/sort.rs | 40 ++++++++++++++---- .../src/sorts/sort_preserving_merge.rs | 14 +++++-- datafusion/physical-plan/src/streaming.rs | 14 ++++++- .../tests/cases/roundtrip_physical_plan.rs | 26 ++++++++++-- 24 files changed, 309 insertions(+), 70 deletions(-) diff --git a/datafusion-examples/examples/custom_data_source/custom_datasource.rs b/datafusion-examples/examples/custom_data_source/custom_datasource.rs index a2d7d7699927f..c88d4af258db1 100644 --- a/datafusion-examples/examples/custom_data_source/custom_datasource.rs +++ b/datafusion-examples/examples/custom_data_source/custom_datasource.rs @@ -35,8 +35,8 @@ use datafusion::physical_expr::EquivalenceProperties; use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType}; use datafusion::physical_plan::memory::MemoryStream; use datafusion::physical_plan::{ - DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, PlanProperties, - SendableRecordBatchStream, project_schema, + ChildrenPropertiesHint, DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, + PlanProperties, SendableRecordBatchStream, project_schema, }; use datafusion::prelude::*; @@ -267,13 +267,21 @@ impl ExecutionPlan for CustomExec { vec![] } - fn with_new_children( + fn replace_children( self: Arc, _: Vec>, + _: ChildrenPropertiesHint, ) -> Result> { Ok(self) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children(children, ChildrenPropertiesHint::Recompute) + } + fn execute( &self, _partition: usize, diff --git a/datafusion-examples/examples/execution_monitoring/memory_pool_execution_plan.rs b/datafusion-examples/examples/execution_monitoring/memory_pool_execution_plan.rs index ca765774d141f..131b03e8dfcc9 100644 --- a/datafusion-examples/examples/execution_monitoring/memory_pool_execution_plan.rs +++ b/datafusion-examples/examples/execution_monitoring/memory_pool_execution_plan.rs @@ -38,7 +38,7 @@ use datafusion::execution::{SendableRecordBatchStream, TaskContext}; use datafusion::logical_expr::LogicalPlanBuilder; use datafusion::physical_plan::stream::RecordBatchStreamAdapter; use datafusion::physical_plan::{ - DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, + ChildrenPropertiesHint, DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, }; use datafusion::prelude::*; use futures::stream::{StreamExt, TryStreamExt}; @@ -236,9 +236,10 @@ impl ExecutionPlan for BufferingExecutionPlan { vec![&self.input] } - fn with_new_children( + fn replace_children( self: Arc, children: Vec>, + _hint: ChildrenPropertiesHint, ) -> Result> { if children.len() == 1 { Ok(Arc::new(BufferingExecutionPlan::new( @@ -250,6 +251,13 @@ impl ExecutionPlan for BufferingExecutionPlan { } } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children(children, ChildrenPropertiesHint::Recompute) + } + fn execute( &self, partition: usize, diff --git a/datafusion-examples/examples/proto/composed_extension_codec.rs b/datafusion-examples/examples/proto/composed_extension_codec.rs index 6077a982c320d..a13528a70a21b 100644 --- a/datafusion-examples/examples/proto/composed_extension_codec.rs +++ b/datafusion-examples/examples/proto/composed_extension_codec.rs @@ -38,6 +38,7 @@ use std::sync::Arc; use datafusion::common::Result; use datafusion::common::internal_err; use datafusion::execution::TaskContext; +use datafusion::physical_plan::ChildrenPropertiesHint; use datafusion::physical_plan::{DisplayAs, ExecutionPlan}; use datafusion::prelude::SessionContext; use datafusion_proto::physical_plan::{ @@ -110,13 +111,21 @@ impl ExecutionPlan for ParentExec { vec![&self.input] } - fn with_new_children( + fn replace_children( self: Arc, - _children: Vec>, + _: Vec>, + _: ChildrenPropertiesHint, ) -> Result> { unreachable!() } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children(children, ChildrenPropertiesHint::Recompute) + } + fn execute( &self, _partition: usize, @@ -188,13 +197,21 @@ impl ExecutionPlan for ChildExec { vec![] } - fn with_new_children( + fn replace_children( self: Arc, - _children: Vec>, + _: Vec>, + _: ChildrenPropertiesHint, ) -> Result> { unreachable!() } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children(children, ChildrenPropertiesHint::Recompute) + } + fn execute( &self, _partition: usize, diff --git a/datafusion-examples/examples/relation_planner/table_sample.rs b/datafusion-examples/examples/relation_planner/table_sample.rs index b0ccff8d10d8c..1de56f650285e 100644 --- a/datafusion-examples/examples/relation_planner/table_sample.rs +++ b/datafusion-examples/examples/relation_planner/table_sample.rs @@ -100,7 +100,6 @@ use futures::{ use rand::{Rng, SeedableRng, rngs::StdRng}; use tonic::async_trait; -use datafusion::optimizer::simplify_expressions::simplify_literal::parse_literal; use datafusion::{ execution::{ RecordBatchStream, SendableRecordBatchStream, SessionState, SessionStateBuilder, @@ -115,6 +114,10 @@ use datafusion::{ physical_planner::{DefaultPhysicalPlanner, ExtensionPlanner, PhysicalPlanner}, prelude::*, }; +use datafusion::{ + optimizer::simplify_expressions::simplify_literal::parse_literal, + physical_plan::ChildrenPropertiesHint, +}; use datafusion_common::{ DFSchemaRef, DataFusionError, Result, Statistics, internal_err, not_impl_err, plan_datafusion_err, plan_err, @@ -697,9 +700,10 @@ impl ExecutionPlan for SampleExec { vec![&self.input] } - fn with_new_children( + fn replace_children( self: Arc, mut children: Vec>, + _: ChildrenPropertiesHint, ) -> Result> { Ok(Arc::new(Self::try_new( children.swap_remove(0), @@ -709,6 +713,13 @@ impl ExecutionPlan for SampleExec { )?)) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children(children, ChildrenPropertiesHint::Recompute) + } + fn execute( &self, partition: usize, diff --git a/datafusion/catalog/src/memory/table.rs b/datafusion/catalog/src/memory/table.rs index 5d07133799ffc..eb0baa6e09925 100644 --- a/datafusion/catalog/src/memory/table.rs +++ b/datafusion/catalog/src/memory/table.rs @@ -44,8 +44,8 @@ use datafusion_physical_expr::{ use datafusion_physical_plan::repartition::RepartitionExec; use datafusion_physical_plan::stream::RecordBatchStreamAdapter; use datafusion_physical_plan::{ - DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, PlanProperties, - collect_partitioned, + ChildrenPropertiesHint, DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, + PlanProperties, collect_partitioned, }; use datafusion_session::Session; @@ -571,13 +571,21 @@ impl ExecutionPlan for DmlResultExec { vec![] } - fn with_new_children( + fn replace_children( self: Arc, - _children: Vec>, + _: Vec>, + _: datafusion_physical_plan::ChildrenPropertiesHint, ) -> Result> { Ok(self) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children(children, ChildrenPropertiesHint::Recompute) + } + fn execute( &self, _partition: usize, diff --git a/datafusion/core/src/physical_planner.rs b/datafusion/core/src/physical_planner.rs index 4e914556b4cc0..49431d8c62a01 100644 --- a/datafusion/core/src/physical_planner.rs +++ b/datafusion/core/src/physical_planner.rs @@ -3389,6 +3389,7 @@ mod tests { use datafusion_functions_aggregate::count::{count_all, count_udaf}; use datafusion_functions_aggregate::expr_fn::sum; use datafusion_physical_expr::EquivalenceProperties; + use datafusion_physical_plan::ChildrenPropertiesHint; use datafusion_physical_plan::execution_plan::{Boundedness, EmissionType}; fn make_session_state() -> SessionState { @@ -4693,9 +4694,10 @@ mod tests { vec![] } - fn with_new_children( + fn replace_children( self: Arc, children: Vec>, + _: ChildrenPropertiesHint, ) -> Result> { if children.is_empty() { Ok(self) @@ -4704,6 +4706,13 @@ mod tests { } } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children(children, ChildrenPropertiesHint::Recompute) + } + fn execute( &self, _partition: usize, @@ -4856,12 +4865,19 @@ digraph { fn name(&self) -> &str { "always ok" } - fn with_new_children( + fn replace_children( self: Arc, children: Vec>, + _: ChildrenPropertiesHint, ) -> Result> { Ok(Arc::new(Self(children))) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children(children, ChildrenPropertiesHint::Recompute) + } fn schema(&self) -> SchemaRef { Arc::new(Schema::empty()) } @@ -4905,12 +4921,19 @@ digraph { fn schema(&self) -> SchemaRef { Arc::new(Schema::empty()) } - fn with_new_children( + fn replace_children( self: Arc, - _children: Vec>, + _: Vec>, + _: ChildrenPropertiesHint, ) -> Result> { unimplemented!() } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children(children, ChildrenPropertiesHint::Recompute) + } fn children(&self) -> Vec<&Arc> { unimplemented!() } @@ -5029,12 +5052,19 @@ digraph { fn schema(&self) -> SchemaRef { Arc::new(Schema::empty()) } - fn with_new_children( + fn replace_children( self: Arc, - _children: Vec>, + _: Vec>, + _: ChildrenPropertiesHint, ) -> Result> { unimplemented!() } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children(children, ChildrenPropertiesHint::Recompute) + } fn children(&self) -> Vec<&Arc> { vec![] } diff --git a/datafusion/core/tests/fuzz_cases/once_exec.rs b/datafusion/core/tests/fuzz_cases/once_exec.rs index 9b57141061518..b6ac0e2f67c8b 100644 --- a/datafusion/core/tests/fuzz_cases/once_exec.rs +++ b/datafusion/core/tests/fuzz_cases/once_exec.rs @@ -16,12 +16,12 @@ // under the License. use arrow_schema::SchemaRef; -use datafusion_common::internal_datafusion_err; +use datafusion_common::{Result, internal_datafusion_err}; use datafusion_execution::{SendableRecordBatchStream, TaskContext}; use datafusion_physical_expr::{EquivalenceProperties, Partitioning}; use datafusion_physical_plan::execution_plan::{Boundedness, EmissionType}; use datafusion_physical_plan::{ - DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, + ChildrenPropertiesHint, DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, }; use std::fmt::{Debug, Formatter}; use std::sync::{Arc, Mutex}; @@ -86,13 +86,21 @@ impl ExecutionPlan for OnceExec { vec![] } - fn with_new_children( + fn replace_children( self: Arc, _: Vec>, - ) -> datafusion_common::Result> { + _: ChildrenPropertiesHint, + ) -> Result> { unimplemented!() } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> datafusion_common::Result> { + self.replace_children(children, ChildrenPropertiesHint::Recompute) + } + /// Returns a stream which yields data fn execute( &self, diff --git a/datafusion/core/tests/user_defined/insert_operation.rs b/datafusion/core/tests/user_defined/insert_operation.rs index f3d3f70bdf925..904dace480c1a 100644 --- a/datafusion/core/tests/user_defined/insert_operation.rs +++ b/datafusion/core/tests/user_defined/insert_operation.rs @@ -27,7 +27,7 @@ use datafusion_catalog::{Session, TableProvider}; use datafusion_common::config::Dialect; use datafusion_expr::{Expr, TableType, dml::InsertOp}; use datafusion_physical_expr::{EquivalenceProperties, Partitioning}; -use datafusion_physical_plan::execution_plan::SchedulingType; +use datafusion_physical_plan::{ChildrenPropertiesHint, execution_plan::SchedulingType}; use datafusion_physical_plan::{ DisplayAs, ExecutionPlan, PlanProperties, execution_plan::{Boundedness, EmissionType}, @@ -161,14 +161,22 @@ impl ExecutionPlan for TestInsertExec { vec![] } - fn with_new_children( + fn replace_children( self: Arc, children: Vec>, + _: ChildrenPropertiesHint, ) -> Result> { assert!(children.is_empty()); Ok(self) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children(children, ChildrenPropertiesHint::Recompute) + } + fn execute( &self, _partition: usize, diff --git a/datafusion/core/tests/user_defined/user_defined_plan.rs b/datafusion/core/tests/user_defined/user_defined_plan.rs index 354a1b3110250..ac9745d9c306c 100644 --- a/datafusion/core/tests/user_defined/user_defined_plan.rs +++ b/datafusion/core/tests/user_defined/user_defined_plan.rs @@ -95,6 +95,7 @@ use datafusion_common::{ScalarValue, assert_eq_or_internal_err, assert_or_intern use datafusion_expr::{FetchType, InvariantLevel, Projection, SortExpr}; use datafusion_optimizer::AnalyzerRule; use datafusion_optimizer::optimizer::ApplyOrder; +use datafusion_physical_plan::ChildrenPropertiesHint; use datafusion_physical_plan::execution_plan::{Boundedness, EmissionType}; use async_trait::async_trait; @@ -722,13 +723,21 @@ impl ExecutionPlan for TopKExec { vec![&self.input] } - fn with_new_children( + fn replace_children( self: Arc, children: Vec>, + _: ChildrenPropertiesHint, ) -> Result> { Ok(Arc::new(TopKExec::new(children[0].clone(), self.k))) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children(children, ChildrenPropertiesHint::Recompute) + } + /// Execute one partition and return an iterator over RecordBatch fn execute( &self, diff --git a/datafusion/ffi/src/execution_plan.rs b/datafusion/ffi/src/execution_plan.rs index 087a351b697cc..ba071944aba8f 100644 --- a/datafusion/ffi/src/execution_plan.rs +++ b/datafusion/ffi/src/execution_plan.rs @@ -24,8 +24,8 @@ use datafusion_common::{DataFusionError, Result, Statistics}; use datafusion_execution::{SendableRecordBatchStream, TaskContext}; use datafusion_physical_expr_common::metrics::MetricsSet; use datafusion_physical_plan::{ - DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, StatisticsArgs, - StatisticsContext, + ChildrenPropertiesHint, DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, + StatisticsArgs, StatisticsContext, }; use stabby::string::String as SString; use stabby::vec::Vec as SVec; @@ -416,9 +416,10 @@ impl ExecutionPlan for ForeignExecutionPlan { self.children.iter().collect() } - fn with_new_children( + fn replace_children( self: Arc, children: Vec>, + _: ChildrenPropertiesHint, ) -> Result> { let children = children .into_iter() @@ -430,6 +431,13 @@ impl ExecutionPlan for ForeignExecutionPlan { (&new_plan).try_into() } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children(children, ChildrenPropertiesHint::Recompute) + } + fn execute( &self, partition: usize, @@ -536,9 +544,10 @@ pub mod tests { self.children.iter().collect() } - fn with_new_children( + fn replace_children( self: Arc, children: Vec>, + _: ChildrenPropertiesHint, ) -> Result> { Ok(Arc::new(EmptyExec { props: Arc::clone(&self.props), @@ -548,6 +557,13 @@ pub mod tests { })) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children(children, ChildrenPropertiesHint::Recompute) + } + fn execute( &self, _partition: usize, diff --git a/datafusion/ffi/src/tests/async_provider.rs b/datafusion/ffi/src/tests/async_provider.rs index 9821c3e501f67..b95b1bf05ecbe 100644 --- a/datafusion/ffi/src/tests/async_provider.rs +++ b/datafusion/ffi/src/tests/async_provider.rs @@ -36,7 +36,7 @@ use datafusion_common::{Result, exec_err}; use datafusion_execution::RecordBatchStream; use datafusion_expr::Expr; use datafusion_physical_expr::{EquivalenceProperties, Partitioning}; -use datafusion_physical_plan::ExecutionPlan; +use datafusion_physical_plan::{ChildrenPropertiesHint, ExecutionPlan}; use datafusion_session::Session; use futures::Stream; use tokio::runtime::Handle; @@ -210,13 +210,21 @@ impl ExecutionPlan for AsyncTestExecutionPlan { Vec::default() } - fn with_new_children( + fn replace_children( self: Arc, - _children: Vec>, + _: Vec>, + _: ChildrenPropertiesHint, ) -> Result> { Ok(self) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children(children, ChildrenPropertiesHint::Recompute) + } + fn execute( &self, _partition: usize, diff --git a/datafusion/physical-optimizer/src/ensure_coop.rs b/datafusion/physical-optimizer/src/ensure_coop.rs index e7aacb2321b67..7593d6ec24cb8 100644 --- a/datafusion/physical-optimizer/src/ensure_coop.rs +++ b/datafusion/physical-optimizer/src/ensure_coop.rs @@ -130,7 +130,9 @@ impl PhysicalOptimizerRule for EnsureCooperative { #[cfg(test)] mod tests { use super::*; - use datafusion_physical_plan::{displayable, test::scan_partitioned}; + use datafusion_physical_plan::{ + ChildrenPropertiesHint, displayable, test::scan_partitioned, + }; use insta::assert_snapshot; #[tokio::test] @@ -327,9 +329,10 @@ mod tests { fn children(&self) -> Vec<&Arc> { vec![&self.input] } - fn with_new_children( + fn replace_children( self: Arc, children: Vec>, + _: ChildrenPropertiesHint, ) -> Result> { Ok(Arc::new(DummyExec::new( &self.name, @@ -338,6 +341,12 @@ mod tests { self.evaluation_type, ))) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children(children, ChildrenPropertiesHint::Recompute) + } fn execute( &self, _: usize, diff --git a/datafusion/physical-optimizer/src/output_requirements.rs b/datafusion/physical-optimizer/src/output_requirements.rs index 0583e5ee885a5..5a727c0db073e 100644 --- a/datafusion/physical-optimizer/src/output_requirements.rs +++ b/datafusion/physical-optimizer/src/output_requirements.rs @@ -40,8 +40,8 @@ use datafusion_physical_plan::scalar_subquery::ScalarSubqueryExec; use datafusion_physical_plan::sorts::sort::SortExec; use datafusion_physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec; use datafusion_physical_plan::{ - ChildStats, DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, - PlanProperties, SendableRecordBatchStream, StatisticsArgs, + ChildStats, ChildrenPropertiesHint, DisplayAs, DisplayFormatType, ExecutionPlan, + ExecutionPlanProperties, PlanProperties, SendableRecordBatchStream, StatisticsArgs, with_new_children_if_necessary, }; @@ -232,9 +232,10 @@ impl ExecutionPlan for OutputRequirementExec { vec![self.order_requirement.clone()] } - fn with_new_children( + fn replace_children( self: Arc, mut children: Vec>, + _: ChildrenPropertiesHint, ) -> Result> { Ok(Arc::new(Self::new( children.remove(0), // has a single child @@ -244,6 +245,13 @@ impl ExecutionPlan for OutputRequirementExec { ))) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children(children, ChildrenPropertiesHint::Recompute) + } + fn execute( &self, _partition: usize, diff --git a/datafusion/physical-plan/src/async_func.rs b/datafusion/physical-plan/src/async_func.rs index 214b7703450e0..a823e61aa6b62 100644 --- a/datafusion/physical-plan/src/async_func.rs +++ b/datafusion/physical-plan/src/async_func.rs @@ -24,8 +24,8 @@ use crate::{ }; use arrow::array::RecordBatch; use arrow_schema::{FieldRef, Fields, Schema, SchemaRef}; +use datafusion_common::Result; use datafusion_common::tree_node::{Transformed, TreeNode, TreeNodeRecursion}; -use datafusion_common::{Result, assert_eq_or_internal_err}; use datafusion_execution::{RecordBatchStream, SendableRecordBatchStream, TaskContext}; use datafusion_physical_expr::ScalarFunctionExpr; use datafusion_physical_expr::async_scalar_function::AsyncFuncExpr; diff --git a/datafusion/physical-plan/src/buffer.rs b/datafusion/physical-plan/src/buffer.rs index 1e82d1f2dddc4..f91f5fd109d87 100644 --- a/datafusion/physical-plan/src/buffer.rs +++ b/datafusion/physical-plan/src/buffer.rs @@ -32,7 +32,7 @@ use crate::{ }; use arrow::array::RecordBatch; use datafusion_common::config::ConfigOptions; -use datafusion_common::{Result, Statistics, internal_err, plan_err}; +use datafusion_common::{Result, Statistics, internal_err}; use datafusion_common_runtime::SpawnedTask; use datafusion_execution::memory_pool::{MemoryConsumer, MemoryReservation}; use datafusion_execution::{SendableRecordBatchStream, TaskContext}; diff --git a/datafusion/physical-plan/src/coop.rs b/datafusion/physical-plan/src/coop.rs index 28eeb61b5cc5f..9f29f8fde87fe 100644 --- a/datafusion/physical-plan/src/coop.rs +++ b/datafusion/physical-plan/src/coop.rs @@ -92,7 +92,7 @@ use crate::{ }; use arrow::record_batch::RecordBatch; use arrow_schema::Schema; -use datafusion_common::{Result, Statistics, assert_eq_or_internal_err}; +use datafusion_common::{Result, Statistics}; use datafusion_execution::TaskContext; use crate::execution_plan::SchedulingType; diff --git a/datafusion/physical-plan/src/limit.rs b/datafusion/physical-plan/src/limit.rs index 431e0803dfe4d..c98581c5a01d1 100644 --- a/datafusion/physical-plan/src/limit.rs +++ b/datafusion/physical-plan/src/limit.rs @@ -35,7 +35,7 @@ use crate::{ use arrow::datatypes::SchemaRef; use arrow::record_batch::RecordBatch; -use datafusion_common::{Result, assert_eq_or_internal_err, internal_err}; +use datafusion_common::{Result, assert_eq_or_internal_err}; use datafusion_execution::TaskContext; use datafusion_physical_expr::LexOrdering; diff --git a/datafusion/physical-plan/src/placeholder_row.rs b/datafusion/physical-plan/src/placeholder_row.rs index 5d71058269f49..dbca18224ff2f 100644 --- a/datafusion/physical-plan/src/placeholder_row.rs +++ b/datafusion/physical-plan/src/placeholder_row.rs @@ -23,8 +23,8 @@ use crate::coop::cooperative; use crate::execution_plan::{Boundedness, EmissionType, SchedulingType}; use crate::memory::MemoryStream; use crate::{ - DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, PlanProperties, - SendableRecordBatchStream, Statistics, common, + ChildrenPropertiesHint, DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, + PlanProperties, SendableRecordBatchStream, Statistics, common, }; use arrow::array::{ArrayRef, NullArray, RecordBatch, RecordBatchOptions}; @@ -136,13 +136,21 @@ impl ExecutionPlan for PlaceholderRowExec { vec![] } - fn with_new_children( + fn replace_children( self: Arc, _: Vec>, + _: ChildrenPropertiesHint, ) -> Result> { Ok(self) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children(children, ChildrenPropertiesHint::Recompute) + } + fn execute( &self, partition: usize, diff --git a/datafusion/physical-plan/src/scalar_subquery.rs b/datafusion/physical-plan/src/scalar_subquery.rs index 2e04b5456bfdd..debc3f1e147e0 100644 --- a/datafusion/physical-plan/src/scalar_subquery.rs +++ b/datafusion/physical-plan/src/scalar_subquery.rs @@ -35,7 +35,9 @@ use crate::execution_plan::{CardinalityEffect, ExecutionPlan, PlanProperties}; use crate::joins::utils::{OnceAsync, OnceFut}; use crate::statistics::{ChildStats, StatisticsArgs}; use crate::stream::RecordBatchStreamAdapter; -use crate::{DisplayAs, DisplayFormatType, SendableRecordBatchStream}; +use crate::{ + ChildrenPropertiesHint, DisplayAs, DisplayFormatType, SendableRecordBatchStream, +}; use futures::StreamExt; use futures::TryStreamExt; @@ -162,9 +164,10 @@ impl ExecutionPlan for ScalarSubqueryExec { children } - fn with_new_children( + fn replace_children( self: Arc, mut children: Vec>, + _: ChildrenPropertiesHint, ) -> Result> { // First child is the main input, the rest are subquery plans. let input = children.remove(0); @@ -184,6 +187,13 @@ impl ExecutionPlan for ScalarSubqueryExec { ))) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children(children, ChildrenPropertiesHint::Recompute) + } + fn reset_state(self: Arc) -> Result> { self.results.clear(); Ok(Arc::new(ScalarSubqueryExec { @@ -381,9 +391,10 @@ mod tests { vec![&self.inner] } - fn with_new_children( + fn replace_children( self: Arc, mut children: Vec>, + _: ChildrenPropertiesHint, ) -> Result> { Ok(Arc::new(Self::new( children.remove(0), @@ -391,6 +402,13 @@ mod tests { ))) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children(children, ChildrenPropertiesHint::Recompute) + } + fn execute( &self, partition: usize, diff --git a/datafusion/physical-plan/src/sorts/partitioned_topk.rs b/datafusion/physical-plan/src/sorts/partitioned_topk.rs index 730440a429c68..b42cf7ef789d8 100644 --- a/datafusion/physical-plan/src/sorts/partitioned_topk.rs +++ b/datafusion/physical-plan/src/sorts/partitioned_topk.rs @@ -44,6 +44,7 @@ use datafusion_physical_expr_common::sort_expr::LexOrdering; use futures::StreamExt; use futures::TryStreamExt; +use crate::ChildrenPropertiesHint; use crate::execution_plan::{Boundedness, EmissionType}; use crate::metrics::ExecutionPlanMetricsSet; use crate::topk::{PartitionedTopK, PartitionedTopKRank, build_sort_fields}; @@ -365,9 +366,10 @@ impl ExecutionPlan for PartitionedTopKExec { vec![&self.input] } - fn with_new_children( + fn replace_children( self: Arc, children: Vec>, + _: ChildrenPropertiesHint, ) -> Result> { assert_eq!(children.len(), 1); Ok(Arc::new(PartitionedTopKExec::try_new( @@ -379,6 +381,13 @@ impl ExecutionPlan for PartitionedTopKExec { )?)) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children(children, ChildrenPropertiesHint::Recompute) + } + fn execute( &self, partition: usize, diff --git a/datafusion/physical-plan/src/sorts/sort.rs b/datafusion/physical-plan/src/sorts/sort.rs index 4b30aede7d02a..f64ed954409aa 100644 --- a/datafusion/physical-plan/src/sorts/sort.rs +++ b/datafusion/physical-plan/src/sorts/sort.rs @@ -51,9 +51,10 @@ use crate::stream::{ObservedStream, RecordBatchStreamAdapter}; use crate::topk::TopK; use crate::topk::TopKDynamicFilters; use crate::{ - DisplayAs, DisplayFormatType, Distribution, EmptyRecordBatchStream, ExecutionPlan, - ExecutionPlanProperties, Partitioning, PlanProperties, SendableRecordBatchStream, - Statistics, + ChildrenPropertiesHint, DisplayAs, DisplayFormatType, Distribution, + EmptyRecordBatchStream, ExecutionPlan, ExecutionPlanProperties, Partitioning, + PlanProperties, SendableRecordBatchStream, Statistics, + with_new_children_if_necessary, }; use arrow::array::{RecordBatch, RecordBatchOptions}; @@ -1278,16 +1279,17 @@ impl ExecutionPlan for SortExec { vec![false] } - fn with_new_children( + fn replace_children( self: Arc, children: Vec>, + hint: ChildrenPropertiesHint, ) -> Result> { let mut new_sort = self.cloned(); assert_eq!(children.len(), 1, "SortExec should have exactly one child"); new_sort.input = Arc::clone(&children[0]); - if !has_same_children_properties(self.as_ref(), &children)? { - // Recompute the properties based on the new input since they may have changed + if hint == ChildrenPropertiesHint::Recompute { + // Recompute the properties based on the new input since they may have changed. let (cache, sort_prefix) = Self::compute_properties( &new_sort.input, new_sort.expr.clone(), @@ -1303,12 +1305,24 @@ impl ExecutionPlan for SortExec { Ok(Arc::new(new_sort)) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + match has_same_children_properties(self.as_ref(), &children)? { + true => { + self.replace_children(children, ChildrenPropertiesHint::SameProperties) + } + false => self.replace_children(children, ChildrenPropertiesHint::Recompute), + } + } + fn reset_state(self: Arc) -> Result> { let children = self.children().into_iter().cloned().collect(); - let new_sort = self.with_new_children(children)?; + let new_sort = with_new_children_if_necessary(self, children)?; let mut new_sort = new_sort .downcast_ref::() - .expect("cloned 1 lines above this line, we know the type") + .expect("rebuilt SortExec with new children") .clone(); // Our dynamic filter and execution metrics are the state we need to reset. new_sort.filter = Some(new_sort.create_filter()); @@ -1753,13 +1767,21 @@ mod tests { vec![] } - fn with_new_children( + fn replace_children( self: Arc, _: Vec>, + _: ChildrenPropertiesHint, ) -> Result> { Ok(self) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children(children, ChildrenPropertiesHint::Recompute) + } + fn execute( &self, _partition: usize, diff --git a/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs b/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs index eb9730ddf1f19..dc8263e9a9ef7 100644 --- a/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs +++ b/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs @@ -28,7 +28,7 @@ use crate::statistics::{ChildStats, StatisticsArgs}; use crate::{ ChildrenPropertiesHint, DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, ExecutionPlanProperties, Partitioning, PlanProperties, SendableRecordBatchStream, - Statistics, validate_child_count, + Statistics, validate_child_count, with_new_children_if_necessary, }; use datafusion_common::{Result, assert_eq_or_internal_err, internal_err}; @@ -249,8 +249,7 @@ impl ExecutionPlan for SortPreservingMergeExec { self.input .with_preserve_order(preserve_order) .and_then(|new_input| { - Arc::new(self.clone()) - .with_new_children(vec![new_input]) + with_new_children_if_necessary(Arc::new(self.clone()), vec![new_input]) .ok() }) } @@ -1584,12 +1583,19 @@ mod tests { fn children(&self) -> Vec<&Arc> { vec![] } - fn with_new_children( + fn replace_children( self: Arc, _: Vec>, + _: ChildrenPropertiesHint, ) -> Result> { Ok(self) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children(children, ChildrenPropertiesHint::Recompute) + } fn execute( &self, partition: usize, diff --git a/datafusion/physical-plan/src/streaming.rs b/datafusion/physical-plan/src/streaming.rs index 61a9b9cc6d0de..df6132a0f03f0 100644 --- a/datafusion/physical-plan/src/streaming.rs +++ b/datafusion/physical-plan/src/streaming.rs @@ -30,7 +30,9 @@ use crate::projection::{ ProjectionExec, all_alias_free_columns, new_projections_for_columns, update_ordering, }; use crate::stream::RecordBatchStreamAdapter; -use crate::{ExecutionPlan, Partitioning, SendableRecordBatchStream}; +use crate::{ + ChildrenPropertiesHint, ExecutionPlan, Partitioning, SendableRecordBatchStream, +}; use arrow::datatypes::{Schema, SchemaRef}; use datafusion_common::{Result, internal_err, plan_err}; @@ -271,9 +273,10 @@ impl ExecutionPlan for StreamingTableExec { vec![] } - fn with_new_children( + fn replace_children( self: Arc, children: Vec>, + _: ChildrenPropertiesHint, ) -> Result> { if children.is_empty() { Ok(self) @@ -282,6 +285,13 @@ impl ExecutionPlan for StreamingTableExec { } } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children(children, ChildrenPropertiesHint::Recompute) + } + fn execute( &self, partition: usize, diff --git a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs index 3d13ffe16e8b9..793fc3b9aaee3 100644 --- a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs +++ b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs @@ -129,6 +129,7 @@ use datafusion_functions_aggregate::string_agg::string_agg_udaf; use datafusion_physical_expr::scalar_subquery::ScalarSubqueryExpr; use datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx; use datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx; +use datafusion_physical_plan::ChildrenPropertiesHint; use datafusion_proto::bytes::{ physical_plan_from_bytes_with_proto_converter, physical_plan_to_bytes_with_proto_converter, @@ -326,14 +327,23 @@ impl ExecutionPlan for DowncastDelegatingExec { self.inner.children() } - fn with_new_children( + fn replace_children( self: Arc, children: Vec>, + _: ChildrenPropertiesHint, ) -> Result> { - let inner = Arc::clone(&self.inner).with_new_children(children)?; + let inner = Arc::clone(&self.inner) + .replace_children(children, ChildrenPropertiesHint::Recompute)?; Ok(Arc::new(Self::new(inner))) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children(children, ChildrenPropertiesHint::Recompute) + } + fn downcast_delegate(&self) -> Option<&dyn ExecutionPlan> { Some(self.inner.as_ref()) } @@ -4600,13 +4610,21 @@ impl ExecutionPlan for CustomExecWithExprs { vec![&self.child] } - fn with_new_children( + fn replace_children( self: Arc, - _children: Vec>, + _: Vec>, + _: ChildrenPropertiesHint, ) -> Result> { unreachable!() } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children(children, ChildrenPropertiesHint::Recompute) + } + fn execute( &self, _partition: usize, From ad3df55f7683a611d9ae29fe991455b1d20c1083 Mon Sep 17 00:00:00 2001 From: Justin O'Dwyer Date: Tue, 28 Jul 2026 15:07:25 -0400 Subject: [PATCH 6/9] Migrate remaining implementations to replace_children. --- .../core/tests/custom_sources_cases/mod.rs | 12 +- .../provider_filter_pushdown.rs | 11 +- .../tests/custom_sources_cases/statistics.rs | 14 ++- .../enforce_distribution.rs | 23 +++- .../physical_optimizer/ensure_requirements.rs | 30 +++-- .../physical_optimizer/join_selection.rs | 22 +++- .../physical_optimizer/pushdown_utils.rs | 11 +- .../tests/physical_optimizer/test_utils.rs | 24 +++- datafusion/datasource/src/sink.rs | 14 ++- datafusion/datasource/src/source.rs | 12 +- .../benches/compute_statistics.rs | 16 ++- .../physical-plan/src/aggregates/mod.rs | 14 ++- datafusion/physical-plan/src/analyze.rs | 12 +- datafusion/physical-plan/src/display.rs | 24 +++- datafusion/physical-plan/src/empty.rs | 15 ++- .../physical-plan/src/execution_plan.rs | 103 +++++++++++++----- datafusion/physical-plan/src/explain.rs | 12 +- .../physical-plan/src/joins/hash_join/exec.rs | 23 +++- datafusion/physical-plan/src/memory.rs | 14 ++- .../src/operator_statistics/mod.rs | 24 +++- .../physical-plan/src/recursive_query.rs | 14 ++- .../physical-plan/src/repartition/mod.rs | 10 +- datafusion/physical-plan/src/test.rs | 12 +- datafusion/physical-plan/src/test/exec.rs | 61 ++++++++++- datafusion/physical-plan/src/work_table.rs | 12 +- 25 files changed, 440 insertions(+), 99 deletions(-) diff --git a/datafusion/core/tests/custom_sources_cases/mod.rs b/datafusion/core/tests/custom_sources_cases/mod.rs index c70722cb2f2ff..f10d9c17aca90 100644 --- a/datafusion/core/tests/custom_sources_cases/mod.rs +++ b/datafusion/core/tests/custom_sources_cases/mod.rs @@ -39,10 +39,10 @@ use datafusion_common::cast::as_primitive_array; use datafusion_common::project_schema; use datafusion_common::stats::Precision; use datafusion_physical_expr::EquivalenceProperties; -use datafusion_physical_plan::PlanProperties; use datafusion_physical_plan::StatisticsArgs; use datafusion_physical_plan::execution_plan::{Boundedness, EmissionType}; use datafusion_physical_plan::placeholder_row::PlaceholderRowExec; +use datafusion_physical_plan::{ChildrenPropertiesHint, PlanProperties}; use async_trait::async_trait; use futures::stream::Stream; @@ -164,13 +164,21 @@ impl ExecutionPlan for CustomExecutionPlan { vec![] } - fn with_new_children( + fn replace_children( self: Arc, _: Vec>, + _: ChildrenPropertiesHint, ) -> Result> { Ok(self) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children(children, ChildrenPropertiesHint::Recompute) + } + fn execute( &self, _partition: usize, diff --git a/datafusion/core/tests/custom_sources_cases/provider_filter_pushdown.rs b/datafusion/core/tests/custom_sources_cases/provider_filter_pushdown.rs index 18695accd0f2e..d6a26a9bc9803 100644 --- a/datafusion/core/tests/custom_sources_cases/provider_filter_pushdown.rs +++ b/datafusion/core/tests/custom_sources_cases/provider_filter_pushdown.rs @@ -39,6 +39,7 @@ use datafusion_common::{DataFusionError, internal_err, not_impl_err}; use datafusion_expr::expr::{BinaryExpr, Cast}; use datafusion_functions_aggregate::expr_fn::count; use datafusion_physical_expr::EquivalenceProperties; +use datafusion_physical_plan::ChildrenPropertiesHint; use datafusion_physical_plan::execution_plan::{Boundedness, EmissionType}; use async_trait::async_trait; @@ -116,9 +117,10 @@ impl ExecutionPlan for CustomPlan { vec![] } - fn with_new_children( + fn replace_children( self: Arc, children: Vec>, + _: ChildrenPropertiesHint, ) -> Result> { // CustomPlan has no children if children.is_empty() { @@ -128,6 +130,13 @@ impl ExecutionPlan for CustomPlan { } } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children(children, ChildrenPropertiesHint::Recompute) + } + fn execute( &self, _partition: usize, diff --git a/datafusion/core/tests/custom_sources_cases/statistics.rs b/datafusion/core/tests/custom_sources_cases/statistics.rs index d289b5c348b3c..22e2bece85106 100644 --- a/datafusion/core/tests/custom_sources_cases/statistics.rs +++ b/datafusion/core/tests/custom_sources_cases/statistics.rs @@ -36,7 +36,9 @@ use datafusion_catalog::Session; use datafusion_common::{project_schema, stats::Precision}; use datafusion_physical_expr::EquivalenceProperties; use datafusion_physical_plan::execution_plan::{Boundedness, EmissionType}; -use datafusion_physical_plan::{StatisticsArgs, StatisticsContext}; +use datafusion_physical_plan::{ + ChildrenPropertiesHint, StatisticsArgs, StatisticsContext, +}; use async_trait::async_trait; @@ -159,13 +161,21 @@ impl ExecutionPlan for StatisticsValidation { vec![] } - fn with_new_children( + fn replace_children( self: Arc, _: Vec>, + _: ChildrenPropertiesHint, ) -> Result> { Ok(self) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children(children, ChildrenPropertiesHint::Recompute) + } + fn execute( &self, _partition: usize, diff --git a/datafusion/core/tests/physical_optimizer/enforce_distribution.rs b/datafusion/core/tests/physical_optimizer/enforce_distribution.rs index 189650fe4afca..dc3b3fbd6ccbf 100644 --- a/datafusion/core/tests/physical_optimizer/enforce_distribution.rs +++ b/datafusion/core/tests/physical_optimizer/enforce_distribution.rs @@ -72,7 +72,8 @@ use datafusion_physical_plan::projection::{ProjectionExec, ProjectionExpr}; use datafusion_physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec; use datafusion_physical_plan::union::UnionExec; use datafusion_physical_plan::{ - DisplayAs, DisplayFormatType, ExecutionPlanProperties, PlanProperties, displayable, + ChildrenPropertiesHint, DisplayAs, DisplayFormatType, ExecutionPlanProperties, + PlanProperties, displayable, }; use insta::Settings; @@ -191,9 +192,10 @@ impl ExecutionPlan for SortRequiredExec { vec![Some(OrderingRequirements::from(self.expr.clone()))] } - fn with_new_children( + fn replace_children( self: Arc, mut children: Vec>, + _: ChildrenPropertiesHint, ) -> Result> { assert_eq!(children.len(), 1); let child = children.pop().unwrap(); @@ -203,6 +205,13 @@ impl ExecutionPlan for SortRequiredExec { ))) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children(children, ChildrenPropertiesHint::Recompute) + } + fn execute( &self, _partition: usize, @@ -281,15 +290,23 @@ impl ExecutionPlan for SinglePartitionMaintainsOrderExec { vec![false] } - fn with_new_children( + fn replace_children( self: Arc, mut children: Vec>, + _: ChildrenPropertiesHint, ) -> Result> { assert_eq!(children.len(), 1); let child = children.pop().unwrap(); Ok(Arc::new(Self::new(child))) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children(children, ChildrenPropertiesHint::Recompute) + } + fn execute( &self, _partition: usize, diff --git a/datafusion/core/tests/physical_optimizer/ensure_requirements.rs b/datafusion/core/tests/physical_optimizer/ensure_requirements.rs index 2c6c46c82985a..5ba949f1343bd 100644 --- a/datafusion/core/tests/physical_optimizer/ensure_requirements.rs +++ b/datafusion/core/tests/physical_optimizer/ensure_requirements.rs @@ -39,8 +39,8 @@ use datafusion_physical_plan::limit::GlobalLimitExec; use datafusion_physical_plan::sorts::sort::SortExec; use datafusion_physical_plan::union::UnionExec; use datafusion_physical_plan::{ - DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, Partitioning, - PlanProperties, SendableRecordBatchStream, + ChildrenPropertiesHint, DisplayAs, DisplayFormatType, ExecutionPlan, + ExecutionPlanProperties, Partitioning, PlanProperties, SendableRecordBatchStream, }; use datafusion_physical_optimizer::output_requirements::OutputRequirementExec; @@ -111,12 +111,19 @@ impl ExecutionPlan for MockMultiPartitionExec { fn children(&self) -> Vec<&Arc> { vec![] } - fn with_new_children( + fn replace_children( self: Arc, - _children: Vec>, + _: Vec>, + _: ChildrenPropertiesHint, ) -> Result> { Ok(self) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children(children, ChildrenPropertiesHint::Recompute) + } fn execute( &self, _partition: usize, @@ -991,17 +998,24 @@ impl ExecutionPlan for MockReqExec { fn maintains_input_order(&self) -> Vec { vec![true] } - fn with_new_children( + fn replace_children( self: Arc, - mut c: Vec>, + mut children: Vec>, + _: ChildrenPropertiesHint, ) -> Result> { - assert_eq!(c.len(), 1); + assert_eq!(children.len(), 1); Ok(Arc::new(MockReqExec::new( - c.pop().expect("1 child"), + children.pop().expect("1 child"), self.dist.clone(), self.ord.clone(), ))) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children(children, ChildrenPropertiesHint::Recompute) + } fn execute( &self, _p: usize, diff --git a/datafusion/core/tests/physical_optimizer/join_selection.rs b/datafusion/core/tests/physical_optimizer/join_selection.rs index 3827e6e98b5e6..c850c59b3986b 100644 --- a/datafusion/core/tests/physical_optimizer/join_selection.rs +++ b/datafusion/core/tests/physical_optimizer/join_selection.rs @@ -38,13 +38,13 @@ use datafusion_physical_expr::{EquivalenceProperties, Partitioning, PhysicalExpr use datafusion_physical_expr_common::sort_expr::PhysicalSortExpr; use datafusion_physical_optimizer::PhysicalOptimizerRule; use datafusion_physical_optimizer::join_selection::JoinSelection; -use datafusion_physical_plan::ExecutionPlanProperties; use datafusion_physical_plan::displayable; use datafusion_physical_plan::joins::utils::ColumnIndex; use datafusion_physical_plan::joins::utils::JoinFilter; use datafusion_physical_plan::joins::{HashJoinExec, NestedLoopJoinExec, PartitionMode}; use datafusion_physical_plan::projection::ProjectionExec; use datafusion_physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec; +use datafusion_physical_plan::{ChildrenPropertiesHint, ExecutionPlanProperties}; use datafusion_physical_plan::{ DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, StatisticsArgs, StatisticsContext, @@ -1107,13 +1107,21 @@ impl ExecutionPlan for UnboundedExec { vec![] } - fn with_new_children( + fn replace_children( self: Arc, _: Vec>, + _: ChildrenPropertiesHint, ) -> Result> { Ok(self) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children(children, ChildrenPropertiesHint::Recompute) + } + fn execute( &self, _partition: usize, @@ -1204,13 +1212,21 @@ impl ExecutionPlan for StatisticsExec { vec![] } - fn with_new_children( + fn replace_children( self: Arc, _: Vec>, + _: ChildrenPropertiesHint, ) -> Result> { Ok(self) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children(children, ChildrenPropertiesHint::Recompute) + } + fn execute( &self, _partition: usize, diff --git a/datafusion/core/tests/physical_optimizer/pushdown_utils.rs b/datafusion/core/tests/physical_optimizer/pushdown_utils.rs index 2ffd1899b3c1d..de548b13c5efc 100644 --- a/datafusion/core/tests/physical_optimizer/pushdown_utils.rs +++ b/datafusion/core/tests/physical_optimizer/pushdown_utils.rs @@ -27,6 +27,7 @@ use datafusion_datasource::{ use datafusion_physical_expr::projection::ProjectionExprs; use datafusion_physical_expr_common::physical_expr::fmt_sql; use datafusion_physical_optimizer::PhysicalOptimizerRule; +use datafusion_physical_plan::ChildrenPropertiesHint; use datafusion_physical_plan::filter::batch_filter; use datafusion_physical_plan::filter_pushdown::{FilterPushdownPhase, PushedDown}; use datafusion_physical_plan::{ @@ -473,9 +474,10 @@ impl ExecutionPlan for TestNode { vec![&self.input] } - fn with_new_children( + fn replace_children( self: Arc, children: Vec>, + _: ChildrenPropertiesHint, ) -> Result> { assert!(children.len() == 1); Ok(Arc::new(TestNode::new( @@ -485,6 +487,13 @@ impl ExecutionPlan for TestNode { ))) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children(children, ChildrenPropertiesHint::Recompute) + } + fn execute( &self, _partition: usize, diff --git a/datafusion/core/tests/physical_optimizer/test_utils.rs b/datafusion/core/tests/physical_optimizer/test_utils.rs index 74230b24e2ab5..026feffdafe3b 100644 --- a/datafusion/core/tests/physical_optimizer/test_utils.rs +++ b/datafusion/core/tests/physical_optimizer/test_utils.rs @@ -68,8 +68,8 @@ use datafusion_physical_plan::tree_node::PlanContext; use datafusion_physical_plan::union::UnionExec; use datafusion_physical_plan::windows::{BoundedWindowAggExec, create_window_expr}; use datafusion_physical_plan::{ - DisplayAs, DisplayFormatType, ExecutionPlan, InputOrderMode, Partitioning, - PlanProperties, SortOrderPushdownResult, StatisticsArgs, displayable, + ChildrenPropertiesHint, DisplayAs, DisplayFormatType, ExecutionPlan, InputOrderMode, + Partitioning, PlanProperties, SortOrderPushdownResult, StatisticsArgs, displayable, }; /// Create a non sorted parquet exec @@ -508,9 +508,10 @@ impl ExecutionPlan for RequirementsTestExec { vec![&self.input] } - fn with_new_children( + fn replace_children( self: Arc, children: Vec>, + _: ChildrenPropertiesHint, ) -> Result> { assert_eq!(children.len(), 1); Ok(RequirementsTestExec::new(Arc::clone(&children[0])) @@ -519,6 +520,13 @@ impl ExecutionPlan for RequirementsTestExec { .into_arc()) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children(children, ChildrenPropertiesHint::Recompute) + } + fn execute( &self, _partition: usize, @@ -999,9 +1007,10 @@ impl ExecutionPlan for TestScan { vec![] } - fn with_new_children( + fn replace_children( self: Arc, children: Vec>, + _: ChildrenPropertiesHint, ) -> Result> { if children.is_empty() { Ok(self) @@ -1010,6 +1019,13 @@ impl ExecutionPlan for TestScan { } } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children(children, ChildrenPropertiesHint::Recompute) + } + fn execute( &self, _partition: usize, diff --git a/datafusion/datasource/src/sink.rs b/datafusion/datasource/src/sink.rs index 18ebe80773e8a..bc5ac39fdf752 100644 --- a/datafusion/datasource/src/sink.rs +++ b/datafusion/datasource/src/sink.rs @@ -31,8 +31,8 @@ use datafusion_physical_expr_common::sort_expr::{LexRequirement, OrderingRequire use datafusion_physical_plan::metrics::MetricsSet; use datafusion_physical_plan::stream::RecordBatchStreamAdapter; use datafusion_physical_plan::{ - DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, - InputDistributionRequirements, Partitioning, PlanProperties, + ChildrenPropertiesHint, DisplayAs, DisplayFormatType, ExecutionPlan, + ExecutionPlanProperties, InputDistributionRequirements, Partitioning, PlanProperties, SendableRecordBatchStream, execute_input_stream, }; @@ -220,9 +220,10 @@ impl ExecutionPlan for DataSinkExec { vec![&self.input] } - fn with_new_children( + fn replace_children( self: Arc, children: Vec>, + _: ChildrenPropertiesHint, ) -> Result> { Ok(Arc::new(Self::new( Arc::clone(&children[0]), @@ -231,6 +232,13 @@ impl ExecutionPlan for DataSinkExec { ))) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children(children, ChildrenPropertiesHint::Recompute) + } + /// Execute the plan and return a stream of `RecordBatch`es for /// the specified partition. fn execute( diff --git a/datafusion/datasource/src/source.rs b/datafusion/datasource/src/source.rs index c280470bb0d0b..1c499f18931df 100644 --- a/datafusion/datasource/src/source.rs +++ b/datafusion/datasource/src/source.rs @@ -33,7 +33,7 @@ use datafusion_physical_plan::metrics::{ use datafusion_physical_plan::projection::ProjectionExec; use datafusion_physical_plan::stream::BatchSplitStream; use datafusion_physical_plan::{ - DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, + ChildrenPropertiesHint, DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, }; use itertools::Itertools; @@ -359,13 +359,21 @@ impl ExecutionPlan for DataSourceExec { Vec::new() } - fn with_new_children( + fn replace_children( self: Arc, _: Vec>, + _: ChildrenPropertiesHint, ) -> Result> { Ok(self) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children(children, ChildrenPropertiesHint::Recompute) + } + /// Implementation of [`ExecutionPlan::repartitioned`] which relies upon the inner [`DataSource::repartitioned`]. /// /// If the data source does not support changing its partitioning, returns `Ok(None)` (the default). Refer diff --git a/datafusion/physical-plan/benches/compute_statistics.rs b/datafusion/physical-plan/benches/compute_statistics.rs index 56a518c95292e..67c736515187e 100644 --- a/datafusion/physical-plan/benches/compute_statistics.rs +++ b/datafusion/physical-plan/benches/compute_statistics.rs @@ -46,8 +46,8 @@ use datafusion_physical_plan::filter::FilterExec; use datafusion_physical_plan::joins::CrossJoinExec; use datafusion_physical_plan::statistics::StatisticsArgs; use datafusion_physical_plan::{ - DisplayAs, DisplayFormatType, Partitioning, SendableRecordBatchStream, - StatisticsContext, + ChildrenPropertiesHint, DisplayAs, DisplayFormatType, Partitioning, + SendableRecordBatchStream, StatisticsContext, }; /// Minimal leaf node for benchmarking @@ -97,13 +97,21 @@ impl ExecutionPlan for BenchLeaf { vec![] } - fn with_new_children( + fn replace_children( self: Arc, - _children: Vec>, + _: Vec>, + _: ChildrenPropertiesHint, ) -> Result> { Ok(self) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children(children, ChildrenPropertiesHint::Recompute) + } + fn execute( &self, _partition: usize, diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index 3197d822bbfdd..c2a66b31f5bae 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -3615,13 +3615,21 @@ mod tests { vec![] } - fn with_new_children( + fn replace_children( self: Arc, _: Vec>, + _: ChildrenPropertiesHint, ) -> Result> { internal_err!("Children cannot be replaced in {self:?}") } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children(children, ChildrenPropertiesHint::Recompute) + } + fn execute( &self, _partition: usize, @@ -4958,8 +4966,8 @@ mod tests { Arc::clone(&blocking_exec) as Arc, schema, )?); - let new_agg = - Arc::clone(&aggregate_exec).with_new_children(vec![blocking_exec])?; + let new_agg = Arc::clone(&aggregate_exec) + .replace_children(vec![blocking_exec], ChildrenPropertiesHint::Recompute)?; assert_eq!(new_agg.schema(), aggregate_exec.schema()); Ok(()) } diff --git a/datafusion/physical-plan/src/analyze.rs b/datafusion/physical-plan/src/analyze.rs index 31e0a27410ff9..bd94b006cb3ca 100644 --- a/datafusion/physical-plan/src/analyze.rs +++ b/datafusion/physical-plan/src/analyze.rs @@ -27,7 +27,7 @@ use super::{ use crate::display::DisplayableExecutionPlan; use crate::execution_plan::EvaluationType; use crate::metrics::{MetricCategory, MetricType}; -use crate::{DisplayFormatType, ExecutionPlan, Partitioning}; +use crate::{ChildrenPropertiesHint, DisplayFormatType, ExecutionPlan, Partitioning}; use arrow::{array::StringBuilder, datatypes::SchemaRef, record_batch::RecordBatch}; use datafusion_common::format::ExplainFormat; @@ -219,9 +219,10 @@ impl ExecutionPlan for AnalyzeExec { ]) } - fn with_new_children( + fn replace_children( self: Arc, mut children: Vec>, + _: ChildrenPropertiesHint, ) -> Result> { Ok(Arc::new( AnalyzeExec::builder( @@ -237,6 +238,13 @@ impl ExecutionPlan for AnalyzeExec { )) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children(children, ChildrenPropertiesHint::Recompute) + } + fn execute( &self, partition: usize, diff --git a/datafusion/physical-plan/src/display.rs b/datafusion/physical-plan/src/display.rs index 6a4d09057bec9..91e79150f1796 100644 --- a/datafusion/physical-plan/src/display.rs +++ b/datafusion/physical-plan/src/display.rs @@ -1455,7 +1455,7 @@ mod tests { use datafusion_execution::{SendableRecordBatchStream, TaskContext}; use crate::statistics::StatisticsArgs; - use crate::{DisplayAs, ExecutionPlan, PlanProperties}; + use crate::{ChildrenPropertiesHint, DisplayAs, ExecutionPlan, PlanProperties}; use super::DisplayableExecutionPlan; @@ -1489,13 +1489,21 @@ mod tests { vec![] } - fn with_new_children( + fn replace_children( self: Arc, _: Vec>, + _: ChildrenPropertiesHint, ) -> Result> { unimplemented!() } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children(children, ChildrenPropertiesHint::Recompute) + } + fn execute( &self, _: usize, @@ -1570,13 +1578,14 @@ mod tests { use insta::assert_snapshot; use super::super::DisplayableExecutionPlan; + use crate::ChildrenPropertiesHint; use crate::empty::EmptyExec; use crate::filter::FilterExec; use crate::projection::ProjectionExec; use datafusion_physical_expr::expressions::{binary, col, lit}; use datafusion_physical_expr::{Partitioning, PhysicalExpr}; - fn sample_plan() -> Arc { + fn sample_plan() -> Arc { let schema = Arc::new(Schema::new(vec![ Field::new("a", DataType::Int32, false), Field::new("b", DataType::Int32, false), @@ -1653,12 +1662,19 @@ mod tests { fn children(&self) -> Vec<&Arc> { vec![&self.inner] } - fn with_new_children( + fn replace_children( self: Arc, _: Vec>, + _: ChildrenPropertiesHint, ) -> Result> { unimplemented!() } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children(children, ChildrenPropertiesHint::Recompute) + } fn execute( &self, _: usize, diff --git a/datafusion/physical-plan/src/empty.rs b/datafusion/physical-plan/src/empty.rs index 3bd38bf238dc1..7308c8efdd94b 100644 --- a/datafusion/physical-plan/src/empty.rs +++ b/datafusion/physical-plan/src/empty.rs @@ -20,7 +20,10 @@ use std::sync::Arc; use crate::memory::MemoryStream; -use crate::{DisplayAs, PlanProperties, SendableRecordBatchStream, Statistics}; +use crate::{ + ChildrenPropertiesHint, DisplayAs, PlanProperties, SendableRecordBatchStream, + Statistics, +}; use crate::{ DisplayFormatType, ExecutionPlan, Partitioning, execution_plan::{Boundedness, EmissionType}, @@ -119,13 +122,21 @@ impl ExecutionPlan for EmptyExec { vec![] } - fn with_new_children( + fn replace_children( self: Arc, _: Vec>, + _: ChildrenPropertiesHint, ) -> Result> { Ok(self) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children(children, ChildrenPropertiesHint::Recompute) + } + fn execute( &self, partition: usize, diff --git a/datafusion/physical-plan/src/execution_plan.rs b/datafusion/physical-plan/src/execution_plan.rs index 372625f3c23da..d6b54546642cb 100644 --- a/datafusion/physical-plan/src/execution_plan.rs +++ b/datafusion/physical-plan/src/execution_plan.rs @@ -1802,13 +1802,21 @@ mod tests { vec![] } - fn with_new_children( + fn replace_children( self: Arc, _: Vec>, + _: ChildrenPropertiesHint, ) -> Result> { unimplemented!() } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children(children, ChildrenPropertiesHint::Recompute) + } + fn execute( &self, _partition: usize, @@ -1865,13 +1873,21 @@ mod tests { vec![] } - fn with_new_children( + fn replace_children( self: Arc, _: Vec>, + _: ChildrenPropertiesHint, ) -> Result> { unimplemented!() } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children(children, ChildrenPropertiesHint::Recompute) + } + fn execute( &self, _partition: usize, @@ -1915,13 +1931,21 @@ mod tests { vec![] } - fn with_new_children( + fn replace_children( self: Arc, _: Vec>, + _: ChildrenPropertiesHint, ) -> Result> { unimplemented!() } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children(children, ChildrenPropertiesHint::Recompute) + } + fn downcast_delegate(&self) -> Option<&dyn ExecutionPlan> { Some(self.0.as_ref()) } @@ -1974,12 +1998,19 @@ mod tests { fn children(&self) -> Vec<&Arc> { vec![] } - fn with_new_children( + fn replace_children( self: Arc, _: Vec>, + _: ChildrenPropertiesHint, ) -> Result> { Ok(self) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children(children, ChildrenPropertiesHint::Recompute) + } fn execute( &self, _partition: usize, @@ -2037,38 +2068,52 @@ mod tests { fn children(&self) -> Vec<&Arc> { vec![&self.input] } - fn with_new_children( + fn replace_children( self: Arc, mut children: Vec>, + hint: ChildrenPropertiesHint, ) -> Result> { - self.recompute_calls - .fetch_add(1, std::sync::atomic::Ordering::SeqCst); - // Full recompute: allocate a fresh `PlanProperties` Arc so this - // path is observable via `Arc::ptr_eq` on properties. - let new_input = children.swap_remove(0); - let cache = Arc::new(PlanProperties::new( - EquivalenceProperties::new(Arc::new(Schema::empty())), - Partitioning::UnknownPartitioning(1), - EmissionType::Final, - Boundedness::Bounded, - )); - Ok(Arc::new(Self { - input: new_input, - cache, - recompute_calls: Arc::clone(&self.recompute_calls), - fast_path_calls: Arc::clone(&self.fast_path_calls), - })) + match hint { + ChildrenPropertiesHint::SameProperties => { + self.fast_path_calls + .fetch_add(1, std::sync::atomic::Ordering::SeqCst); + Ok(Arc::new(Self { + input: children.swap_remove(0), + ..Self::clone(&*self) + })) + } + ChildrenPropertiesHint::Recompute => { + self.recompute_calls + .fetch_add(1, std::sync::atomic::Ordering::SeqCst); + // Full recompute: allocate a fresh `PlanProperties` Arc so this + // path is observable via `Arc::ptr_eq` on properties. + let new_input = children.swap_remove(0); + let cache = Arc::new(PlanProperties::new( + EquivalenceProperties::new(Arc::new(Schema::empty())), + Partitioning::UnknownPartitioning(1), + EmissionType::Final, + Boundedness::Bounded, + )); + Ok(Arc::new(Self { + input: new_input, + cache, + recompute_calls: Arc::clone(&self.recompute_calls), + fast_path_calls: Arc::clone(&self.fast_path_calls), + })) + } + } + } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children(children, ChildrenPropertiesHint::Recompute) } fn with_new_children_and_same_properties( self: Arc, - mut children: Vec>, + children: Vec>, ) -> Result> { - self.fast_path_calls - .fetch_add(1, std::sync::atomic::Ordering::SeqCst); - Ok(Arc::new(Self { - input: children.swap_remove(0), - ..Self::clone(&*self) - })) + self.replace_children(children, ChildrenPropertiesHint::SameProperties) } fn execute( &self, diff --git a/datafusion/physical-plan/src/explain.rs b/datafusion/physical-plan/src/explain.rs index a270a003eba17..9a5db88d5f203 100644 --- a/datafusion/physical-plan/src/explain.rs +++ b/datafusion/physical-plan/src/explain.rs @@ -22,7 +22,7 @@ use std::sync::Arc; use super::{DisplayAs, PlanProperties, SendableRecordBatchStream}; use crate::execution_plan::{Boundedness, EmissionType}; use crate::stream::RecordBatchStreamAdapter; -use crate::{DisplayFormatType, ExecutionPlan, Partitioning}; +use crate::{ChildrenPropertiesHint, DisplayFormatType, ExecutionPlan, Partitioning}; use arrow::{array::StringBuilder, datatypes::SchemaRef, record_batch::RecordBatch}; use datafusion_common::display::StringifiedPlan; @@ -116,13 +116,21 @@ impl ExecutionPlan for ExplainExec { vec![] } - fn with_new_children( + fn replace_children( self: Arc, _: Vec>, + _: ChildrenPropertiesHint, ) -> Result> { Ok(self) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children(children, ChildrenPropertiesHint::Recompute) + } + fn execute( &self, partition: usize, diff --git a/datafusion/physical-plan/src/joins/hash_join/exec.rs b/datafusion/physical-plan/src/joins/hash_join/exec.rs index ccdb050d168e0..190235763beea 100644 --- a/datafusion/physical-plan/src/joins/hash_join/exec.rs +++ b/datafusion/physical-plan/src/joins/hash_join/exec.rs @@ -22,7 +22,6 @@ use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; use std::sync::{Arc, OnceLock}; use std::vec; -use crate::ExecutionPlanProperties; use crate::execution_plan::{ EmissionType, boundedness_from_children, has_same_children_properties, stub_properties, @@ -53,6 +52,7 @@ use crate::projection::{ }; use crate::repartition::REPARTITION_RANDOM_STATE; use crate::statistics::{ChildStats, StatisticsArgs}; +use crate::{ChildrenPropertiesHint, ExecutionPlanProperties}; use crate::{ DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, InputDistributionRequirements, Partitioning, PlanProperties, @@ -1323,13 +1323,21 @@ impl ExecutionPlan for HashJoinExec { /// This method is called during query optimization when the optimizer creates new /// plan nodes. Importantly, it creates a fresh bounds_accumulator via `try_new` /// rather than cloning the existing one because partitioning may have changed. - fn with_new_children( + fn replace_children( self: Arc, children: Vec>, + _: ChildrenPropertiesHint, ) -> Result> { self.builder().with_new_children(children)?.build_exec() } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children(children, ChildrenPropertiesHint::Recompute) + } + fn reset_state(self: Arc) -> Result> { self.builder().reset_state().build_exec() } @@ -2380,6 +2388,7 @@ mod tests { Ok((left_schema, right_schema, on)) } + use crate::ChildrenPropertiesHint; use crate::coalesce_partitions::CoalescePartitionsExec; use crate::execution_plan::Boundedness; use crate::joins::hash_join::stream::lookup_join_hashmap; @@ -2449,13 +2458,21 @@ mod tests { vec![] } - fn with_new_children( + fn replace_children( self: Arc, _: Vec>, + _: ChildrenPropertiesHint, ) -> Result> { Ok(self) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children(children, ChildrenPropertiesHint::Recompute) + } + fn execute( &self, _partition: usize, diff --git a/datafusion/physical-plan/src/memory.rs b/datafusion/physical-plan/src/memory.rs index ad54905f474aa..c11f115e735fc 100644 --- a/datafusion/physical-plan/src/memory.rs +++ b/datafusion/physical-plan/src/memory.rs @@ -26,8 +26,8 @@ use crate::coop::cooperative; use crate::execution_plan::{Boundedness, EmissionType, SchedulingType}; use crate::metrics::{BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet}; use crate::{ - DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, PlanProperties, - RecordBatchStream, SendableRecordBatchStream, + ChildrenPropertiesHint, DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, + PlanProperties, RecordBatchStream, SendableRecordBatchStream, }; use arrow::array::RecordBatch; @@ -311,9 +311,10 @@ impl ExecutionPlan for LazyMemoryExec { vec![] } - fn with_new_children( + fn replace_children( self: Arc, children: Vec>, + _: ChildrenPropertiesHint, ) -> Result> { assert_or_internal_err!( children.is_empty(), @@ -322,6 +323,13 @@ impl ExecutionPlan for LazyMemoryExec { Ok(self) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children(children, ChildrenPropertiesHint::Recompute) + } + fn execute( &self, partition: usize, diff --git a/datafusion/physical-plan/src/operator_statistics/mod.rs b/datafusion/physical-plan/src/operator_statistics/mod.rs index 142768fcf49d2..8b7f1ce064c80 100644 --- a/datafusion/physical-plan/src/operator_statistics/mod.rs +++ b/datafusion/physical-plan/src/operator_statistics/mod.rs @@ -1027,7 +1027,7 @@ mod tests { use crate::filter::FilterExec; use crate::projection::ProjectionExec; use crate::statistics::StatisticsArgs; - use crate::{DisplayAs, DisplayFormatType, PlanProperties}; + use crate::{ChildrenPropertiesHint, DisplayAs, DisplayFormatType, PlanProperties}; use arrow::datatypes::{DataType, Field, Schema}; use datafusion_common::stats::Precision; use datafusion_common::{ColumnStatistics, ScalarValue}; @@ -1106,13 +1106,21 @@ mod tests { vec![] } - fn with_new_children( + fn replace_children( self: Arc, - _children: Vec>, + _: Vec>, + _: ChildrenPropertiesHint, ) -> Result> { Ok(self) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children(children, ChildrenPropertiesHint::Recompute) + } + fn properties(&self) -> &Arc { &self.cache } @@ -1206,15 +1214,23 @@ mod tests { vec![&self.input] } - fn with_new_children( + fn replace_children( self: Arc, children: Vec>, + _: ChildrenPropertiesHint, ) -> Result> { Ok(Arc::new(CustomExec { input: Arc::clone(&children[0]), })) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children(children, ChildrenPropertiesHint::Recompute) + } + fn properties(&self) -> &Arc { self.input.properties() } diff --git a/datafusion/physical-plan/src/recursive_query.rs b/datafusion/physical-plan/src/recursive_query.rs index 00df227cb87db..4798a88d9edd1 100644 --- a/datafusion/physical-plan/src/recursive_query.rs +++ b/datafusion/physical-plan/src/recursive_query.rs @@ -30,8 +30,8 @@ use crate::metrics::{ BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet, RecordOutput, }; use crate::{ - DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, RecordBatchStream, - SendableRecordBatchStream, + ChildrenPropertiesHint, DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, + RecordBatchStream, SendableRecordBatchStream, }; use arrow::array::{BooleanArray, BooleanBuilder}; use arrow::compute::filter_record_batch; @@ -174,9 +174,10 @@ impl ExecutionPlan for RecursiveQueryExec { ]) } - fn with_new_children( + fn replace_children( self: Arc, children: Vec>, + _: ChildrenPropertiesHint, ) -> Result> { RecursiveQueryExec::try_new( self.name.clone(), @@ -188,6 +189,13 @@ impl ExecutionPlan for RecursiveQueryExec { .map(|e| Arc::new(e) as _) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children(children, ChildrenPropertiesHint::Recompute) + } + fn execute( &self, partition: usize, diff --git a/datafusion/physical-plan/src/repartition/mod.rs b/datafusion/physical-plan/src/repartition/mod.rs index b7ce64e9ca9a2..ea1b6be9f6678 100644 --- a/datafusion/physical-plan/src/repartition/mod.rs +++ b/datafusion/physical-plan/src/repartition/mod.rs @@ -3058,13 +3058,21 @@ mod tests { vec![] } - fn with_new_children( + fn replace_children( self: Arc, _: Vec>, + _: ChildrenPropertiesHint, ) -> Result> { Ok(self) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children(children, ChildrenPropertiesHint::Recompute) + } + fn execute( &self, _partition: usize, diff --git a/datafusion/physical-plan/src/test.rs b/datafusion/physical-plan/src/test.rs index e8c775a786578..226c4fc649465 100644 --- a/datafusion/physical-plan/src/test.rs +++ b/datafusion/physical-plan/src/test.rs @@ -24,7 +24,6 @@ use std::pin::Pin; use std::sync::Arc; use std::task::Context; -use crate::ExecutionPlan; use crate::common; use crate::execution_plan::{Boundedness, EmissionType}; use crate::memory::MemoryStream; @@ -32,6 +31,7 @@ use crate::metrics::MetricsSet; use crate::statistics::StatisticsArgs; use crate::stream::RecordBatchStreamAdapter; use crate::streaming::PartitionStream; +use crate::{ChildrenPropertiesHint, ExecutionPlan}; use crate::{DisplayAs, DisplayFormatType, PlanProperties}; use arrow::array::{Array, ArrayRef, Int32Array, RecordBatch}; @@ -138,13 +138,21 @@ impl ExecutionPlan for TestMemoryExec { Vec::new() } - fn with_new_children( + fn replace_children( self: Arc, _: Vec>, + _: ChildrenPropertiesHint, ) -> Result> { Ok(self) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children(children, ChildrenPropertiesHint::Recompute) + } + fn repartitioned( &self, _target_partitions: usize, diff --git a/datafusion/physical-plan/src/test/exec.rs b/datafusion/physical-plan/src/test/exec.rs index b92008c6b219b..ec32f837ab92b 100644 --- a/datafusion/physical-plan/src/test/exec.rs +++ b/datafusion/physical-plan/src/test/exec.rs @@ -17,6 +17,7 @@ //! Simple iterator over batches for use in testing +use crate::ChildrenPropertiesHint; use crate::{ DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, PlanProperties, RecordBatchStream, SendableRecordBatchStream, Statistics, common, @@ -195,13 +196,21 @@ impl ExecutionPlan for MockExec { vec![] } - fn with_new_children( + fn replace_children( self: Arc, _: Vec>, + _: ChildrenPropertiesHint, ) -> Result> { unimplemented!() } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children(children, ChildrenPropertiesHint::Recompute) + } + /// Returns a stream which yields data fn execute( &self, @@ -425,13 +434,21 @@ impl ExecutionPlan for BarrierExec { unimplemented!() } - fn with_new_children( + fn replace_children( self: Arc, _: Vec>, + _: ChildrenPropertiesHint, ) -> Result> { unimplemented!() } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children(children, ChildrenPropertiesHint::Recompute) + } + /// Returns a stream which yields data fn execute( &self, @@ -561,13 +578,21 @@ impl ExecutionPlan for ErrorExec { unimplemented!() } - fn with_new_children( + fn replace_children( self: Arc, _: Vec>, + _: ChildrenPropertiesHint, ) -> Result> { unimplemented!() } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children(children, ChildrenPropertiesHint::Recompute) + } + /// Returns a stream which yields data fn execute( &self, @@ -647,13 +672,21 @@ impl ExecutionPlan for StatisticsExec { vec![] } - fn with_new_children( + fn replace_children( self: Arc, _: Vec>, + _: ChildrenPropertiesHint, ) -> Result> { Ok(self) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children(children, ChildrenPropertiesHint::Recompute) + } + fn execute( &self, _partition: usize, @@ -751,13 +784,21 @@ impl ExecutionPlan for BlockingExec { vec![] } - fn with_new_children( + fn replace_children( self: Arc, _: Vec>, + _: ChildrenPropertiesHint, ) -> Result> { internal_err!("Children cannot be replaced in {self:?}") } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children(children, ChildrenPropertiesHint::Recompute) + } + fn execute( &self, _partition: usize, @@ -893,13 +934,21 @@ impl ExecutionPlan for PanicExec { vec![] } - fn with_new_children( + fn replace_children( self: Arc, _: Vec>, + _: ChildrenPropertiesHint, ) -> Result> { internal_err!("Children cannot be replaced in {:?}", self) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children(children, ChildrenPropertiesHint::Recompute) + } + fn execute( &self, partition: usize, diff --git a/datafusion/physical-plan/src/work_table.rs b/datafusion/physical-plan/src/work_table.rs index c92face1e5404..f8081402d6ba9 100644 --- a/datafusion/physical-plan/src/work_table.rs +++ b/datafusion/physical-plan/src/work_table.rs @@ -25,7 +25,7 @@ use crate::execution_plan::{Boundedness, EmissionType, SchedulingType}; use crate::memory::MemoryStream; use crate::metrics::{ExecutionPlanMetricsSet, MetricsSet}; use crate::{ - DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, + ChildrenPropertiesHint, DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, SendableRecordBatchStream, Statistics, }; @@ -186,13 +186,21 @@ impl ExecutionPlan for WorkTableExec { vec![] } - fn with_new_children( + fn replace_children( self: Arc, _: Vec>, + _: ChildrenPropertiesHint, ) -> Result> { Ok(Arc::clone(&self) as Arc) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children(children, ChildrenPropertiesHint::Recompute) + } + /// Stream the batches that were written to the work table. fn execute( &self, From f3d20d0a5cf8e75550d2cd28b45cbd75d93e29de Mon Sep 17 00:00:00 2001 From: Justin O'Dwyer Date: Wed, 29 Jul 2026 13:17:03 -0400 Subject: [PATCH 7/9] Migrate usages away from with_new_children. --- datafusion/catalog/src/memory/table.rs | 2 +- datafusion/core/src/physical_planner.rs | 42 +++++++++++++------ datafusion/ffi/src/execution_plan.rs | 16 ++++--- datafusion/ffi/tests/ffi_execution_plan.rs | 8 ++-- .../enforce_distribution.rs | 9 ++-- .../physical-optimizer/src/limit_pushdown.rs | 6 ++- .../src/topk_repartition.rs | 6 ++- .../physical-optimizer/src/window_topn.rs | 10 ++--- datafusion/physical-plan/src/async_func.rs | 1 + datafusion/physical-plan/src/buffer.rs | 9 ++-- .../physical-plan/src/coalesce_batches.rs | 9 ++-- .../physical-plan/src/coalesce_partitions.rs | 5 +-- datafusion/physical-plan/src/coop.rs | 15 ++++--- .../physical-plan/src/execution_plan.rs | 2 +- datafusion/physical-plan/src/filter.rs | 7 ++-- datafusion/physical-plan/src/projection.rs | 11 ++--- 16 files changed, 99 insertions(+), 59 deletions(-) diff --git a/datafusion/catalog/src/memory/table.rs b/datafusion/catalog/src/memory/table.rs index eb0baa6e09925..e3bde45f858df 100644 --- a/datafusion/catalog/src/memory/table.rs +++ b/datafusion/catalog/src/memory/table.rs @@ -574,7 +574,7 @@ impl ExecutionPlan for DmlResultExec { fn replace_children( self: Arc, _: Vec>, - _: datafusion_physical_plan::ChildrenPropertiesHint, + _: ChildrenPropertiesHint, ) -> Result> { Ok(self) } diff --git a/datafusion/core/src/physical_planner.rs b/datafusion/core/src/physical_planner.rs index 49431d8c62a01..52806d4dbc9a8 100644 --- a/datafusion/core/src/physical_planner.rs +++ b/datafusion/core/src/physical_planner.rs @@ -4981,10 +4981,16 @@ digraph { // ok plan let ok_node: Arc = Arc::new(OkExtensionNode(vec![])); let child = Arc::clone(&ok_node); - let ok_plan = Arc::clone(&ok_node).with_new_children(vec![ - Arc::clone(&child).with_new_children(vec![Arc::clone(&child)])?, - Arc::clone(&child), - ])?; + let ok_plan = Arc::clone(&ok_node).replace_children( + vec![ + Arc::clone(&child).replace_children( + vec![Arc::clone(&child)], + ChildrenPropertiesHint::Recompute, + )?, + Arc::clone(&child), + ], + ChildrenPropertiesHint::Recompute, + )?; // Test: check should pass with same schema let equal_schema = ok_plan.schema(); @@ -5016,10 +5022,16 @@ digraph { // Test: should fail when descendent extension node fails let failing_node: Arc = Arc::new(InvariantFailsExtensionNode); - let invalid_plan = ok_node.with_new_children(vec![ - Arc::clone(&child).with_new_children(vec![Arc::clone(&failing_node)])?, - Arc::clone(&child), - ])?; + let invalid_plan = ok_node.replace_children( + vec![ + Arc::clone(&child).replace_children( + vec![Arc::clone(&failing_node)], + ChildrenPropertiesHint::Recompute, + )?, + Arc::clone(&child), + ], + ChildrenPropertiesHint::Recompute, + )?; let result = OptimizationInvariantChecker::new(&rule) .check(&invalid_plan, &ok_plan.schema()); if cfg!(debug_assertions) { @@ -5105,10 +5117,16 @@ digraph { let failing_node: Arc = Arc::new(ExecutableInvariantFails); let ok_node: Arc = Arc::new(OkExtensionNode(vec![])); let child = Arc::clone(&ok_node); - let plan = ok_node.with_new_children(vec![ - Arc::clone(&child).with_new_children(vec![Arc::clone(&failing_node)])?, - Arc::clone(&child), - ])?; + let plan = ok_node.replace_children( + vec![ + Arc::clone(&child).replace_children( + vec![Arc::clone(&failing_node)], + ChildrenPropertiesHint::Recompute, + )?, + Arc::clone(&child), + ], + ChildrenPropertiesHint::Recompute, + )?; let expected_err = InvariantChecker(InvariantLevel::Executable) .check(&plan) .unwrap_err(); diff --git a/datafusion/ffi/src/execution_plan.rs b/datafusion/ffi/src/execution_plan.rs index ba071944aba8f..cdc62541be7ba 100644 --- a/datafusion/ffi/src/execution_plan.rs +++ b/datafusion/ffi/src/execution_plan.rs @@ -151,7 +151,9 @@ unsafe extern "C" fn with_new_children_fn_wrapper( .collect(); let children = sresult_return!(children); - let new_plan = sresult_return!(inner_plan.with_new_children(children)); + let new_plan = sresult_return!( + inner_plan.replace_children(children, ChildrenPropertiesHint::Recompute) + ); FFI_Result::Ok(FFI_ExecutionPlan::new(new_plan, runtime)) } @@ -268,7 +270,7 @@ fn pass_runtime_to_children( // If the parent is foreign and the child is local to this library, then when // we called `children()` above we will get something other than a // `ForeignExecutionPlan`. In this case wrap the plan in a `ForeignExecutionPlan` - // because when we call `with_new_children` below it will extract the + // because when we call `replace_children` below it will extract the // FFI plan that does contain the runtime. if plan_is_foreign && !child.is::() { updated_children = true; @@ -281,7 +283,9 @@ fn pass_runtime_to_children( }) .collect::>>()?; if updated_children { - Arc::clone(plan).with_new_children(children).map(Some) + Arc::clone(plan) + .replace_children(children, ChildrenPropertiesHint::Recompute) + .map(Some) } else { Ok(None) } @@ -636,7 +640,8 @@ pub mod tests { assert_eq!(parent_foreign.children().len(), 0); assert_eq!(child_foreign.children().len(), 0); - let parent_foreign = parent_foreign.with_new_children(vec![child_foreign])?; + let parent_foreign = parent_foreign + .replace_children(vec![child_foreign], ChildrenPropertiesHint::Recompute)?; assert_eq!(parent_foreign.children().len(), 1); // Version 2: Adding child to the local plan @@ -646,7 +651,8 @@ pub mod tests { let child_foreign = >::try_from(&child_local)?; let parent_plan = Arc::new(EmptyExec::new(Arc::clone(&schema))); - let parent_plan = parent_plan.with_new_children(vec![child_foreign])?; + let parent_plan = parent_plan + .replace_children(vec![child_foreign], ChildrenPropertiesHint::Recompute)?; let mut parent_local = FFI_ExecutionPlan::new(parent_plan, None); parent_local.library_marker_id = crate::mock_foreign_marker_id; let parent_foreign = >::try_from(&parent_local)?; diff --git a/datafusion/ffi/tests/ffi_execution_plan.rs b/datafusion/ffi/tests/ffi_execution_plan.rs index 7d04e828bd4a5..f308ab169ecb2 100644 --- a/datafusion/ffi/tests/ffi_execution_plan.rs +++ b/datafusion/ffi/tests/ffi_execution_plan.rs @@ -25,7 +25,7 @@ mod tests { use datafusion_ffi::execution_plan::ForeignExecutionPlan; use datafusion_ffi::execution_plan::{ExecutionPlanPrivateData, tests::EmptyExec}; use datafusion_ffi::tests::utils::get_module; - use datafusion_physical_plan::ExecutionPlan; + use datafusion_physical_plan::{ChildrenPropertiesHint, ExecutionPlan}; use std::sync::Arc; #[test] @@ -92,7 +92,8 @@ mod tests { let grandchild_plan = generate_local_plan(); - let child_plan = child_plan.with_new_children(vec![grandchild_plan])?; + let child_plan = child_plan + .replace_children(vec![grandchild_plan], ChildrenPropertiesHint::Recompute)?; unsafe { // Originally the runtime is not set. We go through the unsafe casting @@ -107,7 +108,8 @@ mod tests { assert!((*grandchild_private_data).runtime.is_none()); } - let parent_plan = generate_local_plan().with_new_children(vec![child_plan])?; + let parent_plan = generate_local_plan() + .replace_children(vec![child_plan], ChildrenPropertiesHint::Recompute)?; // Adding the grandchild beneath this FFI plan should get the runtime passed down. let runtime = tokio::runtime::Builder::new_current_thread() diff --git a/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs b/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs index 952aae9846d0f..f801fa70c84e2 100644 --- a/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs +++ b/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs @@ -760,8 +760,10 @@ fn preserving_order_enables_streaming( return Ok(false); } // Build parent with the ordered child - let with_ordered = - Arc::clone(parent).with_new_children(vec![Arc::clone(ordered_child)])?; + let with_ordered = with_new_children_if_necessary( + Arc::clone(parent), + vec![Arc::clone(ordered_child)], + )?; if with_ordered.pipeline_behavior() == EmissionType::Final { // Parent is blocking even with ordering — no benefit return Ok(false); @@ -769,7 +771,8 @@ fn preserving_order_enables_streaming( // Build parent with an unordered child via CoalescePartitionsExec. let unordered_child: Arc = Arc::new(CoalescePartitionsExec::new(Arc::clone(ordered_child))); - let without_ordered = Arc::clone(parent).with_new_children(vec![unordered_child])?; + let without_ordered = + with_new_children_if_necessary(Arc::clone(parent), vec![unordered_child])?; Ok(without_ordered.pipeline_behavior() == EmissionType::Final) } diff --git a/datafusion/physical-optimizer/src/limit_pushdown.rs b/datafusion/physical-optimizer/src/limit_pushdown.rs index 01a288f7f1632..d1b4614faccc4 100644 --- a/datafusion/physical-optimizer/src/limit_pushdown.rs +++ b/datafusion/physical-optimizer/src/limit_pushdown.rs @@ -77,7 +77,9 @@ use datafusion_physical_plan::placeholder_row::PlaceholderRowExec; use datafusion_physical_plan::projection::ProjectionExec; use datafusion_physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec; use datafusion_physical_plan::statistics::{StatisticsArgs, StatisticsContext}; -use datafusion_physical_plan::{ExecutionPlan, ExecutionPlanProperties}; +use datafusion_physical_plan::{ + ExecutionPlan, ExecutionPlanProperties, with_new_children_if_necessary, +}; /// This rule inspects [`ExecutionPlan`]'s and pushes down the fetch limit from /// the parent to the child if applicable. #[derive(Default, Debug)] @@ -403,7 +405,7 @@ pub(crate) fn pushdown_limits( .collect::>()?; if changed { - new_node.data.with_new_children(new_children) + with_new_children_if_necessary(new_node.data, new_children) } else { Ok(new_node.data) } diff --git a/datafusion/physical-optimizer/src/topk_repartition.rs b/datafusion/physical-optimizer/src/topk_repartition.rs index 115bdc3cb535f..1d5c7ca79cd25 100644 --- a/datafusion/physical-optimizer/src/topk_repartition.rs +++ b/datafusion/physical-optimizer/src/topk_repartition.rs @@ -55,7 +55,9 @@ use std::sync::Arc; use datafusion_physical_plan::coalesce_batches::CoalesceBatchesExec; use datafusion_physical_plan::repartition::RepartitionExec; use datafusion_physical_plan::sorts::sort::SortExec; -use datafusion_physical_plan::{ExecutionPlan, Partitioning}; +use datafusion_physical_plan::{ + ExecutionPlan, Partitioning, with_new_children_if_necessary, +}; /// A physical optimizer rule that pushes TopK (Sort with fetch) past /// hash repartition when the partition key is a prefix of the sort key. @@ -151,7 +153,7 @@ impl PhysicalOptimizerRule for TopKRepartition { // Rebuild the tree above the repartition let new_sort_input = if let Some(parent) = repart_parent { - parent.with_new_children(vec![new_repartition])? + with_new_children_if_necessary(parent, vec![new_repartition])? } else { new_repartition }; diff --git a/datafusion/physical-optimizer/src/window_topn.rs b/datafusion/physical-optimizer/src/window_topn.rs index c668608ca241b..8165cec982994 100644 --- a/datafusion/physical-optimizer/src/window_topn.rs +++ b/datafusion/physical-optimizer/src/window_topn.rs @@ -58,7 +58,6 @@ use datafusion_common::{Result, ScalarValue}; use datafusion_expr::Operator; use datafusion_physical_expr::expressions::{BinaryExpr, Column, Literal}; use datafusion_physical_expr::window::StandardWindowExpr; -use datafusion_physical_plan::ExecutionPlan; use datafusion_physical_plan::filter::FilterExec; use datafusion_physical_plan::projection::ProjectionExec; use datafusion_physical_plan::repartition::RepartitionExec; @@ -67,6 +66,7 @@ use datafusion_physical_plan::sorts::partitioned_topk::{ }; use datafusion_physical_plan::sorts::sort::SortExec; use datafusion_physical_plan::windows::{BoundedWindowAggExec, WindowUDFExpr}; +use datafusion_physical_plan::{ExecutionPlan, with_new_children_if_necessary}; /// Physical optimizer rule that converts per-partition `ROW_NUMBER` and /// `RANK` top-K queries into a more efficient plan using @@ -192,13 +192,13 @@ impl WindowTopN { .ok()?; // Step 8: Rebuild window with new child - let mut result = window_exec - .with_new_children(vec![Arc::new(partitioned_topk)]) - .ok()?; + let mut result = + with_new_children_if_necessary(window_exec, vec![Arc::new(partitioned_topk)]) + .ok()?; // Step 9: Rebuild intermediate nodes (ProjectionExec/RepartitionExec) for node in intermediates.into_iter().rev() { - result = node.with_new_children(vec![result]).ok()?; + result = with_new_children_if_necessary(node, vec![result]).ok()?; } Some(result) diff --git a/datafusion/physical-plan/src/async_func.rs b/datafusion/physical-plan/src/async_func.rs index a823e61aa6b62..e225dd68a1cab 100644 --- a/datafusion/physical-plan/src/async_func.rs +++ b/datafusion/physical-plan/src/async_func.rs @@ -297,6 +297,7 @@ impl AsyncFuncExec { node: &datafusion_proto_models::protobuf::PhysicalPlanNode, ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, ) -> Result> { + use datafusion_common::assert_eq_or_internal_err; use datafusion_proto_models::protobuf; let async_func = crate::expect_plan_variant!( node, diff --git a/datafusion/physical-plan/src/buffer.rs b/datafusion/physical-plan/src/buffer.rs index f91f5fd109d87..5de43113f7cdf 100644 --- a/datafusion/physical-plan/src/buffer.rs +++ b/datafusion/physical-plan/src/buffer.rs @@ -28,7 +28,7 @@ use crate::statistics::{ChildStats, StatisticsArgs}; use crate::stream::RecordBatchStreamAdapter; use crate::{ ChildrenPropertiesHint, DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, - SortOrderPushdownResult, validate_child_count, + SortOrderPushdownResult, validate_child_count, with_new_children_if_necessary, }; use arrow::array::RecordBatch; use datafusion_common::config::ConfigOptions; @@ -273,9 +273,10 @@ impl ExecutionPlan for BufferExec { projection: &ProjectionExec, ) -> Result>> { match self.input.try_swapping_with_projection(projection)? { - Some(new_input) => Ok(Some( - Arc::new(self.clone()).with_new_children(vec![new_input])?, - )), + Some(new_input) => Ok(Some(with_new_children_if_necessary( + Arc::new(self.clone()), + vec![new_input], + )?)), None => Ok(None), } } diff --git a/datafusion/physical-plan/src/coalesce_batches.rs b/datafusion/physical-plan/src/coalesce_batches.rs index a4dd239e7906e..efc0c7f8aba03 100644 --- a/datafusion/physical-plan/src/coalesce_batches.rs +++ b/datafusion/physical-plan/src/coalesce_batches.rs @@ -28,7 +28,7 @@ use crate::statistics::{ChildStats, StatisticsArgs}; use crate::stream::EmptyRecordBatchStream; use crate::{ ChildrenPropertiesHint, DisplayFormatType, ExecutionPlan, RecordBatchStream, - SendableRecordBatchStream, validate_child_count, + SendableRecordBatchStream, validate_child_count, with_new_children_if_necessary, }; use arrow::datatypes::SchemaRef; @@ -263,9 +263,10 @@ impl ExecutionPlan for CoalesceBatchesExec { projection: &ProjectionExec, ) -> Result>> { match self.input.try_swapping_with_projection(projection)? { - Some(new_input) => Ok(Some( - Arc::new(self.clone()).with_new_children(vec![new_input])?, - )), + Some(new_input) => Ok(Some(with_new_children_if_necessary( + Arc::new(self.clone()), + vec![new_input], + )?)), None => Ok(None), } } diff --git a/datafusion/physical-plan/src/coalesce_partitions.rs b/datafusion/physical-plan/src/coalesce_partitions.rs index 788a00c7ecb58..99a77bdba7be8 100644 --- a/datafusion/physical-plan/src/coalesce_partitions.rs +++ b/datafusion/physical-plan/src/coalesce_partitions.rs @@ -33,7 +33,7 @@ use crate::sort_pushdown::SortOrderPushdownResult; use crate::statistics::{ChildStats, StatisticsArgs}; use crate::{ ChildrenPropertiesHint, DisplayFormatType, ExecutionPlan, Partitioning, - validate_child_count, + validate_child_count, with_new_children_if_necessary, }; use datafusion_physical_expr_common::sort_expr::PhysicalSortExpr; @@ -312,8 +312,7 @@ impl ExecutionPlan for CoalescePartitionsExec { self.input .with_preserve_order(preserve_order) .and_then(|new_input| { - Arc::new(self.clone()) - .with_new_children(vec![new_input]) + with_new_children_if_necessary(Arc::new(self.clone()), vec![new_input]) .ok() }) } diff --git a/datafusion/physical-plan/src/coop.rs b/datafusion/physical-plan/src/coop.rs index 9f29f8fde87fe..b7e7e05112fe8 100644 --- a/datafusion/physical-plan/src/coop.rs +++ b/datafusion/physical-plan/src/coop.rs @@ -88,7 +88,7 @@ use crate::statistics::{ChildStats, StatisticsArgs}; use crate::{ ChildrenPropertiesHint, DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, RecordBatchStream, SendableRecordBatchStream, SortOrderPushdownResult, - validate_child_count, + validate_child_count, with_new_children_if_necessary, }; use arrow::record_batch::RecordBatch; use arrow_schema::Schema; @@ -333,9 +333,10 @@ impl ExecutionPlan for CooperativeExec { projection: &ProjectionExec, ) -> Result>> { match self.input.try_swapping_with_projection(projection)? { - Some(new_input) => Ok(Some( - Arc::new(self.clone()).with_new_children(vec![new_input])?, - )), + Some(new_input) => Ok(Some(with_new_children_if_necessary( + Arc::new(self.clone()), + vec![new_input], + )?)), None => Ok(None), } } @@ -366,11 +367,13 @@ impl ExecutionPlan for CooperativeExec { match child.try_pushdown_sort(order)? { SortOrderPushdownResult::Exact { inner } => { - let new_exec = Arc::new(self.clone()).with_new_children(vec![inner])?; + let new_exec = + with_new_children_if_necessary(Arc::new(self.clone()), vec![inner])?; Ok(SortOrderPushdownResult::Exact { inner: new_exec }) } SortOrderPushdownResult::Inexact { inner } => { - let new_exec = Arc::new(self.clone()).with_new_children(vec![inner])?; + let new_exec = + with_new_children_if_necessary(Arc::new(self.clone()), vec![inner])?; Ok(SortOrderPushdownResult::Inexact { inner: new_exec }) } SortOrderPushdownResult::Unsupported => { diff --git a/datafusion/physical-plan/src/execution_plan.rs b/datafusion/physical-plan/src/execution_plan.rs index d6b54546642cb..06ab60fd2018f 100644 --- a/datafusion/physical-plan/src/execution_plan.rs +++ b/datafusion/physical-plan/src/execution_plan.rs @@ -312,7 +312,7 @@ pub trait ExecutionPlan: Any + Debug + DisplayAs + Send + Sync { /// [`DynamicFilterPhysicalExpr`]: datafusion_physical_expr::expressions::DynamicFilterPhysicalExpr fn reset_state(self: Arc) -> Result> { let children = self.children().into_iter().cloned().collect(); - self.with_new_children(children) + self.replace_children(children, ChildrenPropertiesHint::Recompute) } /// If supported, attempt to increase the partitioning of this `ExecutionPlan` to diff --git a/datafusion/physical-plan/src/filter.rs b/datafusion/physical-plan/src/filter.rs index 77fac49321c6a..93dc3d7ee9416 100644 --- a/datafusion/physical-plan/src/filter.rs +++ b/datafusion/physical-plan/src/filter.rs @@ -43,7 +43,9 @@ use crate::projection::{ }; use crate::statistics::{ChildStats, StatisticsArgs, StatisticsContext}; use crate::stream::EmptyRecordBatchStream; -use crate::{ChildrenPropertiesHint, validate_child_count}; +use crate::{ + ChildrenPropertiesHint, validate_child_count, with_new_children_if_necessary, +}; use crate::{ DisplayFormatType, ExecutionPlan, metrics::{BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet, RatioMetrics}, @@ -820,8 +822,7 @@ impl ExecutionPlan for FilterExec { self.input .with_preserve_order(preserve_order) .and_then(|new_input| { - Arc::new(self.clone()) - .with_new_children(vec![new_input]) + with_new_children_if_necessary(Arc::new(self.clone()), vec![new_input]) .ok() }) } diff --git a/datafusion/physical-plan/src/projection.rs b/datafusion/physical-plan/src/projection.rs index f830044619466..a7463cdc8462c 100644 --- a/datafusion/physical-plan/src/projection.rs +++ b/datafusion/physical-plan/src/projection.rs @@ -36,7 +36,7 @@ use crate::joins::utils::{ColumnIndex, JoinFilter, JoinOn, JoinOnRef}; use crate::statistics::{ChildStats, StatisticsArgs}; use crate::{ ChildrenPropertiesHint, DisplayFormatType, ExecutionPlan, PhysicalExpr, - validate_child_count, + validate_child_count, with_new_children_if_necessary, }; use std::collections::HashMap; use std::pin::Pin; @@ -489,11 +489,13 @@ impl ExecutionPlan for ProjectionExec { // Recursively push down to child node match child.try_pushdown_sort(&child_order)? { SortOrderPushdownResult::Exact { inner } => { - let new_exec = Arc::new(self.clone()).with_new_children(vec![inner])?; + let new_exec = + with_new_children_if_necessary(Arc::new(self.clone()), vec![inner])?; Ok(SortOrderPushdownResult::Exact { inner: new_exec }) } SortOrderPushdownResult::Inexact { inner } => { - let new_exec = Arc::new(self.clone()).with_new_children(vec![inner])?; + let new_exec = + with_new_children_if_necessary(Arc::new(self.clone()), vec![inner])?; Ok(SortOrderPushdownResult::Inexact { inner: new_exec }) } SortOrderPushdownResult::Unsupported => { @@ -509,8 +511,7 @@ impl ExecutionPlan for ProjectionExec { self.input .with_preserve_order(preserve_order) .and_then(|new_input| { - Arc::new(self.clone()) - .with_new_children(vec![new_input]) + with_new_children_if_necessary(Arc::new(self.clone()), vec![new_input]) .ok() }) } From fc8623e844c17f998214da02b3daa3f626885eed Mon Sep 17 00:00:00 2001 From: Justin O'Dwyer Date: Wed, 29 Jul 2026 15:15:11 -0400 Subject: [PATCH 8/9] Add deprecation markers and fixes for CI. --- datafusion/core/tests/fuzz_cases/once_exec.rs | 4 ++-- datafusion/physical-plan/src/display.rs | 2 +- datafusion/physical-plan/src/execution_plan.rs | 10 ++++++++++ 3 files changed, 13 insertions(+), 3 deletions(-) diff --git a/datafusion/core/tests/fuzz_cases/once_exec.rs b/datafusion/core/tests/fuzz_cases/once_exec.rs index b6ac0e2f67c8b..8834acd3f5531 100644 --- a/datafusion/core/tests/fuzz_cases/once_exec.rs +++ b/datafusion/core/tests/fuzz_cases/once_exec.rs @@ -97,7 +97,7 @@ impl ExecutionPlan for OnceExec { fn with_new_children( self: Arc, children: Vec>, - ) -> datafusion_common::Result> { + ) -> Result> { self.replace_children(children, ChildrenPropertiesHint::Recompute) } @@ -106,7 +106,7 @@ impl ExecutionPlan for OnceExec { &self, partition: usize, _context: Arc, - ) -> datafusion_common::Result { + ) -> Result { assert_eq!(partition, 0); let stream = self.stream.lock().unwrap().take(); diff --git a/datafusion/physical-plan/src/display.rs b/datafusion/physical-plan/src/display.rs index 91e79150f1796..8330f1a918101 100644 --- a/datafusion/physical-plan/src/display.rs +++ b/datafusion/physical-plan/src/display.rs @@ -1578,10 +1578,10 @@ mod tests { use insta::assert_snapshot; use super::super::DisplayableExecutionPlan; - use crate::ChildrenPropertiesHint; use crate::empty::EmptyExec; use crate::filter::FilterExec; use crate::projection::ProjectionExec; + use crate::{ChildrenPropertiesHint, ExecutionPlan}; use datafusion_physical_expr::expressions::{binary, col, lit}; use datafusion_physical_expr::{Partitioning, PhysicalExpr}; diff --git a/datafusion/physical-plan/src/execution_plan.rs b/datafusion/physical-plan/src/execution_plan.rs index 06ab60fd2018f..fe1763d7570bb 100644 --- a/datafusion/physical-plan/src/execution_plan.rs +++ b/datafusion/physical-plan/src/execution_plan.rs @@ -249,6 +249,7 @@ pub trait ExecutionPlan: Any + Debug + DisplayAs + Send + Sync { /// Returns a clone of the existing plan with the children replaced, skipping /// recomputation of plan properties if possible as indicated by the hint. + #[expect(deprecated)] fn replace_children( self: Arc, children: Vec>, @@ -264,6 +265,10 @@ pub trait ExecutionPlan: Any + Debug + DisplayAs + Send + Sync { /// Returns a new `ExecutionPlan` where all existing children were replaced /// by the `children`, in order + #[deprecated( + since = "55.0.0", + note = "Use `ExecutionPlan::replace_children` with ChildrenPropertiesHint::Recompute" + )] fn with_new_children( self: Arc, children: Vec>, @@ -283,6 +288,11 @@ pub trait ExecutionPlan: Any + Debug + DisplayAs + Send + Sync { /// /// Callers should route through [`with_new_children_if_necessary`] and /// not invoke this method directly. + #[deprecated( + since = "55.0.0", + note = "Use `ExecutionPlan::replace_children` with ChildrenPropertiesHint::SameProperties" + )] + #[expect(deprecated)] fn with_new_children_and_same_properties( self: Arc, children: Vec>, From 3c9d8e4ce4c8af93916908360e59510f1ea273ef Mon Sep 17 00:00:00 2001 From: Justin O'Dwyer Date: Wed, 29 Jul 2026 16:15:43 -0400 Subject: [PATCH 9/9] Updating docs. --- .../library-user-guide/upgrading/55.0.0.md | 69 +++++++++++++++++++ 1 file changed, 69 insertions(+) diff --git a/docs/source/library-user-guide/upgrading/55.0.0.md b/docs/source/library-user-guide/upgrading/55.0.0.md index 6ca4d71a7883f..7c58b1f05c193 100644 --- a/docs/source/library-user-guide/upgrading/55.0.0.md +++ b/docs/source/library-user-guide/upgrading/55.0.0.md @@ -552,6 +552,75 @@ See [PR #22733](https://github.com/apache/datafusion/pull/22733) for details, including the per-variant size breakdown and benchmark results. +### `ExecutionPlan::with_new_children` and `ExecutionPlan::with_new_children_and_same_properties` deprecated + +`with_new_children` and `with_new_children_and_same_properties` have been +deprecated. These methods are used to replace the child plans of an +`ExecutionPlan` while leaving the plan otherwise identical. + +As noted [here](https://github.com/apache/datafusion/pull/23332#discussion_r3554897693), +while the addition of `with_new_children_and_same_properties` has the benefit +of skipping potentially expensive computation in the case that replacement children +have the same properties as the original children, it widens the API surface area +of `ExecutionPlan` in a way that could be confusing for users. + +Thus, to rectify this, we unify these methods by introducing `replace_children`. +`replace_children` solves this problem by taking an enum called +`ChildrenPropertiesHint` as an argument. The enum has two variants, `SameProperties` +and `Recompute`, which function as a hint to `replace_children` from the caller +as to whether or not the properties need to be recomputed. + +This method is called from `with_new_children_if_necessary`, which is the +standard entry point that should be used for replacing the children of a node. + +**Migration guide:** + +To migrate from `with_new_children` and `with_new_children_and_same_properties` +to `replace_children`, it is recommended to implement `replace_children` with +a `match` statement matching on the `ChildrenPropertiesHint`. In the case that +the properties match the children, `ChildrenPropertiesHint::SameProperties`, +follow the body of `with_new_children_and_same_properties`. In the case that +the properties do not match the children, `ChildrenPropertiesHint::Recompute`, +follow the body of `with_new_children`. + +For example, take a look at the implementation for `FilterExec`: + +``` + fn replace_children( + self: Arc, + mut children: Vec>, + hint: ChildrenPropertiesHint, + ) -> Result> { + validate_child_count!(self, children); + match hint { + ChildrenPropertiesHint::SameProperties => Ok(Arc::new(Self { + input: children.swap_remove(0), + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(&*self) + })), + ChildrenPropertiesHint::Recompute => { + let new_input = children.swap_remove(0); + FilterExecBuilder::from(&*self) + .with_input(new_input) + .build() + .map(|e| Arc::new(e) as _) + } + } + } +``` + +In the case that the hint suggests that the properties are the same, we can +simply swap the children without having to recompute the properties. In the +other case, we create a new node from scratch. + +To ensure that this works correctly, it is recommended that users also look +through their codebase and ensure that they use `with_new_children_if_necessary` +for these changes — `with_new_children_if_necessary` should be preferred over +manual use of `replace_children`, since `with_new_children_if_necessary` will +call `replace_children` with the correct hint filled in. + +See [PR #23903](https://github.com/apache/datafusion/pull/23903) for details. + ### `ListingOptions::target_partitions` and `collect_stat` removed The `target_partitions` and `collect_stat` fields on