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.

Author: asolimando

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())

datafusion/physical-optimizer/src/aggregate_statistics.rs

@@ -25,14 +25,17 @@ 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, 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()
         }
     }
 

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

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);