Add use_statistics_registry config and registry-aware join selection

asolimando · asolimando · commit a38d3ae336e6 · 2026-04-07T14:38:05.000+02:00
Adds a pluggable statistics path for JoinSelection that uses the
StatisticsRegistry instead of each operator's built-in partition_statistics.

- Add optimizer.use_statistics_registry config flag (default=false)
- Always initialize StatisticsRegistry with built-in providers in
  SessionStateBuilder, so SET via SQL takes effect without rebuild
- Override optimize_with_context in JoinSelection to pass the registry
  to should_swap_join_order when the flag is enabled
- Add statistics_registry.slt demonstrating how the registry produces
  more conservative join estimates for skewed data (10*10=100 cartesian
  fallback vs 10*10/3=33 range-NDV estimate), triggering the correct
  build-side swap that the built-in estimator misses
diff --git a/datafusion/common/src/config.rs b/datafusion/common/src/config.rs
@@ -1100,6 +1100,14 @@ config_namespace! {
         /// query is used.
         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`]: datafusion_physical_plan::operator_statistics::StatisticsRegistry
+        pub use_statistics_registry: bool, default = false
+
         /// When set to true, the physical plan optimizer will prefer HashJoin over SortMergeJoin.
         /// HashJoin can work more efficiently than SortMergeJoin but consumes more memory
         pub prefer_hash_join: bool, default = true
diff --git a/datafusion/core/src/execution/session_state.rs b/datafusion/core/src/execution/session_state.rs
@@ -1598,7 +1598,11 @@ impl SessionStateBuilder {
             runtime_env,
             function_factory,
             cache_factory,
-            statistics_registry,
+            // Always build a default registry with built-in providers when no custom
+            // registry was explicitly set; it is only consulted when
+            // use_statistics_registry = true at optimization time.
+            statistics_registry: statistics_registry
+                .or_else(|| Some(StatisticsRegistry::default_with_builtin_providers())),
             prepared_plans: HashMap::new(),
         };
 
diff --git a/datafusion/physical-optimizer/src/join_selection.rs b/datafusion/physical-optimizer/src/join_selection.rs
@@ -23,9 +23,11 @@
 //! pipeline-friendly ones. To achieve the second goal, it selects the proper
 //! `PartitionMode` and the build side using the available statistics for hash joins.
 
+use crate::optimizer::PhysicalOptimizerContext;
 use crate::PhysicalOptimizerRule;
 use datafusion_common::config::ConfigOptions;
 use datafusion_common::error::Result;
+use datafusion_common::Statistics;
 use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode};
 use datafusion_common::{JoinSide, JoinType, internal_err};
 use datafusion_expr_common::sort_properties::SortProperties;
@@ -37,6 +39,7 @@ use datafusion_physical_plan::joins::{
     CrossJoinExec, HashJoinExec, NestedLoopJoinExec, PartitionMode,
     StreamJoinPartitionMode, SymmetricHashJoinExec,
 };
+use datafusion_physical_plan::operator_statistics::{StatsCache, StatisticsRegistry};
 use datafusion_physical_plan::{ExecutionPlan, ExecutionPlanProperties};
 use std::sync::Arc;
 
@@ -53,36 +56,50 @@ impl JoinSelection {
     }
 }
 
+/// Get statistics for a plan node, using the registry if available.
+fn get_stats(
+    plan: &dyn ExecutionPlan,
+    registry: Option<&StatisticsRegistry>,
+    cache: &mut StatsCache,
+) -> Result<Arc<Statistics>> {
+    if let Some(reg) = registry {
+        reg.compute_cached(plan, cache).map(|s| s.base_arc().clone())
+    } else {
+        plan.partition_statistics(None)
+    }
+}
+
 // 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.
 // TODO: In PrestoSQL, the optimizer flips join sides only if one side is much smaller than the other by more than SIZE_DIFFERENCE_THRESHOLD times, by default is 8 times.
 /// Checks whether join inputs should be swapped using available statistics.
 ///
 /// It follows these steps:
-/// 1. Compare the in-memory sizes of both sides, and place the smaller side on
+/// 1. If a [`StatisticsRegistry`] is provided, use it for cross-operator estimates
+///    (e.g., intermediate join outputs that would otherwise have `Absent` statistics).
+/// 2. Compare the in-memory sizes of both sides, and place the smaller side on
 ///    the left (build) side.
-/// 2. If in-memory byte sizes are unavailable, fall back to row counts.
-/// 3. Do not reorder the join if neither statistic is available, or if
+/// 3. If in-memory byte sizes are unavailable, fall back to row counts.
+/// 4. Do not reorder the join if neither statistic is available, or if
 ///    `datafusion.optimizer.join_reordering` is disabled.
 ///
-///
 /// Used configurations inside arg `config`
 /// - `config.optimizer.join_reordering`: allows or forbids statistics-driven join swapping
 pub(crate) fn should_swap_join_order(
     left: &dyn ExecutionPlan,
     right: &dyn ExecutionPlan,
     config: &ConfigOptions,
+    registry: Option<&StatisticsRegistry>,
+    cache: &mut StatsCache,
 ) -> Result<bool> {
     if !config.optimizer.join_reordering {
         return Ok(false);
     }
 
-    // Get the left and right table's total bytes
-    // If both the left and right tables contain total_byte_size statistics,
-    // use `total_byte_size` to determine `should_swap_join_order`, else use `num_rows`
-    let left_stats = left.partition_statistics(None)?;
-    let right_stats = right.partition_statistics(None)?;
-    // First compare `total_byte_size` of left and right side,
-    // if information in this field is insufficient fallback to the `num_rows`
+    let left_stats = get_stats(left, registry, cache)?;
+    let right_stats = get_stats(right, registry, cache)?;
+
+    // First compare total_byte_size, then fall back to num_rows if byte
+    // sizes are unavailable.
     match (
         left_stats.total_byte_size.get_value(),
         right_stats.total_byte_size.get_value(),
@@ -102,8 +119,10 @@ fn supports_collect_by_thresholds(
     plan: &dyn ExecutionPlan,
     threshold_byte_size: usize,
     threshold_num_rows: usize,
+    registry: Option<&StatisticsRegistry>,
+    cache: &mut StatsCache,
 ) -> bool {
-    let Ok(stats) = plan.partition_statistics(None) else {
+    let Ok(stats) = get_stats(plan, registry, cache) else {
         return false;
     };
 
@@ -149,8 +168,38 @@ impl PhysicalOptimizerRule for JoinSelection {
         //   do not modify join sides.
         // - We will also swap left and right sides for cross joins so that the left
         //   side is the small side.
+        let mut cache = StatsCache::new();
+        new_plan
+            .transform_up(|plan| statistical_join_selection_subrule(plan, config, None, &mut cache))
+            .data()
+    }
+
+    // TODO: this intentionally duplicates the subrule pipeline from optimize()
+    // to reduce the change surface; once optimize_with_context fully replaces
+    // optimize, the duplication goes away.
+    fn optimize_with_context(
+        &self,
+        plan: Arc<dyn ExecutionPlan>,
+        context: &dyn PhysicalOptimizerContext,
+    ) -> Result<Arc<dyn ExecutionPlan>> {
+        let config = context.config_options();
+        let registry = if config.optimizer.use_statistics_registry {
+            context.statistics_registry()
+        } else {
+            None
+        };
+        let subrules: Vec<Box<PipelineFixerSubrule>> = vec![
+            Box::new(hash_join_convert_symmetric_subrule),
+            Box::new(hash_join_swap_subrule),
+        ];
+        let new_plan = plan
+            .transform_up(|p| apply_subrules(p, &subrules, config))
+            .data()?;
+        let mut cache = StatsCache::new();
         new_plan
-            .transform_up(|plan| statistical_join_selection_subrule(plan, config))
+            .transform_up(|plan| {
+                statistical_join_selection_subrule(plan, config, registry, &mut cache)
+            })
             .data()
     }
 
@@ -178,6 +227,8 @@ pub(crate) fn try_collect_left(
     hash_join: &HashJoinExec,
     ignore_threshold: bool,
     config: &ConfigOptions,
+    registry: Option<&StatisticsRegistry>,
+    cache: &mut StatsCache,
 ) -> Result<Option<Arc<dyn ExecutionPlan>>> {
     let left = hash_join.left();
     let right = hash_join.right();
@@ -188,20 +239,24 @@ pub(crate) fn try_collect_left(
             &**left,
             optimizer_config.hash_join_single_partition_threshold,
             optimizer_config.hash_join_single_partition_threshold_rows,
+            registry,
+            cache,
         );
     let right_can_collect = ignore_threshold
         || supports_collect_by_thresholds(
             &**right,
             optimizer_config.hash_join_single_partition_threshold,
             optimizer_config.hash_join_single_partition_threshold_rows,
+            registry,
+            cache,
         );
 
     match (left_can_collect, right_can_collect) {
         (true, true) => {
             // Don't swap null-aware anti joins as they have specific side requirements
             if hash_join.join_type().supports_swap()
                 && !hash_join.null_aware
-                && should_swap_join_order(&**left, &**right, config)?
+                && should_swap_join_order(&**left, &**right, config, registry, cache)?
             {
                 Ok(Some(hash_join.swap_inputs(PartitionMode::CollectLeft)?))
             } else {
@@ -245,13 +300,15 @@ pub(crate) fn try_collect_left(
 pub(crate) fn partitioned_hash_join(
     hash_join: &HashJoinExec,
     config: &ConfigOptions,
+    registry: Option<&StatisticsRegistry>,
+    cache: &mut StatsCache,
 ) -> Result<Arc<dyn ExecutionPlan>> {
     let left = hash_join.left();
     let right = hash_join.right();
     // Don't swap null-aware anti joins as they have specific side requirements
     if hash_join.join_type().supports_swap()
         && !hash_join.null_aware
-        && should_swap_join_order(&**left, &**right, config)?
+        && should_swap_join_order(&**left, &**right, config, registry, cache)?
     {
         hash_join.swap_inputs(PartitionMode::Partitioned)
     } else {
@@ -285,27 +342,31 @@ pub(crate) fn partitioned_hash_join(
 fn statistical_join_selection_subrule(
     plan: Arc<dyn ExecutionPlan>,
     config: &ConfigOptions,
+    registry: Option<&StatisticsRegistry>,
+    cache: &mut StatsCache,
 ) -> Result<Transformed<Arc<dyn ExecutionPlan>>> {
     let transformed =
         if let Some(hash_join) = plan.as_any().downcast_ref::<HashJoinExec>() {
             match hash_join.partition_mode() {
-                PartitionMode::Auto => try_collect_left(hash_join, false, config)?
+                PartitionMode::Auto => try_collect_left(hash_join, false, config, registry, cache)?
                     .map_or_else(
-                        || partitioned_hash_join(hash_join, config).map(Some),
-                        |v| Ok(Some(v)),
-                    )?,
-                PartitionMode::CollectLeft => try_collect_left(hash_join, true, config)?
-                    .map_or_else(
-                        || partitioned_hash_join(hash_join, config).map(Some),
+                        || partitioned_hash_join(hash_join, config, registry, cache).map(Some),
                         |v| Ok(Some(v)),
                     )?,
+                PartitionMode::CollectLeft => {
+                    try_collect_left(hash_join, true, config, registry, cache)?
+                        .map_or_else(
+                            || partitioned_hash_join(hash_join, config, registry, cache).map(Some),
+                            |v| Ok(Some(v)),
+                        )?
+                }
                 PartitionMode::Partitioned => {
                     let left = hash_join.left();
                     let right = hash_join.right();
                     // Don't swap null-aware anti joins as they have specific side requirements
                     if hash_join.join_type().supports_swap()
                         && !hash_join.null_aware
-                        && should_swap_join_order(&**left, &**right, config)?
+                        && should_swap_join_order(&**left, &**right, config, registry, cache)?
                     {
                         hash_join
                             .swap_inputs(PartitionMode::Partitioned)
@@ -318,7 +379,7 @@ fn statistical_join_selection_subrule(
         } else if let Some(cross_join) = plan.as_any().downcast_ref::<CrossJoinExec>() {
             let left = cross_join.left();
             let right = cross_join.right();
-            if should_swap_join_order(&**left, &**right, config)? {
+            if should_swap_join_order(&**left, &**right, config, registry, cache)? {
                 cross_join.swap_inputs().map(Some)?
             } else {
                 None
@@ -327,7 +388,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, config)?
+                && should_swap_join_order(&**left, &**right, config, registry, cache)?
             {
                 nl_join.swap_inputs().map(Some)?
             } else {
diff --git a/datafusion/physical-plan/src/operator_statistics/mod.rs b/datafusion/physical-plan/src/operator_statistics/mod.rs
@@ -947,8 +947,10 @@ impl StatisticsProvider for JoinStatisticsProvider {
             Precision::Inexact(estimated)
         };
 
-        // TODO: once #20184 lands, pass enhanced child_stats to partition_statistics
-        // so column-level stats (NDV, min/max) propagate through the registry walk.
+        // TODO: column-level stats (NDV, min/max) enriched by the registry walk
+        // are lost here because partition_statistics(None) re-fetches raw child
+        // stats internally. Once #20184 lands, pass enhanced child_stats so the
+        // operator's built-in column mapping uses them instead.
         let mut base = Arc::unwrap_or_clone(plan.partition_statistics(None)?);
         rescale_byte_size(&mut base, num_rows);
         Ok(StatisticsResult::Computed(ExtendedStatistics::new(base)))
@@ -1001,8 +1003,10 @@ impl StatisticsProvider for LimitStatisticsProvider {
             },
         };
 
-        // TODO: once #20184 lands, pass enhanced child_stats to partition_statistics
-        // so column-level stats (NDV, min/max) propagate through the registry walk.
+        // TODO: column-level stats (NDV, min/max) enriched by the registry walk
+        // are lost here because partition_statistics(None) re-fetches raw child
+        // stats internally. Once #20184 lands, pass enhanced child_stats so the
+        // operator's built-in column mapping uses them instead.
         let mut base = Arc::unwrap_or_clone(plan.partition_statistics(None)?);
         rescale_byte_size(&mut base, num_rows);
         Ok(StatisticsResult::Computed(ExtendedStatistics::new(base)))
@@ -1044,8 +1048,10 @@ impl StatisticsProvider for UnionStatisticsProvider {
             },
         )?;
 
-        // TODO: once #20184 lands, pass enhanced child_stats to partition_statistics
-        // so column-level stats (NDV, min/max) propagate through the registry walk.
+        // TODO: column-level stats (NDV, min/max) enriched by the registry walk
+        // are lost here because partition_statistics(None) re-fetches raw child
+        // stats internally. Once #20184 lands, pass enhanced child_stats so the
+        // operator's built-in column mapping uses them instead.
         let mut base = Arc::unwrap_or_clone(plan.partition_statistics(None)?);
         rescale_byte_size(&mut base, total);
         Ok(StatisticsResult::Computed(ExtendedStatistics::new(base)))
diff --git a/datafusion/sqllogictest/test_files/information_schema.slt b/datafusion/sqllogictest/test_files/information_schema.slt
@@ -315,6 +315,7 @@ datafusion.optimizer.hash_join_inlist_pushdown_max_size 131072
 datafusion.optimizer.hash_join_single_partition_threshold 1048576
 datafusion.optimizer.hash_join_single_partition_threshold_rows 131072
 datafusion.optimizer.join_reordering true
+datafusion.optimizer.use_statistics_registry false
 datafusion.optimizer.max_passes 3
 datafusion.optimizer.prefer_existing_sort false
 datafusion.optimizer.prefer_existing_union false
@@ -458,6 +459,7 @@ datafusion.optimizer.hash_join_inlist_pushdown_max_size 131072 Maximum size in b
 datafusion.optimizer.hash_join_single_partition_threshold 1048576 The maximum estimated size in bytes for one input side of a HashJoin will be collected into a single partition
 datafusion.optimizer.hash_join_single_partition_threshold_rows 131072 The maximum estimated size in rows for one input side of a HashJoin will be collected into a single partition
 datafusion.optimizer.join_reordering true When set to true, the physical plan optimizer may swap join inputs based on statistics. When set to false, statistics-driven join input reordering is disabled and the original join order in the query is used.
+datafusion.optimizer.use_statistics_registry false 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.
 datafusion.optimizer.max_passes 3 Number of times that the optimizer will attempt to optimize the plan
 datafusion.optimizer.prefer_existing_sort false When true, DataFusion will opportunistically remove sorts when the data is already sorted, (i.e. setting `preserve_order` to true on `RepartitionExec`  and using `SortPreservingMergeExec`) When false, DataFusion will maximize plan parallelism using `RepartitionExec` even if this requires subsequently resorting data using a `SortExec`.
 datafusion.optimizer.prefer_existing_union false When set to true, the optimizer will not attempt to convert Union to Interleave
diff --git a/datafusion/sqllogictest/test_files/statistics_registry.slt b/datafusion/sqllogictest/test_files/statistics_registry.slt
diff --git a/docs/source/user-guide/configs.md b/docs/source/user-guide/configs.md
