perf: reuse one StatisticsContext across ensure_distribution - #25098
perf: reuse one StatisticsContext across ensure_distribution#25098zhuqi-lucas wants to merge 1 commit into
Conversation
There was a problem hiding this comment.
🟡 Changes recommended
It introduces a breaking change to the public ensure_distribution function signature, which is re-exported and may be used by downstream crates.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Pull request overview
This PR improves physical-optimizer planning performance by reusing a single StatisticsContext (and its memoization cache) across an ensure_distribution traversal, avoiding repeated recursive statistics computation on deep/wide plans.
Changes:
- Create one
StatisticsContextfor the distribution-enforcement pass and reset its cache only when a node’s plan pointer actually changes. - Thread
&StatisticsContextthroughensure_distributionand intoget_repartition_requirement_statusso child stats are memoized across the pass. - Add a deep operator-stack regression test asserting the repartition decisions remain unchanged.
File summaries
| File | Description |
|---|---|
| datafusion/physical-optimizer/src/ensure_requirements/mod.rs | Reuses a single StatisticsContext during the bottom-up distribution pass and resets cache on actual plan-pointer rewrites. |
| datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs | Threads &StatisticsContext into distribution enforcement logic to share memoized statistics across children/subtrees. |
| datafusion/core/tests/physical_optimizer/enforce_distribution.rs | Updates call sites for the new ensure_distribution signature and adds a deep-stack test for unchanged repartition behavior. |
Review details
- Files reviewed: 3/3 changed files
- Comments generated: 1
- Review effort level: Lite
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
| pub fn ensure_distribution( | ||
| dist_context: DistributionContext, | ||
| config: &ConfigOptions, | ||
| stats_ctx: &StatisticsContext, | ||
| ) -> Result<Transformed<DistributionContext>> { |
|
Thank you for opening this pull request! Reviewer note: cargo-semver-checks reported the current version number is not SemVer-compatible with the changes in this pull request (compared against the base branch). Details |
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #25098 +/- ##
==========================================
- Coverage 81.80% 81.80% -0.01%
==========================================
Files 1130 1130
Lines 417708 417722 +14
Branches 417708 417722 +14
==========================================
+ Hits 341701 341708 +7
- Misses 55878 55881 +3
- Partials 20129 20133 +4 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
b23e2ce to
a6f3f81
Compare
get_repartition_requirement_status created a fresh StatisticsContext per child. StatisticsContext::compute recurses the child's whole subtree and carries a pointer-keyed memoization cache its docstring describes as a per-call cache meant to be reused across a walk, so allocating a new context per child discards it every time. A single ensure_distribution pass then recomputes shared subtree statistics O(depth) times. Share one StatisticsContext across the whole ensure_distribution transform_up. The cache is keyed by raw node pointer; ensure_distribution reports Transformed::yes unconditionally, so the reset is keyed on whether the node's plan pointer actually changed (a rewrite can free a cached node and a later allocation could reuse its address). Nodes that make no change cannot free anything, so the cache safely persists across no-op nodes, which dominate a deep plan. Adds ensure_distribution_shares_statistics_cache, which counts a leaf's statistics computations under a stack of pass-through operators and asserts the shared cache saves progressively more as the stack deepens (a cache that is not actually shared saves nothing). No plan changes.
a6f3f81 to
a2cb2c8
Compare
|
Hi @zhuqi-lucas, I can confirm that your use-case fits what the cache in Glad to read that you are seeing a drop in planning time of ~10% due to caching! We can probably push this a little further: instead of just passing
I have implemented something similar in #24716 for So, concretely, my suggestion is to:
This would keep the caching fix from this PR, and also allow the rule to benefit from the present and future statistics context (I will be working on #21120 as my next task). Note that there is no behavior change by default: a session without registered providers is unaffected, but it gives room for improvement without further breaking changes. I am off this week with very limited access to a computer, but I can surely offer a review or help with a PR from next week! Regarding the cache reset: your proposed solution seems safe to me, but in the future I'd like to make the cache more robust so consumers don't have to think about it. Either by introducing a unique id per constructed |
Which issue does this PR close?
None filed; small self-contained perf fix. Rationale below.
Rationale for this change
get_repartition_requirement_statuscreates a freshStatisticsContext::new()once per child.StatisticsContext::computerecurses the child's whole subtree and carries a pointer-keyed memoization cache its own docstring describes as a "per-call memoization cache" meant to be reused across a walk. Allocating a new context per child discards that cache every time, so a singleensure_distributionpass recomputes shared subtree statisticsO(depth)times.On a deep/wide plan this is measurable. In our deployment (
EnsureRequirementsruns several times over a ~200-node plan) sharing the cache cut physical planning by ~10% with no plan change.What changes are included in this PR?
StatisticsContextthrough theensure_distributiontransform_up(pass&StatisticsContextintoget_repartition_requirement_status) so each subtree's statistics are computed once per pass.StatsCacheis keyed by raw node pointer, andensure_distributionreturnsTransformed::yesunconditionally, so the cache reset is keyed on whether the node's plan pointer actually changed (Arc::ptr_eqbefore/after). A node that changed may have freed a cached child (which would make a stale pointer key unsafe); a node that made no change cannot, so the cache safely persists across the no-op nodes that dominate a deep plan.(A second per-child
StatisticsContext::new()inPlanSize::from_plan/enforce_distribution_relationshipscan get the same treatment; left as a follow-up to keep this PR focused.)Are these changes tested?
Yes. New test
ensure_distribution_shares_statistics_cacheputs a leaf that counts its own statistics computations under a stack of pass-through operators, runs the distribution pass with a shared context vs a fresh-per-node context, and asserts the shared cache saves progressively more as the stack deepens. A cache that is not actually shared (e.g. reset on every node) saves nothing and fails the test — which a plan-output assertion cannot catch, since the optimized plan is identical either way.Existing suites remain green and unchanged:
datafusion --test core_integration physical_optimizer(569 passed) anddatafusion-physical-planstatistics tests (96 passed).Are there any user-facing changes?
No. Internal physical-optimizer performance only; planner output is identical.