Misc minor optimization in the Physical Optimizer (#21216)

## Which issue does this PR close?

- Closes #.

## Rationale for this change

Similar to #21128, just trying to shave time off the optimizer. Locally,
it improves some sql-planner benchmarks by up to 10% but they seem
relatively noisy on my laptop.

## What changes are included in this PR?

1. Avoid allocation `plan.children()` in a loop in `sort_pushdown.rs`.
2. Try and avoid some expensive tree rewrites in `join_selection.rs`
3. Avoid deep clones of exec limit nodes in `limit_pushdown.rs`, and
only mutate the plan if it was actually changed.
4. Use cheaper code path to change the limit on an `AggregateExec` in
`limited_distinct_aggregation.rs`.
5. Use a read-only traversal in `sanity_checker.rs`. Its read only and
`transform_up` is always more expensive. I've considered extending
`TreeNode` but this seems to be basically the only place in the codebase
that does something like this.

There are a few places where we unconditionally return
`Transformed::yes` which might unintended downstream consequences
because it breaks pointer equality and I think it also just end up
allocating more memory, but they are harder to untangle so I'll try and
do them in followups.

## Are these changes tested?

One new test for limits, otherwise the existing tests.

## Are there any user-facing changes?

Removes the `LimitExec` type, I can't imagine why someone would use it,
and its only used in one place. Happy to bring it back as a deprecated
type.

---------

Signed-off-by: Adam Gutglick <adamgsal@gmail.com>

Author: AdamGS

.gitignore

@@ -75,3 +75,6 @@ rat.txt
 
 # data generated by examples
 datafusion-examples/examples/datafusion-examples/
+
+# Samply profile data
+profile.json.gz

datafusion/core/tests/physical_optimizer/enforce_sorting.rs

@@ -1256,7 +1256,8 @@ async fn test_union_inputs_different_sorted_with_limit() -> Result<()> {
     let physical_plan = sort_preserving_merge_exec(ordering3, union);
 
     let test = EnforceSortingTest::new(physical_plan).with_repartition_sorts(true);
-    // Should not change the unnecessarily fine `SortExec`s because there is `LimitExec`
+    // Should not change the unnecessarily fine `SortExec`s because there are
+    // explicit limit nodes above the second sort.
     assert_snapshot!(test.run(), @r"
     Input Plan:
     SortPreservingMergeExec: [nullable_col@0 ASC]

datafusion/core/tests/physical_optimizer/limit_pushdown.rs

@@ -682,3 +682,38 @@ fn preserves_skip_before_sort() -> Result<()> {
 
     Ok(())
 }
+
+#[test]
+fn no_limit_preserves_plan_identity() -> Result<()> {
+    // When there is no limit in the plan, the optimizer should return the
+    // exact same Arc (pointer-equal) for every node, avoiding unnecessary
+    // plan reconstruction and property recomputation.
+    let schema = create_schema();
+
+    let left = empty_exec(Arc::clone(&schema));
+    let right = empty_exec(Arc::clone(&schema));
+    let on = join_on_columns("c1", "c1");
+    let join = hash_join_exec(left, right, on, None, &JoinType::Inner)?;
+    let plan = filter_exec(Arc::clone(&schema), join)?;
+
+    let optimized =
+        LimitPushdown::new().optimize(Arc::clone(&plan), &ConfigOptions::new())?;
+
+    assert!(
+        Arc::ptr_eq(&plan, &optimized),
+        "Expected optimizer to return the same Arc when no limit is present"
+    );
+
+    let optimized = format_plan(&optimized);
+    insta::assert_snapshot!(
+        optimized,
+        @r"
+    FilterExec: c3@2 > 0
+      HashJoinExec: mode=Partitioned, join_type=Inner, on=[(c1@0, c1@0)]
+        EmptyExec
+        EmptyExec
+    "
+    );
+
+    Ok(())
+}

datafusion/physical-optimizer/src/enforce_sorting/mod.rs

@@ -548,8 +548,9 @@ pub fn ensure_sorting(
 
 /// Analyzes if there are any immediate sort removals by checking the `SortExec`s
 /// and their ordering requirement satisfactions with children
-/// If the sort is unnecessary, either replaces it with [`SortPreservingMergeExec`]/`LimitExec`
-/// or removes the [`SortExec`].
+/// If the sort is unnecessary, either replaces it with
+/// [`SortPreservingMergeExec`] and/or a limit node, or removes the
+/// [`SortExec`].
 /// Otherwise, returns the original plan
 fn analyze_immediate_sort_removal(
     mut node: PlanWithCorrespondingSort,

datafusion/physical-optimizer/src/enforce_sorting/sort_pushdown.rs

@@ -680,8 +680,10 @@ fn handle_custom_pushdown(
     parent_required: OrderingRequirements,
     maintains_input_order: &[bool],
 ) -> Result<Option<Vec<Option<OrderingRequirements>>>> {
+    let plan_children = plan.children();
+
     // If the plan has no children, return early:
-    if plan.children().is_empty() {
+    if plan_children.is_empty() {
         return Ok(None);
     }
 
@@ -699,8 +701,7 @@ fn handle_custom_pushdown(
         .collect();
 
     // Get the number of fields in each child's schema:
-    let children_schema_lengths: Vec<usize> = plan
-        .children()
+    let children_schema_lengths: Vec<usize> = plan_children
         .iter()
         .map(|c| c.schema().fields().len())
         .collect();
@@ -734,7 +735,7 @@ fn handle_custom_pushdown(
         let updated_parent_req = requirement
             .into_iter()
             .map(|req| {
-                let child_schema = plan.children()[maintained_child_idx].schema();
+                let child_schema = plan_children[maintained_child_idx].schema();
                 let updated_columns = req
                     .expr
                     .transform_up(|expr| {
@@ -809,13 +810,15 @@ fn handle_hash_join(
 
     let all_from_right_child = all_indices.iter().all(|i| *i >= len_of_left_fields);
 
+    let plan_children = plan.children();
+
     // If all columns are from the right child, update the parent requirements
     if all_from_right_child {
         // Transform the parent-required expression for the child schema by adjusting columns
         let updated_parent_req = requirement
             .into_iter()
             .map(|req| {
-                let child_schema = plan.children()[1].schema();
+                let child_schema = plan_children[1].schema();
                 let updated_columns = req
                     .expr
                     .transform_up(|expr| {

datafusion/physical-optimizer/src/join_selection.rs

@@ -563,10 +563,14 @@ fn apply_subrules(
     subrules: &Vec<Box<PipelineFixerSubrule>>,
     config_options: &ConfigOptions,
 ) -> Result<Transformed<Arc<dyn ExecutionPlan>>> {
+    let original = Arc::clone(&input);
     for subrule in subrules {
         input = subrule(input, config_options)?;
     }
-    Ok(Transformed::yes(input))
+
+    let transformed = !Arc::ptr_eq(&original, &input);
+
+    Ok(Transformed::new_transformed(input, transformed))
 }
 
 // See tests in datafusion/core/tests/physical_optimizer

datafusion/physical-optimizer/src/limit_pushdown.rs

@@ -19,8 +19,9 @@
 //! data transfer as much as possible.
 //!
 //! # Plan Limit Absorption
-//! In addition to pushing down [`LimitExec`] in the plan, some operators can
-//! "absorb" a limit and stop early during execution.
+//! In addition to pushing down `GlobalLimitExec` and `LocalLimitExec` nodes in
+//! the plan, some operators can "absorb" a limit and stop early during
+//! execution.
 //!
 //! ## Background: vectorized volcano execution model
 //! DataFusion uses a batched volcano model. For most operators, output is
@@ -33,7 +34,7 @@
 //! ## Example
 //! For a join with an expensive, selective predicate:
 //! ```text
-//! LimitExec(fetch=10)
+//! GlobalLimitExec: skip=0, fetch=10
 //! -- NestedLoopJoinExec(on=expr_expensive_and_selective)
 //! --- DataSourceExec()
 //! --- DataSourceExec()
@@ -78,14 +79,13 @@ use datafusion_physical_plan::{ExecutionPlan, ExecutionPlanProperties};
 pub struct LimitPushdown {}
 
 /// This is a "data class" we use within the [`LimitPushdown`] rule to push
-/// down [`LimitExec`] in the plan. GlobalRequirements are hold as a rule-wide state
+/// down limits in the plan. GlobalRequirements are hold as a rule-wide state
 /// and holds the fetch and skip information. The struct also has a field named
 /// satisfied which means if the "current" plan is valid in terms of limits or not.
 ///
 /// For example: If the plan is satisfied with current fetch info, we decide to not add a LocalLimit
 ///
 /// [`LimitPushdown`]: crate::limit_pushdown::LimitPushdown
-/// [`LimitExec`]: crate::limit_pushdown::LimitExec
 #[derive(Default, Clone, Debug)]
 pub struct GlobalRequirements {
     fetch: Option<usize>,
@@ -125,51 +125,11 @@ impl PhysicalOptimizerRule for LimitPushdown {
     }
 }
 
-/// This enumeration makes `skip` and `fetch` calculations easier by providing
-/// a single API for both local and global limit operators.
-#[derive(Debug)]
-pub enum LimitExec {
-    Global(GlobalLimitExec),
-    Local(LocalLimitExec),
-}
-
-impl LimitExec {
-    fn input(&self) -> &Arc<dyn ExecutionPlan> {
-        match self {
-            Self::Global(global) => global.input(),
-            Self::Local(local) => local.input(),
-        }
-    }
-
-    fn fetch(&self) -> Option<usize> {
-        match self {
-            Self::Global(global) => global.fetch(),
-            Self::Local(local) => Some(local.fetch()),
-        }
-    }
-
-    fn skip(&self) -> usize {
-        match self {
-            Self::Global(global) => global.skip(),
-            Self::Local(_) => 0,
-        }
-    }
-
-    fn preserve_order(&self) -> bool {
-        match self {
-            Self::Global(global) => global.required_ordering().is_some(),
-            Self::Local(local) => local.required_ordering().is_some(),
-        }
-    }
-}
-
-impl From<LimitExec> for Arc<dyn ExecutionPlan> {
-    fn from(limit_exec: LimitExec) -> Self {
-        match limit_exec {
-            LimitExec::Global(global) => Arc::new(global),
-            LimitExec::Local(local) => Arc::new(local),
-        }
-    }
+struct LimitInfo {
+    input: Arc<dyn ExecutionPlan>,
+    fetch: Option<usize>,
+    skip: usize,
+    preserve_order: bool,
 }
 
 /// This function is the main helper function of the `LimitPushDown` rule.
@@ -184,26 +144,26 @@ pub fn pushdown_limit_helper(
     mut global_state: GlobalRequirements,
 ) -> Result<(Transformed<Arc<dyn ExecutionPlan>>, GlobalRequirements)> {
     // Extract limit, if exist, and return child inputs.
-    if let Some(limit_exec) = extract_limit(&pushdown_plan) {
+    if let Some(limit_info) = extract_limit(&pushdown_plan) {
         // If we have fetch/skip info in the global state already, we need to
         // decide which one to continue with:
         let (skip, fetch) = combine_limit(
             global_state.skip,
             global_state.fetch,
-            limit_exec.skip(),
-            limit_exec.fetch(),
+            limit_info.skip,
+            limit_info.fetch,
         );
         global_state.skip = skip;
         global_state.fetch = fetch;
-        global_state.preserve_order = limit_exec.preserve_order();
+        global_state.preserve_order = limit_info.preserve_order;
         global_state.satisfied = false;
 
         // Now the global state has the most recent information, we can remove
-        // the `LimitExec` plan. We will decide later if we should add it again
-        // or not.
+        // the limit node. We will decide later if we should add it again or
+        // not.
         return Ok((
             Transformed {
-                data: Arc::clone(limit_exec.input()),
+                data: limit_info.input,
                 transformed: true,
                 tnr: TreeNodeRecursion::Stop,
             },
@@ -253,7 +213,7 @@ pub fn pushdown_limit_helper(
             Ok((Transformed::no(pushdown_plan), global_state))
         } else if let Some(plan_with_fetch) = pushdown_plan.with_fetch(skip_and_fetch) {
             // This plan is combining input partitions, so we need to add the
-            // fetch info to plan if possible. If not, we must add a `LimitExec`
+            // fetch info to plan if possible. If not, we must add a limit node
             // with the information from the global state.
             let mut new_plan = plan_with_fetch;
             // Execution plans can't (yet) handle skip, so if we have one,
@@ -341,32 +301,45 @@ pub(crate) fn pushdown_limits(
 
     // Apply pushdown limits in children
     let children = new_node.data.children();
+    let mut changed = false;
     let new_children = children
         .into_iter()
-        .map(|child| {
-            pushdown_limits(Arc::<dyn ExecutionPlan>::clone(child), global_state.clone())
+        .map(|child: &Arc<dyn ExecutionPlan>| {
+            let new_child = pushdown_limits(
+                Arc::<dyn ExecutionPlan>::clone(child),
+                global_state.clone(),
+            )?;
+            // Tracking if any of the children changed
+            changed |= !Arc::ptr_eq(child, &new_child);
+            Ok(new_child)
         })
         .collect::<Result<_>>()?;
-    new_node.data.with_new_children(new_children)
+
+    if changed {
+        new_node.data.with_new_children(new_children)
+    } else {
+        Ok(new_node.data)
+    }
 }
 
-/// Transforms the [`ExecutionPlan`] into a [`LimitExec`] if it is a
+/// Extracts limit information from the [`ExecutionPlan`] if it is a
 /// [`GlobalLimitExec`] or a [`LocalLimitExec`].
-fn extract_limit(plan: &Arc<dyn ExecutionPlan>) -> Option<LimitExec> {
+fn extract_limit(plan: &Arc<dyn ExecutionPlan>) -> Option<LimitInfo> {
     if let Some(global_limit) = plan.as_any().downcast_ref::<GlobalLimitExec>() {
-        Some(LimitExec::Global(GlobalLimitExec::new(
-            Arc::clone(global_limit.input()),
-            global_limit.skip(),
-            global_limit.fetch(),
-        )))
+        Some(LimitInfo {
+            input: Arc::clone(global_limit.input()),
+            fetch: global_limit.fetch(),
+            skip: global_limit.skip(),
+            preserve_order: global_limit.required_ordering().is_some(),
+        })
     } else {
         plan.as_any()
             .downcast_ref::<LocalLimitExec>()
-            .map(|local_limit| {
-                LimitExec::Local(LocalLimitExec::new(
-                    Arc::clone(local_limit.input()),
-                    local_limit.fetch(),
-                ))
+            .map(|local_limit| LimitInfo {
+                input: Arc::clone(local_limit.input()),
+                fetch: Some(local_limit.fetch()),
+                skip: 0,
+                preserve_order: local_limit.required_ordering().is_some(),
             })
     }
 }

datafusion/physical-optimizer/src/limited_distinct_aggregation.rs

@@ -54,16 +54,8 @@ impl LimitedDistinctAggregation {
         }
 
         // We found what we want: clone, copy the limit down, and return modified node
-        let new_aggr = AggregateExec::try_new(
-            *aggr.mode(),
-            aggr.group_expr().clone(),
-            aggr.aggr_expr().to_vec(),
-            aggr.filter_expr().to_vec(),
-            aggr.input().to_owned(),
-            aggr.input_schema(),
-        )
-        .expect("Unable to copy Aggregate!")
-        .with_limit_options(Some(LimitOptions::new(limit)));
+        let new_aggr = aggr.with_new_limit_options(Some(LimitOptions::new(limit)));
+
         Some(Arc::new(new_aggr))
     }
 

datafusion/physical-optimizer/src/output_requirements.rs

@@ -430,7 +430,14 @@ fn require_top_ordering_helper(
         // be responsible for (i.e. the originator of) the global ordering.
         let (new_child, is_changed) =
             require_top_ordering_helper(Arc::clone(children.swap_remove(0)))?;
-        Ok((plan.with_new_children(vec![new_child])?, is_changed))
+
+        let plan = if is_changed {
+            plan.with_new_children(vec![new_child])?
+        } else {
+            plan
+        };
+
+        Ok((plan, is_changed))
     } else {
         // Stop searching, there is no global ordering desired for the query.
         Ok((plan, false))