[branch-52] perf: skip ensure_distribution rebuild when children are unchanged

[zhuqi-lucas]

Summary

Cherry-pick of upstream apache/datafusion#22521 (approved by 2010YOUY01 and alamb, awaiting maintainer merge) into our branch-52 fork.

ensure_distribution was calling plan.with_new_children(children_plans) unconditionally after collecting the (possibly redistributed) children, even when none of the children were actually replaced. For nodes like ProjectionExec that path recomputes schema / equivalence properties / output ordering / partitioning via try_new — pure overhead when child Arcs are pointer-identical.

Routes the non-UnionExec branch through the existing with_new_children_if_necessary helper, which does the Arc::ptr_eq short-circuit and reuses the input plan when no children changed.

Why pick now

Linear X-2631. Atlas reference-cluster query servers run on branch-52-pinned DataFusion and hit the slow path on every point-query plan node. Profiling on NY5 prod showed ProjectionExec::with_new_children as the single largest cost inside ensure_distribution (1.94s of a 2.87s rule-time slice in a 60s CPU sample).

Adaptations from upstream

• Upstream lives at datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs; branch-52 keeps the flat path datafusion/physical-optimizer/src/enforce_distribution.rs. Same logic, different module path.
• Upstream uses plan.is::<UnionExec>(); branch-52 still has plan.as_any().is::<UnionExec>(). Adjusted.

Affects

• Staging: query servers will see the fast path after a DF rev bump (separate PR).
• Production: same, after staging soak.

Tests

cargo test --release -p datafusion --test core_integration -- physical_optimizer::enforce_distribution — 63/63 pass (62 pre-existing + 1 new regression test ensure_distribution_reuses_plan_arc_when_no_redistribution_needed).

Micro-benchmark from upstream PR

Plan shape: 30-deep ProjectionExec stack over a sorted parquet scan, 5000 iterations.

• Without fix: 852.74 ms total, 170.55 us/call
• With fix: 296.81 ms total, 59.36 us/call
• ~2.87x speedup, -65% CPU per call

[Copilot]

Pull request overview

Cherry-pick of upstream apache#22521 into branch-52 that avoids unnecessary with_new_children rebuilds in ensure_distribution when no child plan was actually replaced, by routing through the existing with_new_children_if_necessary helper (which short-circuits via Arc::ptr_eq). This eliminates redundant schema / equivalence / ordering / partitioning recomputation (notably for ProjectionExec) on plan shapes where no redistribution is needed.

Changes:

• Switch the non-UnionExec/non-InterleaveExec branch in ensure_distribution from plan.with_new_children(children_plans)? to with_new_children_if_necessary(plan, children_plans)?.
• Add the with_new_children_if_necessary import from datafusion_physical_plan.
• Add a regression test verifying that a deep ProjectionExec chain over a single-partition scan returns the same Arc from ensure_distribution.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated no comments.

File	Description
datafusion/physical-optimizer/src/enforce_distribution.rs	Route the common branch through with_new_children_if_necessary so unchanged children short-circuit the rebuild; add corresponding import and explanatory comment.
datafusion/core/tests/physical_optimizer/enforce_distribution.rs	New regression test ensure_distribution_reuses_plan_arc_when_no_redistribution_needed asserting Arc::ptr_eq between input and output for a no-op-distribution projection stack.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.