Add operator_statistics framework for pluggable statistics propagation

Introduces a chain-of-responsibility architecture for statistics
computation on physical plan nodes:

- StatisticsProvider trait: chain element computing stats for operators
- StatisticsRegistry: chains providers, with fast-path for empty config
- ExtendedStatistics: Statistics with type-erased extension map
- DefaultStatisticsProvider: delegates to each operator's partition_statistics()
- PhysicalOptimizerContext trait with optimize_with_context() dispatch
- ConfigOnlyContext for backward-compatible rule invocation
- SessionState integration with statistics_registry field and builder

Author: asolimando

datafusion/core/src/execution/session_state.rs

@@ -68,9 +68,11 @@ use datafusion_optimizer::{
 };
 use datafusion_physical_expr::create_physical_expr;
 use datafusion_physical_expr_common::physical_expr::PhysicalExpr;
+use datafusion_physical_optimizer::PhysicalOptimizerContext;
 use datafusion_physical_optimizer::PhysicalOptimizerRule;
 use datafusion_physical_optimizer::optimizer::PhysicalOptimizer;
 use datafusion_physical_plan::ExecutionPlan;
+use datafusion_physical_plan::operator_statistics::StatisticsRegistry;
 use datafusion_session::Session;
 #[cfg(feature = "sql")]
 use datafusion_sql::{
@@ -191,11 +193,27 @@ pub struct SessionState {
     /// thus, changing dialect o PostgreSql is required
     function_factory: Option<Arc<dyn FunctionFactory>>,
     cache_factory: Option<Arc<dyn CacheFactory>>,
+    /// Optional statistics registry for pluggable statistics providers.
+    ///
+    /// When set, physical optimizer rules can use this registry to obtain
+    /// enhanced statistics (e.g., NDV overrides, histograms) beyond what
+    /// is available from `ExecutionPlan::partition_statistics()`.
+    statistics_registry: Option<StatisticsRegistry>,
     /// Cache logical plans of prepared statements for later execution.
     /// Key is the prepared statement name.
     prepared_plans: HashMap<String, Arc<PreparedPlan>>,
 }
 
+impl PhysicalOptimizerContext for SessionState {
+    fn config_options(&self) -> &ConfigOptions {
+        self.config_options()
+    }
+
+    fn statistics_registry(&self) -> Option<&StatisticsRegistry> {
+        self.statistics_registry.as_ref()
+    }
+}
+
 impl Debug for SessionState {
     /// Prefer having short fields at the top and long vector fields near the end
     /// Group fields by
@@ -817,6 +835,14 @@ impl SessionState {
         self.config.options()
     }
 
+    /// Returns the statistics registry if one is configured.
+    ///
+    /// The registry provides pluggable statistics providers for enhanced
+    /// cardinality estimation (e.g., NDV overrides, histograms).
+    pub fn statistics_registry(&self) -> Option<&StatisticsRegistry> {
+        self.statistics_registry.as_ref()
+    }
+
     /// Mark the start of the execution
     pub fn mark_start_execution(&mut self) {
         let config = Arc::clone(self.config.options());
@@ -1006,6 +1032,7 @@ pub struct SessionStateBuilder {
     runtime_env: Option<Arc<RuntimeEnv>>,
     function_factory: Option<Arc<dyn FunctionFactory>>,
     cache_factory: Option<Arc<dyn CacheFactory>>,
+    statistics_registry: Option<StatisticsRegistry>,
     // fields to support convenience functions
     analyzer_rules: Option<Vec<Arc<dyn AnalyzerRule + Send + Sync>>>,
     optimizer_rules: Option<Vec<Arc<dyn OptimizerRule + Send + Sync>>>,
@@ -1047,6 +1074,7 @@ impl SessionStateBuilder {
             runtime_env: None,
             function_factory: None,
             cache_factory: None,
+            statistics_registry: None,
             // fields to support convenience functions
             analyzer_rules: None,
             optimizer_rules: None,
@@ -1103,6 +1131,7 @@ impl SessionStateBuilder {
             runtime_env: Some(existing.runtime_env),
             function_factory: existing.function_factory,
             cache_factory: existing.cache_factory,
+            statistics_registry: existing.statistics_registry,
             // fields to support convenience functions
             analyzer_rules: None,
             optimizer_rules: None,
@@ -1424,6 +1453,16 @@ impl SessionStateBuilder {
         self
     }
 
+    /// Set a [`StatisticsRegistry`] for pluggable statistics providers.
+    ///
+    /// The registry allows physical optimizer rules to access enhanced statistics
+    /// (e.g., NDV overrides, histograms) beyond what is available from
+    /// `ExecutionPlan::partition_statistics()`.
+    pub fn with_statistics_registry(mut self, registry: StatisticsRegistry) -> Self {
+        self.statistics_registry = Some(registry);
+        self
+    }
+
     /// Register an `ObjectStore` to the [`RuntimeEnv`]. See [`RuntimeEnv::register_object_store`]
     /// for more details.
     ///
@@ -1491,6 +1530,7 @@ impl SessionStateBuilder {
             runtime_env,
             function_factory,
             cache_factory,
+            statistics_registry,
             analyzer_rules,
             optimizer_rules,
             physical_optimizer_rules,
@@ -1531,6 +1571,7 @@ impl SessionStateBuilder {
             runtime_env,
             function_factory,
             cache_factory,
+            statistics_registry,
             prepared_plans: HashMap::new(),
         };
 

datafusion/core/src/physical_planner.rs

@@ -2773,7 +2773,7 @@ impl DefaultPhysicalPlanner {
         for optimizer in optimizers {
             let before_schema = new_plan.schema();
             new_plan = optimizer
-                .optimize(new_plan, session_state.config_options())
+                .optimize_with_context(new_plan, session_state)
                 .map_err(|e| {
                     DataFusionError::Context(optimizer.name().to_string(), Box::new(e))
                 })?;

datafusion/physical-optimizer/src/lib.rs

@@ -47,4 +47,4 @@ pub mod topk_repartition;
 pub mod update_aggr_exprs;
 pub mod utils;
 
-pub use optimizer::PhysicalOptimizerRule;
+pub use optimizer::{ConfigOnlyContext, PhysicalOptimizerContext, PhysicalOptimizerRule};

datafusion/physical-optimizer/src/optimizer.rs

@@ -42,6 +42,48 @@ use crate::pushdown_sort::PushdownSort;
 use datafusion_common::Result;
 use datafusion_common::config::ConfigOptions;
 use datafusion_physical_plan::ExecutionPlan;
+use datafusion_physical_plan::operator_statistics::StatisticsRegistry;
+
+/// Context available to physical optimizer rules.
+///
+/// This trait provides access to configuration options and optional statistics
+/// registry for enhanced statistics lookup. It allows optimizer rules to access
+/// extended context without changing the core [`PhysicalOptimizerRule::optimize`]
+/// signature.
+pub trait PhysicalOptimizerContext: Send + Sync {
+    /// Returns the configuration options.
+    fn config_options(&self) -> &ConfigOptions;
+
+    /// Returns the statistics registry for enhanced statistics lookup.
+    ///
+    /// Returns `None` if no registry is configured, in which case rules
+    /// should fall back to using `ExecutionPlan::partition_statistics()`.
+    fn statistics_registry(&self) -> Option<&StatisticsRegistry> {
+        None
+    }
+}
+
+/// 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.
+pub struct ConfigOnlyContext<'a> {
+    config: &'a ConfigOptions,
+}
+
+impl<'a> ConfigOnlyContext<'a> {
+    /// Create a new context wrapping the given config options.
+    pub fn new(config: &'a ConfigOptions) -> Self {
+        Self { config }
+    }
+}
+
+impl PhysicalOptimizerContext for ConfigOnlyContext<'_> {
+    fn config_options(&self) -> &ConfigOptions {
+        self.config
+    }
+}
 
 /// `PhysicalOptimizerRule` transforms one ['ExecutionPlan'] into another which
 /// computes the same results, but in a potentially more efficient way.
@@ -51,13 +93,29 @@ use datafusion_physical_plan::ExecutionPlan;
 ///
 /// [`SessionState::add_physical_optimizer_rule`]: https://docs.rs/datafusion/latest/datafusion/execution/session_state/struct.SessionState.html#method.add_physical_optimizer_rule
 pub trait PhysicalOptimizerRule: Debug + std::any::Any {
-    /// Rewrite `plan` to an optimized form
+    /// Rewrite `plan` to an optimized form.
+    ///
+    /// This is the primary optimization method. For rules that need access to
+    /// the statistics registry, override [`optimize_with_context`](Self::optimize_with_context) instead.
     fn optimize(
         &self,
         plan: Arc<dyn ExecutionPlan>,
         config: &ConfigOptions,
     ) -> Result<Arc<dyn ExecutionPlan>>;
 
+    /// Rewrite `plan` with access to extended context (statistics registry, etc.).
+    ///
+    /// Override this method if you need access to the statistics registry for
+    /// enhanced statistics lookup. The default implementation simply calls
+    /// [`optimize`](Self::optimize) with the config options from the context.
+    fn optimize_with_context(
+        &self,
+        plan: Arc<dyn ExecutionPlan>,
+        context: &dyn PhysicalOptimizerContext,
+    ) -> Result<Arc<dyn ExecutionPlan>> {
+        self.optimize(plan, context.config_options())
+    }
+
     /// A human readable name for this optimizer rule
     fn name(&self) -> &str;
 

datafusion/physical-plan/src/lib.rs

@@ -87,6 +87,7 @@ pub mod repartition;
 pub mod sort_pushdown;
 pub mod sorts;
 pub mod spill;
+pub mod operator_statistics;
 pub mod stream;
 pub mod streaming;
 pub mod tree_node;