[branch-53] 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-53 fork. Companion to #57 (branch-52).

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. 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-53 keeps the flat path. Same logic, different module path.
• Upstream uses plan.is::<UnionExec>(); branch-53 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

This PR cherry-picks upstream DataFusion change apache#22521 into branch-53 to reduce overhead in the EnforceDistribution optimizer rule by avoiding unnecessary plan rebuilds when child plans are unchanged.

Changes:

• Update ensure_distribution to use with_new_children_if_necessary for non-UnionExec nodes, reusing the original Arc<dyn ExecutionPlan> when children are pointer-identical.
• Add a regression test asserting ensure_distribution returns the exact same plan Arc when no redistribution is required (deep ProjectionExec chain, target_partitions = 1).

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	Uses with_new_children_if_necessary to skip expensive with_new_children rebuilds when children are unchanged.
datafusion/core/tests/physical_optimizer/enforce_distribution.rs	Adds regression coverage ensuring the optimizer reuses the input plan Arc when no redistribution occurs.

---

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