From 3179bd9765b881320564ce270c38241ef2ff44ec Mon Sep 17 00:00:00 2001
From: Alessandro Solimando <asolimando@apache.org>
Date: Tue, 6 Oct 2026 22:14:11 +0200
Subject: [PATCH 1/7] perf: share one StatisticsContext across all physical
 optimizer rules

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.
---
 datafusion/core/src/physical_planner.rs       | 87 ++++++++++++++++-
 .../src/aggregate_statistics.rs               | 32 ++++--
 .../enforce_distribution.rs                   | 97 +++++++++----------
 .../src/ensure_requirements/mod.rs            | 23 ++---
 .../physical-optimizer/src/join_selection.rs  | 36 +++----
 .../physical-optimizer/src/limit_pushdown.rs  | 51 ++++++++--
 .../physical-optimizer/src/optimizer.rs       | 33 ++++++-
 datafusion/physical-plan/src/statistics.rs    | 29 +++---
 datafusion/session/src/physical_optimizer.rs  | 14 +++
 9 files changed, 282 insertions(+), 120 deletions(-)

diff --git a/datafusion/core/src/physical_planner.rs b/datafusion/core/src/physical_planner.rs
index f11b39157ce0f..bad8b8932e657 100644
--- a/datafusion/core/src/physical_planner.rs
+++ b/datafusion/core/src/physical_planner.rs
@@ -106,6 +106,7 @@ use datafusion_physical_plan::joins::PiecewiseMergeJoinExec;
 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};
 
@@ -123,6 +124,20 @@ 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<'_> {
@@ -135,6 +150,10 @@ impl PhysicalOptimizerContext for SessionOptimizerContext<'_> {
     ) -> Option<&datafusion_physical_plan::operator_statistics::StatisticsRegistry> {
         self.session.statistics_registry()
     }
+
+    fn statistics_context(&self) -> Option<&StatisticsContext> {
+        Some(&self.statistics_context)
+    }
 }
 
 /// Default single node physical query planner that converts a
@@ -3119,9 +3138,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);
         for optimizer in optimizers {
             let before_schema = new_plan.schema();
             new_plan = optimizer
@@ -3532,6 +3549,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 _,
@@ -3554,8 +3572,10 @@ 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::statistics::StatisticsArgs;
     use datafusion_physical_plan::{ChildrenPropertiesMode, ReplaceChildrenOptions};
     use datafusion_session::QueryPlanner;
+    use parking_lot::Mutex as SyncMutex;
 
     #[derive(Debug)]
     struct ContextCheckingRule {
@@ -3590,6 +3610,44 @@ mod tests {
         }
     }
 
+    /// Records the root statistics computed with the shared statistics context
+    #[derive(Debug)]
+    struct StatisticsRecordingRule {
+        recorded: Arc<SyncMutex<Vec<Arc<Statistics>>>>,
+    }
+
+    impl PhysicalOptimizerRule for StatisticsRecordingRule {
+        fn optimize(
+            &self,
+            plan: Arc<dyn ExecutionPlan>,
+            _config: &ConfigOptions,
+        ) -> Result<Arc<dyn ExecutionPlan>> {
+            Ok(plan)
+        }
+
+        fn optimize_with_context(
+            &self,
+            plan: Arc<dyn ExecutionPlan>,
+            context: &dyn PhysicalOptimizerContext,
+        ) -> Result<Arc<dyn ExecutionPlan>> {
+            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)
+        }
+
+        fn name(&self) -> &str {
+            "statistics_recording_rule"
+        }
+
+        fn schema_check(&self) -> bool {
+            true
+        }
+    }
+
     #[derive(Debug)]
     struct TestQueryPlanner {
         invoked: Arc<AtomicBool>,
@@ -3839,6 +3897,29 @@ mod tests {
         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(StatisticsRecordingRule {
+                recorded: Arc::clone(&recorded),
+            })
+        };
+        let session_state = SessionStateBuilder::new()
+            .with_default_features()
+            .with_physical_optimizer_rules(vec![rule(), rule()])
+            .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(())
+    }
+
     async fn aggregate_explain(logical_plan: &LogicalPlan) -> Result<String> {
         let physical_plan = plan(logical_plan).await?;
         Ok(displayable(physical_plan.as_ref()).indent(true).to_string())
diff --git a/datafusion/physical-optimizer/src/aggregate_statistics.rs b/datafusion/physical-optimizer/src/aggregate_statistics.rs
index 43b1abb4b68a9..ee8295c2b417e 100644
--- a/datafusion/physical-optimizer/src/aggregate_statistics.rs
+++ b/datafusion/physical-optimizer/src/aggregate_statistics.rs
@@ -25,7 +25,7 @@ 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,
 };
@@ -33,6 +33,9 @@ use datafusion_physical_plan::{ExecutionPlan, expressions};
 use std::sync::Arc;
 
 use crate::PhysicalOptimizerRule;
+use crate::optimizer::{
+    ConfigOnlyContext, PhysicalOptimizerContext, with_statistics_context,
+};
 
 /// Optimizer that uses available statistics for aggregate functions
 #[derive(Default, Debug)]
@@ -46,20 +49,27 @@ 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 = with_statistics_context(context, |stats_ctx| {
+                stats_ctx.compute_arc(partial_agg_exec.input(), &StatisticsArgs::new())
+            })?;
             let mut projections = vec![];
             for expr in partial_agg_exec.aggr_expr() {
                 let field = expr.field();
@@ -92,13 +102,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()
         }
     }
 
diff --git a/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs b/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs
index 3c2426a5ef2b1..31f00fde90ee5 100644
--- a/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs
+++ b/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs
@@ -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()
@@ -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()];
 
@@ -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)> =
+            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 =
@@ -1666,6 +1664,7 @@ pub fn ensure_distribution_with_stats(
         &input_distributions,
         &mut children,
         target_partitions,
+        stats_ctx,
     )?;
 
     let children = children
diff --git a/datafusion/physical-optimizer/src/ensure_requirements/mod.rs b/datafusion/physical-optimizer/src/ensure_requirements/mod.rs
index 3909b56c17e58..56ca430d69973 100644
--- a/datafusion/physical-optimizer/src/ensure_requirements/mod.rs
+++ b/datafusion/physical-optimizer/src/ensure_requirements/mod.rs
@@ -150,13 +150,14 @@ pub mod enforce_sorting;
 use std::sync::Arc;
 
 use crate::PhysicalOptimizerRule;
-use crate::optimizer::{ConfigOnlyContext, PhysicalOptimizerContext};
+use crate::optimizer::{
+    ConfigOnlyContext, PhysicalOptimizerContext, with_statistics_context,
+};
 
 use datafusion_common::Result;
 use datafusion_common::config::ConfigOptions;
 use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode};
 use datafusion_physical_plan::ExecutionPlan;
-use datafusion_physical_plan::statistics::StatisticsContext;
 
 /// Optimizer rule that enforces both distribution and sorting requirements.
 ///
@@ -221,17 +222,13 @@ impl PhysicalOptimizerRule for EnsureRequirements {
 
         // Step 2a: Distribution enforcement (bottom-up)
         let dist_ctx = DistributionContext::new_default(plan);
-        // Share one statistics context across the whole distribution pass so each
-        // subtree's statistics are computed once instead of once per ancestor.
-        // Build it from the session's statistics registry so registered providers
-        // are consulted (an empty registry, the default, is unchanged behavior).
-        let stats_ctx = match context.statistics_registry() {
-            Some(registry) => StatisticsContext::new_with_registry(registry.clone()),
-            None => StatisticsContext::new(),
-        };
-        let dist_ctx = dist_ctx
-            .transform_up(|ctx| ensure_distribution_with_stats(ctx, config, &stats_ctx))
-            .data()?;
+        let dist_ctx = with_statistics_context(context, |stats_ctx| {
+            dist_ctx
+                .transform_up(|ctx| {
+                    ensure_distribution_with_stats(ctx, config, stats_ctx)
+                })
+                .data()
+        })?;
 
         // Step 2b: Sorting enforcement (bottom-up) — runs on distribution-fixed plan
         let sort_ctx = PlanWithCorrespondingSort::new_default(dist_ctx.plan);
diff --git a/datafusion/physical-optimizer/src/join_selection.rs b/datafusion/physical-optimizer/src/join_selection.rs
index 2b342bece7040..54873900444b6 100644
--- a/datafusion/physical-optimizer/src/join_selection.rs
+++ b/datafusion/physical-optimizer/src/join_selection.rs
@@ -24,7 +24,9 @@
 //! `PartitionMode` and the build side using the available statistics for hash joins.
 
 use crate::PhysicalOptimizerRule;
-use crate::optimizer::{ConfigOnlyContext, PhysicalOptimizerContext};
+use crate::optimizer::{
+    ConfigOnlyContext, PhysicalOptimizerContext, with_statistics_context,
+};
 use datafusion_common::Statistics;
 use datafusion_common::config::ConfigOptions;
 use datafusion_common::error::Result;
@@ -39,7 +41,7 @@ use datafusion_physical_plan::joins::{
     CrossJoinExec, HashJoinExec, NestedLoopJoinExec, PartitionMode,
     StreamJoinPartitionMode, SymmetricHashJoinExec,
 };
-use datafusion_physical_plan::statistics::{StatisticsArgs, StatisticsContext};
+use datafusion_physical_plan::statistics::StatisticsArgs;
 use datafusion_physical_plan::{ExecutionPlan, ExecutionPlanProperties};
 use std::sync::Arc;
 
@@ -59,14 +61,10 @@ impl JoinSelection {
 /// Get statistics for a plan node, consulting the registry's providers if one
 /// is available.
 fn get_stats(
-    plan: &dyn ExecutionPlan,
+    plan: &Arc<dyn ExecutionPlan>,
     context: &dyn PhysicalOptimizerContext,
 ) -> Result<Arc<Statistics>> {
-    let ctx = match context.statistics_registry() {
-        Some(reg) => StatisticsContext::new_with_registry(reg.clone()),
-        None => StatisticsContext::new(),
-    };
-    ctx.compute(plan, &StatisticsArgs::new())
+    with_statistics_context(context, |ctx| ctx.compute_arc(plan, &StatisticsArgs::new()))
 }
 
 // TODO: We need some performance test for Right Semi/Right Join swap to Left Semi/Left Join in case that the right side is smaller but not much smaller.
@@ -86,8 +84,8 @@ fn get_stats(
 /// Used configurations
 /// - `optimizer.join_reordering`: allows or forbids statistics-driven join swapping
 pub(crate) fn should_swap_join_order(
-    left: &dyn ExecutionPlan,
-    right: &dyn ExecutionPlan,
+    left: &Arc<dyn ExecutionPlan>,
+    right: &Arc<dyn ExecutionPlan>,
     context: &dyn PhysicalOptimizerContext,
 ) -> Result<bool> {
     if !context.config_options().optimizer.join_reordering {
@@ -115,7 +113,7 @@ pub(crate) fn should_swap_join_order(
 }
 
 fn supports_collect_by_thresholds(
-    plan: &dyn ExecutionPlan,
+    plan: &Arc<dyn ExecutionPlan>,
     threshold_byte_size: usize,
     threshold_num_rows: usize,
     context: &dyn PhysicalOptimizerContext,
@@ -204,14 +202,14 @@ pub(crate) fn try_collect_left(
 
     let left_can_collect = ignore_threshold
         || supports_collect_by_thresholds(
-            &**left,
+            left,
             optimizer_config.hash_join_single_partition_threshold,
             optimizer_config.hash_join_single_partition_threshold_rows,
             context,
         );
     let right_can_collect = ignore_threshold
         || supports_collect_by_thresholds(
-            &**right,
+            right,
             optimizer_config.hash_join_single_partition_threshold,
             optimizer_config.hash_join_single_partition_threshold_rows,
             context,
@@ -221,7 +219,7 @@ pub(crate) fn try_collect_left(
         (true, true) => {
             // For null-aware joins, we only swap `LeftAnti` joins where the left side is > right side
             if can_swap_hash_join(hash_join)
-                && should_swap_join_order(&**left, &**right, context)?
+                && should_swap_join_order(left, right, context)?
             {
                 Ok(Some(hash_join.swap_inputs(PartitionMode::CollectLeft)?))
             } else {
@@ -269,9 +267,7 @@ pub(crate) fn partitioned_hash_join(
     } else {
         PartitionMode::Partitioned
     };
-    if can_swap_hash_join(hash_join)
-        && should_swap_join_order(&**left, &**right, context)?
-    {
+    if can_swap_hash_join(hash_join) && should_swap_join_order(left, right, context)? {
         hash_join.swap_inputs(partition_mode)
     } else {
         // Null-aware anti joins must use CollectLeft mode because they track probe-side state
@@ -316,7 +312,7 @@ fn statistical_join_selection_subrule(
                 let left = hash_join.left();
                 let right = hash_join.right();
                 if can_swap_hash_join(hash_join)
-                    && should_swap_join_order(&**left, &**right, context)?
+                    && should_swap_join_order(left, right, context)?
                 {
                     // Null-aware RightAnti only supports CollectLeft
                     let partition_mode = if hash_join.null_aware {
@@ -333,7 +329,7 @@ fn statistical_join_selection_subrule(
     } else if let Some(cross_join) = plan.downcast_ref::<CrossJoinExec>() {
         let left = cross_join.left();
         let right = cross_join.right();
-        if should_swap_join_order(&**left, &**right, context)? {
+        if should_swap_join_order(left, right, context)? {
             cross_join.swap_inputs().map(Some)?
         } else {
             None
@@ -342,7 +338,7 @@ fn statistical_join_selection_subrule(
         let left = nl_join.left();
         let right = nl_join.right();
         if nl_join.join_type().supports_swap()
-            && should_swap_join_order(&**left, &**right, context)?
+            && should_swap_join_order(left, right, context)?
         {
             nl_join.swap_inputs().map(Some)?
         } else {
diff --git a/datafusion/physical-optimizer/src/limit_pushdown.rs b/datafusion/physical-optimizer/src/limit_pushdown.rs
index f88a2be14e984..8029fd64bae72 100644
--- a/datafusion/physical-optimizer/src/limit_pushdown.rs
+++ b/datafusion/physical-optimizer/src/limit_pushdown.rs
@@ -64,6 +64,9 @@ use std::fmt::Debug;
 use std::sync::Arc;
 
 use crate::PhysicalOptimizerRule;
+use crate::optimizer::{
+    ConfigOnlyContext, PhysicalOptimizerContext, with_statistics_context,
+};
 
 use datafusion_common::config::ConfigOptions;
 use datafusion_common::error::Result;
@@ -111,7 +114,15 @@ impl PhysicalOptimizerRule for LimitPushdown {
     fn optimize(
         &self,
         plan: Arc<dyn ExecutionPlan>,
-        _config: &ConfigOptions,
+        config: &ConfigOptions,
+    ) -> Result<Arc<dyn ExecutionPlan>> {
+        self.optimize_with_context(plan, &ConfigOnlyContext::new(config))
+    }
+
+    fn optimize_with_context(
+        &self,
+        plan: Arc<dyn ExecutionPlan>,
+        context: &dyn PhysicalOptimizerContext,
     ) -> Result<Arc<dyn ExecutionPlan>> {
         let global_state = GlobalRequirements {
             fetch: None,
@@ -119,7 +130,9 @@ impl PhysicalOptimizerRule for LimitPushdown {
             satisfied: false,
             preserve_order: false,
         };
-        pushdown_limits(plan, global_state)
+        with_statistics_context(context, |stats_ctx| {
+            pushdown_limits(plan, global_state, stats_ctx)
+        })
     }
 
     fn name(&self) -> &str {
@@ -145,9 +158,26 @@ struct LimitInfo {
 ///
 /// If a limit is encountered, a [`TreeNodeRecursion::Stop`] is returned. Otherwise,
 /// return a [`TreeNodeRecursion::Continue`].
+///
+/// Computes statistics with a new [`StatisticsContext`]; see
+/// [`pushdown_limit_helper_with_stats`] to share one.
 pub fn pushdown_limit_helper(
+    pushdown_plan: Arc<dyn ExecutionPlan>,
+    global_state: GlobalRequirements,
+) -> Result<(Transformed<Arc<dyn ExecutionPlan>>, GlobalRequirements)> {
+    pushdown_limit_helper_with_stats(
+        pushdown_plan,
+        global_state,
+        &StatisticsContext::new(),
+    )
+}
+
+/// Like [`pushdown_limit_helper`], but computes statistics with `stats_ctx`,
+/// so its cache can be shared with other calls and other rules.
+pub fn pushdown_limit_helper_with_stats(
     mut pushdown_plan: Arc<dyn ExecutionPlan>,
     mut global_state: GlobalRequirements,
+    stats_ctx: &StatisticsContext,
 ) -> Result<(Transformed<Arc<dyn ExecutionPlan>>, GlobalRequirements)> {
     // Extract limit, if exist, and return child inputs.
     if let Some(limit_info) = extract_limit(&pushdown_plan) {
@@ -165,7 +195,7 @@ pub fn pushdown_limit_helper(
         global_state.satisfied = false;
 
         if let Some(fetch) = fetch
-            && limit_satisfied_by_input(&limit_info.input, skip, fetch)?
+            && limit_satisfied_by_input(&limit_info.input, skip, fetch, stats_ctx)?
         {
             // The input already produces at most `fetch` rows, so no new limit
             // node is needed. Mark satisfied so downstream won't re-add one,
@@ -315,6 +345,7 @@ fn limit_satisfied_by_input(
     plan: &Arc<dyn ExecutionPlan>,
     skip: usize,
     fetch: usize,
+    stats_ctx: &StatisticsContext,
 ) -> Result<bool> {
     if skip > 0 {
         return Ok(false);
@@ -324,7 +355,7 @@ fn limit_satisfied_by_input(
         return Ok(false);
     }
 
-    let Some(num_rows) = limit_eliminable_exact_num_rows(plan)? else {
+    let Some(num_rows) = limit_eliminable_exact_num_rows(plan, stats_ctx)? else {
         return Ok(false);
     };
 
@@ -335,6 +366,7 @@ fn limit_satisfied_by_input(
 /// whose row-count guarantees are strong enough to remove a limit.
 fn limit_eliminable_exact_num_rows(
     plan: &Arc<dyn ExecutionPlan>,
+    stats_ctx: &StatisticsContext,
 ) -> Result<Option<usize>> {
     // Unwrap any wrapping ProjectionExec layers; projections preserve row count
     // but may derive statistics in ways that are not trustworthy, so we peek
@@ -353,8 +385,8 @@ fn limit_eliminable_exact_num_rows(
     }
 
     if matches!(
-        StatisticsContext::new()
-            .compute(current.as_ref(), &StatisticsArgs::new())?
+        stats_ctx
+            .compute_arc(current, &StatisticsArgs::new())?
             .num_rows,
         Precision::Exact(0)
     ) {
@@ -368,15 +400,17 @@ fn limit_eliminable_exact_num_rows(
 pub(crate) fn pushdown_limits(
     pushdown_plan: Arc<dyn ExecutionPlan>,
     global_state: GlobalRequirements,
+    stats_ctx: &StatisticsContext,
 ) -> Result<Arc<dyn ExecutionPlan>> {
     // Call pushdown_limit_helper.
     // This will either extract the limit node (returning the child), or apply the limit pushdown.
     let (mut new_node, mut global_state) =
-        pushdown_limit_helper(pushdown_plan, global_state)?;
+        pushdown_limit_helper_with_stats(pushdown_plan, global_state, stats_ctx)?;
 
     // While limits exist, continue combining the global_state.
     while new_node.tnr == TreeNodeRecursion::Stop {
-        (new_node, global_state) = pushdown_limit_helper(new_node.data, global_state)?;
+        (new_node, global_state) =
+            pushdown_limit_helper_with_stats(new_node.data, global_state, stats_ctx)?;
     }
 
     // Once a limit has been materialized above the current node, child
@@ -396,6 +430,7 @@ pub(crate) fn pushdown_limits(
             let new_child = pushdown_limits(
                 Arc::<dyn ExecutionPlan>::clone(child),
                 global_state.clone(),
+                stats_ctx,
             )?;
             // Tracking if any of the children changed
             changed |= !Arc::ptr_eq(child, &new_child);
diff --git a/datafusion/physical-optimizer/src/optimizer.rs b/datafusion/physical-optimizer/src/optimizer.rs
index aed25546cd09b..299546582e0d6 100644
--- a/datafusion/physical-optimizer/src/optimizer.rs
+++ b/datafusion/physical-optimizer/src/optimizer.rs
@@ -40,6 +40,7 @@ use crate::limit_pushdown_past_window::LimitPushPastWindows;
 use crate::pushdown_sort::PushdownSort;
 use crate::window_topn::WindowTopN;
 use datafusion_common::config::ConfigOptions;
+use datafusion_physical_plan::statistics::StatisticsContext;
 
 // Re-export from this module for backwards compatibility.
 pub use datafusion_session::{PhysicalOptimizerContext, PhysicalOptimizerRule};
@@ -47,16 +48,21 @@ pub use datafusion_session::{PhysicalOptimizerContext, PhysicalOptimizerRule};
 /// Simple context wrapping [`ConfigOptions`] for backward compatibility.
 ///
 /// This struct provides a minimal implementation of [`PhysicalOptimizerContext`]
-/// that only supplies configuration options. Used when no statistics registry
-/// is available or needed.
+/// that supplies configuration options and a [`StatisticsContext`] without a
+/// statistics registry. Used when no statistics registry is available or
+/// needed.
 pub struct ConfigOnlyContext<'a> {
     config: &'a ConfigOptions,
+    statistics_context: StatisticsContext,
 }
 
 impl<'a> ConfigOnlyContext<'a> {
     /// Create a new context wrapping the given config options.
     pub fn new(config: &'a ConfigOptions) -> Self {
-        Self { config }
+        Self {
+            config,
+            statistics_context: StatisticsContext::new(),
+        }
     }
 }
 
@@ -64,6 +70,27 @@ impl PhysicalOptimizerContext for ConfigOnlyContext<'_> {
     fn config_options(&self) -> &ConfigOptions {
         self.config
     }
+
+    fn statistics_context(&self) -> Option<&StatisticsContext> {
+        Some(&self.statistics_context)
+    }
+}
+
+/// Calls `f` with the [`StatisticsContext`] shared through `context`, or, if
+/// `context` does not share one, with a new context built from its
+/// statistics registry.
+pub(crate) fn with_statistics_context<R>(
+    context: &dyn PhysicalOptimizerContext,
+    f: impl FnOnce(&StatisticsContext) -> R,
+) -> R {
+    if let Some(shared) = context.statistics_context() {
+        return f(shared);
+    }
+    let local = match context.statistics_registry() {
+        Some(registry) => StatisticsContext::new_with_registry(registry.clone()),
+        None => StatisticsContext::new(),
+    };
+    f(&local)
 }
 
 /// A rule-based physical optimizer.
diff --git a/datafusion/physical-plan/src/statistics.rs b/datafusion/physical-plan/src/statistics.rs
index b4194d61e6407..6e474f8f3c462 100644
--- a/datafusion/physical-plan/src/statistics.rs
+++ b/datafusion/physical-plan/src/statistics.rs
@@ -31,10 +31,9 @@ use datafusion_common::{
     Result, Statistics, assert_eq_or_internal_err, assert_or_internal_err,
 };
 use log::debug;
-use std::cell::RefCell;
+use parking_lot::Mutex;
 use std::collections::HashMap;
 use std::ptr::from_ref;
-use std::rc::Rc;
 use std::sync::Arc;
 
 type CacheKey = (usize, Option<usize>);
@@ -173,7 +172,7 @@ pub enum ChildStats {
 /// [`Self::compute_extended`] observes extensions; [`Self::compute`] returns core
 /// [`Statistics`] only.
 pub struct StatisticsContext {
-    cache: Rc<RefCell<StatsCache>>,
+    cache: Mutex<StatsCache>,
     registry: StatisticsRegistry,
 }
 
@@ -192,7 +191,7 @@ impl StatisticsContext {
     /// Creates a context whose walk consults `registry`'s provider chain.
     pub fn new_with_registry(registry: StatisticsRegistry) -> Self {
         Self {
-            cache: Rc::new(RefCell::new(StatsCache::default())),
+            cache: Mutex::new(StatsCache::default()),
             registry,
         }
     }
@@ -203,7 +202,7 @@ impl StatisticsContext {
     /// node that supplied its pointer key. Use it to bound memory at a logical
     /// lifecycle boundary, such as after an optimizer pass.
     pub fn reset_cache(&self) {
-        let mut cache = self.cache.borrow_mut();
+        let mut cache = self.cache.lock();
         cache.statistics.clear();
         cache.extensions.clear();
     }
@@ -569,7 +568,7 @@ impl StatisticsContext {
         partition: Option<usize>,
     ) -> Option<Arc<Statistics>> {
         self.cache
-            .borrow()
+            .lock()
             .statistics
             .get(&cache_key(plan, partition))
             .map(|entry| Arc::clone(&entry.value))
@@ -599,7 +598,7 @@ impl StatisticsContext {
         statistics: Arc<Statistics>,
     ) {
         Self::store_cache_entry(
-            &mut self.cache.borrow_mut().statistics,
+            &mut self.cache.lock().statistics,
             owner,
             partition,
             statistics,
@@ -612,7 +611,7 @@ impl StatisticsContext {
         partition: Option<usize>,
     ) -> Option<Extensions> {
         self.cache
-            .borrow()
+            .lock()
             .extensions
             .get(&cache_key(plan, partition))
             .map(|entry| entry.value.clone())
@@ -625,7 +624,7 @@ impl StatisticsContext {
         extensions: Extensions,
     ) {
         Self::store_cache_entry(
-            &mut self.cache.borrow_mut().extensions,
+            &mut self.cache.lock().extensions,
             owner,
             partition,
             extensions,
@@ -773,7 +772,7 @@ mod tests {
         let args = StatisticsArgs::new();
 
         let s1 = ctx.compute_arc(&leaf, &args).unwrap();
-        assert!(!ctx.cache.borrow().statistics.is_empty());
+        assert!(!ctx.cache.lock().statistics.is_empty());
 
         let s2 = ctx.compute_arc(&leaf, &args).unwrap();
         assert!(Arc::ptr_eq(&s1, &s2));
@@ -847,7 +846,7 @@ mod tests {
             .compute(parent.as_ref(), &StatisticsArgs::new())
             .unwrap();
         assert!(
-            !ctx.cache.borrow().statistics.contains_key(&parent_key),
+            !ctx.cache.lock().statistics.contains_key(&parent_key),
             "borrowed roots must not be memoized"
         );
 
@@ -879,7 +878,7 @@ mod tests {
         let _ = ctx
             .compute_extended_arc(&leaf, &StatisticsArgs::new())
             .unwrap();
-        ctx.cache.borrow_mut().statistics.clear();
+        ctx.cache.lock().statistics.clear();
         drop(leaf);
         assert!(weak.upgrade().is_some());
 
@@ -892,9 +891,9 @@ mod tests {
         let leaf = make_stats_leaf(10);
         let ctx = StatisticsContext::new();
         let _ = ctx.compute_arc(&leaf, &StatisticsArgs::new()).unwrap();
-        assert!(!ctx.cache.borrow().statistics.is_empty());
+        assert!(!ctx.cache.lock().statistics.is_empty());
         ctx.reset_cache();
-        assert!(ctx.cache.borrow().statistics.is_empty());
+        assert!(ctx.cache.lock().statistics.is_empty());
     }
 
     #[test]
@@ -969,7 +968,7 @@ mod tests {
         assert_eq!(extensions.unwrap().get::<Tag>(), Some(&Tag(7)));
         assert!(
             ctx.cache
-                .borrow()
+                .lock()
                 .extensions
                 .contains_key(&cache_key(leaf.as_ref(), None))
         );
diff --git a/datafusion/session/src/physical_optimizer.rs b/datafusion/session/src/physical_optimizer.rs
index d59dce24a484f..6d9d842b12eef 100644
--- a/datafusion/session/src/physical_optimizer.rs
+++ b/datafusion/session/src/physical_optimizer.rs
@@ -24,6 +24,7 @@ use datafusion_common::Result;
 use datafusion_common::config::ConfigOptions;
 use datafusion_physical_plan::ExecutionPlan;
 use datafusion_physical_plan::operator_statistics::StatisticsRegistry;
+use datafusion_physical_plan::statistics::StatisticsContext;
 
 /// Context available to physical optimizer rules.
 ///
@@ -40,6 +41,19 @@ pub trait PhysicalOptimizerContext: Send + Sync {
     fn statistics_registry(&self) -> Option<&StatisticsRegistry> {
         None
     }
+
+    /// Returns a [`StatisticsContext`] shared by every rule of one optimizer
+    /// run, so statistics computed by one rule are reused by later rules.
+    ///
+    /// The context must be built from [`Self::statistics_registry`]. Its cache
+    /// entries hold the plan nodes they were computed for, so it is safe to
+    /// share across plan rewrites (see [`StatisticsContext`]).
+    ///
+    /// Returns `None` if no shared context is available, in which case rules
+    /// create their own context.
+    fn statistics_context(&self) -> Option<&StatisticsContext> {
+        None
+    }
 }
 
 /// `PhysicalOptimizerRule` transforms one [`ExecutionPlan`] into another which

From 3dd8da10167dbffcd43be012e5138067c2d4a0f4 Mon Sep 17 00:00:00 2001
From: Alessandro Solimando <asolimando@apache.org>
Date: Wed, 7 Oct 2026 07:07:48 +0200
Subject: [PATCH 2/7] test: record shared statistics in the existing
 ContextCheckingRule

---
 datafusion/core/src/physical_planner.rs | 57 +++++--------------------
 1 file changed, 11 insertions(+), 46 deletions(-)

diff --git a/datafusion/core/src/physical_planner.rs b/datafusion/core/src/physical_planner.rs
index bad8b8932e657..794352413f5ae 100644
--- a/datafusion/core/src/physical_planner.rs
+++ b/datafusion/core/src/physical_planner.rs
@@ -3572,14 +3572,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 {
@@ -3597,39 +3599,6 @@ mod tests {
             context: &dyn PhysicalOptimizerContext,
         ) -> Result<Arc<dyn ExecutionPlan>> {
             assert!(context.statistics_registry().is_some());
-            self.invoked.store(true, AtomicOrdering::Relaxed);
-            Ok(plan)
-        }
-
-        fn name(&self) -> &str {
-            "context_checking_rule"
-        }
-
-        fn schema_check(&self) -> bool {
-            true
-        }
-    }
-
-    /// Records the root statistics computed with the shared statistics context
-    #[derive(Debug)]
-    struct StatisticsRecordingRule {
-        recorded: Arc<SyncMutex<Vec<Arc<Statistics>>>>,
-    }
-
-    impl PhysicalOptimizerRule for StatisticsRecordingRule {
-        fn optimize(
-            &self,
-            plan: Arc<dyn ExecutionPlan>,
-            _config: &ConfigOptions,
-        ) -> Result<Arc<dyn ExecutionPlan>> {
-            Ok(plan)
-        }
-
-        fn optimize_with_context(
-            &self,
-            plan: Arc<dyn ExecutionPlan>,
-            context: &dyn PhysicalOptimizerContext,
-        ) -> Result<Arc<dyn ExecutionPlan>> {
             let statistics_context = context
                 .statistics_context()
                 .expect("the planner shares a statistics context");
@@ -3640,7 +3609,7 @@ mod tests {
         }
 
         fn name(&self) -> &str {
-            "statistics_recording_rule"
+            "context_checking_rule"
         }
 
         fn schema_check(&self) -> bool {
@@ -3698,10 +3667,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()
         }
 
@@ -3870,15 +3836,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 {
@@ -3893,7 +3857,7 @@ 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(())
     }
 
@@ -3901,13 +3865,14 @@ mod tests {
     async fn optimizer_rules_share_statistics_context() -> Result<()> {
         let recorded = Arc::new(SyncMutex::new(vec![]));
         let rule = || -> Arc<dyn PhysicalOptimizerRule + Send + Sync> {
-            Arc::new(StatisticsRecordingRule {
+            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()?;

From 784fca08c4e9174419b0cce96c3e082182f7041e Mon Sep 17 00:00:00 2001
From: Alessandro Solimando <asolimando@apache.org>
Date: Wed, 7 Oct 2026 07:47:10 +0200
Subject: [PATCH 3/7] Deprecate pushdown_limit_helper and document the shared
 StatisticsContext

---
 datafusion/core/src/physical_planner.rs       |  8 +++----
 .../physical-optimizer/src/limit_pushdown.rs  |  7 +++---
 .../library-user-guide/upgrading/56.0.0.md    | 22 +++++++++++++++++++
 3 files changed, 30 insertions(+), 7 deletions(-)

diff --git a/datafusion/core/src/physical_planner.rs b/datafusion/core/src/physical_planner.rs
index 794352413f5ae..0b94ca1e4463c 100644
--- a/datafusion/core/src/physical_planner.rs
+++ b/datafusion/core/src/physical_planner.rs
@@ -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};
@@ -103,6 +104,7 @@ 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};
@@ -141,13 +143,11 @@ impl<'a> SessionOptimizerContext<'a> {
 }
 
 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()
     }
 
diff --git a/datafusion/physical-optimizer/src/limit_pushdown.rs b/datafusion/physical-optimizer/src/limit_pushdown.rs
index 8029fd64bae72..69e29a360b97b 100644
--- a/datafusion/physical-optimizer/src/limit_pushdown.rs
+++ b/datafusion/physical-optimizer/src/limit_pushdown.rs
@@ -158,9 +158,10 @@ struct LimitInfo {
 ///
 /// If a limit is encountered, a [`TreeNodeRecursion::Stop`] is returned. Otherwise,
 /// return a [`TreeNodeRecursion::Continue`].
-///
-/// Computes statistics with a new [`StatisticsContext`]; see
-/// [`pushdown_limit_helper_with_stats`] to share one.
+#[deprecated(
+    since = "56.0.0",
+    note = "use `pushdown_limit_helper_with_stats` and share one `StatisticsContext` across calls"
+)]
 pub fn pushdown_limit_helper(
     pushdown_plan: Arc<dyn ExecutionPlan>,
     global_state: GlobalRequirements,
diff --git a/docs/source/library-user-guide/upgrading/56.0.0.md b/docs/source/library-user-guide/upgrading/56.0.0.md
index fceda8f6111a1..902ef85934d45 100644
--- a/docs/source/library-user-guide/upgrading/56.0.0.md
+++ b/docs/source/library-user-guide/upgrading/56.0.0.md
@@ -770,6 +770,28 @@ let plan = DistributionContext::new_default(plan)
 Callers that only need the previous behaviour can keep using the deprecated
 form, or pass a freshly constructed `StatisticsContext` per call.
 
+### Physical optimizer rules share one `StatisticsContext` per query
+
+`PhysicalOptimizerContext` has a new `statistics_context` method. The physical
+planner returns one `StatisticsContext` for all the rules of one
+`optimize_physical_plan` call, so statistics computed by one rule are reused by
+later rules. The method returns `None` by default, so existing implementations
+are not affected. Custom rules that compute statistics can use the shared
+context in `optimize_with_context`:
+
+```rust,ignore
+let statistics = match context.statistics_context() {
+    Some(shared) => shared.compute_arc(&plan, &StatisticsArgs::new())?,
+    None => StatisticsContext::new().compute_arc(&plan, &StatisticsArgs::new())?,
+};
+```
+
+For the same reason,
+`datafusion_physical_optimizer::limit_pushdown::pushdown_limit_helper` is
+deprecated in favour of `pushdown_limit_helper_with_stats`, which takes the
+`StatisticsContext` from the caller. The deprecated form creates a new context
+on every call.
+
 ### `MergeIntoOp` requires the SQL-visible target qualifier
 
 `MergeIntoOp` now stores the SQL-visible target qualifier separately from the

From ae341dc1b326ded47669a3e4c641e04ccb7ad664 Mon Sep 17 00:00:00 2001
From: Alessandro Solimando <asolimando@apache.org>
Date: Wed, 7 Oct 2026 14:04:47 +0200
Subject: [PATCH 4/7] Make with_statistics_context public for custom rules

---
 .../physical-optimizer/src/optimizer.rs       | 26 ++++++++++++++++++-
 .../library-user-guide/upgrading/56.0.0.md    | 18 ++++++++-----
 2 files changed, 37 insertions(+), 7 deletions(-)

diff --git a/datafusion/physical-optimizer/src/optimizer.rs b/datafusion/physical-optimizer/src/optimizer.rs
index 299546582e0d6..2875630a3399f 100644
--- a/datafusion/physical-optimizer/src/optimizer.rs
+++ b/datafusion/physical-optimizer/src/optimizer.rs
@@ -79,7 +79,31 @@ impl PhysicalOptimizerContext for ConfigOnlyContext<'_> {
 /// Calls `f` with the [`StatisticsContext`] shared through `context`, or, if
 /// `context` does not share one, with a new context built from its
 /// statistics registry.
-pub(crate) fn with_statistics_context<R>(
+///
+/// Rules that compute statistics should use this, so they reuse the
+/// statistics computed by earlier rules and consult the same statistics
+/// providers as the built-in rules.
+///
+/// # Example
+///
+/// ```
+/// # use std::sync::Arc;
+/// # use arrow::datatypes::Schema;
+/// # use datafusion_common::config::ConfigOptions;
+/// # use datafusion_physical_optimizer::optimizer::{ConfigOnlyContext, with_statistics_context};
+/// # use datafusion_physical_plan::ExecutionPlan;
+/// # use datafusion_physical_plan::empty::EmptyExec;
+/// # use datafusion_physical_plan::statistics::StatisticsArgs;
+/// let config = ConfigOptions::new();
+/// let context = ConfigOnlyContext::new(&config);
+/// let plan: Arc<dyn ExecutionPlan> = Arc::new(EmptyExec::new(Arc::new(Schema::empty())));
+///
+/// let statistics = with_statistics_context(&context, |stats_ctx| {
+///     stats_ctx.compute_arc(&plan, &StatisticsArgs::new())
+/// })?;
+/// # Ok::<(), datafusion_common::DataFusionError>(())
+/// ```
+pub fn with_statistics_context<R>(
     context: &dyn PhysicalOptimizerContext,
     f: impl FnOnce(&StatisticsContext) -> R,
 ) -> R {
diff --git a/docs/source/library-user-guide/upgrading/56.0.0.md b/docs/source/library-user-guide/upgrading/56.0.0.md
index 902ef85934d45..0e62e1bf85139 100644
--- a/docs/source/library-user-guide/upgrading/56.0.0.md
+++ b/docs/source/library-user-guide/upgrading/56.0.0.md
@@ -776,16 +776,22 @@ form, or pass a freshly constructed `StatisticsContext` per call.
 planner returns one `StatisticsContext` for all the rules of one
 `optimize_physical_plan` call, so statistics computed by one rule are reused by
 later rules. The method returns `None` by default, so existing implementations
-are not affected. Custom rules that compute statistics can use the shared
-context in `optimize_with_context`:
+are not affected. Custom rules that compute statistics can call
+`datafusion_physical_optimizer::optimizer::with_statistics_context` in
+`optimize_with_context`. It uses the shared context, or, if there is none, a new
+context built from the statistics registry:
 
 ```rust,ignore
-let statistics = match context.statistics_context() {
-    Some(shared) => shared.compute_arc(&plan, &StatisticsArgs::new())?,
-    None => StatisticsContext::new().compute_arc(&plan, &StatisticsArgs::new())?,
-};
+let statistics = with_statistics_context(context, |stats_ctx| {
+    stats_ctx.compute_arc(&plan, &StatisticsArgs::new())
+})?;
 ```
 
+`AggregateStatistics`, `LimitPushdown` and the co-partitioning choice in
+`EnsureRequirements` now also consult the session's statistics providers, as
+`JoinSelection` already did. Sessions without registered providers (the
+default) are not affected.
+
 For the same reason,
 `datafusion_physical_optimizer::limit_pushdown::pushdown_limit_helper` is
 deprecated in favour of `pushdown_limit_helper_with_stats`, which takes the

From 5a4626e38cd342ecc57d3b771a18ab26815cf8ec Mon Sep 17 00:00:00 2001
From: Alessandro Solimando <asolimando@apache.org>
Date: Wed, 7 Oct 2026 14:17:25 +0200
Subject: [PATCH 5/7] Add PhysicalOptimizerContext::compute_statistics for
 custom rules

---
 .../src/aggregate_statistics.rs               |  9 +--
 .../physical-optimizer/src/join_selection.rs  |  6 +-
 .../physical-optimizer/src/optimizer.rs       | 42 +-----------
 datafusion/session/src/lib.rs                 |  4 +-
 datafusion/session/src/physical_optimizer.rs  | 65 ++++++++++++++++++-
 .../library-user-guide/upgrading/56.0.0.md    | 15 +++--
 6 files changed, 80 insertions(+), 61 deletions(-)

diff --git a/datafusion/physical-optimizer/src/aggregate_statistics.rs b/datafusion/physical-optimizer/src/aggregate_statistics.rs
index ee8295c2b417e..a1260529431de 100644
--- a/datafusion/physical-optimizer/src/aggregate_statistics.rs
+++ b/datafusion/physical-optimizer/src/aggregate_statistics.rs
@@ -33,9 +33,7 @@ use datafusion_physical_plan::{ExecutionPlan, expressions};
 use std::sync::Arc;
 
 use crate::PhysicalOptimizerRule;
-use crate::optimizer::{
-    ConfigOnlyContext, PhysicalOptimizerContext, with_statistics_context,
-};
+use crate::optimizer::{ConfigOnlyContext, PhysicalOptimizerContext};
 
 /// Optimizer that uses available statistics for aggregate functions
 #[derive(Default, Debug)]
@@ -67,9 +65,8 @@ impl PhysicalOptimizerRule for AggregateStatistics {
             let partial_agg_exec = partial_agg_exec
                 .downcast_ref::<AggregateExec>()
                 .expect("take_optimizable() ensures that this is a AggregateExec");
-            let stats = with_statistics_context(context, |stats_ctx| {
-                stats_ctx.compute_arc(partial_agg_exec.input(), &StatisticsArgs::new())
-            })?;
+            let stats = context
+                .compute_statistics(partial_agg_exec.input(), &StatisticsArgs::new())?;
             let mut projections = vec![];
             for expr in partial_agg_exec.aggr_expr() {
                 let field = expr.field();
diff --git a/datafusion/physical-optimizer/src/join_selection.rs b/datafusion/physical-optimizer/src/join_selection.rs
index 54873900444b6..c7f414d91bfa5 100644
--- a/datafusion/physical-optimizer/src/join_selection.rs
+++ b/datafusion/physical-optimizer/src/join_selection.rs
@@ -24,9 +24,7 @@
 //! `PartitionMode` and the build side using the available statistics for hash joins.
 
 use crate::PhysicalOptimizerRule;
-use crate::optimizer::{
-    ConfigOnlyContext, PhysicalOptimizerContext, with_statistics_context,
-};
+use crate::optimizer::{ConfigOnlyContext, PhysicalOptimizerContext};
 use datafusion_common::Statistics;
 use datafusion_common::config::ConfigOptions;
 use datafusion_common::error::Result;
@@ -64,7 +62,7 @@ fn get_stats(
     plan: &Arc<dyn ExecutionPlan>,
     context: &dyn PhysicalOptimizerContext,
 ) -> Result<Arc<Statistics>> {
-    with_statistics_context(context, |ctx| ctx.compute_arc(plan, &StatisticsArgs::new()))
+    context.compute_statistics(plan, &StatisticsArgs::new())
 }
 
 // TODO: We need some performance test for Right Semi/Right Join swap to Left Semi/Left Join in case that the right side is smaller but not much smaller.
diff --git a/datafusion/physical-optimizer/src/optimizer.rs b/datafusion/physical-optimizer/src/optimizer.rs
index 2875630a3399f..a0de1814a3c24 100644
--- a/datafusion/physical-optimizer/src/optimizer.rs
+++ b/datafusion/physical-optimizer/src/optimizer.rs
@@ -43,6 +43,7 @@ use datafusion_common::config::ConfigOptions;
 use datafusion_physical_plan::statistics::StatisticsContext;
 
 // Re-export from this module for backwards compatibility.
+pub use datafusion_session::with_statistics_context;
 pub use datafusion_session::{PhysicalOptimizerContext, PhysicalOptimizerRule};
 
 /// Simple context wrapping [`ConfigOptions`] for backward compatibility.
@@ -76,47 +77,6 @@ impl PhysicalOptimizerContext for ConfigOnlyContext<'_> {
     }
 }
 
-/// Calls `f` with the [`StatisticsContext`] shared through `context`, or, if
-/// `context` does not share one, with a new context built from its
-/// statistics registry.
-///
-/// Rules that compute statistics should use this, so they reuse the
-/// statistics computed by earlier rules and consult the same statistics
-/// providers as the built-in rules.
-///
-/// # Example
-///
-/// ```
-/// # use std::sync::Arc;
-/// # use arrow::datatypes::Schema;
-/// # use datafusion_common::config::ConfigOptions;
-/// # use datafusion_physical_optimizer::optimizer::{ConfigOnlyContext, with_statistics_context};
-/// # use datafusion_physical_plan::ExecutionPlan;
-/// # use datafusion_physical_plan::empty::EmptyExec;
-/// # use datafusion_physical_plan::statistics::StatisticsArgs;
-/// let config = ConfigOptions::new();
-/// let context = ConfigOnlyContext::new(&config);
-/// let plan: Arc<dyn ExecutionPlan> = Arc::new(EmptyExec::new(Arc::new(Schema::empty())));
-///
-/// let statistics = with_statistics_context(&context, |stats_ctx| {
-///     stats_ctx.compute_arc(&plan, &StatisticsArgs::new())
-/// })?;
-/// # Ok::<(), datafusion_common::DataFusionError>(())
-/// ```
-pub fn with_statistics_context<R>(
-    context: &dyn PhysicalOptimizerContext,
-    f: impl FnOnce(&StatisticsContext) -> R,
-) -> R {
-    if let Some(shared) = context.statistics_context() {
-        return f(shared);
-    }
-    let local = match context.statistics_registry() {
-        Some(registry) => StatisticsContext::new_with_registry(registry.clone()),
-        None => StatisticsContext::new(),
-    };
-    f(&local)
-}
-
 /// A rule-based physical optimizer.
 #[derive(Clone, Debug)]
 pub struct PhysicalOptimizer {
diff --git a/datafusion/session/src/lib.rs b/datafusion/session/src/lib.rs
index 26c675c20f928..0fc99ad66eb79 100644
--- a/datafusion/session/src/lib.rs
+++ b/datafusion/session/src/lib.rs
@@ -56,7 +56,9 @@ pub use crate::catalog::{
     CatalogProvider, CatalogProviderFactory, CatalogProviderList,
     EmptyCatalogProviderList,
 };
-pub use crate::physical_optimizer::{PhysicalOptimizerContext, PhysicalOptimizerRule};
+pub use crate::physical_optimizer::{
+    PhysicalOptimizerContext, PhysicalOptimizerRule, with_statistics_context,
+};
 pub use crate::planner::{
     ExtensionPlanner, PhysicalPlanner, QueryPlanner, UnsupportedQueryPlanner,
 };
diff --git a/datafusion/session/src/physical_optimizer.rs b/datafusion/session/src/physical_optimizer.rs
index 6d9d842b12eef..e0b2b2d2ce0d7 100644
--- a/datafusion/session/src/physical_optimizer.rs
+++ b/datafusion/session/src/physical_optimizer.rs
@@ -20,11 +20,11 @@
 use std::fmt::Debug;
 use std::sync::Arc;
 
-use datafusion_common::Result;
 use datafusion_common::config::ConfigOptions;
+use datafusion_common::{Result, Statistics};
 use datafusion_physical_plan::ExecutionPlan;
 use datafusion_physical_plan::operator_statistics::StatisticsRegistry;
-use datafusion_physical_plan::statistics::StatisticsContext;
+use datafusion_physical_plan::statistics::{StatisticsArgs, StatisticsContext};
 
 /// Context available to physical optimizer rules.
 ///
@@ -54,6 +54,67 @@ pub trait PhysicalOptimizerContext: Send + Sync {
     fn statistics_context(&self) -> Option<&StatisticsContext> {
         None
     }
+
+    /// Computes the statistics of `plan` with [`Self::statistics_context`],
+    /// or, if it is `None`, with a new context built from
+    /// [`Self::statistics_registry`].
+    ///
+    /// Rules that compute statistics should use this, so they reuse the
+    /// statistics computed by earlier rules and consult the same statistics
+    /// providers as the built-in rules.
+    ///
+    /// # Example
+    ///
+    /// ```
+    /// # use std::sync::Arc;
+    /// # use arrow_schema::Schema;
+    /// # use datafusion_common::config::ConfigOptions;
+    /// # use datafusion_physical_plan::ExecutionPlan;
+    /// # use datafusion_physical_plan::empty::EmptyExec;
+    /// # use datafusion_physical_plan::statistics::StatisticsArgs;
+    /// # use datafusion_session::PhysicalOptimizerContext;
+    /// # struct MyContext(ConfigOptions);
+    /// # impl PhysicalOptimizerContext for MyContext {
+    /// #     fn config_options(&self) -> &ConfigOptions {
+    /// #         &self.0
+    /// #     }
+    /// # }
+    /// # let context = MyContext(ConfigOptions::new());
+    /// let plan: Arc<dyn ExecutionPlan> = Arc::new(EmptyExec::new(Arc::new(Schema::empty())));
+    /// let statistics = context.compute_statistics(&plan, &StatisticsArgs::new())?;
+    /// # Ok::<(), datafusion_common::DataFusionError>(())
+    /// ```
+    fn compute_statistics(
+        &self,
+        plan: &Arc<dyn ExecutionPlan>,
+        args: &StatisticsArgs,
+    ) -> Result<Arc<Statistics>> {
+        with_statistics_context(self, |stats_ctx| stats_ctx.compute_arc(plan, args))
+    }
+}
+
+/// Calls `f` with [`PhysicalOptimizerContext::statistics_context`], or, if it
+/// is `None`, with a new context built from
+/// [`PhysicalOptimizerContext::statistics_registry`].
+///
+/// Use this to pass one [`StatisticsContext`] down a traversal. To compute
+/// the statistics of a single plan, use
+/// [`PhysicalOptimizerContext::compute_statistics`].
+pub fn with_statistics_context<C, R>(
+    context: &C,
+    f: impl FnOnce(&StatisticsContext) -> R,
+) -> R
+where
+    C: PhysicalOptimizerContext + ?Sized,
+{
+    if let Some(shared) = context.statistics_context() {
+        return f(shared);
+    }
+    let local = match context.statistics_registry() {
+        Some(registry) => StatisticsContext::new_with_registry(registry.clone()),
+        None => StatisticsContext::new(),
+    };
+    f(&local)
 }
 
 /// `PhysicalOptimizerRule` transforms one [`ExecutionPlan`] into another which
diff --git a/docs/source/library-user-guide/upgrading/56.0.0.md b/docs/source/library-user-guide/upgrading/56.0.0.md
index 0e62e1bf85139..16b77832a9ee1 100644
--- a/docs/source/library-user-guide/upgrading/56.0.0.md
+++ b/docs/source/library-user-guide/upgrading/56.0.0.md
@@ -776,17 +776,18 @@ form, or pass a freshly constructed `StatisticsContext` per call.
 planner returns one `StatisticsContext` for all the rules of one
 `optimize_physical_plan` call, so statistics computed by one rule are reused by
 later rules. The method returns `None` by default, so existing implementations
-are not affected. Custom rules that compute statistics can call
-`datafusion_physical_optimizer::optimizer::with_statistics_context` in
-`optimize_with_context`. It uses the shared context, or, if there is none, a new
-context built from the statistics registry:
+are not affected. Custom rules that compute statistics should call the new
+`PhysicalOptimizerContext::compute_statistics` method in `optimize_with_context`.
+It uses the shared context, or, if there is none, a new context built from the
+statistics registry:
 
 ```rust,ignore
-let statistics = with_statistics_context(context, |stats_ctx| {
-    stats_ctx.compute_arc(&plan, &StatisticsArgs::new())
-})?;
+let statistics = context.compute_statistics(&plan, &StatisticsArgs::new())?;
 ```
 
+To pass one `StatisticsContext` down a traversal, use
+`datafusion_session::with_statistics_context`, which applies the same choice.
+
 `AggregateStatistics`, `LimitPushdown` and the co-partitioning choice in
 `EnsureRequirements` now also consult the session's statistics providers, as
 `JoinSelection` already did. Sessions without registered providers (the

From c101751fcdfdcdb47393ff3e0dfd382dd04b3200 Mon Sep 17 00:00:00 2001
From: Alessandro Solimando <asolimando@apache.org>
Date: Thu, 8 Oct 2026 15:10:40 +0200
Subject: [PATCH 6/7] Address review: reset_cache guidance, provider note on
 the deprecated helper, ConfigOnlyContext invariant

---
 datafusion/physical-optimizer/src/limit_pushdown.rs | 6 ++++++
 datafusion/physical-optimizer/src/optimizer.rs      | 4 +++-
 datafusion/physical-plan/src/statistics.rs          | 9 ++++++---
 docs/source/library-user-guide/upgrading/56.0.0.md  | 5 ++++-
 4 files changed, 19 insertions(+), 5 deletions(-)

diff --git a/datafusion/physical-optimizer/src/limit_pushdown.rs b/datafusion/physical-optimizer/src/limit_pushdown.rs
index 69e29a360b97b..226c41cf8d737 100644
--- a/datafusion/physical-optimizer/src/limit_pushdown.rs
+++ b/datafusion/physical-optimizer/src/limit_pushdown.rs
@@ -158,6 +158,12 @@ struct LimitInfo {
 ///
 /// If a limit is encountered, a [`TreeNodeRecursion::Stop`] is returned. Otherwise,
 /// return a [`TreeNodeRecursion::Continue`].
+///
+/// Computes statistics with a new [`StatisticsContext`] that has no statistics
+/// providers. A context built from a statistics registry, as [`LimitPushdown`]
+/// uses, also consults the registered providers, so switching to
+/// [`pushdown_limit_helper_with_stats`] with such a context can change the
+/// result when providers are registered.
 #[deprecated(
     since = "56.0.0",
     note = "use `pushdown_limit_helper_with_stats` and share one `StatisticsContext` across calls"
diff --git a/datafusion/physical-optimizer/src/optimizer.rs b/datafusion/physical-optimizer/src/optimizer.rs
index a0de1814a3c24..bd8d387657e12 100644
--- a/datafusion/physical-optimizer/src/optimizer.rs
+++ b/datafusion/physical-optimizer/src/optimizer.rs
@@ -43,9 +43,10 @@ use datafusion_common::config::ConfigOptions;
 use datafusion_physical_plan::statistics::StatisticsContext;
 
 // Re-export from this module for backwards compatibility.
-pub use datafusion_session::with_statistics_context;
 pub use datafusion_session::{PhysicalOptimizerContext, PhysicalOptimizerRule};
 
+pub use datafusion_session::with_statistics_context;
+
 /// Simple context wrapping [`ConfigOptions`] for backward compatibility.
 ///
 /// This struct provides a minimal implementation of [`PhysicalOptimizerContext`]
@@ -62,6 +63,7 @@ impl<'a> ConfigOnlyContext<'a> {
     pub fn new(config: &'a ConfigOptions) -> Self {
         Self {
             config,
+            // No providers, matching `statistics_registry()`, which is `None`
             statistics_context: StatisticsContext::new(),
         }
     }
diff --git a/datafusion/physical-plan/src/statistics.rs b/datafusion/physical-plan/src/statistics.rs
index 6e474f8f3c462..c22bc628b44cc 100644
--- a/datafusion/physical-plan/src/statistics.rs
+++ b/datafusion/physical-plan/src/statistics.rs
@@ -157,7 +157,8 @@ pub enum ChildStats {
 /// plan rewrites. Each entry holds a strong reference to the plan node it was
 /// computed for, so cached nodes (and their per-partition statistics) stay
 /// alive until [`Self::reset_cache`] is called or the context is dropped. Reset
-/// a long-lived context at a lifecycle boundary to bound its memory.
+/// a long-lived context at a lifecycle boundary to bound its memory (see
+/// [`Self::reset_cache`]).
 ///
 /// An optional [`StatisticsRegistry`] plugs providers into the walk: at each node
 /// they are consulted before the operator's built-in
@@ -199,8 +200,10 @@ impl StatisticsContext {
     /// Clears the memoization cache and releases its retained plan nodes.
     ///
     /// Resetting is optional for correctness: each cache entry retains the plan
-    /// node that supplied its pointer key. Use it to bound memory at a logical
-    /// lifecycle boundary, such as after an optimizer pass.
+    /// node that supplied its pointer key. Use it to bound memory at a
+    /// lifecycle boundary, such as the end of a query's physical optimization.
+    /// Do not reset a context that several optimizer rules share between those
+    /// rules: that discards the statistics that later rules would reuse.
     pub fn reset_cache(&self) {
         let mut cache = self.cache.lock();
         cache.statistics.clear();
diff --git a/docs/source/library-user-guide/upgrading/56.0.0.md b/docs/source/library-user-guide/upgrading/56.0.0.md
index 16b77832a9ee1..944ba8656ae92 100644
--- a/docs/source/library-user-guide/upgrading/56.0.0.md
+++ b/docs/source/library-user-guide/upgrading/56.0.0.md
@@ -797,7 +797,10 @@ For the same reason,
 `datafusion_physical_optimizer::limit_pushdown::pushdown_limit_helper` is
 deprecated in favour of `pushdown_limit_helper_with_stats`, which takes the
 `StatisticsContext` from the caller. The deprecated form creates a new context
-on every call.
+on every call, without statistics providers. When providers are registered,
+switching to `pushdown_limit_helper_with_stats` with a context built from the
+statistics registry can change the result, because the providers are then
+consulted.
 
 ### `MergeIntoOp` requires the SQL-visible target qualifier
 

From f7587af72580025d7ca6442b204f675f6d8c17af Mon Sep 17 00:00:00 2001
From: Alessandro Solimando <asolimando@apache.org>
Date: Thu, 8 Oct 2026 16:19:52 +0200
Subject: [PATCH 7/7] Document that cached plan nodes must not change their
 statistics in place

---
 datafusion/physical-plan/src/statistics.rs | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/datafusion/physical-plan/src/statistics.rs b/datafusion/physical-plan/src/statistics.rs
index 72f79a3fda98f..c52636b21ff56 100644
--- a/datafusion/physical-plan/src/statistics.rs
+++ b/datafusion/physical-plan/src/statistics.rs
@@ -160,6 +160,12 @@ pub enum ChildStats {
 /// a long-lived context at a lifecycle boundary to bound its memory (see
 /// [`Self::reset_cache`]).
 ///
+/// Cached statistics are served for as long as the context lives, so a plan
+/// node must not change its statistics in place (for example through interior
+/// mutability) while a context holds it: the cache would return stale values.
+/// Optimizer rules satisfy this because they replace nodes instead of changing
+/// them.
+///
 /// An optional [`StatisticsRegistry`] plugs providers into the walk: at each node
 /// they are consulted before the operator's built-in
 /// [`ExecutionPlan::statistics_from_inputs`]. An empty registry is the built-in