Skip to content

perf: share one StatisticsContext across all physical optimizer rules #26094

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
asolimando wants to merge 8 commits into
base: main
Choose a base branch
from
+359 −140
Open

perf: share one StatisticsContext across all physical optimizer rules #26094

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
3179bd9
perf: share one StatisticsContext across all physical optimizer rules
asolimando Oct 6, 2026
3dd8da1
test: record shared statistics in the existing ContextCheckingRule
asolimando Oct 7, 2026
784fca0
Deprecate pushdown_limit_helper and document the shared StatisticsCon…
asolimando Oct 7, 2026
ae341dc
Make with_statistics_context public for custom rules
asolimando Oct 7, 2026
5a4626e
Add PhysicalOptimizerContext::compute_statistics for custom rules
asolimando Oct 7, 2026
c101751
Address review: reset_cache guidance, provider note on the deprecated…
asolimando Oct 8, 2026
63aa427
Merge branch 'main' into asolimando/query-lifetime-stats-cache
zhuqi-lucas Oct 8, 2026
f7587af
Document that cached plan nodes must not change their statistics in p…
asolimando Oct 8, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
84 changes: 65 additions & 19 deletions datafusion/core/src/physical_planner.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -64,6 +64,7 @@ use arrow_schema::Field;
use datafusion_catalog::ScanArgs;
use datafusion_common::Column;
use datafusion_common::HashMap as DFHashMap;
use datafusion_common::config::ConfigOptions;
use datafusion_common::display::ToStringifiedPlan;
use datafusion_common::format::ExplainAnalyzeCategories;
use datafusion_common::tree_node::{Transformed, TreeNode, TreeNodeRecursion};
Expand Down Expand Up @@ -103,9 +104,11 @@ use datafusion_physical_expr::{
use datafusion_physical_plan::empty::EmptyExec;
use datafusion_physical_plan::execution_plan::InvariantLevel;
use datafusion_physical_plan::joins::PiecewiseMergeJoinExec;
use datafusion_physical_plan::operator_statistics::StatisticsRegistry;
use datafusion_physical_plan::placeholder_row::PlaceholderRowExec;
use datafusion_physical_plan::recursive_query::RecursiveQueryExec;
use datafusion_physical_plan::scalar_subquery::{ScalarSubqueryExec, ScalarSubqueryLink};
use datafusion_physical_plan::statistics::StatisticsContext;
use datafusion_physical_plan::unnest::ListUnnest;
use datafusion_session::{PhysicalOptimizerContext, PhysicalOptimizerRule, Session};

Expand All @@ -123,18 +126,34 @@ pub use datafusion_session::{ExtensionPlanner, PhysicalPlanner};

struct SessionOptimizerContext<'a> {
session: &'a dyn Session,
statistics_context: StatisticsContext,
}

impl<'a> SessionOptimizerContext<'a> {
fn new(session: &'a dyn Session) -> Self {
let statistics_context = match session.statistics_registry() {
Some(registry) => StatisticsContext::new_with_registry(registry.clone()),
None => StatisticsContext::new(),
};
Self {
session,
statistics_context,
}
}
}

impl PhysicalOptimizerContext for SessionOptimizerContext<'_> {
fn config_options(&self) -> &datafusion_common::config::ConfigOptions {
fn config_options(&self) -> &ConfigOptions {
self.session.config_options()
}

fn statistics_registry(
&self,
) -> Option<&datafusion_physical_plan::operator_statistics::StatisticsRegistry> {
fn statistics_registry(&self) -> Option<&StatisticsRegistry> {
self.session.statistics_registry()
}

fn statistics_context(&self) -> Option<&StatisticsContext> {
Some(&self.statistics_context)
}
}

/// Default single node physical query planner that converts a
Expand Down Expand Up @@ -3123,9 +3142,7 @@ impl DefaultPhysicalPlanner {
InvariantChecker(InvariantLevel::Always).check(&plan)?;

let mut new_plan = Arc::clone(&plan);
let optimizer_context = SessionOptimizerContext {
session: session_state,
};
let optimizer_context = SessionOptimizerContext::new(session_state);

@zhuqi-lucas zhuqi-lucas Oct 8, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

One more thing this widens, worth stating somewhere: the cache is now exposed to in-place mutation for the whole planning run, not just one rule.

A pointer-keyed hit is only valid because a plan node never changes behind its Arc. That held trivially when each rule had its own context — anything a rule did produced new nodes. Now an entry computed by the first rule is still served to the last one, so a node that mutated its own statistics through interior mutability without changing identity would be read stale.

Nothing does that today (planning-time rules all rebuild nodes, and dynamic filters are updated at execution time, after this context is dropped), so this isn't a bug — but it's an invariant the design now leans on much harder. A line on StatisticsContext saying cached nodes must be immutable for the context's lifetime would make it checkable in review.

@asolimando asolimando Oct 8, 2026 •
edited
Loading

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Agreed, it is an invariant the design now relies on much more. Added a paragraph in f7587af to the StatisticsContext docs: a plan node must not change its statistics in place while a context holds it, otherwise the cache returns stale values; optimizer rules satisfy this because they replace nodes instead of changing them.

for optimizer in optimizers {
let before_schema = new_plan.schema();
new_plan = optimizer
Expand Down Expand Up @@ -3536,6 +3553,7 @@ mod tests {
use arrow::datatypes::{DataType, Field, Int32Type};
use arrow_schema::{FieldRef, SchemaRef};
use datafusion_catalog::CatalogProviderList;
use datafusion_common::Statistics;
use datafusion_common::config::{ConfigOptions, TableOptions};
use datafusion_common::{
DFSchemaRef, ScalarValue, SplitPoint, TableReference, ToDFSchema as _,
Expand All @@ -3558,12 +3576,16 @@ mod tests {
use datafusion_functions_aggregate::expr_fn::sum;
use datafusion_physical_expr::EquivalenceProperties;
use datafusion_physical_plan::execution_plan::{Boundedness, EmissionType};
use datafusion_physical_plan::operator_statistics::StatisticsRegistry;
use datafusion_physical_plan::statistics::StatisticsArgs;
use datafusion_physical_plan::{ChildrenPropertiesMode, ReplaceChildrenOptions};
use datafusion_session::QueryPlanner;
use parking_lot::Mutex as SyncMutex;

/// Records the root statistics computed with the shared statistics context
#[derive(Debug)]
struct ContextCheckingRule {
invoked: Arc<AtomicBool>,
recorded: Arc<SyncMutex<Vec<Arc<Statistics>>>>,
}

impl PhysicalOptimizerRule for ContextCheckingRule {
Expand All @@ -3581,7 +3603,12 @@ mod tests {
context: &dyn PhysicalOptimizerContext,
) -> Result<Arc<dyn ExecutionPlan>> {
assert!(context.statistics_registry().is_some());
self.invoked.store(true, AtomicOrdering::Relaxed);
let statistics_context = context
.statistics_context()
.expect("the planner shares a statistics context");
let statistics =
statistics_context.compute_arc(&plan, &StatisticsArgs::new())?;
self.recorded.lock().push(statistics);
Ok(plan)
}

Expand Down Expand Up @@ -3644,10 +3671,7 @@ mod tests {
self.inner.physical_optimizers()
}

fn statistics_registry(
&self,
) -> Option<&datafusion_physical_plan::operator_statistics::StatisticsRegistry>
{
fn statistics_registry(&self) -> Option<&StatisticsRegistry> {
self.inner.statistics_registry()
}

Expand Down Expand Up @@ -3816,15 +3840,13 @@ mod tests {

#[tokio::test]
async fn plans_with_non_session_state_implementation() -> Result<()> {
let invoked = Arc::new(AtomicBool::new(false));
let recorded = Arc::new(SyncMutex::new(vec![]));
let inner = SessionStateBuilder::new()
.with_default_features()
.with_physical_optimizer_rules(vec![Arc::new(ContextCheckingRule {
invoked: Arc::clone(&invoked),
recorded: Arc::clone(&recorded),
})])
.with_statistics_registry(
datafusion_physical_plan::operator_statistics::StatisticsRegistry::new(),
)
.with_statistics_registry(StatisticsRegistry::new())
.build();
let query_planner_invoked = Arc::new(AtomicBool::new(false));
let session = TestSession {
Expand All @@ -3839,7 +3861,31 @@ mod tests {
let physical_plan = session.create_physical_plan(&logical_plan).await?;
assert!(physical_plan.is::<EmptyExec>());
assert!(query_planner_invoked.load(AtomicOrdering::Relaxed));
assert!(invoked.load(AtomicOrdering::Relaxed));
assert_eq!(recorded.lock().len(), 1);
Ok(())
}

#[tokio::test]
async fn optimizer_rules_share_statistics_context() -> Result<()> {
let recorded = Arc::new(SyncMutex::new(vec![]));
let rule = || -> Arc<dyn PhysicalOptimizerRule + Send + Sync> {
Arc::new(ContextCheckingRule {
recorded: Arc::clone(&recorded),
})
};
let session_state = SessionStateBuilder::new()
.with_default_features()
.with_physical_optimizer_rules(vec![rule(), rule()])
.with_statistics_registry(StatisticsRegistry::new())
.build();

let logical_plan = LogicalPlanBuilder::empty(false).build()?;
session_state.create_physical_plan(&logical_plan).await?;

// The second rule reads the statistics the first rule cached
let recorded = recorded.lock();
assert_eq!(recorded.len(), 2);
assert!(Arc::ptr_eq(&recorded[0], &recorded[1]));
Ok(())
}

Expand Down
29 changes: 20 additions & 9 deletions datafusion/physical-optimizer/src/aggregate_statistics.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,14 +25,15 @@ use datafusion_physical_plan::aggregates::{
};
use datafusion_physical_plan::placeholder_row::PlaceholderRowExec;
use datafusion_physical_plan::projection::{ProjectionExec, ProjectionExpr};
use datafusion_physical_plan::statistics::{StatisticsArgs, StatisticsContext};
use datafusion_physical_plan::statistics::StatisticsArgs;
use datafusion_physical_plan::udaf::{
AggregateFunctionExpr, StatisticsArgs as PlanStatisticsArgs,
};
use datafusion_physical_plan::{ExecutionPlan, expressions};
use std::sync::Arc;

use crate::PhysicalOptimizerRule;
use crate::optimizer::{ConfigOnlyContext, PhysicalOptimizerContext};

/// Optimizer that uses available statistics for aggregate functions
#[derive(Default, Debug)]
Expand All @@ -46,20 +47,26 @@ impl AggregateStatistics {
}

impl PhysicalOptimizerRule for AggregateStatistics {
#[cfg_attr(feature = "recursive_protection", recursive::recursive)]
#[expect(clippy::allow_attributes)] // See https://github.com/apache/datafusion/issues/18881#issuecomment-3621545670
#[allow(clippy::only_used_in_recursion)] // See https://github.com/rust-lang/rust-clippy/issues/14566
fn optimize(
&self,
plan: Arc<dyn ExecutionPlan>,
config: &ConfigOptions,
) -> Result<Arc<dyn ExecutionPlan>> {
self.optimize_with_context(plan, &ConfigOnlyContext::new(config))
}

#[cfg_attr(feature = "recursive_protection", recursive::recursive)]
fn optimize_with_context(
&self,
plan: Arc<dyn ExecutionPlan>,
context: &dyn PhysicalOptimizerContext,
) -> Result<Arc<dyn ExecutionPlan>> {
if let Some(partial_agg_exec) = take_optimizable(&plan) {
let partial_agg_exec = partial_agg_exec
.downcast_ref::<AggregateExec>()
.expect("take_optimizable() ensures that this is a AggregateExec");
let stats = StatisticsContext::new()
.compute(partial_agg_exec.input().as_ref(), &StatisticsArgs::new())?;
let stats = context

@jayzhan211 jayzhan211 Oct 9, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This lets registered providers decide query results, not just plan choices: AggregateStatistics folds COUNT(*)/MIN/MAX from a provider's Exact values, and LimitPushdown (limit_pushdown.rs:398) drops a limit over a provider's Exact(0). No test covers it. The only provider in any test is the Inexact HashJoinExec one in statistics_registry.slt. With a provider claiming Exact(42) on a 3-row table, SELECT COUNT(*) returns 42 on this branch. If this line goes back to main's StatisticsContext::new(), it returns 3, and every existing test still passes.

Test for physical_planner.rs tests (passes here, fails with the line reverted):

/// `AggregateStatistics` trusts an `Exact` row count from a registered
/// provider: it answers `COUNT(*)` without scanning
#[tokio::test]
async fn aggregate_statistics_consults_statistics_providers() -> Result<()> {
    use datafusion_common::stats::Precision;
    use datafusion_datasource::source::DataSourceExec;
    use datafusion_physical_plan::operator_statistics::{
        ClosureStatisticsProvider, StatisticsResult,
    };

    let provider = ClosureStatisticsProvider::with_matches(
        |plan| plan.is::<DataSourceExec>(),
        |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(42);
            Ok(StatisticsResult::Computed(stats.into()))
        },
    );
    let state = SessionStateBuilder::new()
        .with_default_features()
        .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 COUNT(*) FROM t").await?.collect().await?;
    assert_batches_eq!(
        &[
            "+----------+",
            "| count(*) |",
            "+----------+",
            "| 42       |",
            "+----------+"
        ],
        &batches
    );
    Ok(())
}

State the contract where provider authors will read it (operator_statistics/mod.rs:257):

 /// Implementations can handle specific operator types or override default
 /// estimation logic. The chain of providers is traversed until one returns
 /// [`StatisticsResult::Computed`].
+///
+/// `Exact` statistics are trusted for correctness, not just cost: optimizer
+/// rules use them to answer `COUNT(*)`, `MIN` and `MAX` without scanning and
+/// to remove limits. Return `Inexact` for anything that is an estimate.

And in 56.0.0.md:

 `JoinSelection` already did. Sessions without registered providers (the
-default) are not affected.
+default) are not affected. Provider statistics can now change query results:
+`AggregateStatistics` answers `COUNT(*)`, `MIN` and `MAX` from `Exact` values,
+and `LimitPushdown` removes a limit over an `Exact(0)` input, so a provider
+must return `Exact` only for guaranteed values.

.compute_statistics(partial_agg_exec.input(), &StatisticsArgs::new())?;
let mut projections = vec![];
for expr in partial_agg_exec.aggr_expr() {
let field = expr.field();
Expand Down Expand Up @@ -92,13 +99,17 @@ impl PhysicalOptimizerRule for AggregateStatistics {
)?))
} else {
plan.map_children(|child| {
self.optimize(child, config).map(Transformed::yes)
self.optimize_with_context(child, context)
.map(Transformed::yes)
})
.data()
}
} else {
plan.map_children(|child| self.optimize(child, config).map(Transformed::yes))
.data()
plan.map_children(|child| {
self.optimize_with_context(child, context)
.map(Transformed::yes)
})
.data()
}
}

Expand Down
97 changes: 48 additions & 49 deletions datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1063,10 +1063,8 @@ struct PlanSize {
}

impl PlanSize {
fn from_plan(plan: &dyn ExecutionPlan) -> Self {
let stats = StatisticsContext::new()
.compute(plan, &StatisticsArgs::new())
.ok();
fn from_plan(plan: &Arc<dyn ExecutionPlan>, stats_ctx: &StatisticsContext) -> Self {
let stats = stats_ctx.compute_arc(plan, &StatisticsArgs::new()).ok();
Self {
byte_size: stats
.as_ref()
Expand Down Expand Up @@ -1124,6 +1122,7 @@ fn enforce_distribution_relationships(
input_distributions: &InputDistributionRequirements,
children: &mut [DistributionChildState],
target_partitions: usize,
stats_ctx: &StatisticsContext,
) -> Result<()> {
let mut repartitioned_for_relationship = vec![false; children.len()];

Expand Down Expand Up @@ -1162,53 +1161,52 @@ fn enforce_distribution_relationships(
}
}

let best_satisfied_child: Option<(usize, Partitioning)> = match satisfied_children
.len()
{
0 => None,
1 => satisfied_children
.into_iter()
.next()
.map(|(i, p, _)| (i, p)),
_ => {
// Prefer native partitioned children over newly repartitioned ones
let native_children: Vec<_> = satisfied_children
.iter()
.filter(|(_, _, is_native)| *is_native)
.collect();
if native_children.len() == 1 {
let (i, p, _) = native_children[0];
Some((*i, p.clone()))
} else {
let pool = if !native_children.is_empty() {
native_children
} else {
satisfied_children.iter().collect()
};
let candidates: Vec<_> = pool
.into_iter()
.map(|(idx, part, _)| {
let size =
PlanSize::from_plan(children[*idx].context.plan.as_ref());
(size, *idx, part.clone())
})
.collect();

// Prefer a unique, strictly larger winner (`size_a > size_b`).
// Otherwise, fall back to standard distribution rather
// than choosing an arbitrary reference.
candidates
let best_satisfied_child: Option<(usize, Partitioning)> =

@asolimando asolimando Oct 7, 2026

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Note to reviewers: re-indented by rustfmt, the real change is passing stats_ctx to PlanSize::from_plan, but the diff can be confusing.

match satisfied_children.len() {
0 => None,
1 => satisfied_children
.into_iter()
.next()
.map(|(i, p, _)| (i, p)),
_ => {
// Prefer native partitioned children over newly repartitioned ones
let native_children: Vec<_> = satisfied_children
.iter()
.find(|(size_a, idx_a, _)| {
size_a.is_known()
&& candidates.iter().all(|(size_b, idx_b, _)| {
idx_a == idx_b || size_a > size_b
})
})
.map(|(_, idx, part)| (*idx, part.clone()))
.filter(|(_, _, is_native)| *is_native)
.collect();
if native_children.len() == 1 {
let (i, p, _) = native_children[0];
Some((*i, p.clone()))
} else {
let pool = if !native_children.is_empty() {
native_children
} else {
satisfied_children.iter().collect()
};
let candidates: Vec<_> = pool
.into_iter()
.map(|(idx, part, _)| {
let plan = &children[*idx].context.plan;
let size = PlanSize::from_plan(plan, stats_ctx);
(size, *idx, part.clone())
})
.collect();

// Prefer a unique, strictly larger winner (`size_a > size_b`).
// Otherwise, fall back to standard distribution rather
// than choosing an arbitrary reference.
candidates
.iter()
.find(|(size_a, idx_a, _)| {
size_a.is_known()
&& candidates.iter().all(|(size_b, idx_b, _)| {
idx_a == idx_b || size_a > size_b
})
})
.map(|(_, idx, part)| (*idx, part.clone()))
}
}
}
};
};

// Validate that best_satisfied_child can be adapted across all unsatisfied children in the group
let best_satisfied_child =
Expand Down Expand Up @@ -1666,6 +1664,7 @@ pub fn ensure_distribution_with_stats(
&input_distributions,
&mut children,
target_partitions,
stats_ctx,
)?;

let children = children
Expand Down
Loading
Loading