diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index 11446137f3ca1..714a533402d5c 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -1610,17 +1610,6 @@ impl AggregateExec { _ => Precision::Absent, } } - - fn with_new_children_and_same_properties( - &self, - mut children: Vec>, - ) -> Self { - Self { - input: children.swap_remove(0), - metrics: ExecutionPlanMetricsSet::new(), - ..Self::clone(self) - } - } } impl DisplayAs for AggregateExec { @@ -1838,6 +1827,17 @@ impl ExecutionPlan for AggregateExec { Ok(Arc::new(me)) } + fn with_new_children_and_same_properties( + self: Arc, + mut children: Vec>, + ) -> Result> { + Ok(Arc::new(Self { + input: children.swap_remove(0), + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(&*self) + })) + } + fn execute( &self, partition: usize, diff --git a/datafusion/physical-plan/src/async_func.rs b/datafusion/physical-plan/src/async_func.rs index efac83bbbe5ba..5a65c9aedc2f1 100644 --- a/datafusion/physical-plan/src/async_func.rs +++ b/datafusion/physical-plan/src/async_func.rs @@ -113,17 +113,6 @@ impl AsyncFuncExec { pub fn input(&self) -> &Arc { &self.input } - - fn with_new_children_and_same_properties( - &self, - mut children: Vec>, - ) -> Self { - Self { - input: children.swap_remove(0), - metrics: ExecutionPlanMetricsSet::new(), - ..Self::clone(self) - } - } } impl DisplayAs for AsyncFuncExec { @@ -180,6 +169,17 @@ impl ExecutionPlan for AsyncFuncExec { )?)) } + fn with_new_children_and_same_properties( + self: Arc, + mut children: Vec>, + ) -> Result> { + Ok(Arc::new(Self { + input: children.swap_remove(0), + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(&*self) + })) + } + fn execute( &self, partition: usize, diff --git a/datafusion/physical-plan/src/buffer.rs b/datafusion/physical-plan/src/buffer.rs index 5e220b7e48544..4e88daae73d18 100644 --- a/datafusion/physical-plan/src/buffer.rs +++ b/datafusion/physical-plan/src/buffer.rs @@ -123,17 +123,6 @@ impl BufferExec { pub fn capacity(&self) -> usize { self.capacity } - - fn with_new_children_and_same_properties( - &self, - mut children: Vec>, - ) -> Self { - Self { - input: children.swap_remove(0), - metrics: ExecutionPlanMetricsSet::new(), - ..Self::clone(self) - } - } } impl DisplayAs for BufferExec { @@ -181,6 +170,17 @@ impl ExecutionPlan for BufferExec { Ok(Arc::new(Self::new(children.swap_remove(0), self.capacity))) } + fn with_new_children_and_same_properties( + self: Arc, + mut children: Vec>, + ) -> Result> { + Ok(Arc::new(Self { + input: children.swap_remove(0), + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(&*self) + })) + } + fn execute( &self, partition: usize, diff --git a/datafusion/physical-plan/src/coalesce_batches.rs b/datafusion/physical-plan/src/coalesce_batches.rs index 59b3138b55430..fc0fae6cc34c2 100644 --- a/datafusion/physical-plan/src/coalesce_batches.rs +++ b/datafusion/physical-plan/src/coalesce_batches.rs @@ -117,17 +117,6 @@ impl CoalesceBatchesExec { input.boundedness(), ) } - - fn with_new_children_and_same_properties( - &self, - mut children: Vec>, - ) -> Self { - Self { - input: children.swap_remove(0), - metrics: ExecutionPlanMetricsSet::new(), - ..Self::clone(self) - } - } } #[expect(deprecated)] @@ -195,6 +184,17 @@ impl ExecutionPlan for CoalesceBatchesExec { )) } + fn with_new_children_and_same_properties( + self: Arc, + mut children: Vec>, + ) -> Result> { + Ok(Arc::new(Self { + input: children.swap_remove(0), + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(&*self) + })) + } + fn execute( &self, partition: usize, diff --git a/datafusion/physical-plan/src/coalesce_partitions.rs b/datafusion/physical-plan/src/coalesce_partitions.rs index 0a8c5f78882c5..a858b1cd1b487 100644 --- a/datafusion/physical-plan/src/coalesce_partitions.rs +++ b/datafusion/physical-plan/src/coalesce_partitions.rs @@ -100,17 +100,6 @@ impl CoalescePartitionsExec { .with_evaluation_type(drive) .with_scheduling_type(scheduling) } - - fn with_new_children_and_same_properties( - &self, - mut children: Vec>, - ) -> Self { - Self { - input: children.swap_remove(0), - metrics: ExecutionPlanMetricsSet::new(), - ..Self::clone(self) - } - } } impl DisplayAs for CoalescePartitionsExec { @@ -164,6 +153,17 @@ impl ExecutionPlan for CoalescePartitionsExec { Ok(Arc::new(plan)) } + fn with_new_children_and_same_properties( + self: Arc, + mut children: Vec>, + ) -> Result> { + Ok(Arc::new(Self { + input: children.swap_remove(0), + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(&*self) + })) + } + fn execute( &self, partition: usize, diff --git a/datafusion/physical-plan/src/coop.rs b/datafusion/physical-plan/src/coop.rs index 7bd84a3a6b392..46141b7e7a213 100644 --- a/datafusion/physical-plan/src/coop.rs +++ b/datafusion/physical-plan/src/coop.rs @@ -234,16 +234,6 @@ impl CooperativeExec { pub fn input(&self) -> &Arc { &self.input } - - fn with_new_children_and_same_properties( - &self, - mut children: Vec>, - ) -> Self { - Self { - input: children.swap_remove(0), - ..Self::clone(self) - } - } } impl DisplayAs for CooperativeExec { @@ -290,6 +280,16 @@ impl ExecutionPlan for CooperativeExec { Ok(Arc::new(CooperativeExec::new(children.swap_remove(0)))) } + fn with_new_children_and_same_properties( + self: Arc, + mut children: Vec>, + ) -> Result> { + Ok(Arc::new(Self { + input: children.swap_remove(0), + ..Self::clone(&*self) + })) + } + fn execute( &self, partition: usize, diff --git a/datafusion/physical-plan/src/execution_plan.rs b/datafusion/physical-plan/src/execution_plan.rs index 76abf73e0ebbe..5837e4d07b913 100644 --- a/datafusion/physical-plan/src/execution_plan.rs +++ b/datafusion/physical-plan/src/execution_plan.rs @@ -235,6 +235,27 @@ pub trait ExecutionPlan: Any + Debug + DisplayAs + Send + Sync { children: Vec>, ) -> Result>; + /// Fast-path used by [`with_new_children_if_necessary`] when the new + /// `children` are known to have the same [`PlanProperties`] as the current + /// children. Implementations should swap the children in without + /// recomputing this plan's `PlanProperties` (typically by cloning `self` + /// and replacing the child pointers). + /// + /// The default implementation falls back to + /// [`ExecutionPlan::with_new_children`] which is always correct but + /// forfeits the fast-path: implementations that own an expensive + /// `PlanProperties` (e.g. projection mapping, complex equivalence + /// classes) should override this method. + /// + /// Callers should route through [`with_new_children_if_necessary`] and + /// not invoke this method directly. + fn with_new_children_and_same_properties( + self: Arc, + children: Vec>, + ) -> Result> { + self.with_new_children(children) + } + /// Reset any internal state within this [`ExecutionPlan`]. /// /// This method is called when an [`ExecutionPlan`] needs to be re-executed, @@ -1267,7 +1288,25 @@ pub fn need_data_exchange(plan: Arc) -> bool { } } -/// Returns a copy of this plan if we change any child according to the pointer comparison. +/// Returns a plan with the given children, skipping as much work as possible. +/// +/// This helper is the single entry point for "rebuild a plan from new +/// children" and applies three layers of short-circuits, from cheapest to +/// most expensive: +/// +/// 1. **Same child pointers** — if every `children[i]` is `Arc::ptr_eq` to the +/// corresponding existing child, the original `plan` is returned +/// unchanged (no allocation, no [`ExecutionPlan::with_new_children`] +/// call). +/// 2. **Same child properties** — if the children's `PlanProperties` Arcs +/// match (via [`has_same_children_properties`]), the plan's own +/// `PlanProperties` cache can be reused. This calls +/// [`ExecutionPlan::with_new_children_and_same_properties`], which +/// swaps the child pointers without recomputing `PlanProperties`. +/// 3. **Full recompute** — otherwise, delegate to +/// [`ExecutionPlan::with_new_children`], which recomputes +/// `PlanProperties` from scratch. +/// /// The size of `children` must be equal to the size of `ExecutionPlan::children()`. pub fn with_new_children_if_necessary( plan: Arc, @@ -1279,16 +1318,22 @@ pub fn with_new_children_if_necessary( old_children.len(), "Wrong number of children" ); - if children.is_empty() - || children + if !children.is_empty() { + // Layer 1: same child pointers → return the plan unchanged. + if children .iter() .zip(old_children.iter()) - .any(|(c1, c2)| !Arc::ptr_eq(c1, c2)) - { - plan.with_new_children(children) - } else { - Ok(plan) + .all(|(c1, c2)| Arc::ptr_eq(c1, c2)) + { + return Ok(plan); + } + // 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); + } } + // Layer 3: full recompute. + plan.with_new_children(children) } /// Return a [`DisplayableExecutionPlan`] wrapper around an @@ -1531,7 +1576,7 @@ pub fn reset_plan_states(plan: Arc) -> Result], ) -> Result { let old_children = plan.children(); @@ -1551,6 +1596,11 @@ pub fn has_same_children_properties( /// Helper macro to avoid properties re-computation if passed children properties /// the same as plan already has. Could be used to implement fast-path for method /// [`ExecutionPlan::with_new_children`]. +/// +/// New call sites should route through [`with_new_children_if_necessary`], +/// which applies this check together with the child-pointer short-circuit +/// (see [`with_new_children_if_necessary`] for the layered policy). This +/// macro remains for direct-caller sites that have not been migrated yet. #[macro_export] macro_rules! check_if_same_properties { ($plan: expr, $children: expr) => { @@ -1558,8 +1608,8 @@ macro_rules! check_if_same_properties { $plan.as_ref(), &$children, )? { - let plan = $plan.with_new_children_and_same_properties($children); - return Ok(::std::sync::Arc::new(plan)); + return ::std::sync::Arc::clone(&$plan) + .with_new_children_and_same_properties($children); } }; } @@ -1779,6 +1829,334 @@ mod tests { unimplemented!() } } + /// Test leaf plan with a real [`PlanProperties`] cache. Different instances + /// can share the same cache Arc by cloning `cache`. + #[derive(Debug, Clone)] + struct WithChildrenTestLeaf { + cache: Arc, + } + + impl WithChildrenTestLeaf { + fn new(cache: Arc) -> Self { + Self { cache } + } + } + + impl DisplayAs for WithChildrenTestLeaf { + fn fmt_as( + &self, + _t: DisplayFormatType, + _f: &mut std::fmt::Formatter, + ) -> std::fmt::Result { + unimplemented!() + } + } + + impl ExecutionPlan for WithChildrenTestLeaf { + fn name(&self) -> &'static str { + "WithChildrenTestLeaf" + } + fn properties(&self) -> &Arc { + &self.cache + } + fn children(&self) -> Vec<&Arc> { + vec![] + } + fn with_new_children( + self: Arc, + _: Vec>, + ) -> Result> { + Ok(self) + } + fn execute( + &self, + _partition: usize, + _context: Arc, + ) -> Result { + unimplemented!() + } + } + + /// Test unary plan that counts which of `with_new_children` (full + /// recompute) vs `with_new_children_and_same_properties` (fast path) is + /// taken. + #[derive(Debug, Clone)] + struct WithChildrenTestParent { + input: Arc, + cache: Arc, + recompute_calls: Arc, + fast_path_calls: Arc, + } + + impl WithChildrenTestParent { + fn new(input: Arc) -> Self { + let cache = Arc::new(PlanProperties::new( + EquivalenceProperties::new(Arc::new(Schema::empty())), + Partitioning::UnknownPartitioning(1), + EmissionType::Final, + Boundedness::Bounded, + )); + Self { + input, + cache, + recompute_calls: Arc::new(std::sync::atomic::AtomicUsize::new(0)), + fast_path_calls: Arc::new(std::sync::atomic::AtomicUsize::new(0)), + } + } + } + + impl DisplayAs for WithChildrenTestParent { + fn fmt_as( + &self, + _t: DisplayFormatType, + _f: &mut std::fmt::Formatter, + ) -> std::fmt::Result { + unimplemented!() + } + } + + impl ExecutionPlan for WithChildrenTestParent { + fn name(&self) -> &'static str { + "WithChildrenTestParent" + } + fn properties(&self) -> &Arc { + &self.cache + } + fn children(&self) -> Vec<&Arc> { + vec![&self.input] + } + fn with_new_children( + self: Arc, + mut children: Vec>, + ) -> 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), + })) + } + fn with_new_children_and_same_properties( + self: Arc, + mut 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) + })) + } + fn execute( + &self, + _partition: usize, + _context: Arc, + ) -> Result { + unimplemented!() + } + } + + /// Test unary plan that does **not** override + /// `with_new_children_and_same_properties`. Used to verify the default + /// trait fallback still routes through `with_new_children` (which is + /// the semantics-preserving path for downstream / external + /// `ExecutionPlan` implementations that haven't opted into the + /// fast path yet). + #[derive(Debug, Clone)] + struct WithChildrenTestParentDefault { + input: Arc, + cache: Arc, + recompute_calls: Arc, + } + + impl WithChildrenTestParentDefault { + fn new(input: Arc) -> Self { + let cache = Arc::new(PlanProperties::new( + EquivalenceProperties::new(Arc::new(Schema::empty())), + Partitioning::UnknownPartitioning(1), + EmissionType::Final, + Boundedness::Bounded, + )); + Self { + input, + cache, + recompute_calls: Arc::new(std::sync::atomic::AtomicUsize::new(0)), + } + } + } + + impl DisplayAs for WithChildrenTestParentDefault { + fn fmt_as( + &self, + _t: DisplayFormatType, + _f: &mut std::fmt::Formatter, + ) -> std::fmt::Result { + unimplemented!() + } + } + + impl ExecutionPlan for WithChildrenTestParentDefault { + fn name(&self) -> &'static str { + "WithChildrenTestParentDefault" + } + fn properties(&self) -> &Arc { + &self.cache + } + fn children(&self) -> Vec<&Arc> { + vec![&self.input] + } + fn with_new_children( + self: Arc, + mut children: Vec>, + ) -> Result> { + self.recompute_calls + .fetch_add(1, std::sync::atomic::Ordering::SeqCst); + 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), + })) + } + // Intentionally does **not** override + // `with_new_children_and_same_properties` — relies on the trait + // default that falls back to `with_new_children`. + fn execute( + &self, + _partition: usize, + _context: Arc, + ) -> Result { + unimplemented!() + } + } + + /// Cover the three short-circuit layers of + /// [`with_new_children_if_necessary`]. + #[test] + fn test_with_new_children_if_necessary_layers() -> Result<()> { + use std::sync::atomic::Ordering; + + // Two leaves that share the same `PlanProperties` Arc but sit behind + // distinct `Arc` pointers. + let leaf_props = Arc::new(PlanProperties::new( + EquivalenceProperties::new(Arc::new(Schema::empty())), + Partitioning::UnknownPartitioning(1), + EmissionType::Final, + Boundedness::Bounded, + )); + let leaf_a: Arc = + Arc::new(WithChildrenTestLeaf::new(Arc::clone(&leaf_props))); + let leaf_b: Arc = + Arc::new(WithChildrenTestLeaf::new(Arc::clone(&leaf_props))); + // A third leaf with a *different* `PlanProperties` Arc — for layer 3. + let leaf_c_props = Arc::new(PlanProperties::new( + EquivalenceProperties::new(Arc::new(Schema::empty())), + Partitioning::UnknownPartitioning(1), + EmissionType::Final, + Boundedness::Bounded, + )); + let leaf_c: Arc = + Arc::new(WithChildrenTestLeaf::new(leaf_c_props)); + + let parent = Arc::new(WithChildrenTestParent::new(Arc::clone(&leaf_a))); + let parent_dyn: Arc = Arc::clone(&parent) as _; + let orig_props = Arc::clone(parent.properties()); + + // Layer 1: same child pointer → returns the original plan Arc verbatim. + 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); + assert_eq!(parent.fast_path_calls.load(Ordering::SeqCst), 0); + + // Layer 2: distinct child Arc, but children share the same + // `PlanProperties` Arc → fast path, parent's `PlanProperties` cache + // Arc is reused (not reallocated). + assert!(!Arc::ptr_eq(&leaf_a, &leaf_b)); + assert!(Arc::ptr_eq(leaf_a.properties(), leaf_b.properties())); + let out = with_new_children_if_necessary( + Arc::clone(&parent_dyn), + vec![Arc::clone(&leaf_b)], + )?; + assert!(Arc::ptr_eq(out.properties(), &orig_props)); + assert_eq!(parent.recompute_calls.load(Ordering::SeqCst), 0); + assert_eq!(parent.fast_path_calls.load(Ordering::SeqCst), 1); + + // Layer 3: child's `PlanProperties` Arc differs → full recompute. + assert!(!Arc::ptr_eq(leaf_a.properties(), leaf_c.properties())); + let out = with_new_children_if_necessary( + Arc::clone(&parent_dyn), + vec![Arc::clone(&leaf_c)], + )?; + assert!(!Arc::ptr_eq(out.properties(), &orig_props)); + assert_eq!(parent.recompute_calls.load(Ordering::SeqCst), 1); + assert_eq!(parent.fast_path_calls.load(Ordering::SeqCst), 1); + + Ok(()) + } + + /// A plan that does not override `with_new_children_and_same_properties` + /// (per @kosiew's review on #23332) must still be routed through + /// `with_new_children` when the helper hits the "same properties" + /// branch. The default trait implementation forwards to + /// `with_new_children`, so downstream / external `ExecutionPlan` + /// implementations keep the semantics-preserving path. + #[test] + fn test_with_new_children_if_necessary_default_fallback() -> Result<()> { + use std::sync::atomic::Ordering; + + let leaf_props = Arc::new(PlanProperties::new( + EquivalenceProperties::new(Arc::new(Schema::empty())), + Partitioning::UnknownPartitioning(1), + EmissionType::Final, + Boundedness::Bounded, + )); + let leaf_a: Arc = + Arc::new(WithChildrenTestLeaf::new(Arc::clone(&leaf_props))); + let leaf_b: Arc = + Arc::new(WithChildrenTestLeaf::new(Arc::clone(&leaf_props))); + assert!(!Arc::ptr_eq(&leaf_a, &leaf_b)); + assert!(Arc::ptr_eq(leaf_a.properties(), leaf_b.properties())); + + 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 + // enters the "same properties" branch and calls the trait method, + // whose default forwards to `with_new_children`. + let out = with_new_children_if_necessary( + Arc::clone(&parent_dyn), + vec![Arc::clone(&leaf_b)], + )?; + // `with_new_children` was invoked exactly once via the default. + assert_eq!(parent.recompute_calls.load(Ordering::SeqCst), 1); + // The returned plan has a freshly-recomputed `PlanProperties` Arc, + // so it differs from the parent's original cache. This confirms + // the fallback ran and did not short-circuit. + assert!(!Arc::ptr_eq(out.properties(), parent.properties())); + + Ok(()) + } + #[test] fn test_execution_plan_name() { let schema1 = Arc::new(Schema::empty()); diff --git a/datafusion/physical-plan/src/filter.rs b/datafusion/physical-plan/src/filter.rs index d23dd380423d1..9c09ff6f4f7fd 100644 --- a/datafusion/physical-plan/src/filter.rs +++ b/datafusion/physical-plan/src/filter.rs @@ -470,17 +470,6 @@ impl FilterExec { input.boundedness(), )) } - - fn with_new_children_and_same_properties( - &self, - mut children: Vec>, - ) -> Self { - Self { - input: children.swap_remove(0), - metrics: ExecutionPlanMetricsSet::new(), - ..Self::clone(self) - } - } } impl DisplayAs for FilterExec { @@ -559,6 +548,17 @@ impl ExecutionPlan for FilterExec { .map(|e| Arc::new(e) as _) } + fn with_new_children_and_same_properties( + self: Arc, + mut children: Vec>, + ) -> Result> { + Ok(Arc::new(Self { + input: children.swap_remove(0), + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(&*self) + })) + } + fn execute( &self, partition: usize, diff --git a/datafusion/physical-plan/src/joins/cross_join.rs b/datafusion/physical-plan/src/joins/cross_join.rs index 79295ba2fb556..f8a9b80179102 100644 --- a/datafusion/physical-plan/src/joins/cross_join.rs +++ b/datafusion/physical-plan/src/joins/cross_join.rs @@ -194,23 +194,6 @@ impl CrossJoinExec { &self.right.schema(), ) } - - fn with_new_children_and_same_properties( - &self, - mut children: Vec>, - ) -> Self { - let left = children.swap_remove(0); - let right = children.swap_remove(0); - - Self { - left, - right, - metrics: ExecutionPlanMetricsSet::new(), - left_fut: Default::default(), - cache: Arc::clone(&self.cache), - schema: Arc::clone(&self.schema), - } - } } /// Asynchronously collect the result of the left child @@ -294,6 +277,23 @@ impl ExecutionPlan for CrossJoinExec { ))) } + fn with_new_children_and_same_properties( + self: Arc, + mut 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), + })) + } + fn reset_state(self: Arc) -> Result> { let new_exec = CrossJoinExec { left: Arc::clone(&self.left), diff --git a/datafusion/physical-plan/src/joins/nested_loop_join.rs b/datafusion/physical-plan/src/joins/nested_loop_join.rs index 3cae05a3a815a..c120654319a63 100644 --- a/datafusion/physical-plan/src/joins/nested_loop_join.rs +++ b/datafusion/physical-plan/src/joins/nested_loop_join.rs @@ -489,28 +489,6 @@ impl NestedLoopJoinExec { Ok(plan) } - - fn with_new_children_and_same_properties( - &self, - mut children: Vec>, - ) -> Self { - let left = children.swap_remove(0); - let right = children.swap_remove(0); - - 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(), - } - } } impl DisplayAs for NestedLoopJoinExec { @@ -597,6 +575,28 @@ impl ExecutionPlan for NestedLoopJoinExec { )) } + fn with_new_children_and_same_properties( + self: Arc, + mut 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(), + })) + } + fn execute( &self, partition: usize, 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 50e9252a21131..bec91fdb62ff6 100644 --- a/datafusion/physical-plan/src/joins/piecewise_merge_join/exec.rs +++ b/datafusion/physical-plan/src/joins/piecewise_merge_join/exec.rs @@ -467,31 +467,6 @@ impl PiecewiseMergeJoinExec { pub fn swap_inputs(&self) -> Result> { todo!() } - - fn with_new_children_and_same_properties( - &self, - mut children: Vec>, - ) -> Self { - let buffered = children.swap_remove(0); - let streamed = children.swap_remove(0); - 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(), - } - } } impl ExecutionPlan for PiecewiseMergeJoinExec { @@ -550,11 +525,35 @@ impl ExecutionPlan for PiecewiseMergeJoinExec { } } + fn with_new_children_and_same_properties( + self: Arc, + mut 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(), + })) + } + fn reset_state(self: Arc) -> Result> { - Ok(Arc::new(self.with_new_children_and_same_properties(vec![ - Arc::clone(&self.buffered), - Arc::clone(&self.streamed), - ]))) + let buffered = Arc::clone(&self.buffered); + let streamed = Arc::clone(&self.streamed); + self.with_new_children_and_same_properties(vec![buffered, streamed]) } 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 2dc7065eee04e..3c0743ad894d3 100644 --- a/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs +++ b/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs @@ -344,20 +344,6 @@ impl SortMergeJoinExec { reorder_output_after_swap(Arc::new(new_join), &left.schema(), &right.schema()) } } - - fn with_new_children_and_same_properties( - &self, - mut children: Vec>, - ) -> Self { - let left = children.swap_remove(0); - let right = children.swap_remove(0); - Self { - left, - right, - metrics: ExecutionPlanMetricsSet::new(), - ..Self::clone(self) - } - } } impl DisplayAs for SortMergeJoinExec { @@ -469,6 +455,20 @@ impl ExecutionPlan for SortMergeJoinExec { } } + fn with_new_children_and_same_properties( + self: Arc, + mut 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) + })) + } + fn execute( &self, partition: usize, diff --git a/datafusion/physical-plan/src/joins/symmetric_hash_join.rs b/datafusion/physical-plan/src/joins/symmetric_hash_join.rs index 52a1aa056d244..0cef4690718d7 100644 --- a/datafusion/physical-plan/src/joins/symmetric_hash_join.rs +++ b/datafusion/physical-plan/src/joins/symmetric_hash_join.rs @@ -360,20 +360,6 @@ impl SymmetricHashJoinExec { } Ok(false) } - - fn with_new_children_and_same_properties( - &self, - mut children: Vec>, - ) -> Self { - let left = children.swap_remove(0); - let right = children.swap_remove(0); - Self { - left, - right, - metrics: ExecutionPlanMetricsSet::new(), - ..Self::clone(self) - } - } } impl DisplayAs for SymmetricHashJoinExec { @@ -477,6 +463,20 @@ impl ExecutionPlan for SymmetricHashJoinExec { )?)) } + fn with_new_children_and_same_properties( + self: Arc, + mut 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) + })) + } + fn metrics(&self) -> Option { Some(self.metrics.clone_inner()) } diff --git a/datafusion/physical-plan/src/limit.rs b/datafusion/physical-plan/src/limit.rs index 1e4b5e5bb6426..4f1abf288db05 100644 --- a/datafusion/physical-plan/src/limit.rs +++ b/datafusion/physical-plan/src/limit.rs @@ -109,17 +109,6 @@ impl GlobalLimitExec { pub fn set_required_ordering(&mut self, required_ordering: Option) { self.required_ordering = required_ordering; } - - fn with_new_children_and_same_properties( - &self, - mut children: Vec>, - ) -> Self { - Self { - input: children.swap_remove(0), - metrics: ExecutionPlanMetricsSet::new(), - ..Self::clone(self) - } - } } impl DisplayAs for GlobalLimitExec { @@ -186,6 +175,17 @@ impl ExecutionPlan for GlobalLimitExec { ))) } + fn with_new_children_and_same_properties( + self: Arc, + mut children: Vec>, + ) -> Result> { + Ok(Arc::new(Self { + input: children.swap_remove(0), + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(&*self) + })) + } + fn execute( &self, partition: usize, @@ -294,17 +294,6 @@ impl LocalLimitExec { pub fn set_required_ordering(&mut self, required_ordering: Option) { self.required_ordering = required_ordering; } - - fn with_new_children_and_same_properties( - &self, - mut children: Vec>, - ) -> Self { - Self { - input: children.swap_remove(0), - metrics: ExecutionPlanMetricsSet::new(), - ..Self::clone(self) - } - } } impl DisplayAs for LocalLimitExec { @@ -360,6 +349,17 @@ impl ExecutionPlan for LocalLimitExec { } } + fn with_new_children_and_same_properties( + self: Arc, + mut children: Vec>, + ) -> Result> { + Ok(Arc::new(Self { + input: children.swap_remove(0), + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(&*self) + })) + } + fn execute( &self, partition: usize, diff --git a/datafusion/physical-plan/src/projection.rs b/datafusion/physical-plan/src/projection.rs index 16b0a5ad7e4b5..18f9e8d938c59 100644 --- a/datafusion/physical-plan/src/projection.rs +++ b/datafusion/physical-plan/src/projection.rs @@ -222,17 +222,6 @@ impl ProjectionExec { } Ok(alias_map) } - - fn with_new_children_and_same_properties( - &self, - mut children: Vec>, - ) -> Self { - Self { - input: children.swap_remove(0), - metrics: ExecutionPlanMetricsSet::new(), - ..Self::clone(self) - } - } } impl DisplayAs for ProjectionExec { @@ -325,6 +314,17 @@ impl ExecutionPlan for ProjectionExec { .map(|p| Arc::new(p) as _) } + fn with_new_children_and_same_properties( + self: Arc, + mut children: Vec>, + ) -> Result> { + Ok(Arc::new(Self { + input: children.swap_remove(0), + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(&*self) + })) + } + fn execute( &self, partition: usize, diff --git a/datafusion/physical-plan/src/repartition/mod.rs b/datafusion/physical-plan/src/repartition/mod.rs index 1617af3a68baa..67dabd7c725f7 100644 --- a/datafusion/physical-plan/src/repartition/mod.rs +++ b/datafusion/physical-plan/src/repartition/mod.rs @@ -1117,18 +1117,6 @@ impl RepartitionExec { pub fn name(&self) -> &str { "RepartitionExec" } - - fn with_new_children_and_same_properties( - &self, - mut children: Vec>, - ) -> Self { - Self { - input: children.swap_remove(0), - metrics: ExecutionPlanMetricsSet::new(), - state: Default::default(), - ..Self::clone(self) - } - } } impl DisplayAs for RepartitionExec { @@ -1207,6 +1195,18 @@ impl ExecutionPlan for RepartitionExec { Ok(Arc::new(repartition)) } + fn with_new_children_and_same_properties( + self: Arc, + mut children: Vec>, + ) -> Result> { + Ok(Arc::new(Self { + input: children.swap_remove(0), + metrics: ExecutionPlanMetricsSet::new(), + state: Default::default(), + ..Self::clone(&*self) + })) + } + fn benefits_from_input_partitioning(&self) -> Vec { vec![matches!(self.partitioning(), Partitioning::Hash(_, _))] } diff --git a/datafusion/physical-plan/src/sorts/partial_sort.rs b/datafusion/physical-plan/src/sorts/partial_sort.rs index d215e5296f91d..95c47c5d01a8c 100644 --- a/datafusion/physical-plan/src/sorts/partial_sort.rs +++ b/datafusion/physical-plan/src/sorts/partial_sort.rs @@ -235,17 +235,6 @@ impl PartialSortExec { input.boundedness(), )) } - - fn with_new_children_and_same_properties( - &self, - mut children: Vec>, - ) -> Self { - Self { - input: children.swap_remove(0), - metrics_set: ExecutionPlanMetricsSet::new(), - ..Self::clone(self) - } - } } impl DisplayAs for PartialSortExec { @@ -330,6 +319,17 @@ impl ExecutionPlan for PartialSortExec { Ok(Arc::new(new_partial_sort)) } + fn with_new_children_and_same_properties( + self: Arc, + mut children: Vec>, + ) -> Result> { + Ok(Arc::new(Self { + input: children.swap_remove(0), + metrics_set: ExecutionPlanMetricsSet::new(), + ..Self::clone(&*self) + })) + } + 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 77a7d8f8f2e11..6d377e64a6cdd 100644 --- a/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs +++ b/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs @@ -181,17 +181,6 @@ impl SortPreservingMergeExec { .with_evaluation_type(drive) .with_scheduling_type(scheduling) } - - fn with_new_children_and_same_properties( - &self, - mut children: Vec>, - ) -> Self { - Self { - input: children.swap_remove(0), - metrics: ExecutionPlanMetricsSet::new(), - ..Self::clone(self) - } - } } impl DisplayAs for SortPreservingMergeExec { @@ -297,6 +286,17 @@ impl ExecutionPlan for SortPreservingMergeExec { )) } + fn with_new_children_and_same_properties( + self: Arc, + mut children: Vec>, + ) -> Result> { + Ok(Arc::new(Self { + input: children.swap_remove(0), + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(&*self) + })) + } + fn execute( &self, partition: usize, diff --git a/datafusion/physical-plan/src/union.rs b/datafusion/physical-plan/src/union.rs index d6f664c0059bc..1074ed1f55063 100644 --- a/datafusion/physical-plan/src/union.rs +++ b/datafusion/physical-plan/src/union.rs @@ -172,17 +172,6 @@ impl UnionExec { boundedness_from_children(inputs), )) } - - fn with_new_children_and_same_properties( - &self, - children: Vec>, - ) -> Self { - Self { - inputs: children, - metrics: ExecutionPlanMetricsSet::new(), - ..Self::clone(self) - } - } } impl DisplayAs for UnionExec { @@ -259,6 +248,17 @@ impl ExecutionPlan for UnionExec { UnionExec::try_new(children) } + fn with_new_children_and_same_properties( + self: Arc, + children: Vec>, + ) -> Result> { + Ok(Arc::new(Self { + inputs: children, + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(&*self) + })) + } + fn execute( &self, mut partition: usize, @@ -527,17 +527,6 @@ impl InterleaveExec { boundedness_from_children(inputs), )) } - - fn with_new_children_and_same_properties( - &self, - children: Vec>, - ) -> Self { - Self { - inputs: children, - metrics: ExecutionPlanMetricsSet::new(), - ..Self::clone(self) - } - } } impl DisplayAs for InterleaveExec { @@ -586,6 +575,17 @@ impl ExecutionPlan for InterleaveExec { Ok(Arc::new(InterleaveExec::try_new(children)?)) } + fn with_new_children_and_same_properties( + self: Arc, + children: Vec>, + ) -> Result> { + Ok(Arc::new(Self { + inputs: children, + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(&*self) + })) + } + fn execute( &self, partition: usize, diff --git a/datafusion/physical-plan/src/unnest.rs b/datafusion/physical-plan/src/unnest.rs index c31d0dd23fa68..632306b80c334 100644 --- a/datafusion/physical-plan/src/unnest.rs +++ b/datafusion/physical-plan/src/unnest.rs @@ -195,17 +195,6 @@ impl UnnestExec { pub fn options(&self) -> &UnnestOptions { &self.options } - - fn with_new_children_and_same_properties( - &self, - mut children: Vec>, - ) -> Self { - Self { - input: children.swap_remove(0), - metrics: ExecutionPlanMetricsSet::new(), - ..Self::clone(self) - } - } } impl DisplayAs for UnnestExec { @@ -252,6 +241,17 @@ impl ExecutionPlan for UnnestExec { )?)) } + fn with_new_children_and_same_properties( + self: Arc, + mut children: Vec>, + ) -> Result> { + Ok(Arc::new(Self { + input: children.swap_remove(0), + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(&*self) + })) + } + fn required_input_distribution(&self) -> Vec { vec![Distribution::UnspecifiedDistribution] } 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 a9d580f4c687d..f0e18ee818ca3 100644 --- a/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs +++ b/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs @@ -250,17 +250,6 @@ impl BoundedWindowAggExec { total_byte_size: Precision::Absent, }) } - - fn with_new_children_and_same_properties( - &self, - mut children: Vec>, - ) -> Self { - Self { - input: children.swap_remove(0), - metrics: ExecutionPlanMetricsSet::new(), - ..Self::clone(self) - } - } } impl DisplayAs for BoundedWindowAggExec { @@ -357,6 +346,17 @@ impl ExecutionPlan for BoundedWindowAggExec { )?)) } + fn with_new_children_and_same_properties( + self: Arc, + mut children: Vec>, + ) -> Result> { + Ok(Arc::new(Self { + input: children.swap_remove(0), + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(&*self) + })) + } + fn execute( &self, partition: usize, diff --git a/datafusion/physical-plan/src/windows/window_agg_exec.rs b/datafusion/physical-plan/src/windows/window_agg_exec.rs index 72474c6a55483..0a9916414476d 100644 --- a/datafusion/physical-plan/src/windows/window_agg_exec.rs +++ b/datafusion/physical-plan/src/windows/window_agg_exec.rs @@ -159,17 +159,6 @@ impl WindowAggExec { .unwrap_or_else(Vec::new) } } - - fn with_new_children_and_same_properties( - &self, - mut children: Vec>, - ) -> Self { - Self { - input: children.swap_remove(0), - metrics: ExecutionPlanMetricsSet::new(), - ..Self::clone(self) - } - } } impl DisplayAs for WindowAggExec { @@ -260,6 +249,17 @@ impl ExecutionPlan for WindowAggExec { )?)) } + fn with_new_children_and_same_properties( + self: Arc, + mut children: Vec>, + ) -> Result> { + Ok(Arc::new(Self { + input: children.swap_remove(0), + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(&*self) + })) + } + fn execute( &self, partition: usize,