Repository navigation
perf: share one StatisticsContext across all physical optimizer rules - #26094
Merged
zhuqi-lucas merged 12 commits intoOct 13, 2026
Merged
Conversation
Create one StatisticsContext per optimize_physical_plan call and expose it through the new PhysicalOptimizerContext::statistics_context, so statistics computed by one rule are reused by later rules. JoinSelection, EnsureRequirements, AggregateStatistics and LimitPushdown use it. StatisticsContext stores its cache in a parking_lot::Mutex so it is Send + Sync, as PhysicalOptimizerContext requires.
github-actions
Bot
added
optimizer
Member
Author
|
@kosiew @zhuqi-lucas, could you run I figure you'd be interested in this PR as it builds on #25929, and it's the query-lifetime follow-up of what #25098 did for EnsureDistribution. |
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #26094 +/- ##
==========================================
- Coverage 82.80% 82.80% -0.01%
==========================================
Files 1149 1149
Lines 452234 452326 +92
Branches 452234 452326 +92
==========================================
+ Hits 374484 374558 +74
- Misses 54952 54965 +13
- Partials 22798 22803 +5 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
Contributor
|
run benchmark sql_planner |
Contributor
@asolimando triggered now |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing asolimando/query-lifetime-stats-cache (3179bd9) to db83fcc (merge-base) diff Run configurationrun benchmark sql_plannerResults will be posted here when complete File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Instance: Comparing asolimando/query-lifetime-stats-cache (3179bd9) to db83fcc (merge-base) diff Run configurationrun benchmark sql_plannerCPU Details (lscpu)Details
Resource Usagesql_planner — base (merge-base)
sql_planner — branch
File an issue against this benchmark runner |
Member
Author
|
Thanks @zhuqi-lucas! My understanding of the benchmark run:
The +5% on The gain is little smaller than the sf1 numbers I had locally (6%, as reported in the description) because the benchmark tables are empty, so each statistics computation is cheap. If you think this is interesting I can do another self-review pass on the PR and remove it from draft, wdyt? |
Contributor
|
Thanks for the detailed breakdown, your reading matches mine. This looks valuable, please go ahead with the self-review and take it out of draft. |
|
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 |
Member
Author
|
@zhuqi-lucas thanks a lot for your feedback, after self-review I have pushed some test refactoring, doc improvements and improved API ergonomics (added I noticed the coverage warning, I could fix that easily but I don't see anything that genuinely needs more coverage, happy to add more tests if you see fit. OT: is there a process to get into the allow-list to run benchmarks? What are the criteria for eligibility? |
asolimando
marked this pull request as ready for review
October 7, 2026 13:05
asolimando
commented
Oct 7, 2026
| /// [`Statistics`] only. | ||
| pub struct StatisticsContext { | ||
| cache: Rc<RefCell<StatsCache>>, | ||
| cache: Mutex<StatsCache>, |
Member
Author
There was a problem hiding this comment.
Note for reviewers: the Mutex is needed because PhysicalOptimizerContext: Send + Sync, and the shared context is held inside it and returned as &StatisticsContext. RefCell is not Sync (and the Rc was never cloned, so it was dropped rather than replaced).
The optimizer uses the context from one thread, so the lock is never contended, and it is briefly held (single map lookup/insert), never across the recursive walk. The compute_statistics benchmark shows no measurable difference from RefCell (all cases within ±4% of main, in both directions).
Alternatives I considered: a RefCell restricted to the creating thread (needs unsafe impl Sync), a thread-local cache (hidden global state), or removing Sync from PhysicalOptimizerContext (a breaking change). None of them seemed worth it for a lock that costs nothing measurable.
asolimando
commented
Oct 7, 2026
| // Otherwise, fall back to standard distribution rather | ||
| // than choosing an arbitrary reference. | ||
| candidates | ||
| let best_satisfied_child: Option<(usize, Partitioning)> = |
Member
Author
There was a problem hiding this comment.
Note to reviewers: re-indented by rustfmt, the real change is passing stats_ctx to PlanSize::from_plan, but the diff can be confusing.
zhuqi-lucas
reviewed
Oct 8, 2026
zhuqi-lucas
left a comment
zhuqi-lucas
left a comment
Contributor
There was a problem hiding this comment.
Nice work — this is the piece #25098 and #25929 were building toward, and the numbers are convincing.
I checked the two things that worried me most and both hold up:
- No lock held across the recursive walk. Every
borrow()→lock()site is a single lookup or insert whose guard dies at the end of the statement, so the non-reentrantparking_lot::Mutexcan't self-deadlock. - Pointer keys stay safe over the much longer sharing window.
store_cache_entryis the only insertion point and always clones the owner intoCacheEntry::_plan, so an address can't be recycled while its entry lives. That invariant is what makes widening the scope from one rule to the whole run sound — worth keeping in mind for anyone who later adds a second insertion path.
Also worth calling out, since it's easy to miss in the diff: the compute(x.as_ref()) → compute_arc(x) switches are not cosmetic. Borrowed roots are deliberately not memoized, so those lines are what actually lets the root of each traversal land in the shared cache.
Three small notes inline, none blocking.
… helper, ConfigOnlyContext invariant
github-actions
Bot
added
the
documentation
Contributor
|
run benchmark sql_planner |
zhuqi-lucas
approved these changes
Oct 8, 2026
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing asolimando/query-lifetime-stats-cache (63aa427) to 3ed377a (merge-base) diff Run configurationrun benchmark sql_plannerResults will be posted here when complete File an issue against this benchmark runner |
zhuqi-lucas
reviewed
Oct 8, 2026
|
🤖 Benchmark completed (GKE) | trigger Instance: Comparing asolimando/query-lifetime-stats-cache (63aa427) to 3ed377a (merge-base) diff Run configurationrun benchmark sql_plannerCPU Details (lscpu)Details
Resource Usagesql_planner — base (merge-base)
sql_planner — branch
File an issue against this benchmark runner |
Member
Author
|
Thanks @zhuqi-lucas for re-running the benchmark! Here is how I read the second run:
On TPC-DS we save less time than in the first run (13 ms instead of 25 ms). I checked the commits merged to WDYT? |
jayzhan211
self-requested a review
October 9, 2026 00:16
jayzhan211
reviewed
Oct 9, 2026
jayzhan211
left a comment
jayzhan211
left a comment
Contributor
There was a problem hiding this comment.
Thanks @asolimando , 1 suggestion
…tract for providers
Contributor
|
Thanks for the /// `LimitPushdown` trusts an `Exact(0)` row count from a registered
/// provider: it removes a limit over that input
#[tokio::test]
async fn limit_pushdown_consults_statistics_providers() -> Result<()> {
let provider = ClosureStatisticsProvider::with_matches(
|plan| plan.is::<AggregateExec>(),
|plan, child_stats| {
let child_stats = child_stats
.iter()
.map(|c| Arc::clone(c.base_arc()))
.collect::<Vec<_>>();
let mut stats = Arc::unwrap_or_clone(
plan.statistics_from_inputs(&child_stats, &StatisticsArgs::new())?,
);
stats.num_rows = Precision::Exact(0);
Ok(StatisticsResult::Computed(stats.into()))
},
);
let state = SessionStateBuilder::new()
.with_default_features()
.with_config(SessionConfig::new().with_target_partitions(1))
.with_statistics_registry(StatisticsRegistry::with_providers(vec![Arc::new(
provider,
)]))
.build();
let ctx = SessionContext::new_with_state(state);
ctx.sql("CREATE TABLE t AS VALUES (1), (2), (3)")
.await?
.collect()
.await?;
let batches = ctx
.sql("SELECT column1, count(*) FROM t GROUP BY column1 LIMIT 1")
.await?
.collect()
.await?;
assert_eq!(batches.iter().map(|b| b.num_rows()).sum::<usize>(), 3);
Ok(())
} |
Member
Author
Thanks @jayzhan211! I have added your I hope you don't mind me moving both provider tests next to the other tests of each rule ( |
Contributor
|
run benchmark sql_planner |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing asolimando/query-lifetime-stats-cache (247230b) to b40ef06 (merge-base) diff Run configurationrun benchmark sql_plannerResults will be posted here when complete File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Instance: Comparing asolimando/query-lifetime-stats-cache (247230b) to b40ef06 (merge-base) diff Run configurationrun benchmark sql_plannerCPU Details (lscpu)Details
Resource Usagesql_planner — base (merge-base)
sql_planner — branch
File an issue against this benchmark runner |
Contributor
|
run benchmark sql_planner |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing asolimando/query-lifetime-stats-cache (6493ea3) to de06656 (merge-base) diff Run configurationrun benchmark sql_plannerResults will be posted here when complete File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Instance: Comparing asolimando/query-lifetime-stats-cache (6493ea3) to de06656 (merge-base) diff Run configurationrun benchmark sql_plannerCPU Details (lscpu)Details
Resource Usagesql_planner — base (merge-base)
sql_planner — branch
File an issue against this benchmark runner |
Member
Author
|
Thanks @zhuqi-lucas for this extra round of benchmarking. TPC-DS planning has shown improvements in all runs ranging from 2% to 4%, which is good but it is also likely a lower bound, because the benchmark tables are empty (each statistics computation is cheap, as discussed above). The other benchmarks aren't really affected and are within noise within noise: the logical planning benchmarks, which this PR does not touch, move by a similar amount. Unless @jayzhan211 has further comments, I think this is good to go, but happy to improve it further if needed. |
Contributor
|
Thanks @asolimando , i will merge this PR if no new comment in one or two days. |
jayzhan211
approved these changes
Oct 12, 2026
github-merge-queue
Bot
removed this pull request from the merge queue due to no response for status checks
Oct 12, 2026
Contributor
|
Thanks @asolimando and @jayzhan211 , merged now. |
jayzhan211
added a commit
to CuteChuanChuan/datafusion
that referenced
this pull request
Oct 31, 2026
…pache#26660) ## Which issue does this PR close? - No issue. This is a refactor; query results and plans of DataFusion's own operators are unchanged. ## Rationale for this change Four optimizer rules move a `LIMIT`/`OFFSET` (or a fetch) through the physical plan: `LimitPushdown`, `LimitPushPastWindows`, the sort pushdown inside `EnsureRequirements`, and `PushdownSort`. Each rule answers the same three questions with its own code: - **How do `skip` and `fetch` combine?** Each rule adds `skip + fetch` itself. `LimitPushdown` adds them unchecked, while the sort pushdown in `EnsureRequirements` and `LimitPushPastWindows` saturate (the latter since apache#26662). - **Can a limit go below this node?** There are three different checks: - `supports_limit_pushdown()` alone. - The same plus `CardinalityEffect::Equal`. - The same plus `Equal | GreaterEqual`, with a hard-coded type list. They differ because a rule can *copy* a limit below a node as an early-stop hint while keeping it above, or *move* it below and drop it above. Moving needs more from the node than copying. Nothing in the code names that difference. - **What does `LimitPushdown` do at each node?** It removes every limit node and carries `(fetch, skip, satisfied)` down. It decides later whether to re-add the limit, and clears `skip` before recursing. The decision is spread over about 150 lines of nested `if`s. Several recent fixes came from one of these places disagreeing with another: apache#24958 (skip ignored when a sort moved below a limit), apache#22525 (OFFSET applied twice), apache#21744 and apache#24809 (fetch lost when an operator was removed). ## What changes are included in this PR? Three main commits, each of which builds and passes the tests on its own, followed by review fixes: 1. **`RowBound`** (crate-private, `physical-optimizer/src/row_bound.rs`). This is `OFFSET skip LIMIT fetch` as a value: combining two bounds, `skip + fetch` (saturating), the tighter of two fetches, and the one rule for choosing between a `GlobalLimitExec` and a `LocalLimitExec`. The sort pushdown in `EnsureRequirements` uses it instead of its own arithmetic. `LimitPushPastWindows` uses it to read limit nodes. 2. **`LimitPushdown` as one walk.** The bound a node receives from its ancestors is a small enum: - `Nothing`. - `Open(bound)`: no ancestor enforces it yet, so it must be placed at or above the first node it can't pass. - `Enforced(fetch)`: an ancestor enforces it, so nodes below may stop early. Each node is one of `Limit`, `Pass`, `Merge` or `Stop`. The rule is a single `match` over the two (`apply_bound` in `limit_pushdown.rs`), which replaces the `satisfied` flag and the skip-clearing step. `GlobalRequirements`, `pushdown_limit_helper` and `pushdown_limit_helper_with_stats` are deprecated; they now run on the new walk. 3. **`ExecutionPlan::limit_behavior()`.** Operators declare what a limit on their output means for their input: - `Stops`. - `Limit { skip, fetch }`. - `CopyToChildren` (Union). - `MoveToInput` (Projection). - `Merge` (CoalescePartitions, SortPreservingMerge). - `Window(lookahead)`. The copy/move difference is documented once, on the enum. So is the rule that a merge with its own fetch (`ExecutionPlan::fetch()`) already selects rows: a limit above it can't move below it. Before, that rule was only implied by the merge's `cardinality_effect()` returning `LowerEqual` when it has a fetch. The three rules match on it instead of combining `supports_limit_pushdown()` with `cardinality_effect()` and type checks. The window's lookahead (frame end plus `LEAD` offsets) is now computed by `BoundedWindowAggExec` rather than by the rule. `supports_limit_pushdown` is deprecated. The default `limit_behavior()` is derived from it, so operators outside DataFusion keep their current behavior. `LimitPushdown` moves a limit below a `CopyToChildren` or `MoveToInput` node only when each output partition is one input partition. A node whose children have more partitions in total than it outputs takes the limit like a merge: the limit stays above it and is copied below. No DataFusion operator is such a node, so this only fixes custom operators (see below). The fetch handling in `PushdownSort` and the other "a fetched operator is a row-selection boundary" checks in `EnsureRequirements` are unchanged. ## What is the testing strategy for this PR? The goal is no change in results or plans, so the existing tests are the check. All of these pass with **no change to expected output**: - every sqllogictest file - `core_integration` `physical_optimizer::` (663 tests: limit pushdown, sort pushdown, enforce sorting, window) - the `limit_fuzz`, `window_fuzz`, `sort_fuzz`, `sort_query_fuzz` and `topk_filter_pushdown` fuzz tests - `datafusion-physical-optimizer` and `datafusion-physical-plan` unit tests New tests: - `row_bound.rs` unit tests: combining bounds, saturation at `usize::MAX`, and the Global/Local choice. - `test_fetched_merge_keeps_stricter_topk_above` in `core/tests/physical_optimizer/ensure_requirements.rs`: a TopK over a `SortPreservingMergeExec` that has its own fetch keeps its fetch above the merge. It passes on main, and fails if a fetched merge is treated as `MoveToInput`. - Three tests in `core/tests/physical_optimizer/limit_pushdown.rs` with a custom operator that combines partitions into one: - declared `Merge` but can't take a fetch: under `OFFSET 5 LIMIT 3`, its input gets `fetch=8`, not 3; - declared `CopyToChildren` with two inputs, or `MoveToInput` with one 8-partition input: the limit stays above it. Benchmarks: the plans are identical, so only planning time can change. Measured on Apple silicon against the merge base: - EXPLAIN output for LIMIT/OFFSET, window, UNION and TopK queries is identical, and planning them takes the same time (1.01x). - `sql_planner` `physical_plan_*` (TPC-H, TPC-DS, ClickBench; 84 benchmarks): median +0.7% on this branch, while re-running main against itself gave +3.2%, so the difference is within noise. TPC-H q3 and q4, the largest differences, were 3-15% faster on this branch when re-run. ## Are there any user-facing changes? No change to query results or plans for DataFusion's own operators. API (upgrade guide updated): - New `ExecutionPlan::limit_behavior()`, `LimitBehavior` and `WindowLookahead`. - `ExecutionPlan::supports_limit_pushdown()` is deprecated. Only the default `limit_behavior()` reads it. DataFusion's own operators keep returning the same values, so a wrapper that forwards it to its input keeps working. - `limit_pushdown::{GlobalRequirements, pushdown_limit_helper}` are deprecated. They now run on the new walk. A caller that drives them one node at a time no longer applies an `OFFSET` twice when a limit node with only an `OFFSET` sits above another limit. - `limit_pushdown::pushdown_limit_helper_with_stats` is removed. It was added in this release cycle (apache#26094) and is in no release. - Two differences for custom operators that still implement `supports_limit_pushdown()`: - A single-input operator that returns `true` and whose `cardinality_effect()` is `Unknown` or `LowerEqual` can now have a window's grown limit pushed below it. `Unknown` is the default, so this includes a wrapper that forwards `supports_limit_pushdown()` and does not override `cardinality_effect()`. `LimitPushdown` already pushes plain limits through such operators. - An operator that returns `true` and combines its input partitions (for example several inputs into one partition) now keeps a limit above it. On main, `LimitPushdown` moved the limit below it, which kept that many rows of each input partition. - `skip + fetch` now saturates at `usize::MAX` in `LimitPushdown`, as it already does in the other two rules. SQL can't reach the overflow there, because `LIMIT` and `OFFSET` are each at most `i64::MAX`.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Which issue does this PR close?
Rationale for this change
Each physical optimizer rule that reads statistics creates its own
StatisticsContext, so the statistics of the same plan nodes are computed again in every rule.JoinSelectioncreates a new context for everyget_statscall, so it has no caching at all, even within its own pass.#25098 showed the gain of sharing one context inside
EnsureDistribution. #25929 made cache entries keep the plan node they were computed for, so a context can now be shared safely across plan rewrites. This PR uses that to share oneStatisticsContextacross all the rules of oneoptimize_physical_plancall.On TPC-DS (98 queries) and TPC-H (21 queries), sf1 Parquet, planned with a local harness (not part of this PR):
mainThe largest gain is TPC-DS q64: cache misses go from 3,785 to 531 and planning time drops by about 21%.
What changes are included in this PR?
StatisticsContextstores its cache in aparking_lot::Mutexinstead ofRc<RefCell>, so it isSend + Sync. This is needed becausePhysicalOptimizerContext: Send + Sync. The lock is held only for single map lookups and inserts, never across the recursive walk. Thecompute_statisticsbenchmark shows no difference frommain(all cases within ±4%, in both directions).PhysicalOptimizerContext::statistics_context(), which returnsNoneby default.DefaultPhysicalPlannercreates one context peroptimize_physical_plancall, built from the session's statistics registry, and returns it to every rule.ConfigOnlyContextalso owns aStatisticsContext, so a rule called throughoptimize()shares one cache for its whole pass.JoinSelection,EnsureRequirements(includingPlanSize::from_plan),AggregateStatisticsandLimitPushdownuse the shared context. When a context does not share one (for example, a context received through FFI), the rules create a new context from the statistics registry, as they did before.PhysicalOptimizerContext::compute_statistics, which custom rules can call on the context they already receive. Without it, a custom rule would have to repeat the fallback itself (use the shared context, otherwise build one fromstatistics_registry()), and the obvious shortcut,StatisticsContext::new(), does not consult the registered statistics providers. With it, custom rules behave exactly like the built-in ones.with_statistics_context. Rules that walk the plan (EnsureRequirements,LimitPushdown) use it to get oneStatisticsContextfor the whole traversal, so a fallback context still caches within the pass.pushdown_limit_helperis deprecated in favour of the newpushdown_limit_helper_with_stats, which takes a&StatisticsContext. This is the same pattern asensure_distribution_with_stats.optimize_physical_planreturns, including nodes that a later rule replaces. Only the rules that read statistics add entries, and the memory is released at the end of physical planning.What is the testing strategy for this PR?
optimizer_rules_share_statistics_contexttest inphysical_planner.rs: two rules compute the root statistics through the shared context, and the second one gets theArccached by the first.Are there any user-facing changes?
No breaking changes.
PhysicalOptimizerContext::statistics_context()andPhysicalOptimizerContext::compute_statistics().datafusion_session::with_statistics_contextandpushdown_limit_helper_with_stats;pushdown_limit_helperis deprecated. All are described in the 56.0.0 upgrade guide.StatisticsContextis nowSend + Sync.AggregateStatistics,LimitPushdownandPlanSize::from_plannow consult the session's statistics providers, asJoinSelectionandEnsureRequirementsalready do. Without registered providers (the default), their results are unchanged.Disclaimer: I used AI to assist in the code generation, I have manually reviewed the output and it matches my intention and understanding.