Inject ExpressionAnalyzerRegistry into exec nodes and re-inject via optimizer loop

Add trait methods on ExecutionPlan for expression-level statistics injection
(uses_expression_level_statistics, with_expression_analyzer_registry,
expression_analyzer_registry). The physical planner injects the registry
after plan creation and re-injects after each optimizer rule that modifies
the plan, gated by the use_expression_analyzer config flag.

Author: asolimando

datafusion/common/src/config.rs

@@ -1113,11 +1113,9 @@ config_namespace! {
         /// So if you disable `enable_topk_dynamic_filter_pushdown`, then enable `enable_dynamic_filter_pushdown`, the `enable_topk_dynamic_filter_pushdown` will be overridden.
         pub enable_dynamic_filter_pushdown: bool, default = true
 
-        /// When set to true, the physical planner uses the ExpressionAnalyzer
-        /// framework for expression-level statistics estimation (NDV, selectivity,
-        /// min/max, null fraction). When `use_statistics_registry` is also enabled,
-        /// the registry providers (filters, projections) also use it.
-        /// When false, existing behavior is unchanged.
+        /// When set to true, the pluggable `ExpressionAnalyzerRegistry` from
+        /// `SessionState` is used for expression-level statistics estimation
+        /// (NDV, selectivity, min/max, null fraction) in physical plan operators.
         pub use_expression_analyzer: bool, default = false
 
         /// When set to true, the optimizer will insert filters before a join between
@@ -1252,9 +1250,11 @@ config_namespace! {
         pub join_reordering: bool, default = true
 
         /// When set to true, the physical plan optimizer uses the pluggable
-        /// `StatisticsRegistry` for statistics propagation across operators.
-        /// This enables more accurate cardinality estimates compared to each
-        /// operator's built-in `partition_statistics`.
+        /// `StatisticsRegistry` for a bottom-up statistics walk across operators,
+        /// enabling more accurate cardinality estimates. Enabling
+        /// `use_expression_analyzer` alongside this flag gives built-in
+        /// providers access to custom expression-level analyzers (NDV,
+        /// selectivity) for the operators they process.
         pub use_statistics_registry: bool, default = false
 
         /// When set to true, the physical plan optimizer will prefer HashJoin over SortMergeJoin.

datafusion/core/src/execution/session_state.rs

@@ -190,8 +190,6 @@ pub struct SessionState {
     table_factories: HashMap<String, Arc<dyn TableProviderFactory>>,
     /// Runtime environment
     runtime_env: Arc<RuntimeEnv>,
-    /// Registry for expression-level statistics analyzers (NDV, selectivity, etc.)
-    expression_analyzer_registry: Arc<ExpressionAnalyzerRegistry>,
     /// [FunctionFactory] to support pluggable user defined function handler.
     ///
     /// It will be invoked on `CREATE FUNCTION` statements.
@@ -217,6 +215,13 @@ impl PhysicalOptimizerContext for SessionState {
     fn statistics_registry(&self) -> Option<&StatisticsRegistry> {
         self.statistics_registry.as_ref()
     }
+
+    fn expression_analyzer_registry(&self) -> Option<&Arc<ExpressionAnalyzerRegistry>> {
+        self.config_options()
+            .optimizer
+            .use_expression_analyzer
+            .then_some(&self.expression_analyzer_registry)
+    }
 }
 
 impl Debug for SessionState {

datafusion/core/src/physical_planner.rs

@@ -1111,11 +1111,6 @@ impl DefaultPhysicalPlanner {
                             physical_input,
                         )
                         .with_batch_size(session_state.config().batch_size());
-                        let builder =
-                            match Self::expression_analyzer_registry(session_state) {
-                                Some(r) => builder.with_expression_analyzer_registry(r),
-                                None => builder,
-                            };
                         builder.build()?
                     }
                     PlanAsyncExpr::Async(
@@ -1126,7 +1121,7 @@ impl DefaultPhysicalPlanner {
                             async_map.async_exprs,
                             physical_input,
                         )?;
-                        let builder = FilterExecBuilder::new(
+                        FilterExecBuilder::new(
                             Arc::clone(&runtime_expr[0]),
                             Arc::new(async_exec),
                         )
@@ -1135,13 +1130,8 @@ impl DefaultPhysicalPlanner {
                         .apply_projection(Some(
                             (0..input.schema().fields().len()).collect::<Vec<_>>(),
                         ))?
-                        .with_batch_size(session_state.config().batch_size());
-                        let builder =
-                            match Self::expression_analyzer_registry(session_state) {
-                                Some(r) => builder.with_expression_analyzer_registry(r),
-                                None => builder,
-                            };
-                        builder.build()?
+                        .with_batch_size(session_state.config().batch_size())
+                        .build()?
                     }
                     _ => {
                         return internal_err!(
@@ -1664,15 +1654,16 @@ impl DefaultPhysicalPlanner {
                 {
                     // Use SortMergeJoin if hash join is not preferred
                     let join_on_len = join_on.len();
-                    Arc::new(SortMergeJoinExec::try_new(
+                    let exec = SortMergeJoinExec::try_new(
                         physical_left,
                         physical_right,
                         join_on,
                         join_filter,
                         *join_type,
                         vec![SortOptions::default(); join_on_len],
                         *null_equality,
-                    )?)
+                    )?;
+                    Arc::new(exec)
                 } else if session_state.config().target_partitions() > 1
                     && session_state.config().repartition_joins()
                     && prefer_hash_join
@@ -2780,15 +2771,30 @@ impl DefaultPhysicalPlanner {
         // to verify that the plan fulfills the base requirements.
         InvariantChecker(InvariantLevel::Always).check(&plan)?;
 
+        let use_expression_analyzer = session_state
+            .config_options()
+            .optimizer
+            .use_expression_analyzer;
         let mut new_plan = Arc::clone(&plan);
         for optimizer in optimizers {
             let before_schema = new_plan.schema();
+            let plan_before_rule = Arc::clone(&new_plan);
             new_plan = optimizer
                 .optimize_with_context(new_plan, session_state)
                 .map_err(|e| {
                     DataFusionError::Context(optimizer.name().to_string(), Box::new(e))
                 })?;
 
+            // Re-inject ExpressionAnalyzer registry into any exec nodes created or replaced by
+            // this rule. Skip if the rule returned the same plan unchanged to
+            // avoid an O(nodes) walk for no-op rules.
+            if use_expression_analyzer && !Arc::ptr_eq(&plan_before_rule, &new_plan) {
+                new_plan = Self::inject_expression_analyzer(
+                    new_plan,
+                    session_state.expression_analyzer_registry(),
+                )?;
+            }
+
             // This only checks the schema in release build, and performs additional checks in debug mode.
             OptimizationInvariantChecker::new(optimizer)
                 .check(&new_plan, &before_schema)?;
@@ -2858,15 +2864,22 @@ impl DefaultPhysicalPlanner {
         Ok(mem_exec)
     }
 
-    /// Returns the expression analyzer registry if the config option is enabled.
-    fn expression_analyzer_registry(
-        session_state: &SessionState,
-    ) -> Option<Arc<ExpressionAnalyzerRegistry>> {
-        session_state
-            .config_options()
-            .optimizer
-            .use_expression_analyzer
-            .then(|| Arc::clone(session_state.expression_analyzer_registry()))
+    /// Walks `plan` and injects `registry` into every exec node that accepts it,
+    /// skipping nodes that already have a registry set.
+    fn inject_expression_analyzer(
+        plan: Arc<dyn ExecutionPlan>,
+        registry: &Arc<ExpressionAnalyzerRegistry>,
+    ) -> Result<Arc<dyn ExecutionPlan>> {
+        use datafusion_common::tree_node::{Transformed, TreeNode};
+
+        plan.transform_up(|node| {
+            if let Some(updated) = node.with_expression_analyzer_registry(registry) {
+                Ok(Transformed::yes(updated))
+            } else {
+                Ok(Transformed::no(node))
+            }
+        })
+        .map(|t| t.data)
     }
 
     fn create_project_physical_exec(
@@ -2928,11 +2941,7 @@ impl DefaultPhysicalPlanner {
                     .into_iter()
                     .map(|(expr, alias)| ProjectionExpr { expr, alias })
                     .collect();
-                let mut proj_exec = ProjectionExec::try_new(proj_exprs, input_exec)?;
-                if let Some(r) = Self::expression_analyzer_registry(session_state) {
-                    proj_exec = proj_exec.with_expression_analyzer_registry(r);
-                }
-                Ok(Arc::new(proj_exec))
+                Ok(Arc::new(ProjectionExec::try_new(proj_exprs, input_exec)?))
             }
             PlanAsyncExpr::Async(
                 async_map,
@@ -2944,18 +2953,10 @@ impl DefaultPhysicalPlanner {
                     .into_iter()
                     .map(|(expr, alias)| ProjectionExpr { expr, alias })
                     .collect();
-                let mut new_proj_exec =
-                    ProjectionExec::try_new(proj_exprs, Arc::new(async_exec))?;
-                if session_state
-                    .config_options()
-                    .optimizer
-                    .use_expression_analyzer
-                {
-                    new_proj_exec = new_proj_exec.with_expression_analyzer_registry(
-                        Arc::clone(session_state.expression_analyzer_registry()),
-                    );
-                }
-                Ok(Arc::new(new_proj_exec))
+                Ok(Arc::new(ProjectionExec::try_new(
+                    proj_exprs,
+                    Arc::new(async_exec),
+                )?))
             }
             _ => internal_err!("Unexpected PlanAsyncExpressions variant"),
         }

datafusion/core/tests/physical_optimizer/expression_analyzer_injection.rs

@@ -0,0 +1,142 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Tests that ExpressionAnalyzerRegistry is preserved through optimizer rules.
+//!
+//! Optimizer rules like CombinePartialFinalAggregate, EnforceDistribution, and
+//! ProjectionPushdown rebuild exec nodes internally. The post-rule injection loop
+//! in the physical planner must re-inject the ExpressionAnalyzer registry into those new nodes.
+
+#[cfg(test)]
+mod test {
+    use std::sync::Arc;
+
+    use arrow::array::{Int64Array, RecordBatch};
+    use arrow_schema::{DataType, Field, Schema};
+    use datafusion::prelude::SessionContext;
+    use datafusion_common::Result;
+    use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion};
+    use datafusion_execution::config::SessionConfig;
+    use datafusion_physical_plan::ExecutionPlan;
+
+    /// Creates a SessionContext with `use_expression_analyzer = true`.
+    fn make_ctx() -> SessionContext {
+        let mut config = SessionConfig::new();
+        config.options_mut().optimizer.use_expression_analyzer = true;
+        SessionContext::new_with_config(config)
+    }
+
+    /// Registers an in-memory table with a single int64 column.
+    async fn register_table(ctx: &SessionContext, name: &str) -> Result<()> {
+        let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)]));
+        let batch = RecordBatch::try_new(
+            Arc::clone(&schema),
+            vec![Arc::new(Int64Array::from(vec![1i64, 2, 3]))],
+        )?;
+        let table = datafusion::datasource::MemTable::try_new(schema, vec![vec![batch]])?;
+        ctx.register_table(name, Arc::new(table))?;
+        Ok(())
+    }
+
+    /// Asserts that every node in `plan` that declares
+    /// `uses_expression_level_statistics() == true` has the registry set.
+    fn assert_expression_analyzer_injected(plan: &Arc<dyn ExecutionPlan>) -> Result<()> {
+        plan.clone().apply(|node| {
+            if node.uses_expression_level_statistics() {
+                assert!(
+                    node.expression_analyzer_registry().is_some(),
+                    "{} declares uses_expression_level_statistics but has no registry after optimization",
+                    node.name()
+                );
+            }
+            Ok(TreeNodeRecursion::Continue)
+        })?;
+        Ok(())
+    }
+
+    /// Filter and aggregate go through CombinePartialFinalAggregate which rebuilds
+    /// the partial/final AggregateExec pair; both must end up with the ExpressionAnalyzer registry injected.
+    #[tokio::test]
+    async fn test_expression_analyzer_injected_through_aggregate_rewrite() -> Result<()> {
+        let ctx = make_ctx();
+        register_table(&ctx, "t").await?;
+
+        let df = ctx
+            .sql("SELECT a, COUNT(*) FROM t WHERE a > 1 GROUP BY a")
+            .await?;
+        let plan = df.create_physical_plan().await?;
+
+        assert_expression_analyzer_injected(&plan)
+    }
+
+    /// A simple filter query - FilterExec must have the registry even without complex rewrites.
+    #[tokio::test]
+    async fn test_expression_analyzer_injected_filter() -> Result<()> {
+        let ctx = make_ctx();
+        register_table(&ctx, "t").await?;
+
+        let df = ctx.sql("SELECT a FROM t WHERE a > 1").await?;
+        let plan = df.create_physical_plan().await?;
+
+        assert_expression_analyzer_injected(&plan)
+    }
+
+    /// A join query - HashJoinExec must have the registry after join selection rewrites it.
+    #[tokio::test]
+    async fn test_expression_analyzer_injected_join() -> Result<()> {
+        let ctx = make_ctx();
+        register_table(&ctx, "t1").await?;
+        register_table(&ctx, "t2").await?;
+
+        let df = ctx
+            .sql("SELECT t1.a FROM t1 JOIN t2 ON t1.a = t2.a WHERE t1.a > 0")
+            .await?;
+        let plan = df.create_physical_plan().await?;
+
+        assert_expression_analyzer_injected(&plan)
+    }
+
+    /// Disabled by default: verify that the flag gates the injection walk.
+    /// When `use_expression_analyzer = false`, exec nodes must NOT have the registry set.
+    #[tokio::test]
+    async fn test_expression_analyzer_not_injected_when_disabled() -> Result<()> {
+        let ctx = SessionContext::new(); // default config, ExpressionAnalyzer disabled
+
+        let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)]));
+        let batch = RecordBatch::try_new(
+            Arc::clone(&schema),
+            vec![Arc::new(Int64Array::from(vec![1i64, 2, 3]))],
+        )?;
+        let table = datafusion::datasource::MemTable::try_new(schema, vec![vec![batch]])?;
+        ctx.register_table("t", Arc::new(table))?;
+
+        let df = ctx.sql("SELECT a FROM t WHERE a > 1").await?;
+        let plan = df.create_physical_plan().await?;
+
+        plan.clone().apply(|node| {
+            if node.uses_expression_level_statistics() {
+                assert!(
+                    node.expression_analyzer_registry().is_none(),
+                    "{} should NOT have the registry when use_expression_analyzer=false",
+                    node.name()
+                );
+            }
+            Ok(TreeNodeRecursion::Continue)
+        })?;
+        Ok(())
+    }
+}

datafusion/core/tests/physical_optimizer/mod.rs

@@ -38,4 +38,5 @@ mod sanity_checker;
 mod test_utils;
 mod window_optimize;
 
+mod expression_analyzer_injection;
 mod pushdown_utils;