From 7452b167b0fd40addb601a9be964f8e062c64004 Mon Sep 17 00:00:00 2001 From: Alessandro Solimando Date: Fri, 17 Jul 2026 00:52:01 +0200 Subject: [PATCH 01/15] feat: make StatisticsRegistry the default operator-statistics path StatisticsContext (the single statistics walk, #23051) now consults the StatisticsRegistry (#21483) at each node before falling back to the operator's statistics_from_inputs, using the provider result's base Statistics. The registry is always-on: datafusion.optimizer.use_statistics_registry is deprecated and ignored (register providers on the SessionState instead). The now-redundant DefaultStatisticsProvider is deprecated, since the walk falls back natively. --- datafusion/common/src/config.rs | 9 +- .../physical-optimizer/src/join_selection.rs | 25 +- .../src/operator_statistics/mod.rs | 331 ++++++++---------- datafusion/physical-plan/src/statistics.rs | 161 ++++++++- datafusion/sqllogictest/src/test_context.rs | 10 + .../test_files/information_schema.slt | 2 +- .../test_files/statistics_registry.slt | 52 +-- .../library-user-guide/upgrading/55.0.0.md | 37 ++ docs/source/user-guide/configs.md | 2 +- 9 files changed, 366 insertions(+), 263 deletions(-) diff --git a/datafusion/common/src/config.rs b/datafusion/common/src/config.rs index b649ecad570d2..42994d60b3ef0 100644 --- a/datafusion/common/src/config.rs +++ b/datafusion/common/src/config.rs @@ -1568,11 +1568,10 @@ 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`. - pub use_statistics_registry: bool, default = false + /// (Deprecated) Ignored: the physical plan optimizer always consults the + /// session's pluggable `StatisticsRegistry` (register providers on the + /// `SessionState`; with none it is a no-op). + pub use_statistics_registry: bool, warn = "`use_statistics_registry` is deprecated and ignored; the StatisticsRegistry is always consulted, register providers on the SessionState", 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 diff --git a/datafusion/physical-optimizer/src/join_selection.rs b/datafusion/physical-optimizer/src/join_selection.rs index 42736f8205089..947fcf855df19 100644 --- a/datafusion/physical-optimizer/src/join_selection.rs +++ b/datafusion/physical-optimizer/src/join_selection.rs @@ -57,17 +57,17 @@ impl JoinSelection { } } -/// Get statistics for a plan node, using the registry if available. +/// Get statistics for a plan node, consulting the registry's providers if one +/// is available. fn get_stats( plan: &dyn ExecutionPlan, registry: Option<&StatisticsRegistry>, ) -> Result> { - if let Some(reg) = registry { - reg.compute(plan) - .map(|s| Arc::::clone(s.base_arc())) - } else { - StatisticsContext::new().compute(plan, &StatisticsArgs::new()) - } + let ctx = match registry { + Some(reg) => StatisticsContext::new_with_registry(reg.clone()), + None => StatisticsContext::new(), + }; + ctx.compute(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. @@ -153,16 +153,7 @@ impl PhysicalOptimizerRule for JoinSelection { context: &dyn PhysicalOptimizerContext, ) -> Result> { let config = context.config_options(); - let mut default_registry = None; - let registry: Option<&StatisticsRegistry> = - if config.optimizer.use_statistics_registry { - Some(context.statistics_registry().unwrap_or_else(|| { - default_registry - .insert(StatisticsRegistry::default_with_builtin_providers()) - })) - } else { - None - }; + let registry = context.statistics_registry(); let subrules: Vec> = vec![ Box::new(hash_join_convert_symmetric_subrule), Box::new(hash_join_swap_subrule), diff --git a/datafusion/physical-plan/src/operator_statistics/mod.rs b/datafusion/physical-plan/src/operator_statistics/mod.rs index 142768fcf49d2..9ad3340e419c2 100644 --- a/datafusion/physical-plan/src/operator_statistics/mod.rs +++ b/datafusion/physical-plan/src/operator_statistics/mod.rs @@ -23,9 +23,9 @@ //! //! # Overview //! -//! The default implementation delegates to each operator's built-in -//! `partition_statistics`. Users can register custom [`StatisticsProvider`] -//! implementations to: +//! With no providers registered, statistics come from each operator's built-in +//! [`ExecutionPlan::statistics_from_inputs`]. Users can register custom +//! [`StatisticsProvider`] implementations to: //! //! 1. Provide statistics for custom [`ExecutionPlan`] implementations //! 2. Override default estimation with advanced approaches (e.g., histograms) @@ -48,35 +48,23 @@ //! 5. [`JoinStatisticsProvider`] - NDV-based join output estimation (hash, sort-merge, cross) //! 6. [`LimitStatisticsProvider`] - caps output at the fetch limit (local and global) //! 7. [`UnionStatisticsProvider`] - sums input row counts -//! 8. [`DefaultStatisticsProvider`] - fallback to `partition_statistics(None)` //! -//! # Relationship to [#20184](https://github.com/apache/datafusion/issues/20184) +//! # Statistics walk //! -//! This module performs its own bottom-up tree walk in [`StatisticsRegistry::compute`], -//! separate from the walk optimizer rules do via `transform_up`. This means existing -//! rules that call `partition_statistics` directly bypass the registry. +//! Providers plug into the single [`StatisticsContext`] walk: at each node the +//! chain is consulted before falling back to the operator's built-in +//! [`ExecutionPlan::statistics_from_inputs`], with children already resolved by +//! that same walk. The first provider that returns a computed result sets the +//! node's statistics. //! -//! [#20184](https://github.com/apache/datafusion/issues/20184) adds a `child_stats` -//! parameter to `partition_statistics`. Once it lands, the registry can feed enriched -//! **base** [`Statistics`] into operators' built-in `partition_statistics` calls, -//! removing redundancy for the base-stats path (row counts, column stats). However, -//! the separate registry walk is still required for [`ExtendedStatistics`] extension -//! propagation: `partition_statistics` returns `Arc`, so extensions -//! (histograms, sketches, etc.) are stripped at that boundary and can only flow -//! through the registry walk. -//! -//! If [`Statistics`] itself were extended to carry a type-erased extension map -//! (similar to [`ExtendedStatistics`]), the registry walk could be dropped entirely: -//! extensions would flow naturally through `partition_statistics(child_stats)` and -//! the registry would become a pure chain-of-responsibility on top of the existing -//! traversal with no separate walk needed. +//! [`StatisticsContext`]: crate::statistics::StatisticsContext //! //! # Example //! //! ```ignore //! use datafusion_physical_plan::operator_statistics::*; //! -//! // Create registry with default provider +//! // Create an empty registry //! let mut registry = StatisticsRegistry::new(); //! //! // Register custom provider (higher priority) @@ -239,7 +227,10 @@ pub enum StatisticsResult { /// } /// ``` pub trait StatisticsProvider: Debug + Send + Sync { - /// Compute statistics for an [`ExecutionPlan`] node. + /// Compute *overall* (all-partitions) statistics for a node. + /// + /// The partition-agnostic entry point; per-partition requests go through + /// [`Self::compute_statistics_with_args`] (whose default delegates per partition). /// /// # Arguments /// * `plan` - The execution plan node to compute statistics for @@ -253,21 +244,56 @@ pub trait StatisticsProvider: Debug + Send + Sync { &self, plan: &dyn ExecutionPlan, child_stats: &[ExtendedStatistics], - ) -> Result; + ) -> Result { + let _ = (plan, child_stats); + Ok(StatisticsResult::Delegate) + } + + /// Compute statistics for `args.partition()` (or overall when `None`). + /// + /// The default is overall-only: forwards to [`Self::compute_statistics`] for + /// `None`, and returns `Delegate` for a specific partition so the operator's + /// own per-partition statistics are used. Override to provide per-partition stats. + fn compute_statistics_with_args( + &self, + plan: &dyn ExecutionPlan, + child_stats: &[ExtendedStatistics], + args: &StatisticsArgs, + ) -> Result { + if args.partition().is_some() { + Ok(StatisticsResult::Delegate) + } else { + self.compute_statistics(plan, child_stats) + } + } } -/// Default statistics provider that delegates to each operator's built-in -/// `partition_statistics` implementation. +/// Deprecated statistics provider that delegates to each operator's built-in +/// [`ExecutionPlan::statistics_from_inputs`]. +/// +/// Redundant now that the statistics walk falls back to `statistics_from_inputs` +/// natively when the provider chain delegates or is empty, so a terminal +/// "default" provider is unnecessary. Kept for backward compatibility; do not +/// add it to a provider chain. +#[deprecated( + since = "55.0.0", + note = "redundant: the statistics walk falls back to `statistics_from_inputs` when the provider chain delegates or is empty; register no terminal provider" +)] #[derive(Debug, Default)] pub struct DefaultStatisticsProvider; +#[expect(deprecated)] impl StatisticsProvider for DefaultStatisticsProvider { fn compute_statistics( &self, plan: &dyn ExecutionPlan, - _child_stats: &[ExtendedStatistics], + child_stats: &[ExtendedStatistics], ) -> Result { - let base = StatisticsContext::new().compute(plan, &StatisticsArgs::new())?; + let child_base: Vec> = child_stats + .iter() + .map(|c| Arc::clone(c.base_arc())) + .collect(); + let base = plan.statistics_from_inputs(&child_base, &StatisticsArgs::new())?; Ok(StatisticsResult::Computed(ExtendedStatistics::new_arc( base, ))) @@ -299,7 +325,7 @@ impl StatisticsRegistry { /// Create a new empty registry. /// /// With no providers, `compute()` falls back to each plan node's - /// built-in `partition_statistics()`. Register providers to enhance + /// built-in `statistics_from_inputs()`. Register providers to enhance /// statistics (e.g., inject NDV, use histograms). pub fn new() -> Self { Self { @@ -322,7 +348,6 @@ impl StatisticsRegistry { /// 5. [`JoinStatisticsProvider`] /// 6. [`LimitStatisticsProvider`] /// 7. [`UnionStatisticsProvider`] - /// 8. [`DefaultStatisticsProvider`] pub fn default_with_builtin_providers() -> Self { Self::with_providers(vec![ Arc::new(FilterStatisticsProvider), @@ -332,7 +357,6 @@ impl StatisticsRegistry { Arc::new(JoinStatisticsProvider), Arc::new(LimitStatisticsProvider), Arc::new(UnionStatisticsProvider), - Arc::new(DefaultStatisticsProvider), ]) } @@ -346,44 +370,14 @@ impl StatisticsRegistry { &self.providers } - /// Compute extended statistics for a plan through the provider chain. + /// Compute extended statistics for `plan` through the provider chain. /// - /// Performs a bottom-up tree walk: child statistics are computed recursively - /// and passed to providers, mirroring how `partition_statistics` composes - /// operators. Once [#20184](https://github.com/apache/datafusion/issues/20184) - /// lands, the registry can feed enriched base stats directly into - /// `partition_statistics(child_stats)`, removing the need for a separate walk. - /// - /// If no providers are registered, falls back to the plan's built-in - /// `partition_statistics(None)` with no overhead. + /// Thin wrapper over the single [`StatisticsContext`] walk (the same path the + /// optimizer and EXPLAIN use), so there is one traversal implementation. The + /// walk carries base [`Statistics`] only; extensions are not propagated. pub fn compute(&self, plan: &dyn ExecutionPlan) -> Result { - // Fast path: no providers registered, skip the walk entirely - if self.providers.is_empty() { - let base = StatisticsContext::new().compute(plan, &StatisticsArgs::new())?; - return Ok(ExtendedStatistics::new_arc(base)); - } - - let children = plan.children(); - - // For leaf nodes, try providers with empty child stats. - // For non-leaf nodes, recursively compute enhanced child stats first. - let child_stats: Vec = if children.is_empty() { - Vec::new() - } else { - children - .iter() - .map(|child| self.compute(child.as_ref())) - .collect::>>()? - }; - - for provider in &self.providers { - match provider.compute_statistics(plan, &child_stats)? { - StatisticsResult::Computed(stats) => return Ok(stats), - StatisticsResult::Delegate => continue, - } - } - // Fallback: use plan's built-in stats - let base = StatisticsContext::new().compute(plan, &StatisticsArgs::new())?; + let base = StatisticsContext::new_with_registry(self.clone()) + .compute(plan, &StatisticsArgs::new())?; Ok(ExtendedStatistics::new_arc(base)) } @@ -471,7 +465,7 @@ pub fn ndv_after_selectivity( /// Rescale `total_byte_size` proportionally after overriding `num_rows`. /// /// When a provider replaces `num_rows` but keeps the rest of the stats from -/// `partition_statistics`, the original `total_byte_size` becomes inconsistent. +/// `statistics_from_inputs`, the original `total_byte_size` becomes inconsistent. /// This function adjusts it by the ratio `new_rows / old_rows`, preserving the /// average bytes-per-row from the original estimate. fn rescale_byte_size(stats: &mut Statistics, new_num_rows: Precision) { @@ -496,18 +490,20 @@ fn rescale_byte_size(stats: &mut Statistics, new_num_rows: Precision) { }; } -/// Fetches base statistics from the operator's built-in `partition_statistics`, -/// overrides `num_rows` with the registry-computed estimate, and rescales -/// `total_byte_size` proportionally. -/// -/// Used by providers that compute a better row count but cannot yet propagate -/// column-level stats (NDV, min/max) through the operator — pending #20184. +/// Overrides the operator's built-in output statistics (computed from the +/// pre-resolved `child_stats`) with `num_rows`, rescaling `total_byte_size` +/// proportionally. Used by providers that refine only the row count. fn computed_with_row_count( plan: &dyn ExecutionPlan, num_rows: Precision, + child_stats: &[ExtendedStatistics], ) -> Result { + let child_base: Vec> = child_stats + .iter() + .map(|c| Arc::clone(c.base_arc())) + .collect(); let mut base = Arc::unwrap_or_clone( - StatisticsContext::new().compute(plan, &StatisticsArgs::new())?, + plan.statistics_from_inputs(&child_base, &StatisticsArgs::new())?, ); rescale_byte_size(&mut base, num_rows); Ok(StatisticsResult::Computed(ExtendedStatistics::new(base))) @@ -516,7 +512,7 @@ fn computed_with_row_count( /// Statistics provider for [`FilterExec`](crate::filter::FilterExec) that uses /// pre-computed enhanced child statistics from the registry walk. /// -/// Unlike the default provider (which calls `partition_statistics` and gets raw +/// Unlike the built-in fallback (which calls `statistics_from_inputs` and gets raw /// child stats), this provider receives enhanced child stats that may include /// NDV overrides injected at the scan level. It applies the same selectivity /// estimation logic as `FilterExec::statistics_helper`, then additionally @@ -612,7 +608,7 @@ impl StatisticsProvider for ProjectionStatisticsProvider { /// /// These operators (Sort, Repartition, CoalescePartitions, etc.) don't /// transform statistics, so we pass through the enhanced child stats directly. -/// This avoids the fallback calling `partition_statistics(None)` which would +/// This avoids the fallback calling `statistics_from_inputs` (overall) which would /// trigger a redundant internal recursion with raw (non-enhanced) stats. #[derive(Debug, Default)] pub struct PassthroughStatisticsProvider; @@ -654,7 +650,7 @@ impl StatisticsProvider for PassthroughStatisticsProvider { /// produce overestimates. /// /// For GROUPING SETS / CUBE / ROLLUP, delegates to the built-in -/// `partition_statistics`, which handles per-set NDV estimation correctly. +/// `statistics_from_inputs`, which handles per-set NDV estimation correctly. /// /// Delegates when: /// - The plan is not an `AggregateExec` @@ -732,7 +728,7 @@ impl StatisticsProvider for AggregateStatisticsProvider { let num_rows = Precision::Inexact(estimate); - computed_with_row_count(plan, num_rows) + computed_with_row_count(plan, num_rows, child_stats) } } @@ -867,7 +863,7 @@ impl StatisticsProvider for JoinStatisticsProvider { Precision::Inexact(estimated) }; - computed_with_row_count(plan, num_rows) + computed_with_row_count(plan, num_rows, child_stats) } } @@ -917,7 +913,7 @@ impl StatisticsProvider for LimitStatisticsProvider { }, }; - computed_with_row_count(plan, num_rows) + computed_with_row_count(plan, num_rows, child_stats) } } @@ -956,7 +952,7 @@ impl StatisticsProvider for UnionStatisticsProvider { }, )?; - computed_with_row_count(plan, total) + computed_with_row_count(plan, total, child_stats) } } @@ -1159,7 +1155,7 @@ mod tests { let custom_only = StatisticsRegistry::with_providers(vec![Arc::new(CustomStatisticsProvider)]); // CustomStatisticsProvider only handles CustomExec, delegates for others - // With no default provider, filter returns fallback statistics + // With no provider handling FilterExec, it returns the built-in fallback statistics let filter: Arc = Arc::new(FilterExec::try_new(lit(true), Arc::clone(&source))?); let stats = custom_only.compute(filter.as_ref())?; @@ -1206,6 +1202,13 @@ mod tests { vec![&self.input] } + fn child_stats_requests( + &self, + partition: Option, + ) -> Vec { + vec![crate::statistics::ChildStats::At(partition)] + } + fn with_new_children( self: Arc, children: Vec>, @@ -1363,10 +1366,8 @@ mod tests { let filter: Arc = Arc::new(FilterExec::try_new(predicate, source)?); - let registry = StatisticsRegistry::with_providers(vec![ - Arc::new(FilterStatisticsProvider), - Arc::new(DefaultStatisticsProvider), - ]); + let registry = + StatisticsRegistry::with_providers(vec![Arc::new(FilterStatisticsProvider)]); let stats = registry.compute(filter.as_ref())?; let output_ndv_a = stats.base.column_statistics[0] @@ -1553,10 +1554,9 @@ mod tests { )]); let agg = make_aggregate(source, group_by)?; - let registry = StatisticsRegistry::with_providers(vec![ - Arc::new(AggregateStatisticsProvider), - Arc::new(DefaultStatisticsProvider), - ]); + let registry = StatisticsRegistry::with_providers(vec![Arc::new( + AggregateStatisticsProvider, + )]); let stats = registry.compute(agg.as_ref())?; assert_eq!(stats.base.num_rows, Precision::Inexact(10)); Ok(()) @@ -1571,10 +1571,9 @@ mod tests { ]); let agg = make_aggregate(source, group_by)?; - let registry = StatisticsRegistry::with_providers(vec![ - Arc::new(AggregateStatisticsProvider), - Arc::new(DefaultStatisticsProvider), - ]); + let registry = StatisticsRegistry::with_providers(vec![Arc::new( + AggregateStatisticsProvider, + )]); let stats = registry.compute(agg.as_ref())?; // 10 * 5 = 50 assert_eq!(stats.base.num_rows, Precision::Inexact(50)); @@ -1591,10 +1590,9 @@ mod tests { ]); let agg = make_aggregate(source, group_by)?; - let registry = StatisticsRegistry::with_providers(vec![ - Arc::new(AggregateStatisticsProvider), - Arc::new(DefaultStatisticsProvider), - ]); + let registry = StatisticsRegistry::with_providers(vec![Arc::new( + AggregateStatisticsProvider, + )]); let stats = registry.compute(agg.as_ref())?; assert_eq!(stats.base.num_rows, Precision::Inexact(500)); Ok(()) @@ -1610,12 +1608,11 @@ mod tests { )]); let agg = make_aggregate(source, group_by)?; - let registry = StatisticsRegistry::with_providers(vec![ - Arc::new(AggregateStatisticsProvider), - Arc::new(DefaultStatisticsProvider), - ]); + let registry = StatisticsRegistry::with_providers(vec![Arc::new( + AggregateStatisticsProvider, + )]); let stats = registry.compute(agg.as_ref())?; - // Delegates to DefaultStatisticsProvider, which calls partition_statistics + // Delegates, falling back to the operator's built-in statistics_from_inputs assert!( stats.base.num_rows.get_value().is_some() || matches!(stats.base.num_rows, Precision::Absent) @@ -1635,10 +1632,9 @@ mod tests { let group_by = PhysicalGroupBy::new_single(vec![(expr, "sum".to_string())]); let agg = make_aggregate(source, group_by)?; - let registry = StatisticsRegistry::with_providers(vec![ - Arc::new(AggregateStatisticsProvider), - Arc::new(DefaultStatisticsProvider), - ]); + let registry = StatisticsRegistry::with_providers(vec![Arc::new( + AggregateStatisticsProvider, + )]); let stats = registry.compute(agg.as_ref())?; // Should delegate (expression is not a Column) assert!( @@ -1676,13 +1672,12 @@ mod tests { ); let agg = make_aggregate(source, group_by)?; - let registry = StatisticsRegistry::with_providers(vec![ - Arc::new(AggregateStatisticsProvider), - Arc::new(DefaultStatisticsProvider), - ]); + let registry = StatisticsRegistry::with_providers(vec![Arc::new( + AggregateStatisticsProvider, + )]); let stats = registry.compute(agg.as_ref())?; - // Multiple grouping sets: provider delegates to DefaultStatisticsProvider, - // which calls the built-in partition_statistics for correct per-set + // Multiple grouping sets: provider delegates, so the operator's built-in + // statistics_from_inputs computes the correct per-set // NDV estimation. The exact value depends on the built-in implementation. assert!( stats.base.num_rows.get_value().is_some() @@ -1709,12 +1704,11 @@ mod tests { source.schema(), )?); - let registry = StatisticsRegistry::with_providers(vec![ - Arc::new(AggregateStatisticsProvider), - Arc::new(DefaultStatisticsProvider), - ]); + let registry = StatisticsRegistry::with_providers(vec![Arc::new( + AggregateStatisticsProvider, + )]); let stats = registry.compute(agg.as_ref())?; - // Should fall through to DefaultStatisticsProvider (partition_statistics). + // Should fall through to the operator's built-in statistics_from_inputs. // The exact value depends on the built-in implementation. assert!( stats.base.num_rows.get_value().is_some() @@ -1782,10 +1776,8 @@ mod tests { let right = make_source_with_ndv_2col(500, Some(50)); let join = make_hash_join(left, right)?; - let registry = StatisticsRegistry::with_providers(vec![ - Arc::new(JoinStatisticsProvider), - Arc::new(DefaultStatisticsProvider), - ]); + let registry = + StatisticsRegistry::with_providers(vec![Arc::new(JoinStatisticsProvider)]); let stats = registry.compute(join.as_ref())?; assert_eq!(stats.base.num_rows, Precision::Inexact(5000)); Ok(()) @@ -1836,10 +1828,8 @@ mod tests { false, )?); - let registry = StatisticsRegistry::with_providers(vec![ - Arc::new(JoinStatisticsProvider), - Arc::new(DefaultStatisticsProvider), - ]); + let registry = + StatisticsRegistry::with_providers(vec![Arc::new(JoinStatisticsProvider)]); let stats = registry.compute(join.as_ref())?; assert_eq!(stats.base.num_rows, Precision::Inexact(10_000)); Ok(()) @@ -1898,10 +1888,8 @@ mod tests { false, )?); - let registry = StatisticsRegistry::with_providers(vec![ - Arc::new(JoinStatisticsProvider), - Arc::new(DefaultStatisticsProvider), - ]); + let registry = + StatisticsRegistry::with_providers(vec![Arc::new(JoinStatisticsProvider)]); let stats = registry.compute(join.as_ref())?; assert_eq!(stats.base.num_rows, Precision::Inexact(250)); Ok(()) @@ -1914,10 +1902,8 @@ mod tests { let right = make_source_with_ndv_2col(200, None); let join = make_hash_join(left, right)?; - let registry = StatisticsRegistry::with_providers(vec![ - Arc::new(JoinStatisticsProvider), - Arc::new(DefaultStatisticsProvider), - ]); + let registry = + StatisticsRegistry::with_providers(vec![Arc::new(JoinStatisticsProvider)]); let stats = registry.compute(join.as_ref())?; assert_eq!(stats.base.num_rows, Precision::Inexact(20_000)); Ok(()) @@ -1939,12 +1925,10 @@ mod tests { None, )?); - let registry = StatisticsRegistry::with_providers(vec![ - Arc::new(JoinStatisticsProvider), - Arc::new(DefaultStatisticsProvider), - ]); + let registry = + StatisticsRegistry::with_providers(vec![Arc::new(JoinStatisticsProvider)]); let stats = registry.compute(join.as_ref())?; - // Provider delegates; result comes from built-in partition_statistics. + // Provider delegates; result comes from built-in statistics_from_inputs. assert!( stats.base.num_rows.get_value().is_some() || matches!(stats.base.num_rows, Precision::Absent) @@ -1984,10 +1968,8 @@ mod tests { let left = make_source_with_ndv_2col(left_rows, left_ndv); let right = make_source_with_ndv_2col(right_rows, right_ndv); let join = make_hash_join_typed(left, right, join_type)?; - let registry = StatisticsRegistry::with_providers(vec![ - Arc::new(JoinStatisticsProvider), - Arc::new(DefaultStatisticsProvider), - ]); + let registry = + StatisticsRegistry::with_providers(vec![Arc::new(JoinStatisticsProvider)]); Ok(registry.compute(join.as_ref())?.base.num_rows) } @@ -2071,10 +2053,8 @@ mod tests { let right = make_source(200); let join: Arc = Arc::new(CrossJoinExec::new(left, right)); - let registry = StatisticsRegistry::with_providers(vec![ - Arc::new(JoinStatisticsProvider), - Arc::new(DefaultStatisticsProvider), - ]); + let registry = + StatisticsRegistry::with_providers(vec![Arc::new(JoinStatisticsProvider)]); let stats = registry.compute(join.as_ref())?; // Both inputs have Exact row counts -> result is also Exact assert_eq!(stats.base.num_rows, Precision::Exact(20_000)); @@ -2093,10 +2073,8 @@ mod tests { let source = make_source(1000); let limit: Arc = Arc::new(LocalLimitExec::new(source, 100)); - let registry = StatisticsRegistry::with_providers(vec![ - Arc::new(LimitStatisticsProvider), - Arc::new(DefaultStatisticsProvider), - ]); + let registry = + StatisticsRegistry::with_providers(vec![Arc::new(LimitStatisticsProvider)]); let stats = registry.compute(limit.as_ref())?; assert_eq!(stats.base.num_rows, Precision::Exact(100)); Ok(()) @@ -2108,10 +2086,8 @@ mod tests { let source = make_source(50); let limit: Arc = Arc::new(LocalLimitExec::new(source, 200)); - let registry = StatisticsRegistry::with_providers(vec![ - Arc::new(LimitStatisticsProvider), - Arc::new(DefaultStatisticsProvider), - ]); + let registry = + StatisticsRegistry::with_providers(vec![Arc::new(LimitStatisticsProvider)]); let stats = registry.compute(limit.as_ref())?; assert_eq!(stats.base.num_rows, Precision::Exact(50)); Ok(()) @@ -2124,10 +2100,8 @@ mod tests { let limit: Arc = Arc::new(GlobalLimitExec::new(source, 200, Some(100))); - let registry = StatisticsRegistry::with_providers(vec![ - Arc::new(LimitStatisticsProvider), - Arc::new(DefaultStatisticsProvider), - ]); + let registry = + StatisticsRegistry::with_providers(vec![Arc::new(LimitStatisticsProvider)]); let stats = registry.compute(limit.as_ref())?; assert_eq!(stats.base.num_rows, Precision::Exact(100)); Ok(()) @@ -2140,10 +2114,8 @@ mod tests { let limit: Arc = Arc::new(GlobalLimitExec::new(source, 200, Some(50))); - let registry = StatisticsRegistry::with_providers(vec![ - Arc::new(LimitStatisticsProvider), - Arc::new(DefaultStatisticsProvider), - ]); + let registry = + StatisticsRegistry::with_providers(vec![Arc::new(LimitStatisticsProvider)]); let stats = registry.compute(limit.as_ref())?; assert_eq!(stats.base.num_rows, Precision::Exact(0)); Ok(()) @@ -2156,10 +2128,8 @@ mod tests { let source = make_source_with_precision(Precision::Inexact(1000)); let limit: Arc = Arc::new(LocalLimitExec::new(source, 100)); - let registry = StatisticsRegistry::with_providers(vec![ - Arc::new(LimitStatisticsProvider), - Arc::new(DefaultStatisticsProvider), - ]); + let registry = + StatisticsRegistry::with_providers(vec![Arc::new(LimitStatisticsProvider)]); let stats = registry.compute(limit.as_ref())?; assert_eq!(stats.base.num_rows, Precision::Inexact(100)); Ok(()) @@ -2179,10 +2149,8 @@ mod tests { fn test_union_provider_sums_rows() -> Result<()> { let union = UnionExec::try_new(vec![make_source(300), make_source(700)])?; - let registry = StatisticsRegistry::with_providers(vec![ - Arc::new(UnionStatisticsProvider), - Arc::new(DefaultStatisticsProvider), - ]); + let registry = + StatisticsRegistry::with_providers(vec![Arc::new(UnionStatisticsProvider)]); let stats = registry.compute(union.as_ref())?; assert_eq!(stats.base.num_rows, Precision::Exact(1000)); Ok(()) @@ -2196,10 +2164,8 @@ mod tests { make_source(300), ])?; - let registry = StatisticsRegistry::with_providers(vec![ - Arc::new(UnionStatisticsProvider), - Arc::new(DefaultStatisticsProvider), - ]); + let registry = + StatisticsRegistry::with_providers(vec![Arc::new(UnionStatisticsProvider)]); let stats = registry.compute(union.as_ref())?; assert_eq!(stats.base.num_rows, Precision::Exact(600)); Ok(()) @@ -2213,10 +2179,8 @@ mod tests { make_source_with_precision(Precision::Absent), ])?; - let registry = StatisticsRegistry::with_providers(vec![ - Arc::new(UnionStatisticsProvider), - Arc::new(DefaultStatisticsProvider), - ]); + let registry = + StatisticsRegistry::with_providers(vec![Arc::new(UnionStatisticsProvider)]); let stats = registry.compute(union.as_ref())?; assert_eq!(stats.base.num_rows, Precision::Absent); Ok(()) @@ -2243,10 +2207,7 @@ mod tests { } }); - let registry = StatisticsRegistry::with_providers(vec![ - Arc::new(provider), - Arc::new(DefaultStatisticsProvider), - ]); + let registry = StatisticsRegistry::with_providers(vec![Arc::new(provider)]); let source = make_source(1000); let filter: Arc = diff --git a/datafusion/physical-plan/src/statistics.rs b/datafusion/physical-plan/src/statistics.rs index 9246d7d9f5a9c..ae84cec4c58a1 100644 --- a/datafusion/physical-plan/src/statistics.rs +++ b/datafusion/physical-plan/src/statistics.rs @@ -21,6 +21,9 @@ //! [`ExecutionPlan::statistics_from_inputs`]. use crate::ExecutionPlan; +use crate::operator_statistics::{ + ExtendedStatistics, StatisticsRegistry, StatisticsResult, +}; use datafusion_common::{ Result, Statistics, assert_eq_or_internal_err, assert_or_internal_err, }; @@ -118,8 +121,14 @@ pub enum ChildStats { /// Owns the bottom-up traversal and per-walk memoization cache for statistics /// computation. Call [`StatisticsContext::compute`] to walk a plan tree. +/// +/// 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 +/// computation. pub struct StatisticsContext { cache: Rc>, + registry: StatisticsRegistry, } impl Default for StatisticsContext { @@ -129,10 +138,17 @@ impl Default for StatisticsContext { } impl StatisticsContext { - /// Creates a context with an empty cache. + /// Creates a context with an empty cache and no statistics providers. pub fn new() -> Self { + Self::new_with_registry(StatisticsRegistry::new()) + } + + /// Creates a context whose walk consults `registry`'s provider chain before + /// falling back to each operator's built-in statistics. + pub fn new_with_registry(registry: StatisticsRegistry) -> Self { Self { cache: Rc::new(RefCell::new(StatsCache::default())), + registry, } } @@ -172,6 +188,30 @@ impl StatisticsContext { return Ok(Arc::clone(cached)); } + let child_stats = self.resolve_children(plan, partition)?; + + let result = match self.try_provider_stats(plan, &child_stats, args)? { + Some(stats) => stats, + None => plan.statistics_from_inputs(&child_stats, args)?, + }; + self.cache + .borrow_mut() + .insert(plan, partition, Arc::clone(&result)); + Ok(result) + } + + /// Resolves each child's statistics following the operator's + /// [`ExecutionPlan::child_stats_requests`] mapping: `At(p)` computes the child + /// at partition `p`, `Skip` supplies a [`Statistics::new_unknown`] placeholder. + /// + /// A provider computes the same node's statistics, so it depends on the same + /// children as the node itself: an operator whose statistics (built-in or via + /// a provider) depend on a child must declare `At` for it. + fn resolve_children( + &self, + plan: &dyn ExecutionPlan, + partition: Option, + ) -> Result>> { let children = plan.children(); let requests = plan.child_stats_requests(partition); assert_eq_or_internal_err!( @@ -182,7 +222,7 @@ impl StatisticsContext { requests.len(), children.len() ); - let child_stats = children + children .iter() .zip(requests) .map(|(child, directive)| match directive { @@ -193,13 +233,39 @@ impl StatisticsContext { Ok(Arc::new(Statistics::new_unknown(child.schema().as_ref()))) } }) - .collect::>>()?; + .collect() + } - let result = plan.statistics_from_inputs(&child_stats, args)?; - self.cache - .borrow_mut() - .insert(plan, partition, Arc::clone(&result)); - Ok(result) + /// Runs the provider chain, returning the first `Computed` result's base + /// statistics, or `None` if all delegate or the chain is empty. A + /// partition-blind provider applies only to overall stats (its default + /// `compute_statistics_with_args` delegates per partition). + /// + /// Providers may attach custom extensions to their [`ExtendedStatistics`] + /// result, but the walk uses only the base [`Statistics`]. + fn try_provider_stats( + &self, + plan: &dyn ExecutionPlan, + child_stats: &[Arc], + args: &StatisticsArgs, + ) -> Result>> { + let providers = self.registry.providers(); + if providers.is_empty() { + return Ok(None); + } + // Providers take `&[ExtendedStatistics]`; wrap the resolved base stats. + let child_ext: Vec = child_stats + .iter() + .map(|s| ExtendedStatistics::new_arc(Arc::clone(s))) + .collect(); + for provider in providers { + if let StatisticsResult::Computed(stats) = + provider.compute_statistics_with_args(plan, &child_ext, args)? + { + return Ok(Some(Arc::clone(stats.base_arc()))); + } + } + Ok(None) } } @@ -207,10 +273,50 @@ impl StatisticsContext { mod tests { use super::*; use crate::coalesce_partitions::CoalescePartitionsExec; + use crate::operator_statistics::StatisticsProvider; use crate::test::exec::StatisticsExec; + use crate::union::UnionExec; use arrow::datatypes::{DataType, Field, Schema}; use datafusion_common::{ColumnStatistics, stats::Precision}; + /// Overall-only provider: sets a fixed row count for any node. + #[derive(Debug)] + struct FixedRowCountProvider(usize); + impl StatisticsProvider for FixedRowCountProvider { + fn compute_statistics( + &self, + plan: &dyn ExecutionPlan, + _child_stats: &[ExtendedStatistics], + ) -> Result { + let mut stats = Statistics::new_unknown(&plan.schema()); + stats.num_rows = Precision::Exact(self.0); + Ok(StatisticsResult::Computed(ExtendedStatistics::new(stats))) + } + } + + /// Partition-aware provider: encodes the requested partition in the row count. + #[derive(Debug)] + struct PartitionRowCountProvider; + impl StatisticsProvider for PartitionRowCountProvider { + fn compute_statistics_with_args( + &self, + plan: &dyn ExecutionPlan, + _child_stats: &[ExtendedStatistics], + args: &StatisticsArgs, + ) -> Result { + let marker = 700 + args.partition().map_or(0, |p| p + 1); + let mut stats = Statistics::new_unknown(&plan.schema()); + stats.num_rows = Precision::Exact(marker); + Ok(StatisticsResult::Computed(ExtendedStatistics::new(stats))) + } + } + + fn ctx_with(provider: Arc) -> StatisticsContext { + StatisticsContext::new_with_registry(StatisticsRegistry::with_providers(vec![ + provider, + ])) + } + fn make_stats_leaf(num_rows: usize) -> Arc { let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]); let col_stats = vec![ColumnStatistics { @@ -271,4 +377,43 @@ mod tests { ctx.reset_cache(); assert!(ctx.cache.borrow().0.is_empty()); } + + #[test] + fn partition_aware_provider_applies_per_partition() { + let leaf = make_stats_leaf(10); + let ctx = ctx_with(Arc::new(PartitionRowCountProvider)); + + let per_part = ctx + .compute( + leaf.as_ref(), + &StatisticsArgs::new().with_partition(Some(0)), + ) + .unwrap(); + assert_eq!(per_part.num_rows, Precision::Exact(701)); + } + + #[test] + fn per_partition_union_with_registry_no_out_of_bounds() { + // Two 2-partition inputs -> 4 output partitions. Union owns output + // partition 3 via its second input (owning_input(3) = (1, 1)); the first + // input is Skipped, so the walk supplies a placeholder for it (never + // resolving it at partition 3, which is out of that input's 0..2 range). + // An overall-only provider delegates for a specific partition, so p3 keeps + // the operator's honest per-partition answer while the overall request + // picks up the provider's row count. + let union = + UnionExec::try_new(vec![make_stats_leaf(10), make_stats_leaf(20)]).unwrap(); + let ctx = ctx_with(Arc::new(FixedRowCountProvider(999))); + + let p3 = ctx + .compute( + union.as_ref(), + &StatisticsArgs::new().with_partition(Some(3)), + ) + .unwrap(); + assert_eq!(p3.num_rows, Precision::Absent); + + let overall = ctx.compute(union.as_ref(), &StatisticsArgs::new()).unwrap(); + assert_eq!(overall.num_rows, Precision::Exact(999)); + } } diff --git a/datafusion/sqllogictest/src/test_context.rs b/datafusion/sqllogictest/src/test_context.rs index e0aaa91ef6369..927f8239c202c 100644 --- a/datafusion/sqllogictest/src/test_context.rs +++ b/datafusion/sqllogictest/src/test_context.rs @@ -60,6 +60,7 @@ use async_trait::async_trait; use datafusion::common::cast::as_float64_array; use datafusion::execution::SessionStateBuilder; use datafusion::execution::runtime_env::RuntimeEnv; +use datafusion::physical_plan::operator_statistics::StatisticsRegistry; use log::info; use sqlparser::ast; use tempfile::TempDir; @@ -130,6 +131,15 @@ impl TestContext { state_builder.with_type_planner(Arc::new(SqlLogicTestTypePlanner)); } + if matches!( + relative_path.file_name().and_then(|name| name.to_str()), + Some("statistics_registry.slt") + ) { + state_builder = state_builder.with_statistics_registry( + StatisticsRegistry::default_with_builtin_providers(), + ); + } + let state = state_builder.build(); let mut test_ctx = TestContext::new(SessionContext::new_with_state(state)); diff --git a/datafusion/sqllogictest/test_files/information_schema.slt b/datafusion/sqllogictest/test_files/information_schema.slt index 1adf98f67ff99..2e0afa11b2a76 100644 --- a/datafusion/sqllogictest/test_files/information_schema.slt +++ b/datafusion/sqllogictest/test_files/information_schema.slt @@ -499,7 +499,7 @@ datafusion.optimizer.repartition_windows true Should DataFusion repartition data datafusion.optimizer.skip_failed_rules false When set to true, the logical plan optimizer will produce warning messages if any optimization rules produce errors and then proceed to the next rule. When set to false, any rules that produce errors will cause the query to fail datafusion.optimizer.subset_repartition_threshold 4 Partition count threshold for subset satisfaction optimization. When the current partition count is >= this threshold, DataFusion will skip repartitioning if the required partitioning expression is a subset of the current partition expression such as Hash(a) satisfies Hash(a, b). When the current partition count is < this threshold, DataFusion will repartition to increase parallelism even when subset satisfaction applies. Set to 0 to always repartition (disable subset satisfaction optimization). Set to a high value to always use subset satisfaction. Example (subset_repartition_threshold = 4): ```text Hash([a]) satisfies Hash([a, b]) because (Hash([a, b]) is subset of Hash([a]) If current partitions (3) < threshold (4), repartition: AggregateExec: mode=FinalPartitioned, gby=[a, b], aggr=[SUM(x)] RepartitionExec: partitioning=Hash([a, b], 8), input_partitions=3 AggregateExec: mode=Partial, gby=[a, b], aggr=[SUM(x)] DataSourceExec: file_groups={...}, output_partitioning=Hash([a], 3) If current partitions (8) >= threshold (4), use subset satisfaction: AggregateExec: mode=SinglePartitioned, gby=[a, b], aggr=[SUM(x)] DataSourceExec: file_groups={...}, output_partitioning=Hash([a], 8) ``` datafusion.optimizer.top_down_join_key_reordering true When set to true, the physical plan optimizer will run a top down process to reorder the join keys -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.use_statistics_registry false (Deprecated) Ignored: the physical plan optimizer always consults the session's pluggable `StatisticsRegistry` (register providers on the `SessionState`; with none it is a no-op). datafusion.runtime.file_statistics_cache_limit 20M Maximum memory to use for file statistics cache. Supports suffixes K (kilobytes), M (megabytes), and G (gigabytes) or '0' for 0. Example: '2G' for 2 gigabytes. datafusion.runtime.list_files_cache_limit 1M Maximum memory to use for list files cache. Supports suffixes K (kilobytes), M (megabytes), and G (gigabytes) or '0' for 0. Example: '2G' for 2 gigabytes. datafusion.runtime.list_files_cache_ttl NULL TTL (time-to-live) of the entries in the list file cache. Supports units m (minutes), and s (seconds). Example: '2m' for 2 minutes. diff --git a/datafusion/sqllogictest/test_files/statistics_registry.slt b/datafusion/sqllogictest/test_files/statistics_registry.slt index 89258bec299c1..ff4a6ae67bde5 100644 --- a/datafusion/sqllogictest/test_files/statistics_registry.slt +++ b/datafusion/sqllogictest/test_files/statistics_registry.slt @@ -22,8 +22,9 @@ # orders (10 rows, same distribution) # dim_small (50 rows) # -# Built-in: 10*10 / NDV(3) = 33 < 50 -> keeps inner join on build side (wrong; actual = 66) -# Registry: 10*10 = 100 > 50 -> swaps dim_small to build side (correct) +# Built-in operator stats give 10*10 / NDV(3) = 33 < 50, which would keep the +# inner join on the build side (wrong; actual output is 66). The registry's +# conservative estimate 10*10 = 100 > 50 swaps dim_small to the build side. # # Parquet files written by COPY TO carry min/max stats (NDV=3 via range) but no # distinct_count, so the registry falls back to the cartesian product upper bound. @@ -87,32 +88,8 @@ CREATE EXTERNAL TABLE dim_small STORED AS PARQUET LOCATION 'test_files/scratch/statistics_registry/dim_small.parquet'; -# -- Without registry -------------------------------------------------------- -# Built-in estimate 33 < 50: inner join stays on left (suboptimal; actual output is 66 rows) - -statement ok -set datafusion.optimizer.use_statistics_registry = false; - -query TT -EXPLAIN SELECT o.order_id, c.region_id, d.label -FROM customers c -JOIN orders o ON c.customer_id = o.customer_id -JOIN dim_small d ON o.small_id = d.small_id; ----- -physical_plan -01)HashJoinExec: mode=Partitioned, join_type=Inner, on=[(small_id@2, small_id@0)], projection=[order_id@1, region_id@0, label@4] -02)--RepartitionExec: partitioning=Hash([small_id@2], 4), input_partitions=1, maintains_sort_order=true -03)----HashJoinExec: mode=Partitioned, join_type=Inner, on=[(customer_id@0, customer_id@1)], projection=[region_id@1, order_id@2, small_id@4] -04)------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/statistics_registry/customers.parquet]]}, projection=[customer_id, region_id], output_ordering=[region_id@1 ASC NULLS LAST], file_type=parquet -05)------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/statistics_registry/orders.parquet]]}, projection=[order_id, customer_id, small_id], output_ordering=[order_id@0 ASC NULLS LAST], file_type=parquet, predicate=DynamicFilter [ empty ], dynamic_rg_pruning=eligible -06)--RepartitionExec: partitioning=Hash([small_id@0], 4), input_partitions=1, maintains_sort_order=true -07)----DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/statistics_registry/dim_small.parquet]]}, projection=[small_id, label], output_ordering=[small_id@0 ASC NULLS LAST], file_type=parquet, predicate=DynamicFilter [ empty ], dynamic_rg_pruning=eligible - -# -- With registry ----------------------------------------------------------- -# Conservative estimate 100 > 50: dim_small correctly swapped to build side - -statement ok -set datafusion.optimizer.use_statistics_registry = true; +# -- Registry-driven join ordering ------------------------------------------- +# Conservative estimate 100 > 50: dim_small correctly swapped to the build side. query TT EXPLAIN SELECT o.order_id, c.region_id, d.label @@ -129,21 +106,7 @@ physical_plan 06)------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/statistics_registry/customers.parquet]]}, projection=[customer_id, region_id], output_ordering=[region_id@1 ASC NULLS LAST], file_type=parquet 07)------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/statistics_registry/orders.parquet]]}, projection=[order_id, customer_id, small_id], output_ordering=[order_id@0 ASC NULLS LAST], file_type=parquet, predicate=DynamicFilter [ empty ] AND DynamicFilter [ empty ], dynamic_rg_pruning=eligible -# -- Verify results are identical regardless of join order -------------------- - -statement ok -set datafusion.optimizer.use_statistics_registry = false; - -query I -SELECT count(*) -FROM customers c -JOIN orders o ON c.customer_id = o.customer_id -JOIN dim_small d ON o.small_id = d.small_id; ----- -66 - -statement ok -set datafusion.optimizer.use_statistics_registry = true; +# -- Correctness ------------------------------------------------------------- query I SELECT count(*) @@ -158,9 +121,6 @@ JOIN dim_small d ON o.small_id = d.small_id; statement ok set datafusion.explain.physical_plan_only = false; -statement ok -set datafusion.optimizer.use_statistics_registry = false; - statement ok set datafusion.optimizer.hash_join_single_partition_threshold = 1048576; diff --git a/docs/source/library-user-guide/upgrading/55.0.0.md b/docs/source/library-user-guide/upgrading/55.0.0.md index f7d2745549c99..63ee5b26b3779 100644 --- a/docs/source/library-user-guide/upgrading/55.0.0.md +++ b/docs/source/library-user-guide/upgrading/55.0.0.md @@ -374,6 +374,43 @@ let stats = plan.partition_statistics(None)?; let stats = StatisticsContext::new().compute(plan.as_ref(), &StatisticsArgs::new())?; ``` +### `datafusion.optimizer.use_statistics_registry` is deprecated and ignored + +The `datafusion.optimizer.use_statistics_registry` config flag is deprecated and +now ignored (setting it emits a deprecation warning). The pluggable +`StatisticsRegistry` is always consulted during the single statistics walk: +providers registered on the session take effect directly. With no providers +registered the registry is a no-op, so default behavior is unchanged from the +previous default (`use_statistics_registry = false`). + +**Who is affected:** + +- Anyone who set `datafusion.optimizer.use_statistics_registry`. The setting is + still accepted (with a deprecation warning) but has no effect, and will be + removed in a future release. + +**Migration guide:** + +- Drop any `use_statistics_registry` setting. +- To use statistics providers, register them on the session with + `SessionStateBuilder::with_statistics_registry(...)`. To keep the built-in + NDV-aware providers the flag previously enabled, register + `StatisticsRegistry::default_with_builtin_providers()`. +- To opt out, register nothing (the default). + +### `DefaultStatisticsProvider` is deprecated + +`datafusion_physical_plan::operator_statistics::DefaultStatisticsProvider` is +deprecated. It is redundant: the statistics walk (`StatisticsContext`) falls back +to each operator's `statistics_from_inputs` when the provider chain delegates or +is empty, so a terminal "default" provider is no longer needed. It is no longer +part of `StatisticsRegistry::default_with_builtin_providers()`. + +**Migration guide:** + +- Remove `DefaultStatisticsProvider` from any custom provider chain; register no + terminal provider instead (the walk falls back on its own). + ### `DdlStatement::CreateExternalTable` and `CreateFunction` are now boxed The two largest variants of `datafusion_expr::DdlStatement` are now diff --git a/docs/source/user-guide/configs.md b/docs/source/user-guide/configs.md index f6e072b59bceb..7c3af9aefd61e 100644 --- a/docs/source/user-guide/configs.md +++ b/docs/source/user-guide/configs.md @@ -169,7 +169,7 @@ The following configuration settings are available: | datafusion.optimizer.max_passes | 3 | Number of times that the optimizer will attempt to optimize the plan | | datafusion.optimizer.top_down_join_key_reordering | true | When set to true, the physical plan optimizer will run a top down process to reorder the join keys | | 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.use_statistics_registry | false | (Deprecated) Ignored: the physical plan optimizer always consults the session's pluggable `StatisticsRegistry` (register providers on the `SessionState`; with none it is a no-op). | | datafusion.optimizer.prefer_hash_join | true | When set to true, the physical plan optimizer will prefer HashJoin over SortMergeJoin. HashJoin can work more efficiently than SortMergeJoin but consumes more memory | | datafusion.optimizer.enable_piecewise_merge_join | false | When set to true, piecewise merge join is enabled. PiecewiseMergeJoin is currently experimental. Physical planner will opt for PiecewiseMergeJoin when there is only one range filter. | | 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 | From bfad5815a7d91e032785492eaf49b5bcab05ed2b Mon Sep 17 00:00:00 2001 From: Alessandro Solimando Date: Fri, 17 Jul 2026 10:10:18 +0200 Subject: [PATCH 02/15] feat: propagate provider extensions through the statistics walk The walk keeps core `Statistics` as its currency; a per-walk side cache carries provider `Extensions` separately, so a walk with no providers has no per-node overhead. A provider's `Computed` extensions are cached and reach parent-node providers as their `child_stats`; the built-in `statistics_from_inputs` fallback yields none, so extensions survive only across an unbroken chain of provider-handled nodes. Adds `StatisticsContext::compute_extended` (statistics + extensions) alongside `compute` (core statistics only). `StatisticsRegistry::compute` preserves extensions again; it and `compute_base` are deprecated in favor of `StatisticsContext`. --- .../src/operator_statistics/mod.rs | 152 +++++--- datafusion/physical-plan/src/statistics.rs | 353 ++++++++++++++---- 2 files changed, 377 insertions(+), 128 deletions(-) diff --git a/datafusion/physical-plan/src/operator_statistics/mod.rs b/datafusion/physical-plan/src/operator_statistics/mod.rs index 9ad3340e419c2..5ba68f11b147e 100644 --- a/datafusion/physical-plan/src/operator_statistics/mod.rs +++ b/datafusion/physical-plan/src/operator_statistics/mod.rs @@ -55,7 +55,8 @@ //! chain is consulted before falling back to the operator's built-in //! [`ExecutionPlan::statistics_from_inputs`], with children already resolved by //! that same walk. The first provider that returns a computed result sets the -//! node's statistics. +//! node's statistics. The walk carries [`ExtendedStatistics`]; see +//! [`StatisticsContext`] for how a provider's extensions propagate. //! //! [`StatisticsContext`]: crate::statistics::StatisticsContext //! @@ -71,7 +72,8 @@ //! registry.register(Arc::new(MyHistogramProvider)); //! //! // Compute statistics through the chain -//! let stats = registry.compute(plan.as_ref())?; +//! let stats = StatisticsContext::new_with_registry(registry) +//! .compute_extended(plan.as_ref(), &StatisticsArgs::new())?; //! ``` use std::fmt::{self, Debug}; @@ -165,6 +167,19 @@ impl ExtendedStatistics { pub fn merge_extensions(&mut self, other: &ExtendedStatistics) { self.extensions.merge(&other.extensions); } + + /// Create from base statistics plus an existing extension map. + pub(crate) fn new_with_extensions( + base: Arc, + extensions: Extensions, + ) -> Self { + Self { base, extensions } + } + + /// Returns the extension map. + pub(crate) fn extensions(&self) -> &Extensions { + &self.extensions + } } impl From for ExtendedStatistics { @@ -373,19 +388,36 @@ impl StatisticsRegistry { /// Compute extended statistics for `plan` through the provider chain. /// /// Thin wrapper over the single [`StatisticsContext`] walk (the same path the - /// optimizer and EXPLAIN use), so there is one traversal implementation. The - /// walk carries base [`Statistics`] only; extensions are not propagated. + /// optimizer and EXPLAIN use), so there is one traversal implementation. + /// Provider extensions are preserved; see [`StatisticsContext`] for how they + /// propagate up the tree. + /// + /// Children are resolved via [`ExecutionPlan::child_stats_requests`] (default + /// `Skip`), so a custom operator whose provider reads `child_stats` must + /// declare `At` for those children to receive their computed statistics. + #[deprecated( + since = "55.0.0", + note = "use `StatisticsContext::new_with_registry(registry).compute_extended(plan, &StatisticsArgs::new())`" + )] pub fn compute(&self, plan: &dyn ExecutionPlan) -> Result { - let base = StatisticsContext::new_with_registry(self.clone()) - .compute(plan, &StatisticsArgs::new())?; - Ok(ExtendedStatistics::new_arc(base)) + Ok(StatisticsContext::new_with_registry(self.clone()) + .compute_extended(plan, &StatisticsArgs::new())? + .as_ref() + .clone()) } /// Compute statistics and return only the base Statistics (no extensions). /// /// Convenience method for callers that don't need extensions. + #[deprecated( + since = "55.0.0", + note = "use `StatisticsContext::new_with_registry(registry).compute(plan, &StatisticsArgs::new())`" + )] pub fn compute_base(&self, plan: &dyn ExecutionPlan) -> Result { - Ok(self.compute(plan)?.base().clone()) + Ok(StatisticsContext::new_with_registry(self.clone()) + .compute(plan, &StatisticsArgs::new())? + .as_ref() + .clone()) } } @@ -1022,7 +1054,7 @@ mod tests { use super::*; use crate::filter::FilterExec; use crate::projection::ProjectionExec; - use crate::statistics::StatisticsArgs; + use crate::statistics::{StatisticsArgs, StatisticsContext}; use crate::{DisplayAs, DisplayFormatType, PlanProperties}; use arrow::datatypes::{DataType, Field, Schema}; use datafusion_common::stats::Precision; @@ -1035,6 +1067,17 @@ mod tests { use crate::execution_plan::{Boundedness, EmissionType}; + /// Compute statistics via [`StatisticsContext`] (replaces the deprecated + /// `compute`/`compute_base`). + fn compute( + registry: &StatisticsRegistry, + plan: &dyn ExecutionPlan, + ) -> Result { + Ok((*StatisticsContext::new_with_registry(registry.clone()) + .compute_extended(plan, &StatisticsArgs::new())?) + .clone()) + } + fn make_schema() -> Arc { Arc::new(Schema::new(vec![ Field::new("a", DataType::Int32, false), @@ -1142,11 +1185,30 @@ mod tests { let engine = StatisticsRegistry::new(); let source = make_source(1000); - let stats = engine.compute(source.as_ref())?; + let stats = compute(&engine, source.as_ref())?; assert!(matches!(stats.base.num_rows, Precision::Exact(1000))); Ok(()) } + #[test] + fn test_compute_preserves_provider_extensions() -> Result<()> { + #[derive(Debug, Clone, PartialEq)] + struct Sketch(u32); + + let source = make_source(1000); + let provider = ClosureStatisticsProvider::new(|plan, _child_stats| { + let mut ext = + ExtendedStatistics::new(Statistics::new_unknown(&plan.schema())); + ext.set_extension(Sketch(42)); + Ok(StatisticsResult::Computed(ext)) + }); + let registry = StatisticsRegistry::with_providers(vec![Arc::new(provider)]); + + let stats = compute(®istry, source.as_ref())?; + assert_eq!(stats.get_extension::(), Some(&Sketch(42))); + Ok(()) + } + #[test] fn test_custom_chain_configuration() -> Result<()> { let source = make_source(1000); @@ -1158,7 +1220,7 @@ mod tests { // With no provider handling FilterExec, it returns the built-in fallback statistics let filter: Arc = Arc::new(FilterExec::try_new(lit(true), Arc::clone(&source))?); - let stats = custom_only.compute(filter.as_ref())?; + let stats = compute(&custom_only, filter.as_ref())?; // Falls back to plan.statistics() since no provider handles it assert!(stats.base.num_rows.get_value().is_some()); @@ -1169,7 +1231,7 @@ mod tests { }) as Arc]); // OverrideFilterProvider handles filters, built-in fallback handles the rest - let stats = with_override.compute(filter.as_ref())?; + let stats = compute(&with_override, filter.as_ref())?; assert!(matches!(stats.base.num_rows, Precision::Inexact(250))); // Verify chain inspection @@ -1256,7 +1318,7 @@ mod tests { let source = make_source(1000); let custom: Arc = Arc::new(CustomExec { input: source }); - let stats = engine.compute(custom.as_ref())?; + let stats = compute(&engine, custom.as_ref())?; assert!(matches!(stats.base.num_rows, Precision::Exact(1000))); Ok(()) } @@ -1305,7 +1367,7 @@ mod tests { let filter: Arc = Arc::new(FilterExec::try_new(lit(true), source)?); - let stats = engine.compute(filter.as_ref())?; + let stats = compute(&engine, filter.as_ref())?; assert!(matches!(stats.base.num_rows, Precision::Inexact(100))); Ok(()) } @@ -1318,7 +1380,7 @@ mod tests { let filter: Arc = Arc::new(FilterExec::try_new(predicate, source)?); - let stats = engine.compute(filter.as_ref())?; + let stats = compute(&engine, filter.as_ref())?; assert!(stats.base.num_rows.get_value().unwrap_or(&0) <= &1000); Ok(()) } @@ -1368,7 +1430,7 @@ mod tests { let registry = StatisticsRegistry::with_providers(vec![Arc::new(FilterStatisticsProvider)]); - let stats = registry.compute(filter.as_ref())?; + let stats = compute(®istry, filter.as_ref())?; let output_ndv_a = stats.base.column_statistics[0] .distinct_count @@ -1406,7 +1468,7 @@ mod tests { source, )?); - let stats = engine.compute(proj.as_ref())?; + let stats = compute(&engine, proj.as_ref())?; assert!(matches!(stats.base.num_rows, Precision::Exact(1000))); Ok(()) } @@ -1420,7 +1482,7 @@ mod tests { let coalesce: Arc = Arc::new(CoalescePartitionsExec::new(source)); - let stats = engine.compute(coalesce.as_ref())?; + let stats = compute(&engine, coalesce.as_ref())?; // PassthroughStatisticsProvider should propagate child row count unchanged assert_eq!(stats.base.num_rows, Precision::Exact(1000)); Ok(()) @@ -1440,13 +1502,13 @@ mod tests { let custom: Arc = Arc::new(CustomExec { input: Arc::clone(&source), }); - let stats = engine.compute(custom.as_ref())?; + let stats = compute(&engine, custom.as_ref())?; assert!(matches!(stats.base.num_rows, Precision::Exact(1000))); // FilterExec: CustomStatisticsProvider delegates, OverrideFilterProvider handles let filter: Arc = Arc::new(FilterExec::try_new(lit(true), source)?); - let stats = engine.compute(filter.as_ref())?; + let stats = compute(&engine, filter.as_ref())?; assert!(matches!(stats.base.num_rows, Precision::Inexact(500))); Ok(()) @@ -1557,7 +1619,7 @@ mod tests { let registry = StatisticsRegistry::with_providers(vec![Arc::new( AggregateStatisticsProvider, )]); - let stats = registry.compute(agg.as_ref())?; + let stats = compute(®istry, agg.as_ref())?; assert_eq!(stats.base.num_rows, Precision::Inexact(10)); Ok(()) } @@ -1574,7 +1636,7 @@ mod tests { let registry = StatisticsRegistry::with_providers(vec![Arc::new( AggregateStatisticsProvider, )]); - let stats = registry.compute(agg.as_ref())?; + let stats = compute(®istry, agg.as_ref())?; // 10 * 5 = 50 assert_eq!(stats.base.num_rows, Precision::Inexact(50)); Ok(()) @@ -1593,7 +1655,7 @@ mod tests { let registry = StatisticsRegistry::with_providers(vec![Arc::new( AggregateStatisticsProvider, )]); - let stats = registry.compute(agg.as_ref())?; + let stats = compute(®istry, agg.as_ref())?; assert_eq!(stats.base.num_rows, Precision::Inexact(500)); Ok(()) } @@ -1611,7 +1673,7 @@ mod tests { let registry = StatisticsRegistry::with_providers(vec![Arc::new( AggregateStatisticsProvider, )]); - let stats = registry.compute(agg.as_ref())?; + let stats = compute(®istry, agg.as_ref())?; // Delegates, falling back to the operator's built-in statistics_from_inputs assert!( stats.base.num_rows.get_value().is_some() @@ -1635,7 +1697,7 @@ mod tests { let registry = StatisticsRegistry::with_providers(vec![Arc::new( AggregateStatisticsProvider, )]); - let stats = registry.compute(agg.as_ref())?; + let stats = compute(®istry, agg.as_ref())?; // Should delegate (expression is not a Column) assert!( stats.base.num_rows.get_value().is_some() @@ -1675,7 +1737,7 @@ mod tests { let registry = StatisticsRegistry::with_providers(vec![Arc::new( AggregateStatisticsProvider, )]); - let stats = registry.compute(agg.as_ref())?; + let stats = compute(®istry, agg.as_ref())?; // Multiple grouping sets: provider delegates, so the operator's built-in // statistics_from_inputs computes the correct per-set // NDV estimation. The exact value depends on the built-in implementation. @@ -1707,7 +1769,7 @@ mod tests { let registry = StatisticsRegistry::with_providers(vec![Arc::new( AggregateStatisticsProvider, )]); - let stats = registry.compute(agg.as_ref())?; + let stats = compute(®istry, agg.as_ref())?; // Should fall through to the operator's built-in statistics_from_inputs. // The exact value depends on the built-in implementation. assert!( @@ -1778,7 +1840,7 @@ mod tests { let registry = StatisticsRegistry::with_providers(vec![Arc::new(JoinStatisticsProvider)]); - let stats = registry.compute(join.as_ref())?; + let stats = compute(®istry, join.as_ref())?; assert_eq!(stats.base.num_rows, Precision::Inexact(5000)); Ok(()) } @@ -1830,7 +1892,7 @@ mod tests { let registry = StatisticsRegistry::with_providers(vec![Arc::new(JoinStatisticsProvider)]); - let stats = registry.compute(join.as_ref())?; + let stats = compute(®istry, join.as_ref())?; assert_eq!(stats.base.num_rows, Precision::Inexact(10_000)); Ok(()) } @@ -1890,7 +1952,7 @@ mod tests { let registry = StatisticsRegistry::with_providers(vec![Arc::new(JoinStatisticsProvider)]); - let stats = registry.compute(join.as_ref())?; + let stats = compute(®istry, join.as_ref())?; assert_eq!(stats.base.num_rows, Precision::Inexact(250)); Ok(()) } @@ -1904,7 +1966,7 @@ mod tests { let registry = StatisticsRegistry::with_providers(vec![Arc::new(JoinStatisticsProvider)]); - let stats = registry.compute(join.as_ref())?; + let stats = compute(®istry, join.as_ref())?; assert_eq!(stats.base.num_rows, Precision::Inexact(20_000)); Ok(()) } @@ -1927,7 +1989,7 @@ mod tests { let registry = StatisticsRegistry::with_providers(vec![Arc::new(JoinStatisticsProvider)]); - let stats = registry.compute(join.as_ref())?; + let stats = compute(®istry, join.as_ref())?; // Provider delegates; result comes from built-in statistics_from_inputs. assert!( stats.base.num_rows.get_value().is_some() @@ -1970,7 +2032,7 @@ mod tests { let join = make_hash_join_typed(left, right, join_type)?; let registry = StatisticsRegistry::with_providers(vec![Arc::new(JoinStatisticsProvider)]); - Ok(registry.compute(join.as_ref())?.base.num_rows) + Ok(compute(®istry, join.as_ref())?.base.num_rows) } #[test] @@ -2055,7 +2117,7 @@ mod tests { let registry = StatisticsRegistry::with_providers(vec![Arc::new(JoinStatisticsProvider)]); - let stats = registry.compute(join.as_ref())?; + let stats = compute(®istry, join.as_ref())?; // Both inputs have Exact row counts -> result is also Exact assert_eq!(stats.base.num_rows, Precision::Exact(20_000)); Ok(()) @@ -2075,7 +2137,7 @@ mod tests { let registry = StatisticsRegistry::with_providers(vec![Arc::new(LimitStatisticsProvider)]); - let stats = registry.compute(limit.as_ref())?; + let stats = compute(®istry, limit.as_ref())?; assert_eq!(stats.base.num_rows, Precision::Exact(100)); Ok(()) } @@ -2088,7 +2150,7 @@ mod tests { let registry = StatisticsRegistry::with_providers(vec![Arc::new(LimitStatisticsProvider)]); - let stats = registry.compute(limit.as_ref())?; + let stats = compute(®istry, limit.as_ref())?; assert_eq!(stats.base.num_rows, Precision::Exact(50)); Ok(()) } @@ -2102,7 +2164,7 @@ mod tests { let registry = StatisticsRegistry::with_providers(vec![Arc::new(LimitStatisticsProvider)]); - let stats = registry.compute(limit.as_ref())?; + let stats = compute(®istry, limit.as_ref())?; assert_eq!(stats.base.num_rows, Precision::Exact(100)); Ok(()) } @@ -2116,7 +2178,7 @@ mod tests { let registry = StatisticsRegistry::with_providers(vec![Arc::new(LimitStatisticsProvider)]); - let stats = registry.compute(limit.as_ref())?; + let stats = compute(®istry, limit.as_ref())?; assert_eq!(stats.base.num_rows, Precision::Exact(0)); Ok(()) } @@ -2130,7 +2192,7 @@ mod tests { let registry = StatisticsRegistry::with_providers(vec![Arc::new(LimitStatisticsProvider)]); - let stats = registry.compute(limit.as_ref())?; + let stats = compute(®istry, limit.as_ref())?; assert_eq!(stats.base.num_rows, Precision::Inexact(100)); Ok(()) } @@ -2151,7 +2213,7 @@ mod tests { let registry = StatisticsRegistry::with_providers(vec![Arc::new(UnionStatisticsProvider)]); - let stats = registry.compute(union.as_ref())?; + let stats = compute(®istry, union.as_ref())?; assert_eq!(stats.base.num_rows, Precision::Exact(1000)); Ok(()) } @@ -2166,7 +2228,7 @@ mod tests { let registry = StatisticsRegistry::with_providers(vec![Arc::new(UnionStatisticsProvider)]); - let stats = registry.compute(union.as_ref())?; + let stats = compute(®istry, union.as_ref())?; assert_eq!(stats.base.num_rows, Precision::Exact(600)); Ok(()) } @@ -2181,7 +2243,7 @@ mod tests { let registry = StatisticsRegistry::with_providers(vec![Arc::new(UnionStatisticsProvider)]); - let stats = registry.compute(union.as_ref())?; + let stats = compute(®istry, union.as_ref())?; assert_eq!(stats.base.num_rows, Precision::Absent); Ok(()) } @@ -2212,7 +2274,7 @@ mod tests { let source = make_source(1000); let filter: Arc = Arc::new(FilterExec::try_new(lit(true), source)?); - let stats = registry.compute(filter.as_ref())?; + let stats = compute(®istry, filter.as_ref())?; assert_eq!(stats.base.num_rows, Precision::Inexact(42)); Ok(()) } @@ -2253,8 +2315,8 @@ mod tests { let filter_b: Arc = Arc::new(FilterExec::try_new(lit(true), make_source(200))?); - let stats_a = registry.compute(filter_a.as_ref())?; - let stats_b = registry.compute(filter_b.as_ref())?; + let stats_a = compute(®istry, filter_a.as_ref())?; + let stats_b = compute(®istry, filter_b.as_ref())?; assert_eq!(stats_a.base.num_rows, Precision::Inexact(100)); assert_eq!(stats_b.base.num_rows, Precision::Inexact(50)); diff --git a/datafusion/physical-plan/src/statistics.rs b/datafusion/physical-plan/src/statistics.rs index ae84cec4c58a1..47499d82a9422 100644 --- a/datafusion/physical-plan/src/statistics.rs +++ b/datafusion/physical-plan/src/statistics.rs @@ -24,6 +24,7 @@ use crate::ExecutionPlan; use crate::operator_statistics::{ ExtendedStatistics, StatisticsRegistry, StatisticsResult, }; +use datafusion_common::extensions::Extensions; use datafusion_common::{ Result, Statistics, assert_eq_or_internal_err, assert_or_internal_err, }; @@ -32,6 +33,15 @@ use std::collections::HashMap; use std::rc::Rc; use std::sync::Arc; +type CacheKey = (usize, Option); + +fn cache_key(plan: &dyn ExecutionPlan, partition: Option) -> CacheKey { + ( + plan as *const dyn ExecutionPlan as *const () as usize, + partition, + ) +} + /// Per-call memoization cache for statistics computation. /// /// Keyed by `(plan node pointer address, partition)`. Shared across @@ -40,34 +50,15 @@ use std::sync::Arc; /// The pointer-based key is safe within a single synchronous walk: /// all `Arc` nodes are held by the plan tree for /// the duration of the walk, so addresses cannot be reused. +/// +/// Core statistics and provider extensions are cached separately: the +/// `statistics` map is the hot path (populated on every walk); the `extensions` +/// map is populated only when a provider returns non-empty extensions, so a walk +/// with no providers never touches it. #[derive(Debug, Default)] -struct StatsCache(HashMap<(usize, Option), Arc>); - -impl StatsCache { - fn get( - &self, - plan: &dyn ExecutionPlan, - partition: Option, - ) -> Option<&Arc> { - let key = ( - plan as *const dyn ExecutionPlan as *const () as usize, - partition, - ); - self.0.get(&key) - } - - fn insert( - &mut self, - plan: &dyn ExecutionPlan, - partition: Option, - stats: Arc, - ) { - let key = ( - plan as *const dyn ExecutionPlan as *const () as usize, - partition, - ); - self.0.insert(key, stats); - } +struct StatsCache { + statistics: HashMap>, + extensions: HashMap, } /// Arguments passed to [`ExecutionPlan::statistics_from_inputs`] carrying @@ -126,6 +117,14 @@ pub enum ChildStats { /// they are consulted before the operator's built-in /// [`ExecutionPlan::statistics_from_inputs`]. An empty registry is the built-in /// computation. +/// +/// The walk carries [`ExtendedStatistics`]. A node has extensions only if a +/// provider `Computed` them for it; a node that falls back to the built-in +/// [`ExecutionPlan::statistics_from_inputs`] has none. So extensions propagate +/// upward only through an unbroken chain of provider-handled nodes: a single +/// built-in node yields no extensions and hides those of everything beneath it. +/// [`Self::compute_extended`] observes extensions; [`Self::compute`] returns core +/// [`Statistics`] only. pub struct StatisticsContext { cache: Rc>, registry: StatisticsRegistry, @@ -143,8 +142,7 @@ impl StatisticsContext { Self::new_with_registry(StatisticsRegistry::new()) } - /// Creates a context whose walk consults `registry`'s provider chain before - /// falling back to each operator's built-in statistics. + /// 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())), @@ -159,15 +157,49 @@ impl StatisticsContext { /// (which rewrite the plan) when reusing one context across them, so stale /// pointer keys cannot collide. pub fn reset_cache(&self) { - self.cache.borrow_mut().0.clear(); + let mut cache = self.cache.borrow_mut(); + cache.statistics.clear(); + cache.extensions.clear(); + } + + /// Computes the core [`Statistics`] for `plan`, discarding any + /// provider-supplied extensions (see [`Self::compute_extended`]). + /// + /// With no providers registered this is the plain built-in walk: only the + /// `statistics` cache is touched, so it carries no extension overhead. + pub fn compute( + &self, + plan: &dyn ExecutionPlan, + args: &StatisticsArgs, + ) -> Result> { + self.compute_base(plan, args) + } + + /// Computes the [`ExtendedStatistics`] for `plan`: the core statistics plus + /// any extensions a provider attached to this node (see the type-level docs + /// for how extensions propagate up the tree). + pub fn compute_extended( + &self, + plan: &dyn ExecutionPlan, + args: &StatisticsArgs, + ) -> Result> { + let statistics = self.compute_base(plan, args)?; + let extensions = self + .cached_extensions(plan, args.partition()) + .unwrap_or_default(); + Ok(Arc::new(ExtendedStatistics::new_with_extensions( + statistics, extensions, + ))) } - /// Computes statistics for `plan`, resolving children first and passing - /// the results to [`ExecutionPlan::statistics_from_inputs`]. + /// Bottom-up walk producing the node's core statistics, resolving children + /// first and consulting the provider chain before the operator's built-in + /// [`ExecutionPlan::statistics_from_inputs`]. Any extensions a provider + /// attaches are recorded in the extension cache for [`Self::compute_extended`]. /// /// When `args.partition()` is `Some(idx)`, `idx` is validated against the /// plan's partition count. - pub fn compute( + fn compute_base( &self, plan: &dyn ExecutionPlan, args: &StatisticsArgs, @@ -184,34 +216,10 @@ impl StatisticsContext { ); } - if let Some(cached) = self.cache.borrow().get(plan, partition) { - return Ok(Arc::clone(cached)); + if let Some(cached) = self.cached_statistics(plan, partition) { + return Ok(cached); } - let child_stats = self.resolve_children(plan, partition)?; - - let result = match self.try_provider_stats(plan, &child_stats, args)? { - Some(stats) => stats, - None => plan.statistics_from_inputs(&child_stats, args)?, - }; - self.cache - .borrow_mut() - .insert(plan, partition, Arc::clone(&result)); - Ok(result) - } - - /// Resolves each child's statistics following the operator's - /// [`ExecutionPlan::child_stats_requests`] mapping: `At(p)` computes the child - /// at partition `p`, `Skip` supplies a [`Statistics::new_unknown`] placeholder. - /// - /// A provider computes the same node's statistics, so it depends on the same - /// children as the node itself: an operator whose statistics (built-in or via - /// a provider) depend on a child must declare `At` for it. - fn resolve_children( - &self, - plan: &dyn ExecutionPlan, - partition: Option, - ) -> Result>> { let children = plan.children(); let requests = plan.child_stats_requests(partition); assert_eq_or_internal_err!( @@ -222,51 +230,148 @@ impl StatisticsContext { requests.len(), children.len() ); - children + let child_statistics: Vec> = children .iter() - .zip(requests) + .zip(&requests) .map(|(child, directive)| match directive { - ChildStats::At(p) => { - self.compute(child.as_ref(), &StatisticsArgs::new().with_partition(p)) - } + ChildStats::At(p) => self.compute_base( + child.as_ref(), + &StatisticsArgs::new().with_partition(*p), + ), ChildStats::Skip => { Ok(Arc::new(Statistics::new_unknown(child.schema().as_ref()))) } }) - .collect() + .collect::>()?; + + let statistics = match self.try_provider_stats( + plan, + &children, + &requests, + &child_statistics, + args, + )? { + Some(statistics) => statistics, + None => plan.statistics_from_inputs(&child_statistics, args)?, + }; + self.store_statistics(plan, partition, Arc::clone(&statistics)); + Ok(statistics) } - /// Runs the provider chain, returning the first `Computed` result's base - /// statistics, or `None` if all delegate or the chain is empty. A - /// partition-blind provider applies only to overall stats (its default - /// `compute_statistics_with_args` delegates per partition). + /// Runs the provider chain, returning the first `Computed` result's core + /// statistics (and recording its extensions), or `None` if the chain is empty + /// or all delegate. A partition-blind provider applies only to overall stats + /// (its default `compute_statistics_with_args` delegates per partition). /// - /// Providers may attach custom extensions to their [`ExtendedStatistics`] - /// result, but the walk uses only the base [`Statistics`]. + /// Assembles each child's [`ExtendedStatistics`] (statistics plus cached + /// extensions) only here, so a walk with no providers pays no extension cost. fn try_provider_stats( &self, plan: &dyn ExecutionPlan, - child_stats: &[Arc], + children: &[&Arc], + requests: &[ChildStats], + child_statistics: &[Arc], args: &StatisticsArgs, ) -> Result>> { let providers = self.registry.providers(); if providers.is_empty() { return Ok(None); } - // Providers take `&[ExtendedStatistics]`; wrap the resolved base stats. - let child_ext: Vec = child_stats - .iter() - .map(|s| ExtendedStatistics::new_arc(Arc::clone(s))) - .collect(); + let child_extended = + self.child_extended_stats(children, requests, child_statistics); for provider in providers { - if let StatisticsResult::Computed(stats) = - provider.compute_statistics_with_args(plan, &child_ext, args)? + if let StatisticsResult::Computed(computed) = + provider.compute_statistics_with_args(plan, &child_extended, args)? { - return Ok(Some(Arc::clone(stats.base_arc()))); + if !computed.extensions().is_empty() { + self.store_extensions( + plan, + args.partition(), + computed.extensions().clone(), + ); + } + return Ok(Some(Arc::clone(computed.base_arc()))); } } Ok(None) } + + /// Pairs each child's core statistics with any extensions cached for it, + /// producing the [`ExtendedStatistics`] the provider chain consumes. Called + /// only when providers exist, so an empty registry pays no extension cost. + fn child_extended_stats( + &self, + children: &[&Arc], + requests: &[ChildStats], + child_statistics: &[Arc], + ) -> Vec { + children + .iter() + .zip(requests) + .zip(child_statistics) + .map(|((child, directive), statistics)| { + let extensions = match directive { + ChildStats::At(p) => self.cached_extensions(child.as_ref(), *p), + ChildStats::Skip => None, + }; + match extensions { + Some(extensions) => ExtendedStatistics::new_with_extensions( + Arc::clone(statistics), + extensions, + ), + None => ExtendedStatistics::new_arc(Arc::clone(statistics)), + } + }) + .collect() + } + + fn cached_statistics( + &self, + plan: &dyn ExecutionPlan, + partition: Option, + ) -> Option> { + self.cache + .borrow() + .statistics + .get(&cache_key(plan, partition)) + .cloned() + } + + fn store_statistics( + &self, + plan: &dyn ExecutionPlan, + partition: Option, + statistics: Arc, + ) { + self.cache + .borrow_mut() + .statistics + .insert(cache_key(plan, partition), statistics); + } + + fn cached_extensions( + &self, + plan: &dyn ExecutionPlan, + partition: Option, + ) -> Option { + self.cache + .borrow() + .extensions + .get(&cache_key(plan, partition)) + .cloned() + } + + fn store_extensions( + &self, + plan: &dyn ExecutionPlan, + partition: Option, + extensions: Extensions, + ) { + self.cache + .borrow_mut() + .extensions + .insert(cache_key(plan, partition), extensions); + } } #[cfg(all(test, feature = "test_utils"))] @@ -311,6 +416,53 @@ mod tests { } } + #[derive(Debug, Clone, PartialEq)] + struct Tag(u32); + + /// Leaf provider: sets a row count and attaches a `Tag` extension. + #[derive(Debug)] + struct TagLeafProvider { + rows: usize, + tag: u32, + } + impl StatisticsProvider for TagLeafProvider { + fn compute_statistics( + &self, + plan: &dyn ExecutionPlan, + child_stats: &[ExtendedStatistics], + ) -> Result { + if !child_stats.is_empty() { + return Ok(StatisticsResult::Delegate); + } + let mut stats = Statistics::new_unknown(&plan.schema()); + stats.num_rows = Precision::Exact(self.rows); + let mut extended = ExtendedStatistics::new(stats); + extended.set_extension(Tag(self.tag)); + Ok(StatisticsResult::Computed(extended)) + } + } + + /// Non-leaf provider: re-emits a `Tag` doubled from the first child's `Tag`, + /// proving the child's extension reached this provider. + #[derive(Debug)] + struct TagDoublingProvider; + impl StatisticsProvider for TagDoublingProvider { + fn compute_statistics( + &self, + plan: &dyn ExecutionPlan, + child_stats: &[ExtendedStatistics], + ) -> Result { + let Some(Tag(v)) = child_stats.first().and_then(|c| c.get_extension::()) + else { + return Ok(StatisticsResult::Delegate); + }; + let mut extended = + ExtendedStatistics::new(Statistics::new_unknown(&plan.schema())); + extended.set_extension(Tag(v * 2)); + Ok(StatisticsResult::Computed(extended)) + } + } + fn ctx_with(provider: Arc) -> StatisticsContext { StatisticsContext::new_with_registry(StatisticsRegistry::with_providers(vec![ provider, @@ -362,7 +514,7 @@ mod tests { let args = StatisticsArgs::new(); let s1 = ctx.compute(leaf.as_ref(), &args).unwrap(); - assert!(!ctx.cache.borrow().0.is_empty()); + assert!(!ctx.cache.borrow().statistics.is_empty()); let s2 = ctx.compute(leaf.as_ref(), &args).unwrap(); assert!(Arc::ptr_eq(&s1, &s2)); @@ -373,9 +525,9 @@ mod tests { let leaf = make_stats_leaf(10); let ctx = StatisticsContext::new(); let _ = ctx.compute(leaf.as_ref(), &StatisticsArgs::new()).unwrap(); - assert!(!ctx.cache.borrow().0.is_empty()); + assert!(!ctx.cache.borrow().statistics.is_empty()); ctx.reset_cache(); - assert!(ctx.cache.borrow().0.is_empty()); + assert!(ctx.cache.borrow().statistics.is_empty()); } #[test] @@ -392,6 +544,41 @@ mod tests { assert_eq!(per_part.num_rows, Precision::Exact(701)); } + #[test] + fn extensions_reach_parent_provider() { + let leaf = make_stats_leaf(100); + let parent: Arc = Arc::new(CoalescePartitionsExec::new(leaf)); + let ctx = StatisticsContext::new_with_registry( + StatisticsRegistry::with_providers(vec![ + Arc::new(TagLeafProvider { rows: 100, tag: 7 }), + Arc::new(TagDoublingProvider), + ]), + ); + let extended = ctx + .compute_extended(parent.as_ref(), &StatisticsArgs::new()) + .unwrap(); + assert_eq!(extended.get_extension::(), Some(&Tag(14))); + } + + #[test] + fn builtin_fallback_drops_extensions() { + let leaf = make_stats_leaf(100); + let parent: Arc = + Arc::new(CoalescePartitionsExec::new(Arc::clone(&leaf))); + let ctx = ctx_with(Arc::new(TagLeafProvider { rows: 100, tag: 7 })); + + let leaf_extended = ctx + .compute_extended(leaf.as_ref(), &StatisticsArgs::new()) + .unwrap(); + assert_eq!(leaf_extended.get_extension::(), Some(&Tag(7))); + + let parent_extended = ctx + .compute_extended(parent.as_ref(), &StatisticsArgs::new()) + .unwrap(); + assert_eq!(parent_extended.get_extension::(), None); + assert_eq!(parent_extended.base().num_rows, Precision::Exact(100)); + } + #[test] fn per_partition_union_with_registry_no_out_of_bounds() { // Two 2-partition inputs -> 4 output partitions. Union owns output From 2a920928edbdd44bb6ad182663ca81c3b19c432d Mon Sep 17 00:00:00 2001 From: Alessandro Solimando Date: Fri, 17 Jul 2026 00:52:02 +0200 Subject: [PATCH 03/15] feat: reflect the StatisticsRegistry in EXPLAIN statistics EXPLAIN ... show_statistics computed statistics with an empty registry, so registered providers were not reflected in the statistics column. Thread the session's StatisticsRegistry into the displayable plan; the column now reflects provider-supplied estimates when providers are registered, and is unchanged when none are. --- datafusion/core/src/physical_planner.rs | 16 ++++++++ datafusion/physical-plan/src/analyze.rs | 18 +++++++++ datafusion/physical-plan/src/display.rs | 26 ++++++++++++- .../test_files/statistics_registry.slt | 39 +++++++++++++++++++ 4 files changed, 97 insertions(+), 2 deletions(-) diff --git a/datafusion/core/src/physical_planner.rs b/datafusion/core/src/physical_planner.rs index b6d28e7b21c79..0a71c77cd5e59 100644 --- a/datafusion/core/src/physical_planner.rs +++ b/datafusion/core/src/physical_planner.rs @@ -2607,6 +2607,10 @@ impl DefaultPhysicalPlanner { let explain_format = &e.explain_format; // Statement-level override wins over session config for show_statistics. let show_statistics = e.show_statistics.unwrap_or(config.show_statistics); + let statistics_registry = session_state + .statistics_registry() + .cloned() + .unwrap_or_default(); if !e.logical_optimization_succeeded { return Ok(Arc::new(ExplainExec::new( @@ -2680,6 +2684,7 @@ impl DefaultPhysicalPlanner { InitialPhysicalPlan, displayable(input.as_ref()) .set_show_statistics(show_statistics) + .set_statistics_registry(statistics_registry.clone()) .set_show_schema(config.show_schema) .indent(e.verbose) .to_string(), @@ -2693,6 +2698,7 @@ impl DefaultPhysicalPlanner { InitialPhysicalPlanWithStats, displayable(input.as_ref()) .set_show_statistics(true) + .set_statistics_registry(statistics_registry.clone()) .indent(e.verbose) .to_string(), )); @@ -2718,6 +2724,7 @@ impl DefaultPhysicalPlanner { plan_type, displayable(plan) .set_show_statistics(show_statistics) + .set_statistics_registry(statistics_registry.clone()) .set_show_schema(config.show_schema) .indent(e.verbose) .to_string(), @@ -2731,6 +2738,7 @@ impl DefaultPhysicalPlanner { FinalPhysicalPlan, displayable(input.as_ref()) .set_show_statistics(show_statistics) + .set_statistics_registry(statistics_registry.clone()) .set_show_schema(config.show_schema) .indent(e.verbose) .to_string(), @@ -2744,6 +2752,9 @@ impl DefaultPhysicalPlanner { FinalPhysicalPlanWithStats, displayable(input.as_ref()) .set_show_statistics(true) + .set_statistics_registry( + statistics_registry.clone(), + ) .indent(e.verbose) .to_string(), )); @@ -2809,11 +2820,16 @@ impl DefaultPhysicalPlanner { ExplainAnalyzeCategories::All => None, ExplainAnalyzeCategories::Only(cats) => Some(cats), }; + let statistics_registry = session_state + .statistics_registry() + .cloned() + .unwrap_or_default(); Ok(Arc::new( AnalyzeExec::builder(a.verbose, show_statistics, input, schema) .with_metric_types(metric_types) .with_metric_categories(metric_categories) .with_format(a.format.clone()) + .with_statistics_registry(statistics_registry) .build(), )) } diff --git a/datafusion/physical-plan/src/analyze.rs b/datafusion/physical-plan/src/analyze.rs index 72cd24ef95673..4bda4bc405620 100644 --- a/datafusion/physical-plan/src/analyze.rs +++ b/datafusion/physical-plan/src/analyze.rs @@ -27,6 +27,7 @@ use super::{ use crate::display::DisplayableExecutionPlan; use crate::execution_plan::EvaluationType; use crate::metrics::{MetricCategory, MetricType}; +use crate::operator_statistics::StatisticsRegistry; use crate::{DisplayFormatType, ExecutionPlan, Partitioning}; use arrow::{array::StringBuilder, datatypes::SchemaRef, record_batch::RecordBatch}; @@ -54,6 +55,9 @@ pub struct AnalyzeExec { metric_categories: Option>, /// Output format for the rendered plan + metrics. format: ExplainFormat, + /// Registry consulted when rendering displayed statistics, so `EXPLAIN + /// ANALYZE` reflects the same provider-refined stats as plain `EXPLAIN`. + statistics_registry: StatisticsRegistry, /// The input plan (the plan being analyzed) pub(crate) input: Arc, /// The output schema for RecordBatches of this exec node @@ -72,6 +76,7 @@ pub struct AnalyzeExecBuilder { metric_types: Vec, metric_categories: Option>, format: ExplainFormat, + statistics_registry: StatisticsRegistry, } impl AnalyzeExecBuilder { @@ -89,6 +94,7 @@ impl AnalyzeExecBuilder { metric_types: vec![MetricType::Summary, MetricType::Dev], metric_categories: None, format: ExplainFormat::Indent, + statistics_registry: StatisticsRegistry::new(), } } @@ -110,6 +116,11 @@ impl AnalyzeExecBuilder { self } + pub fn with_statistics_registry(mut self, registry: StatisticsRegistry) -> Self { + self.statistics_registry = registry; + self + } + pub fn build(self) -> AnalyzeExec { let cache = AnalyzeExec::compute_properties(&self.input, Arc::clone(&self.schema)); @@ -119,6 +130,7 @@ impl AnalyzeExecBuilder { metric_types: self.metric_types, metric_categories: self.metric_categories, format: self.format, + statistics_registry: self.statistics_registry, input: self.input, schema: self.schema, cache: Arc::new(cache), @@ -233,6 +245,7 @@ impl ExecutionPlan for AnalyzeExec { .with_metric_types(self.metric_types.clone()) .with_metric_categories(self.metric_categories.clone()) .with_format(self.format.clone()) + .with_statistics_registry(self.statistics_registry.clone()) .build(), )) } @@ -272,6 +285,7 @@ impl ExecutionPlan for AnalyzeExec { let metric_types = self.metric_types.clone(); let metric_categories = self.metric_categories.clone(); let format = self.format.clone(); + let statistics_registry = self.statistics_registry.clone(); // future that gathers the results from all the tasks in the // JoinSet that computes the overall row count and final @@ -295,6 +309,7 @@ impl ExecutionPlan for AnalyzeExec { &metric_types, metric_categories.as_deref(), &format, + &statistics_registry, ) }; @@ -317,6 +332,7 @@ fn create_output_batch( metric_types: &[MetricType], metric_categories: Option<&[MetricCategory]>, format: &ExplainFormat, + statistics_registry: &StatisticsRegistry, ) -> Result { let mut type_builder = StringBuilder::with_capacity(1, 1024); let mut plan_builder = StringBuilder::with_capacity(1, 1024); @@ -329,6 +345,7 @@ fn create_output_batch( .set_metric_types(metric_types.to_vec()) .set_metric_categories(metric_categories.map(|c| c.to_vec())) .set_show_statistics(show_statistics) + .set_statistics_registry(statistics_registry.clone()) .indent(verbose) .to_string(); plan_builder.append_value(annotated_plan); @@ -341,6 +358,7 @@ fn create_output_batch( .set_metric_types(metric_types.to_vec()) .set_metric_categories(metric_categories.map(|c| c.to_vec())) .set_show_statistics(show_statistics) + .set_statistics_registry(statistics_registry.clone()) .indent(verbose) .to_string(); plan_builder.append_value(annotated_plan); diff --git a/datafusion/physical-plan/src/display.rs b/datafusion/physical-plan/src/display.rs index 6a4d09057bec9..92ae46f275707 100644 --- a/datafusion/physical-plan/src/display.rs +++ b/datafusion/physical-plan/src/display.rs @@ -32,6 +32,7 @@ use datafusion_physical_expr::LexOrdering; use crate::metrics::{MetricCategory, MetricType, MetricValue}; use crate::render_tree::RenderTree; +use crate::operator_statistics::StatisticsRegistry; use crate::statistics::{StatisticsArgs, StatisticsContext}; use super::{ExecutionPlan, ExecutionPlanVisitor, accept}; @@ -122,6 +123,7 @@ pub struct DisplayableExecutionPlan<'a> { show_metrics: ShowMetrics, /// If statistics should be displayed show_statistics: bool, + registry: StatisticsRegistry, /// If schema should be displayed. See [`Self::set_show_schema`] show_schema: bool, /// Which metric categories should be included when rendering @@ -154,6 +156,7 @@ impl<'a> DisplayableExecutionPlan<'a> { pub fn new(inner: &'a dyn ExecutionPlan) -> Self { Self { inner, + registry: StatisticsRegistry::new(), show_metrics: ShowMetrics::None, show_statistics: false, show_schema: false, @@ -170,6 +173,7 @@ impl<'a> DisplayableExecutionPlan<'a> { pub fn with_metrics(inner: &'a dyn ExecutionPlan) -> Self { Self { inner, + registry: StatisticsRegistry::new(), show_metrics: ShowMetrics::Aggregated, show_statistics: false, show_schema: false, @@ -186,6 +190,7 @@ impl<'a> DisplayableExecutionPlan<'a> { pub fn with_full_metrics(inner: &'a dyn ExecutionPlan) -> Self { Self { inner, + registry: StatisticsRegistry::new(), show_metrics: ShowMetrics::Full, show_statistics: false, show_schema: false, @@ -211,6 +216,12 @@ impl<'a> DisplayableExecutionPlan<'a> { self } + /// Set the [`StatisticsRegistry`] consulted when computing displayed statistics. + pub fn set_statistics_registry(mut self, registry: StatisticsRegistry) -> Self { + self.registry = registry; + self + } + /// Specify which metric types should be rendered alongside the plan pub fn set_metric_types(mut self, metric_types: Vec) -> Self { self.metric_types = metric_types; @@ -276,6 +287,7 @@ impl<'a> DisplayableExecutionPlan<'a> { plan: &'a dyn ExecutionPlan, show_metrics: ShowMetrics, show_statistics: bool, + registry: StatisticsRegistry, show_schema: bool, metric_types: Vec, metric_categories: Option>, @@ -288,6 +300,7 @@ impl<'a> DisplayableExecutionPlan<'a> { indent: 0, show_metrics: self.show_metrics, show_statistics: self.show_statistics, + registry: self.registry.clone(), show_schema: self.show_schema, metric_types: &self.metric_types, metric_categories: self.metric_categories.as_deref(), @@ -300,6 +313,7 @@ impl<'a> DisplayableExecutionPlan<'a> { plan: self.inner, show_metrics: self.show_metrics, show_statistics: self.show_statistics, + registry: self.registry.clone(), show_schema: self.show_schema, metric_types: self.metric_types.clone(), metric_categories: self.metric_categories.clone(), @@ -322,6 +336,7 @@ impl<'a> DisplayableExecutionPlan<'a> { plan: &'a dyn ExecutionPlan, show_metrics: ShowMetrics, show_statistics: bool, + registry: StatisticsRegistry, metric_types: Vec, metric_categories: Option>, } @@ -334,6 +349,7 @@ impl<'a> DisplayableExecutionPlan<'a> { t, show_metrics: self.show_metrics, show_statistics: self.show_statistics, + registry: self.registry.clone(), metric_types: &self.metric_types, metric_categories: self.metric_categories.as_deref(), graphviz_builder: GraphvizBuilder::default(), @@ -353,6 +369,7 @@ impl<'a> DisplayableExecutionPlan<'a> { plan: self.inner, show_metrics: self.show_metrics, show_statistics: self.show_statistics, + registry: self.registry.clone(), metric_types: self.metric_types.clone(), metric_categories: self.metric_categories.clone(), } @@ -457,6 +474,7 @@ impl<'a> DisplayableExecutionPlan<'a> { plan: &'a dyn ExecutionPlan, show_metrics: ShowMetrics, show_statistics: bool, + registry: StatisticsRegistry, show_schema: bool, metric_types: Vec, metric_categories: Option>, @@ -470,6 +488,7 @@ impl<'a> DisplayableExecutionPlan<'a> { indent: 0, show_metrics: self.show_metrics, show_statistics: self.show_statistics, + registry: self.registry.clone(), show_schema: self.show_schema, metric_types: &self.metric_types, metric_categories: self.metric_categories.as_deref(), @@ -483,6 +502,7 @@ impl<'a> DisplayableExecutionPlan<'a> { plan: self.inner, show_metrics: self.show_metrics, show_statistics: self.show_statistics, + registry: self.registry.clone(), show_schema: self.show_schema, metric_types: self.metric_types.clone(), metric_categories: self.metric_categories.clone(), @@ -538,6 +558,7 @@ struct IndentVisitor<'a, 'b> { show_metrics: ShowMetrics, /// If statistics should be displayed show_statistics: bool, + registry: StatisticsRegistry, /// If schema should be displayed show_schema: bool, /// Which metric types should be rendered @@ -581,7 +602,7 @@ impl ExecutionPlanVisitor for IndentVisitor<'_, '_> { } } if self.show_statistics { - let stats = StatisticsContext::new() + let stats = StatisticsContext::new_with_registry(self.registry.clone()) .compute(plan, &StatisticsArgs::new()) .map_err(|_e| fmt::Error)?; write!(self.f, ", statistics=[{stats}]")?; @@ -612,6 +633,7 @@ struct GraphvizVisitor<'a, 'b> { show_metrics: ShowMetrics, /// If statistics should be displayed show_statistics: bool, + registry: StatisticsRegistry, /// Which metric types should be rendered metric_types: &'a [MetricType], /// Optional filter by semantic category @@ -679,7 +701,7 @@ impl ExecutionPlanVisitor for GraphvizVisitor<'_, '_> { }; let statistics = if self.show_statistics { - let stats = StatisticsContext::new() + let stats = StatisticsContext::new_with_registry(self.registry.clone()) .compute(plan, &StatisticsArgs::new()) .map_err(|_e| fmt::Error)?; format!("statistics=[{stats}]") diff --git a/datafusion/sqllogictest/test_files/statistics_registry.slt b/datafusion/sqllogictest/test_files/statistics_registry.slt index ff4a6ae67bde5..9e8dbc8293c90 100644 --- a/datafusion/sqllogictest/test_files/statistics_registry.slt +++ b/datafusion/sqllogictest/test_files/statistics_registry.slt @@ -106,6 +106,45 @@ physical_plan 06)------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/statistics_registry/customers.parquet]]}, projection=[customer_id, region_id], output_ordering=[region_id@1 ASC NULLS LAST], file_type=parquet 07)------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/statistics_registry/orders.parquet]]}, projection=[order_id, customer_id, small_id], output_ordering=[order_id@0 ASC NULLS LAST], file_type=parquet, predicate=DynamicFilter [ empty ] AND DynamicFilter [ empty ], dynamic_rg_pruning=eligible +# -- Registry reflected in EXPLAIN statistics -------------------------------- +# With show_statistics on, the displayed statistics column reflects the +# registry's estimates (the top join's conservative Rows), not the built-in +# empty-registry path. + +statement ok +set datafusion.explain.show_statistics = true; + +query TT +EXPLAIN SELECT o.order_id, c.region_id, d.label +FROM customers c +JOIN orders o ON c.customer_id = o.customer_id +JOIN dim_small d ON o.small_id = d.small_id; +---- +physical_plan +01)HashJoinExec: mode=Partitioned, join_type=Inner, on=[(small_id@0, small_id@2)], projection=[order_id@3, region_id@2, label@1], statistics=[Rows=Inexact(5000), Bytes=Absent, [(Col[0]: Min=Inexact(Int32(1)) Max=Inexact(Int32(10)) Null=Inexact(0) ScanBytes=Inexact(40)),(Col[1]: Min=Exact(Int32(1)) Max=Exact(Int32(10)) Null=Exact(0) ScanBytes=Exact(40)),(Col[2]: Min=Exact(Int32(1)) Max=Exact(Int32(50)) Null=Exact(0) ScanBytes=Exact(200))]] +02)--RepartitionExec: partitioning=Hash([small_id@0], 4), input_partitions=1, maintains_sort_order=true, statistics=[Rows=Exact(50), Bytes=Exact(400), [(Col[0]: Min=Exact(Int32(1)) Max=Exact(Int32(50)) Null=Exact(0) ScanBytes=Exact(200)),(Col[1]: Min=Exact(Int32(1)) Max=Exact(Int32(50)) Null=Exact(0) ScanBytes=Exact(200))]] +03)----DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/statistics_registry/dim_small.parquet]]}, projection=[small_id, label], output_ordering=[small_id@0 ASC NULLS LAST], file_type=parquet, statistics=[Rows=Exact(50), Bytes=Exact(400), [(Col[0]: Min=Exact(Int32(1)) Max=Exact(Int32(50)) Null=Exact(0) ScanBytes=Exact(200)),(Col[1]: Min=Exact(Int32(1)) Max=Exact(Int32(50)) Null=Exact(0) ScanBytes=Exact(200))]] +04)--RepartitionExec: partitioning=Hash([small_id@2], 4), input_partitions=1, maintains_sort_order=true, statistics=[Rows=Inexact(100), Bytes=Absent, [(Col[0]: Min=Exact(Int32(1)) Max=Exact(Int32(10)) Null=Exact(0) ScanBytes=Exact(40)),(Col[1]: Min=Inexact(Int32(1)) Max=Inexact(Int32(10)) Null=Inexact(0) ScanBytes=Inexact(40)),(Col[2]: Min=Inexact(Int32(1)) Max=Inexact(Int32(10)) Null=Inexact(0) ScanBytes=Inexact(40))]] +05)----HashJoinExec: mode=Partitioned, join_type=Inner, on=[(customer_id@0, customer_id@1)], projection=[region_id@1, order_id@2, small_id@4], statistics=[Rows=Inexact(100), Bytes=Absent, [(Col[0]: Min=Exact(Int32(1)) Max=Exact(Int32(10)) Null=Exact(0) ScanBytes=Exact(40)),(Col[1]: Min=Inexact(Int32(1)) Max=Inexact(Int32(10)) Null=Inexact(0) ScanBytes=Inexact(40)),(Col[2]: Min=Inexact(Int32(1)) Max=Inexact(Int32(10)) Null=Inexact(0) ScanBytes=Inexact(40))]] +06)------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/statistics_registry/customers.parquet]]}, projection=[customer_id, region_id], output_ordering=[region_id@1 ASC NULLS LAST], file_type=parquet, statistics=[Rows=Exact(10), Bytes=Exact(80), [(Col[0]: Min=Exact(Int32(1)) Max=Exact(Int32(3)) Null=Exact(0) ScanBytes=Exact(40)),(Col[1]: Min=Exact(Int32(1)) Max=Exact(Int32(10)) Null=Exact(0) ScanBytes=Exact(40))]] +07)------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/statistics_registry/orders.parquet]]}, projection=[order_id, customer_id, small_id], output_ordering=[order_id@0 ASC NULLS LAST], file_type=parquet, predicate=DynamicFilter [ empty ] AND DynamicFilter [ empty ], dynamic_rg_pruning=eligible, statistics=[Rows=Inexact(10), Bytes=Inexact(120), [(Col[0]: Min=Inexact(Int32(1)) Max=Inexact(Int32(10)) Null=Inexact(0) ScanBytes=Inexact(40)),(Col[1]: Min=Inexact(Int32(1)) Max=Inexact(Int32(3)) Null=Inexact(0) ScanBytes=Inexact(40)),(Col[2]: Min=Inexact(Int32(1)) Max=Inexact(Int32(10)) Null=Inexact(0) ScanBytes=Inexact(40))]] + +# -- Registry reflected in EXPLAIN ANALYZE statistics ------------------------ +# EXPLAIN ANALYZE renders the same registry-refined statistics as plain EXPLAIN +# (the top join's conservative Rows=Inexact(5000)) alongside runtime metrics. +# Timings are volatile, so match only the deterministic statistics fragment. + +query TT +EXPLAIN ANALYZE SELECT o.order_id, c.region_id, d.label +FROM customers c +JOIN orders o ON c.customer_id = o.customer_id +JOIN dim_small d ON o.small_id = d.small_id; +---- +Plan with Metricsstatistics=[Rows=Inexact(5000) + +statement ok +set datafusion.explain.show_statistics = false; + # -- Correctness ------------------------------------------------------------- query I From ded429241cb1d217c5367f878af625155695eb53 Mon Sep 17 00:00:00 2001 From: Alessandro Solimando Date: Fri, 17 Jul 2026 00:52:02 +0200 Subject: [PATCH 04/15] feat(examples): add StatisticsRegistry example A runnable example showing a StatisticsProvider that supplies and refines column statistics (post-filter distinct count via the survival formula) to flip a join's build side, with a before/after EXPLAIN. --- datafusion-examples/README.md | 10 + .../examples/statistics/join_reorder.rs | 182 ++++++++++++++++++ .../examples/statistics/main.rs | 84 ++++++++ 3 files changed, 276 insertions(+) create mode 100644 datafusion-examples/examples/statistics/join_reorder.rs create mode 100644 datafusion-examples/examples/statistics/main.rs diff --git a/datafusion-examples/README.md b/datafusion-examples/README.md index 4746ac9114733..7aec562640d7b 100644 --- a/datafusion-examples/README.md +++ b/datafusion-examples/README.md @@ -213,6 +213,16 @@ cargo run --example dataframe -- dataframe | frontend | [`sql_ops/frontend.rs`](examples/sql_ops/frontend.rs) | Build LogicalPlans from SQL | | query | [`sql_ops/query.rs`](examples/sql_ops/query.rs) | Query data using SQL | +## Statistics Examples + +### Group: `statistics` + +#### Category: Single Process + +| Subcommand | File Path | Description | +| ------------ | ------------------------------------------------------------------- | ----------------------------------------------------------------------- | +| join_reorder | [`statistics/join_reorder.rs`](examples/statistics/join_reorder.rs) | Supply and refine column statistics via a provider to flip a join order | + ## UDF Examples ### Group: `udf` diff --git a/datafusion-examples/examples/statistics/join_reorder.rs b/datafusion-examples/examples/statistics/join_reorder.rs new file mode 100644 index 0000000000000..4a689a89587f1 --- /dev/null +++ b/datafusion-examples/examples/statistics/join_reorder.rs @@ -0,0 +1,182 @@ +// 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. + +//! Plug refined cardinality estimation into the optimizer via the `StatisticsRegistry`. +//! +//! DataFusion's built-in statistics are intentionally simple defaults; the registry +//! is the seam for plugging in more refined estimation (formulas or heuristics +//! suited to your system or data) without changing the core. +//! +//! `(SELECT user_id FROM events WHERE amount < 50 GROUP BY user_id) JOIN dims`: a +//! provider supplies the column stats a catalog knows (`amount` range, `user_id` +//! distinct count); the built-in `FilterStatisticsProvider` then refines the +//! post-filter distinct count with the survival formula +//! `NDV * (1 - (1 - selectivity)^(rows / NDV))` (Yao/Cardenas) to ~32, below `dims` +//! (48), flipping the join build side. Core's simpler `min(NDV, rows)` cap would +//! give 50 (> 48) and keep the other order; the refinement is the point. The +//! ground-truth query prints the true surviving distinct count (below 48), +//! confirming the flip. + +use std::sync::Arc; + +use datafusion::arrow::array::Int32Array; +use datafusion::arrow::datatypes::{DataType, Field, Schema}; +use datafusion::arrow::record_batch::RecordBatch; +use datafusion::arrow::util::pretty::pretty_format_batches; +use datafusion::catalog::MemTable; +use datafusion::common::Result; +use datafusion::common::ScalarValue; +use datafusion::common::stats::Precision; +use datafusion::execution::SessionStateBuilder; +use datafusion::physical_plan::ExecutionPlan; +use datafusion::physical_plan::operator_statistics::{ + ClosureStatisticsProvider, ExtendedStatistics, StatisticsRegistry, StatisticsResult, +}; +use datafusion::physical_plan::statistics::StatisticsArgs; +use datafusion::prelude::*; +use rand::{Rng, SeedableRng, rngs::StdRng}; + +/// Supplies the catalog-known statistics of the `events` scan that the in-memory +/// table does not carry: the `amount` range (for filter selectivity) and the +/// `user_id` distinct count (for post-filter NDV estimation). Only the base +/// `events` scan (a leaf carrying both columns) is matched; everything else +/// delegates to the rest of the chain. +fn catalog_stats( + plan: &dyn ExecutionPlan, + _child_stats: &[ExtendedStatistics], +) -> Result { + let schema = plan.schema(); + let (Ok(user_id), Ok(amount)) = + (schema.index_of("user_id"), schema.index_of("amount")) + else { + return Ok(StatisticsResult::Delegate); + }; + if !plan.children().is_empty() { + return Ok(StatisticsResult::Delegate); + } + let mut stats = (*plan.statistics_from_inputs(&[], &StatisticsArgs::new())?).clone(); + stats.column_statistics[amount].min_value = + Precision::Inexact(ScalarValue::Int32(Some(0))); + stats.column_statistics[amount].max_value = + Precision::Inexact(ScalarValue::Int32(Some(999))); + stats.column_statistics[user_id].distinct_count = Precision::Inexact(100); + Ok(StatisticsResult::Computed(ExtendedStatistics::new(stats))) +} + +fn int_col(values: &[i32]) -> Arc { + Arc::new(Int32Array::from_iter_values(values.iter().copied())) +} + +fn mem_table(fields: &[(&str, Arc)]) -> Result> { + let schema = Arc::new(Schema::new( + fields + .iter() + .map(|(name, _)| Field::new(*name, DataType::Int32, false)) + .collect::>(), + )); + let cols = fields.iter().map(|(_, col)| Arc::clone(col) as _).collect(); + let batch = RecordBatch::try_new(Arc::clone(&schema), cols)?; + Ok(Arc::new(MemTable::try_new(schema, vec![vec![batch]])?)) +} + +async fn build_ctx(with_registry: bool) -> Result { + let config = SessionConfig::new() + .with_target_partitions(4) + .set_bool("datafusion.explain.physical_plan_only", true) + .set_bool("datafusion.explain.show_statistics", true) + // Force Partitioned hash joins so statistics alone drive the build side. + .set_usize( + "datafusion.optimizer.hash_join_single_partition_threshold", + 1, + ) + .set_usize( + "datafusion.optimizer.hash_join_single_partition_threshold_rows", + 1, + ); + + let mut builder = SessionStateBuilder::new() + .with_config(config) + .with_default_features(); + if with_registry { + let mut registry = StatisticsRegistry::default_with_builtin_providers(); + registry.register(Arc::new(ClosureStatisticsProvider::new(catalog_stats))); + builder = builder.with_statistics_registry(registry); + } + let ctx = SessionContext::new_with_state(builder.build()); + + let n = 1000i32; + let user_ids: Vec = (0..n).map(|v| v % 100).collect(); + // `amount` independent of `user_id` so the filter keeps a representative sample. + let mut rng = StdRng::seed_from_u64(2024); + let amounts: Vec = (0..n).map(|_| rng.random_range(0..1000)).collect(); + ctx.register_table( + "events", + mem_table(&[ + ("user_id", int_col(&user_ids)), + ("amount", int_col(&amounts)), + ])?, + )?; + ctx.register_table( + "dims", + mem_table(&[ + ("user_id", int_col(&(0..48).collect::>())), + ("label", int_col(&(0..48).collect::>())), + ])?, + )?; + Ok(ctx) +} + +const QUERY: &str = "SELECT e.user_id, d.label \ + FROM (SELECT user_id FROM events WHERE amount < 50 GROUP BY user_id) e \ + JOIN dims d ON e.user_id = d.user_id"; + +async fn explain(ctx: &SessionContext) -> Result { + let batches = ctx + .sql(&format!("EXPLAIN {QUERY}")) + .await? + .collect() + .await?; + Ok(pretty_format_batches(&batches)?.to_string()) +} + +pub async fn join_reorder() -> Result<()> { + let truth_query = "SELECT count(DISTINCT user_id) AS true_distinct_users \ + FROM events WHERE amount < 50"; + println!("-- Ground truth --\n{truth_query}\n"); + let truth = build_ctx(false) + .await? + .sql(truth_query) + .await? + .collect() + .await?; + println!("{}\n", pretty_format_batches(&truth)?); + + println!("-- Query --\n{QUERY}\n"); + println!( + "A hash join builds its in-memory hash table from one input and probes with\n\ + the other, so the smaller input should be the build side. Default estimation\n\ + sizes the grouped `events` at 1000 rows and builds from `dims`; the\n\ + registry's refined ~32 estimate is below `dims` (48 rows), so it flips the\n\ + build side to `events`. The ground-truth count above (also below 48)\n\ + confirms `events` really is the smaller, cheaper side.\n" + ); + println!("-- Without the registry (default estimation) --"); + println!("{}\n", explain(&build_ctx(false).await?).await?); + println!("-- With the registry (built-in refinement) --"); + println!("{}", explain(&build_ctx(true).await?).await?); + Ok(()) +} diff --git a/datafusion-examples/examples/statistics/main.rs b/datafusion-examples/examples/statistics/main.rs new file mode 100644 index 0000000000000..313e18645d50f --- /dev/null +++ b/datafusion-examples/examples/statistics/main.rs @@ -0,0 +1,84 @@ +// 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. + +//! # Pluggable operator statistics (`StatisticsRegistry`) +//! +//! These examples show how a `StatisticsProvider` registered on the session can +//! change a physical plan by supplying statistics the built-in estimators +//! cannot express. +//! +//! ## Usage +//! ```bash +//! cargo run --example statistics -- [all|join_reorder] +//! ``` +//! +//! Each subcommand runs a corresponding example: +//! - `all`: run all examples included in this module +//! +//! - `join_reorder` +//! (file: join_reorder.rs, desc: Supply and refine column statistics via a provider to flip a join order) + +mod join_reorder; + +use datafusion::error::{DataFusionError, Result}; +use strum::{IntoEnumIterator, VariantNames}; +use strum_macros::{Display, EnumIter, EnumString, VariantNames}; + +#[derive(EnumIter, EnumString, Display, VariantNames)] +#[strum(serialize_all = "snake_case")] +enum ExampleKind { + All, + JoinReorder, +} + +impl ExampleKind { + const EXAMPLE_NAME: &str = "statistics"; + + fn runnable() -> impl Iterator { + ExampleKind::iter().filter(|v| !matches!(v, ExampleKind::All)) + } + + async fn run(&self) -> Result<()> { + match self { + ExampleKind::All => { + for example in ExampleKind::runnable() { + println!("Running example: {example}"); + Box::pin(example.run()).await?; + } + Ok(()) + } + ExampleKind::JoinReorder => join_reorder::join_reorder().await, + } + } +} + +#[tokio::main] +async fn main() -> Result<()> { + let usage = format!( + "Usage: cargo run --example {} -- [{}]", + ExampleKind::EXAMPLE_NAME, + ExampleKind::VARIANTS.join("|") + ); + + let example: ExampleKind = std::env::args() + .nth(1) + .unwrap_or_else(|| ExampleKind::All.to_string()) + .parse() + .map_err(|_| DataFusionError::Execution(format!("Unknown example. {usage}")))?; + + example.run().await +} From dd79d6f3d44717e943e2ab3be5d3e789a527b2d4 Mon Sep 17 00:00:00 2001 From: Alessandro Solimando Date: Mon, 20 Jul 2026 20:56:40 +0200 Subject: [PATCH 05/15] refactor: dedupe EXPLAIN display setup --- datafusion/core/src/physical_planner.rs | 54 +++++++++++-------------- datafusion/physical-plan/src/display.rs | 33 +++++---------- 2 files changed, 34 insertions(+), 53 deletions(-) diff --git a/datafusion/core/src/physical_planner.rs b/datafusion/core/src/physical_planner.rs index 0a71c77cd5e59..7ce30d6854df3 100644 --- a/datafusion/core/src/physical_planner.rs +++ b/datafusion/core/src/physical_planner.rs @@ -2674,6 +2674,15 @@ impl DefaultPhysicalPlanner { } if !config.logical_plan_only && e.logical_optimization_succeeded { + let render_indent = + |plan: &dyn ExecutionPlan, show_statistics: bool, show_schema: bool| { + displayable(plan) + .set_show_statistics(show_statistics) + .set_statistics_registry(statistics_registry.clone()) + .set_show_schema(show_schema) + .indent(e.verbose) + .to_string() + }; match self .create_initial_plan(e.plan.as_ref(), session_state) .await @@ -2682,12 +2691,11 @@ impl DefaultPhysicalPlanner { // Include statistics / schema if enabled stringified_plans.push(StringifiedPlan::new( InitialPhysicalPlan, - displayable(input.as_ref()) - .set_show_statistics(show_statistics) - .set_statistics_registry(statistics_registry.clone()) - .set_show_schema(config.show_schema) - .indent(e.verbose) - .to_string(), + render_indent( + input.as_ref(), + show_statistics, + config.show_schema, + ), )); // Show statistics + schema in verbose output even if not @@ -2696,16 +2704,14 @@ impl DefaultPhysicalPlanner { if !show_statistics { stringified_plans.push(StringifiedPlan::new( InitialPhysicalPlanWithStats, - displayable(input.as_ref()) - .set_show_statistics(true) - .set_statistics_registry(statistics_registry.clone()) - .indent(e.verbose) - .to_string(), + render_indent(input.as_ref(), true, false), )); } if !config.show_schema { stringified_plans.push(StringifiedPlan::new( InitialPhysicalPlanWithSchema, + // Schema only: statistics are off, so this + // renders without the registry. displayable(input.as_ref()) .set_show_schema(true) .indent(e.verbose) @@ -2722,12 +2728,7 @@ impl DefaultPhysicalPlanner { let plan_type = OptimizedPhysicalPlan { optimizer_name }; stringified_plans.push(StringifiedPlan::new( plan_type, - displayable(plan) - .set_show_statistics(show_statistics) - .set_statistics_registry(statistics_registry.clone()) - .set_show_schema(config.show_schema) - .indent(e.verbose) - .to_string(), + render_indent(plan, show_statistics, config.show_schema), )); }, ); @@ -2736,12 +2737,11 @@ impl DefaultPhysicalPlanner { // This plan will includes statistics if show_statistics is on stringified_plans.push(StringifiedPlan::new( FinalPhysicalPlan, - displayable(input.as_ref()) - .set_show_statistics(show_statistics) - .set_statistics_registry(statistics_registry.clone()) - .set_show_schema(config.show_schema) - .indent(e.verbose) - .to_string(), + render_indent( + input.as_ref(), + show_statistics, + config.show_schema, + ), )); // Show statistics + schema in verbose output even if not @@ -2750,13 +2750,7 @@ impl DefaultPhysicalPlanner { if !show_statistics { stringified_plans.push(StringifiedPlan::new( FinalPhysicalPlanWithStats, - displayable(input.as_ref()) - .set_show_statistics(true) - .set_statistics_registry( - statistics_registry.clone(), - ) - .indent(e.verbose) - .to_string(), + render_indent(input.as_ref(), true, false), )); } if !config.show_schema { diff --git a/datafusion/physical-plan/src/display.rs b/datafusion/physical-plan/src/display.rs index 92ae46f275707..0f20ecd3984ef 100644 --- a/datafusion/physical-plan/src/display.rs +++ b/datafusion/physical-plan/src/display.rs @@ -154,44 +154,31 @@ impl<'a> DisplayableExecutionPlan<'a> { /// Create a wrapper around an [`ExecutionPlan`] which can be /// pretty printed in a variety of ways pub fn new(inner: &'a dyn ExecutionPlan) -> Self { - Self { - inner, - registry: StatisticsRegistry::new(), - show_metrics: ShowMetrics::None, - show_statistics: false, - show_schema: false, - metric_types: Self::default_metric_types(), - metric_categories: None, - tree_maximum_render_width: 240, - summary: None, - } + Self::with_show_metrics(inner, ShowMetrics::None) } /// Create a wrapper around an [`ExecutionPlan`] which can be /// pretty printed in a variety of ways that also shows aggregated /// metrics pub fn with_metrics(inner: &'a dyn ExecutionPlan) -> Self { - Self { - inner, - registry: StatisticsRegistry::new(), - show_metrics: ShowMetrics::Aggregated, - show_statistics: false, - show_schema: false, - metric_types: Self::default_metric_types(), - metric_categories: None, - tree_maximum_render_width: 240, - summary: None, - } + Self::with_show_metrics(inner, ShowMetrics::Aggregated) } /// Create a wrapper around an [`ExecutionPlan`] which can be /// pretty printed in a variety of ways that also shows all low /// level metrics pub fn with_full_metrics(inner: &'a dyn ExecutionPlan) -> Self { + Self::with_show_metrics(inner, ShowMetrics::Full) + } + + fn with_show_metrics( + inner: &'a dyn ExecutionPlan, + show_metrics: ShowMetrics, + ) -> Self { Self { inner, registry: StatisticsRegistry::new(), - show_metrics: ShowMetrics::Full, + show_metrics, show_statistics: false, show_schema: false, metric_types: Self::default_metric_types(), From e4df6301d27f7aaa89add5470760c69b39f723d5 Mon Sep 17 00:00:00 2001 From: Alessandro Solimando Date: Wed, 22 Jul 2026 13:01:34 +0200 Subject: [PATCH 06/15] feat: let StatisticsProvider declare its child stat needs --- .../src/operator_statistics/mod.rs | 78 ++++++++++++++++--- datafusion/physical-plan/src/statistics.rs | 59 +++++++------- .../library-user-guide/upgrading/55.0.0.md | 9 +++ 3 files changed, 110 insertions(+), 36 deletions(-) diff --git a/datafusion/physical-plan/src/operator_statistics/mod.rs b/datafusion/physical-plan/src/operator_statistics/mod.rs index 5ba68f11b147e..68fcb881d9fad 100644 --- a/datafusion/physical-plan/src/operator_statistics/mod.rs +++ b/datafusion/physical-plan/src/operator_statistics/mod.rs @@ -84,7 +84,7 @@ use datafusion_common::stats::Precision; use datafusion_common::{Result, Statistics}; use crate::ExecutionPlan; -use crate::statistics::{StatisticsArgs, StatisticsContext}; +use crate::statistics::{ChildStats, StatisticsArgs, StatisticsContext}; // ============================================================================ // ExtendedStatistics: Statistics with type-safe extensions @@ -281,6 +281,32 @@ pub trait StatisticsProvider: Debug + Send + Sync { self.compute_statistics(plan, child_stats) } } + + /// Which child statistics this provider needs for `plan`, one entry per + /// child in `plan.children()` order. + /// + /// Resolved independently of the operator's own + /// [`ExecutionPlan::child_stats_requests`], so a provider can refine an + /// operator that itself declares [`ChildStats::Skip`]. Results are memoized + /// per `(node, partition)`, so requesting a child already resolved for the + /// operator costs nothing. + /// + /// The default requests each child's overall (`None`) statistics, so + /// refining from child statistics works without modifying the operator. + /// Override to [`ChildStats::Skip`] a child the provider does not need. A + /// provider that computes per-partition statistics (overriding + /// [`Self::compute_statistics_with_args`]) must also override this to request + /// the matching partition, which the default ignores. + fn child_stats_requests( + &self, + plan: &dyn ExecutionPlan, + _partition: Option, + ) -> Vec { + plan.children() + .iter() + .map(|_| ChildStats::At(None)) + .collect() + } } /// Deprecated statistics provider that delegates to each operator's built-in @@ -1054,7 +1080,7 @@ mod tests { use super::*; use crate::filter::FilterExec; use crate::projection::ProjectionExec; - use crate::statistics::{StatisticsArgs, StatisticsContext}; + use crate::statistics::{ChildStats, StatisticsArgs, StatisticsContext}; use crate::{DisplayAs, DisplayFormatType, PlanProperties}; use arrow::datatypes::{DataType, Field, Schema}; use datafusion_common::stats::Precision; @@ -1264,13 +1290,6 @@ mod tests { vec![&self.input] } - fn child_stats_requests( - &self, - partition: Option, - ) -> Vec { - vec![crate::statistics::ChildStats::At(partition)] - } - fn with_new_children( self: Arc, children: Vec>, @@ -1323,6 +1342,47 @@ mod tests { Ok(()) } + /// A provider that overrides `child_stats_requests` to need no children. + #[derive(Debug)] + struct NoChildStatsProvider; + + impl StatisticsProvider for NoChildStatsProvider { + fn child_stats_requests( + &self, + plan: &dyn ExecutionPlan, + _partition: Option, + ) -> Vec { + plan.children().iter().map(|_| ChildStats::Skip).collect() + } + + fn compute_statistics( + &self, + plan: &dyn ExecutionPlan, + child_stats: &[ExtendedStatistics], + ) -> Result { + if plan.downcast_ref::().is_some() { + Ok(StatisticsResult::Computed(child_stats[0].clone())) + } else { + Ok(StatisticsResult::Delegate) + } + } + } + + #[test] + fn test_provider_opts_out_of_child_stats() -> Result<()> { + let mut engine = StatisticsRegistry::new(); + engine.register(Arc::new(NoChildStatsProvider)); + + let source = make_source(1000); + let custom: Arc = Arc::new(CustomExec { input: source }); + + // The provider requested no children, so it sees the unknown placeholder + // rather than the source's row count. + let stats = compute(&engine, custom.as_ref())?; + assert!(stats.base.num_rows.get_value().is_none()); + Ok(()) + } + #[derive(Debug)] struct OverrideFilterProvider { fixed_selectivity: f64, diff --git a/datafusion/physical-plan/src/statistics.rs b/datafusion/physical-plan/src/statistics.rs index 47499d82a9422..ca16d25a0f95b 100644 --- a/datafusion/physical-plan/src/statistics.rs +++ b/datafusion/physical-plan/src/statistics.rs @@ -222,6 +222,26 @@ impl StatisticsContext { let children = plan.children(); let requests = plan.child_stats_requests(partition); + // Resolved for the built-in `statistics_from_inputs` fallback below. + let child_statistics = self.resolve_children(plan, &children, &requests)?; + + let statistics = match self.try_provider_stats(plan, &children, args)? { + Some(statistics) => statistics, + None => plan.statistics_from_inputs(&child_statistics, args)?, + }; + self.store_statistics(plan, partition, Arc::clone(&statistics)); + Ok(statistics) + } + + /// Resolves each child's core statistics per `requests`: computes the child + /// at the requested partition (memoized), or supplies a + /// [`Statistics::new_unknown`] placeholder for [`ChildStats::Skip`]. + fn resolve_children( + &self, + plan: &dyn ExecutionPlan, + children: &[&Arc], + requests: &[ChildStats], + ) -> Result>> { assert_eq_or_internal_err!( requests.len(), children.len(), @@ -230,9 +250,9 @@ impl StatisticsContext { requests.len(), children.len() ); - let child_statistics: Vec> = children + children .iter() - .zip(&requests) + .zip(requests) .map(|(child, directive)| match directive { ChildStats::At(p) => self.compute_base( child.as_ref(), @@ -242,20 +262,7 @@ impl StatisticsContext { Ok(Arc::new(Statistics::new_unknown(child.schema().as_ref()))) } }) - .collect::>()?; - - let statistics = match self.try_provider_stats( - plan, - &children, - &requests, - &child_statistics, - args, - )? { - Some(statistics) => statistics, - None => plan.statistics_from_inputs(&child_statistics, args)?, - }; - self.store_statistics(plan, partition, Arc::clone(&statistics)); - Ok(statistics) + .collect() } /// Runs the provider chain, returning the first `Computed` result's core @@ -263,32 +270,30 @@ impl StatisticsContext { /// or all delegate. A partition-blind provider applies only to overall stats /// (its default `compute_statistics_with_args` delegates per partition). /// - /// Assembles each child's [`ExtendedStatistics`] (statistics plus cached - /// extensions) only here, so a walk with no providers pays no extension cost. + /// Each provider's child statistics come from its own + /// [`child_stats_requests`](crate::operator_statistics::StatisticsProvider::child_stats_requests) + /// and are memoized, so a walk with no providers pays nothing. fn try_provider_stats( &self, plan: &dyn ExecutionPlan, children: &[&Arc], - requests: &[ChildStats], - child_statistics: &[Arc], args: &StatisticsArgs, ) -> Result>> { let providers = self.registry.providers(); if providers.is_empty() { return Ok(None); } - let child_extended = - self.child_extended_stats(children, requests, child_statistics); + let partition = args.partition(); for provider in providers { + let requests = provider.child_stats_requests(plan, partition); + let child_statistics = self.resolve_children(plan, children, &requests)?; + let child_extended = + self.child_extended_stats(children, &requests, &child_statistics); if let StatisticsResult::Computed(computed) = provider.compute_statistics_with_args(plan, &child_extended, args)? { if !computed.extensions().is_empty() { - self.store_extensions( - plan, - args.partition(), - computed.extensions().clone(), - ); + self.store_extensions(plan, partition, computed.extensions().clone()); } return Ok(Some(Arc::clone(computed.base_arc()))); } diff --git a/docs/source/library-user-guide/upgrading/55.0.0.md b/docs/source/library-user-guide/upgrading/55.0.0.md index 63ee5b26b3779..24e1b821f0797 100644 --- a/docs/source/library-user-guide/upgrading/55.0.0.md +++ b/docs/source/library-user-guide/upgrading/55.0.0.md @@ -411,6 +411,15 @@ part of `StatisticsRegistry::default_with_builtin_providers()`. - Remove `DefaultStatisticsProvider` from any custom provider chain; register no terminal provider instead (the walk falls back on its own). +### `StatisticsRegistry::compute` and `compute_base` are deprecated + +Use the walk instead: + +```rust +StatisticsContext::new_with_registry(registry) + .compute_extended(plan, &StatisticsArgs::new())?; // or .compute(...) for core Statistics +``` + ### `DdlStatement::CreateExternalTable` and `CreateFunction` are now boxed The two largest variants of `datafusion_expr::DdlStatement` are now From 5f2b7f1ed8299fd4f122c1a27199ebaa7cf3f828 Mon Sep 17 00:00:00 2001 From: Alessandro Solimando Date: Wed, 22 Jul 2026 13:02:42 +0200 Subject: [PATCH 07/15] refactor: rename `engine` to `registry` in operator_statistics tests --- .../src/operator_statistics/mod.rs | 44 +++++++++---------- 1 file changed, 22 insertions(+), 22 deletions(-) diff --git a/datafusion/physical-plan/src/operator_statistics/mod.rs b/datafusion/physical-plan/src/operator_statistics/mod.rs index 68fcb881d9fad..575eb608ad6ed 100644 --- a/datafusion/physical-plan/src/operator_statistics/mod.rs +++ b/datafusion/physical-plan/src/operator_statistics/mod.rs @@ -1208,10 +1208,10 @@ mod tests { #[test] fn test_default_provider() -> Result<()> { - let engine = StatisticsRegistry::new(); + let registry = StatisticsRegistry::new(); let source = make_source(1000); - let stats = compute(&engine, source.as_ref())?; + let stats = compute(®istry, source.as_ref())?; assert!(matches!(stats.base.num_rows, Precision::Exact(1000))); Ok(()) } @@ -1331,13 +1331,13 @@ mod tests { #[test] fn test_custom_provider_for_custom_exec() -> Result<()> { - let mut engine = StatisticsRegistry::new(); - engine.register(Arc::new(CustomStatisticsProvider)); + let mut registry = StatisticsRegistry::new(); + registry.register(Arc::new(CustomStatisticsProvider)); let source = make_source(1000); let custom: Arc = Arc::new(CustomExec { input: source }); - let stats = compute(&engine, custom.as_ref())?; + let stats = compute(®istry, custom.as_ref())?; assert!(matches!(stats.base.num_rows, Precision::Exact(1000))); Ok(()) } @@ -1370,15 +1370,15 @@ mod tests { #[test] fn test_provider_opts_out_of_child_stats() -> Result<()> { - let mut engine = StatisticsRegistry::new(); - engine.register(Arc::new(NoChildStatsProvider)); + let mut registry = StatisticsRegistry::new(); + registry.register(Arc::new(NoChildStatsProvider)); let source = make_source(1000); let custom: Arc = Arc::new(CustomExec { input: source }); // The provider requested no children, so it sees the unknown placeholder // rather than the source's row count. - let stats = compute(&engine, custom.as_ref())?; + let stats = compute(®istry, custom.as_ref())?; assert!(stats.base.num_rows.get_value().is_none()); Ok(()) } @@ -1418,8 +1418,8 @@ mod tests { #[test] fn test_override_builtin_operator() -> Result<()> { - let mut engine = StatisticsRegistry::new(); - engine.register(Arc::new(OverrideFilterProvider { + let mut registry = StatisticsRegistry::new(); + registry.register(Arc::new(OverrideFilterProvider { fixed_selectivity: 0.1, })); @@ -1427,20 +1427,20 @@ mod tests { let filter: Arc = Arc::new(FilterExec::try_new(lit(true), source)?); - let stats = compute(&engine, filter.as_ref())?; + let stats = compute(®istry, filter.as_ref())?; assert!(matches!(stats.base.num_rows, Precision::Inexact(100))); Ok(()) } #[test] fn test_filter_statistics_propagation() -> Result<()> { - let engine = StatisticsRegistry::new(); + let registry = StatisticsRegistry::new(); let source = make_source(1000); let predicate = lit(true); let filter: Arc = Arc::new(FilterExec::try_new(predicate, source)?); - let stats = compute(&engine, filter.as_ref())?; + let stats = compute(®istry, filter.as_ref())?; assert!(stats.base.num_rows.get_value().unwrap_or(&0) <= &1000); Ok(()) } @@ -1520,7 +1520,7 @@ mod tests { #[test] fn test_projection_statistics_propagation() -> Result<()> { - let engine = StatisticsRegistry::new(); + let registry = StatisticsRegistry::new(); let source = make_source(1000); let schema = make_schema(); let proj: Arc = Arc::new(ProjectionExec::try_new( @@ -1528,7 +1528,7 @@ mod tests { source, )?); - let stats = compute(&engine, proj.as_ref())?; + let stats = compute(®istry, proj.as_ref())?; assert!(matches!(stats.base.num_rows, Precision::Exact(1000))); Ok(()) } @@ -1537,12 +1537,12 @@ mod tests { fn test_passthrough_statistics_propagation() -> Result<()> { use crate::coalesce_partitions::CoalescePartitionsExec; - let engine = StatisticsRegistry::new(); + let registry = StatisticsRegistry::new(); let source = make_source(1000); let coalesce: Arc = Arc::new(CoalescePartitionsExec::new(source)); - let stats = compute(&engine, coalesce.as_ref())?; + let stats = compute(®istry, coalesce.as_ref())?; // PassthroughStatisticsProvider should propagate child row count unchanged assert_eq!(stats.base.num_rows, Precision::Exact(1000)); Ok(()) @@ -1550,11 +1550,11 @@ mod tests { #[test] fn test_chain_priority() -> Result<()> { - let mut engine = StatisticsRegistry::new(); - engine.register(Arc::new(OverrideFilterProvider { + let mut registry = StatisticsRegistry::new(); + registry.register(Arc::new(OverrideFilterProvider { fixed_selectivity: 0.5, })); - engine.register(Arc::new(CustomStatisticsProvider)); + registry.register(Arc::new(CustomStatisticsProvider)); let source = make_source(1000); @@ -1562,13 +1562,13 @@ mod tests { let custom: Arc = Arc::new(CustomExec { input: Arc::clone(&source), }); - let stats = compute(&engine, custom.as_ref())?; + let stats = compute(®istry, custom.as_ref())?; assert!(matches!(stats.base.num_rows, Precision::Exact(1000))); // FilterExec: CustomStatisticsProvider delegates, OverrideFilterProvider handles let filter: Arc = Arc::new(FilterExec::try_new(lit(true), source)?); - let stats = compute(&engine, filter.as_ref())?; + let stats = compute(®istry, filter.as_ref())?; assert!(matches!(stats.base.num_rows, Precision::Inexact(500))); Ok(()) From 94f320bda83278e12c8578b86a091ebb944cf3c9 Mon Sep 17 00:00:00 2001 From: Alessandro Solimando Date: Mon, 27 Jul 2026 11:53:51 +0200 Subject: [PATCH 08/15] refactor: use imports instead of fully qualified names in operator_statistics tests --- .../src/operator_statistics/mod.rs | 22 +++++++++++-------- 1 file changed, 13 insertions(+), 9 deletions(-) diff --git a/datafusion/physical-plan/src/operator_statistics/mod.rs b/datafusion/physical-plan/src/operator_statistics/mod.rs index 575eb608ad6ed..9d44befff52f7 100644 --- a/datafusion/physical-plan/src/operator_statistics/mod.rs +++ b/datafusion/physical-plan/src/operator_statistics/mod.rs @@ -1079,12 +1079,16 @@ impl StatisticsProvider for ClosureStatisticsProvider { mod tests { use super::*; use crate::filter::FilterExec; + use crate::joins::JoinOn; use crate::projection::ProjectionExec; use crate::statistics::{ChildStats, StatisticsArgs, StatisticsContext}; - use crate::{DisplayAs, DisplayFormatType, PlanProperties}; + use crate::{ + DisplayAs, DisplayFormatType, PlanProperties, SendableRecordBatchStream, + }; use arrow::datatypes::{DataType, Field, Schema}; use datafusion_common::stats::Precision; use datafusion_common::{ColumnStatistics, ScalarValue}; + use datafusion_execution::TaskContext; use datafusion_expr::Operator; use datafusion_physical_expr::PhysicalExpr; use datafusion_physical_expr::expressions::{BinaryExpr, Column, Literal, col, lit}; @@ -1185,8 +1189,8 @@ mod tests { fn execute( &self, _partition: usize, - _context: Arc, - ) -> Result { + _context: Arc, + ) -> Result { unimplemented!() } @@ -1306,8 +1310,8 @@ mod tests { fn execute( &self, _partition: usize, - _context: Arc, - ) -> Result { + _context: Arc, + ) -> Result { unimplemented!() } } @@ -1873,7 +1877,7 @@ mod tests { right: Arc, ) -> Result> { let _schema = make_schema(); - let on: crate::joins::JoinOn = vec![( + let on: JoinOn = vec![( Arc::new(Column::new("a", 0)) as Arc, Arc::new(Column::new("a", 0)) as Arc, )]; @@ -1934,7 +1938,7 @@ mod tests { let right = make_source_ndv_b(500, 25); // Join on column "b" (index 1) - let on: crate::joins::JoinOn = vec![( + let on: JoinOn = vec![( Arc::new(Column::new("b", 1)) as Arc, Arc::new(Column::new("b", 1)) as Arc, )]; @@ -1988,7 +1992,7 @@ mod tests { let left = make_source_2ndv(1000, 100, 20); let right = make_source_2ndv(500, 50, 10); - let on: crate::joins::JoinOn = vec![ + let on: JoinOn = vec![ ( Arc::new(Column::new("a", 0)) as Arc, Arc::new(Column::new("a", 0)) as Arc, @@ -2063,7 +2067,7 @@ mod tests { right: Arc, join_type: JoinType, ) -> Result> { - let on: crate::joins::JoinOn = vec![( + let on: JoinOn = vec![( Arc::new(Column::new("a", 0)) as Arc, Arc::new(Column::new("a", 0)) as Arc, )]; From 0b6a0213ccd93f7948360036855e04a420831028 Mon Sep 17 00:00:00 2001 From: Alessandro Solimando Date: Mon, 27 Jul 2026 10:49:26 +0200 Subject: [PATCH 09/15] fix: try statistics providers before resolving the operator's child stats --- .../src/operator_statistics/mod.rs | 109 +++++++++++++++++- datafusion/physical-plan/src/statistics.rs | 13 ++- 2 files changed, 113 insertions(+), 9 deletions(-) diff --git a/datafusion/physical-plan/src/operator_statistics/mod.rs b/datafusion/physical-plan/src/operator_statistics/mod.rs index 9d44befff52f7..0913b01049600 100644 --- a/datafusion/physical-plan/src/operator_statistics/mod.rs +++ b/datafusion/physical-plan/src/operator_statistics/mod.rs @@ -418,9 +418,10 @@ impl StatisticsRegistry { /// Provider extensions are preserved; see [`StatisticsContext`] for how they /// propagate up the tree. /// - /// Children are resolved via [`ExecutionPlan::child_stats_requests`] (default - /// `Skip`), so a custom operator whose provider reads `child_stats` must - /// declare `At` for those children to receive their computed statistics. + /// Each provider's `child_stats` come from its own + /// [`StatisticsProvider::child_stats_requests`], so a provider can read child + /// statistics for an existing operator without that operator declaring + /// anything. #[deprecated( since = "55.0.0", note = "use `StatisticsContext::new_with_registry(registry).compute_extended(plan, &StatisticsArgs::new())`" @@ -1087,7 +1088,7 @@ mod tests { }; use arrow::datatypes::{DataType, Field, Schema}; use datafusion_common::stats::Precision; - use datafusion_common::{ColumnStatistics, ScalarValue}; + use datafusion_common::{ColumnStatistics, ScalarValue, internal_err}; use datafusion_execution::TaskContext; use datafusion_expr::Operator; use datafusion_physical_expr::PhysicalExpr; @@ -1387,6 +1388,106 @@ mod tests { Ok(()) } + /// Errors when its statistics are computed, so resolving it is detectable. + #[derive(Debug)] + struct ErroringExec { + input: Arc, + } + + impl DisplayAs for ErroringExec { + fn fmt_as(&self, _t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result { + write!(f, "ErroringExec") + } + } + + impl ExecutionPlan for ErroringExec { + fn name(&self) -> &str { + "ErroringExec" + } + + fn children(&self) -> Vec<&Arc> { + vec![&self.input] + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + Ok(Arc::new(ErroringExec { + input: Arc::clone(&children[0]), + })) + } + + fn properties(&self) -> &Arc { + self.input.properties() + } + + fn statistics_from_inputs( + &self, + _input_stats: &[Arc], + _args: &StatisticsArgs, + ) -> Result> { + internal_err!("child statistics should not be computed") + } + + fn execute( + &self, + _partition: usize, + _context: Arc, + ) -> Result { + unimplemented!() + } + } + + /// Computes a fixed row count for `ProjectionExec`, requesting no children. + #[derive(Debug)] + struct SkipComputeProvider; + + impl StatisticsProvider for SkipComputeProvider { + fn child_stats_requests( + &self, + plan: &dyn ExecutionPlan, + _partition: Option, + ) -> Vec { + plan.children().iter().map(|_| ChildStats::Skip).collect() + } + + fn compute_statistics( + &self, + plan: &dyn ExecutionPlan, + _child_stats: &[ExtendedStatistics], + ) -> Result { + if plan.downcast_ref::().is_some() { + let mut stats = Statistics::new_unknown(plan.schema().as_ref()); + stats.num_rows = Precision::Exact(7); + Ok(StatisticsResult::Computed(ExtendedStatistics::from(stats))) + } else { + Ok(StatisticsResult::Delegate) + } + } + } + + #[test] + fn test_provider_override_skips_operator_child_walk() -> Result<()> { + // ProjectionExec requests At for a child that errors when computed; the + // overriding provider requests Skip, so the child is never resolved and + // its error never surfaces. + let schema = make_schema(); + let child: Arc = Arc::new(ErroringExec { + input: make_source(1000), + }); + let parent: Arc = Arc::new(ProjectionExec::try_new( + vec![(col("a", &schema)?, "a".to_string())], + child, + )?); + let mut registry = StatisticsRegistry::new(); + registry.register(Arc::new(SkipComputeProvider)); + + let stats = compute(®istry, parent.as_ref())?; + assert!(matches!(stats.base.num_rows, Precision::Exact(7))); + Ok(()) + } + #[derive(Debug)] struct OverrideFilterProvider { fixed_selectivity: f64, diff --git a/datafusion/physical-plan/src/statistics.rs b/datafusion/physical-plan/src/statistics.rs index ca16d25a0f95b..10d9bfdf1a939 100644 --- a/datafusion/physical-plan/src/statistics.rs +++ b/datafusion/physical-plan/src/statistics.rs @@ -221,13 +221,16 @@ impl StatisticsContext { } let children = plan.children(); - let requests = plan.child_stats_requests(partition); - // Resolved for the built-in `statistics_from_inputs` fallback below. - let child_statistics = self.resolve_children(plan, &children, &requests)?; - + // Try providers before resolving the operator's own children, so a + // provider that overrides this node is not blocked by the fallback walk. let statistics = match self.try_provider_stats(plan, &children, args)? { Some(statistics) => statistics, - None => plan.statistics_from_inputs(&child_statistics, args)?, + None => { + let requests = plan.child_stats_requests(partition); + let child_statistics = + self.resolve_children(plan, &children, &requests)?; + plan.statistics_from_inputs(&child_statistics, args)? + } }; self.store_statistics(plan, partition, Arc::clone(&statistics)); Ok(statistics) From b2548db8fff63fb08a6da349dcbaa130f3d67655 Mon Sep 17 00:00:00 2001 From: Alessandro Solimando Date: Mon, 27 Jul 2026 16:56:27 +0200 Subject: [PATCH 10/15] fix: drop unused async in join_reorder example --- .../examples/statistics/join_reorder.rs | 13 ++++--------- 1 file changed, 4 insertions(+), 9 deletions(-) diff --git a/datafusion-examples/examples/statistics/join_reorder.rs b/datafusion-examples/examples/statistics/join_reorder.rs index 4a689a89587f1..ed7b41228d55d 100644 --- a/datafusion-examples/examples/statistics/join_reorder.rs +++ b/datafusion-examples/examples/statistics/join_reorder.rs @@ -93,7 +93,7 @@ fn mem_table(fields: &[(&str, Arc)]) -> Result> { Ok(Arc::new(MemTable::try_new(schema, vec![vec![batch]])?)) } -async fn build_ctx(with_registry: bool) -> Result { +fn build_ctx(with_registry: bool) -> Result { let config = SessionConfig::new() .with_target_partitions(4) .set_bool("datafusion.explain.physical_plan_only", true) @@ -157,12 +157,7 @@ pub async fn join_reorder() -> Result<()> { let truth_query = "SELECT count(DISTINCT user_id) AS true_distinct_users \ FROM events WHERE amount < 50"; println!("-- Ground truth --\n{truth_query}\n"); - let truth = build_ctx(false) - .await? - .sql(truth_query) - .await? - .collect() - .await?; + let truth = build_ctx(false)?.sql(truth_query).await?.collect().await?; println!("{}\n", pretty_format_batches(&truth)?); println!("-- Query --\n{QUERY}\n"); @@ -175,8 +170,8 @@ pub async fn join_reorder() -> Result<()> { confirms `events` really is the smaller, cheaper side.\n" ); println!("-- Without the registry (default estimation) --"); - println!("{}\n", explain(&build_ctx(false).await?).await?); + println!("{}\n", explain(&build_ctx(false)?).await?); println!("-- With the registry (built-in refinement) --"); - println!("{}", explain(&build_ctx(true).await?).await?); + println!("{}", explain(&build_ctx(true)?).await?); Ok(()) } From 9a416b847552df3baf1bba8a4932d56c8c869df4 Mon Sep 17 00:00:00 2001 From: Alessandro Solimando Date: Tue, 28 Jul 2026 12:08:23 +0200 Subject: [PATCH 11/15] refactor: hoist provider exec-type imports to module level --- .../src/operator_statistics/mod.rs | 23 ++++++------------- 1 file changed, 7 insertions(+), 16 deletions(-) diff --git a/datafusion/physical-plan/src/operator_statistics/mod.rs b/datafusion/physical-plan/src/operator_statistics/mod.rs index 0913b01049600..0f729cdc953c1 100644 --- a/datafusion/physical-plan/src/operator_statistics/mod.rs +++ b/datafusion/physical-plan/src/operator_statistics/mod.rs @@ -84,7 +84,14 @@ use datafusion_common::stats::Precision; use datafusion_common::{Result, Statistics}; use crate::ExecutionPlan; +use crate::aggregates::{AggregateExec, AggregateMode}; +use crate::execution_plan::CardinalityEffect; +use crate::filter::FilterExec; +use crate::joins::{CrossJoinExec, HashJoinExec, JoinOnRef, SortMergeJoinExec}; +use crate::limit::{GlobalLimitExec, LocalLimitExec}; +use crate::projection::ProjectionExec; use crate::statistics::{ChildStats, StatisticsArgs, StatisticsContext}; +use crate::union::UnionExec; // ============================================================================ // ExtendedStatistics: Statistics with type-safe extensions @@ -586,8 +593,6 @@ impl StatisticsProvider for FilterStatisticsProvider { plan: &dyn ExecutionPlan, child_stats: &[ExtendedStatistics], ) -> Result { - use crate::filter::FilterExec; - let Some(filter) = plan.downcast_ref::() else { return Ok(StatisticsResult::Delegate); }; @@ -642,8 +647,6 @@ impl StatisticsProvider for ProjectionStatisticsProvider { plan: &dyn ExecutionPlan, child_stats: &[ExtendedStatistics], ) -> Result { - use crate::projection::ProjectionExec; - let Some(proj) = plan.downcast_ref::() else { return Ok(StatisticsResult::Delegate); }; @@ -678,8 +681,6 @@ impl StatisticsProvider for PassthroughStatisticsProvider { plan: &dyn ExecutionPlan, child_stats: &[ExtendedStatistics], ) -> Result { - use crate::execution_plan::CardinalityEffect; - if child_stats.len() != 1 || !matches!(plan.cardinality_effect(), CardinalityEffect::Equal) { @@ -726,11 +727,8 @@ impl StatisticsProvider for AggregateStatisticsProvider { plan: &dyn ExecutionPlan, child_stats: &[ExtendedStatistics], ) -> Result { - use crate::aggregates::AggregateExec; use datafusion_physical_expr::expressions::Column; - use crate::aggregates::AggregateMode; - let Some(agg) = plan.downcast_ref::() else { return Ok(StatisticsResult::Delegate); }; @@ -818,7 +816,6 @@ impl StatisticsProvider for JoinStatisticsProvider { plan: &dyn ExecutionPlan, child_stats: &[ExtendedStatistics], ) -> Result { - use crate::joins::{CrossJoinExec, HashJoinExec, SortMergeJoinExec}; use datafusion_common::JoinType; use datafusion_physical_expr::expressions::Column; @@ -835,8 +832,6 @@ impl StatisticsProvider for JoinStatisticsProvider { return Ok(StatisticsResult::Delegate); }; - use crate::joins::JoinOnRef; - /// Estimate equi-join output using NDV of join key columns: /// left_rows * right_rows / product(max(left_ndv_i, right_ndv_i)) /// Falls back to Cartesian product if any key lacks NDV on both sides. @@ -940,8 +935,6 @@ impl StatisticsProvider for LimitStatisticsProvider { plan: &dyn ExecutionPlan, child_stats: &[ExtendedStatistics], ) -> Result { - use crate::limit::{GlobalLimitExec, LocalLimitExec}; - if child_stats.is_empty() { return Ok(StatisticsResult::Delegate); } @@ -988,8 +981,6 @@ impl StatisticsProvider for UnionStatisticsProvider { plan: &dyn ExecutionPlan, child_stats: &[ExtendedStatistics], ) -> Result { - use crate::union::UnionExec; - if plan.downcast_ref::().is_none() { return Ok(StatisticsResult::Delegate); } From 83b5e8f2c4bcfea0a205993c1a292a78cb7ada5f Mon Sep 17 00:00:00 2001 From: Alessandro Solimando Date: Tue, 28 Jul 2026 12:13:10 +0200 Subject: [PATCH 12/15] feat: gate statistics providers with a matches() hook --- .../examples/statistics/join_reorder.rs | 30 ++-- .../src/operator_statistics/mod.rs | 148 ++++++++++++++---- datafusion/physical-plan/src/statistics.rs | 3 + 3 files changed, 140 insertions(+), 41 deletions(-) diff --git a/datafusion-examples/examples/statistics/join_reorder.rs b/datafusion-examples/examples/statistics/join_reorder.rs index ed7b41228d55d..0eccc93630af2 100644 --- a/datafusion-examples/examples/statistics/join_reorder.rs +++ b/datafusion-examples/examples/statistics/join_reorder.rs @@ -50,24 +50,23 @@ use datafusion::physical_plan::statistics::StatisticsArgs; use datafusion::prelude::*; use rand::{Rng, SeedableRng, rngs::StdRng}; -/// Supplies the catalog-known statistics of the `events` scan that the in-memory -/// table does not carry: the `amount` range (for filter selectivity) and the -/// `user_id` distinct count (for post-filter NDV estimation). Only the base -/// `events` scan (a leaf carrying both columns) is matched; everything else -/// delegates to the rest of the chain. +/// Matches the base `events` scan: a leaf carrying both `user_id` and `amount`. +fn catalog_matches(plan: &dyn ExecutionPlan) -> bool { + let schema = plan.schema(); + plan.children().is_empty() + && schema.index_of("user_id").is_ok() + && schema.index_of("amount").is_ok() +} + +/// Injects the catalog-known `amount` range (for filter selectivity) and +/// `user_id` distinct count (for post-filter NDV) the in-memory table lacks. fn catalog_stats( plan: &dyn ExecutionPlan, _child_stats: &[ExtendedStatistics], ) -> Result { let schema = plan.schema(); - let (Ok(user_id), Ok(amount)) = - (schema.index_of("user_id"), schema.index_of("amount")) - else { - return Ok(StatisticsResult::Delegate); - }; - if !plan.children().is_empty() { - return Ok(StatisticsResult::Delegate); - } + let user_id = schema.index_of("user_id")?; + let amount = schema.index_of("amount")?; let mut stats = (*plan.statistics_from_inputs(&[], &StatisticsArgs::new())?).clone(); stats.column_statistics[amount].min_value = Precision::Inexact(ScalarValue::Int32(Some(0))); @@ -113,7 +112,10 @@ fn build_ctx(with_registry: bool) -> Result { .with_default_features(); if with_registry { let mut registry = StatisticsRegistry::default_with_builtin_providers(); - registry.register(Arc::new(ClosureStatisticsProvider::new(catalog_stats))); + registry.register(Arc::new(ClosureStatisticsProvider::with_matches( + catalog_matches, + catalog_stats, + ))); builder = builder.with_statistics_registry(registry); } let ctx = SessionContext::new_with_state(builder.build()); diff --git a/datafusion/physical-plan/src/operator_statistics/mod.rs b/datafusion/physical-plan/src/operator_statistics/mod.rs index 0f729cdc953c1..c16a4a23703db 100644 --- a/datafusion/physical-plan/src/operator_statistics/mod.rs +++ b/datafusion/physical-plan/src/operator_statistics/mod.rs @@ -233,18 +233,17 @@ pub enum StatisticsResult { /// struct MyStatisticsProvider; /// /// impl StatisticsProvider for MyStatisticsProvider { +/// fn matches(&self, plan: &dyn ExecutionPlan) -> bool { +/// plan.downcast_ref::().is_some() +/// } +/// /// fn compute_statistics( /// &self, /// plan: &dyn ExecutionPlan, /// child_stats: &[ExtendedStatistics], /// ) -> Result { -/// if let Some(my_exec) = plan.downcast_ref::() { -/// // Custom logic for MyCustomExec -/// Ok(StatisticsResult::Computed(/* ... */)) -/// } else { -/// // Let next provider handle it -/// Ok(StatisticsResult::Delegate) -/// } +/// // matches() guaranteed this node is a MyCustomExec +/// Ok(StatisticsResult::Computed(/* ... */)) /// } /// } /// ``` @@ -314,6 +313,13 @@ pub trait StatisticsProvider: Debug + Send + Sync { .map(|_| ChildStats::At(None)) .collect() } + + /// Whether this provider handles `plan`. Checked before its + /// [`Self::child_stats_requests`] are resolved, so a non-matching provider + /// never triggers child statistics computation. + fn matches(&self, _plan: &dyn ExecutionPlan) -> bool { + true + } } /// Deprecated statistics provider that delegates to each operator's built-in @@ -588,6 +594,10 @@ fn computed_with_row_count( pub struct FilterStatisticsProvider; impl StatisticsProvider for FilterStatisticsProvider { + fn matches(&self, plan: &dyn ExecutionPlan) -> bool { + plan.downcast_ref::().is_some() + } + fn compute_statistics( &self, plan: &dyn ExecutionPlan, @@ -642,6 +652,10 @@ impl StatisticsProvider for FilterStatisticsProvider { pub struct ProjectionStatisticsProvider; impl StatisticsProvider for ProjectionStatisticsProvider { + fn matches(&self, plan: &dyn ExecutionPlan) -> bool { + plan.downcast_ref::().is_some() + } + fn compute_statistics( &self, plan: &dyn ExecutionPlan, @@ -676,6 +690,11 @@ impl StatisticsProvider for ProjectionStatisticsProvider { pub struct PassthroughStatisticsProvider; impl StatisticsProvider for PassthroughStatisticsProvider { + fn matches(&self, plan: &dyn ExecutionPlan) -> bool { + plan.children().len() == 1 + && matches!(plan.cardinality_effect(), CardinalityEffect::Equal) + } + fn compute_statistics( &self, plan: &dyn ExecutionPlan, @@ -722,6 +741,10 @@ impl StatisticsProvider for PassthroughStatisticsProvider { pub struct AggregateStatisticsProvider; impl StatisticsProvider for AggregateStatisticsProvider { + fn matches(&self, plan: &dyn ExecutionPlan) -> bool { + plan.downcast_ref::().is_some() + } + fn compute_statistics( &self, plan: &dyn ExecutionPlan, @@ -811,6 +834,12 @@ impl StatisticsProvider for AggregateStatisticsProvider { pub struct JoinStatisticsProvider; impl StatisticsProvider for JoinStatisticsProvider { + fn matches(&self, plan: &dyn ExecutionPlan) -> bool { + plan.downcast_ref::().is_some() + || plan.downcast_ref::().is_some() + || plan.downcast_ref::().is_some() + } + fn compute_statistics( &self, plan: &dyn ExecutionPlan, @@ -930,6 +959,11 @@ impl StatisticsProvider for JoinStatisticsProvider { pub struct LimitStatisticsProvider; impl StatisticsProvider for LimitStatisticsProvider { + fn matches(&self, plan: &dyn ExecutionPlan) -> bool { + plan.downcast_ref::().is_some() + || plan.downcast_ref::().is_some() + } + fn compute_statistics( &self, plan: &dyn ExecutionPlan, @@ -976,6 +1010,10 @@ impl StatisticsProvider for LimitStatisticsProvider { pub struct UnionStatisticsProvider; impl StatisticsProvider for UnionStatisticsProvider { + fn matches(&self, plan: &dyn ExecutionPlan) -> bool { + plan.downcast_ref::().is_some() + } + fn compute_statistics( &self, plan: &dyn ExecutionPlan, @@ -1010,6 +1048,8 @@ type ProviderFn = dyn Fn(&dyn ExecutionPlan, &[ExtendedStatistics]) -> Result bool + Send + Sync; + /// A [`StatisticsProvider`] backed by a user-supplied closure. /// /// Useful for injecting custom statistics in tests or for cardinality feedback @@ -1024,18 +1064,18 @@ type ProviderFn = dyn Fn(&dyn ExecutionPlan, &[ExtendedStatistics]) -> Result().is_some() { +/// let provider = ClosureStatisticsProvider::with_matches( +/// |plan| plan.downcast_ref::().is_some(), +/// |plan, child_stats| { /// Ok(StatisticsResult::Computed(ExtendedStatistics::from(Statistics { /// num_rows: Precision::Inexact(42), /// ..Statistics::new_unknown(plan.schema().as_ref()) /// }))) -/// } else { -/// Ok(StatisticsResult::Delegate) -/// } -/// }); +/// }, +/// ); /// ``` pub struct ClosureStatisticsProvider { + matches_fn: Box, f: Box, } @@ -1047,7 +1087,24 @@ impl ClosureStatisticsProvider { + Sync + 'static, ) -> Self { - Self { f: Box::new(f) } + Self { + matches_fn: Box::new(|_| true), + f: Box::new(f), + } + } + + /// Like [`Self::new`] but applies only where `matches` returns true. + pub fn with_matches( + matches: impl Fn(&dyn ExecutionPlan) -> bool + Send + Sync + 'static, + f: impl Fn(&dyn ExecutionPlan, &[ExtendedStatistics]) -> Result + + Send + + Sync + + 'static, + ) -> Self { + Self { + matches_fn: Box::new(matches), + f: Box::new(f), + } } } @@ -1058,6 +1115,10 @@ impl Debug for ClosureStatisticsProvider { } impl StatisticsProvider for ClosureStatisticsProvider { + fn matches(&self, plan: &dyn ExecutionPlan) -> bool { + (self.matches_fn)(plan) + } + fn compute_statistics( &self, plan: &dyn ExecutionPlan, @@ -1435,6 +1496,10 @@ mod tests { struct SkipComputeProvider; impl StatisticsProvider for SkipComputeProvider { + fn matches(&self, plan: &dyn ExecutionPlan) -> bool { + plan.downcast_ref::().is_some() + } + fn child_stats_requests( &self, plan: &dyn ExecutionPlan, @@ -1448,29 +1513,43 @@ mod tests { plan: &dyn ExecutionPlan, _child_stats: &[ExtendedStatistics], ) -> Result { - if plan.downcast_ref::().is_some() { - let mut stats = Statistics::new_unknown(plan.schema().as_ref()); - stats.num_rows = Precision::Exact(7); - Ok(StatisticsResult::Computed(ExtendedStatistics::from(stats))) - } else { - Ok(StatisticsResult::Delegate) + if plan.downcast_ref::().is_none() { + return Ok(StatisticsResult::Delegate); } + let mut stats = Statistics::new_unknown(plan.schema().as_ref()); + stats.num_rows = Precision::Exact(7); + Ok(StatisticsResult::Computed(ExtendedStatistics::from(stats))) } } - #[test] - fn test_provider_override_skips_operator_child_walk() -> Result<()> { - // ProjectionExec requests At for a child that errors when computed; the - // overriding provider requests Skip, so the child is never resolved and - // its error never surfaces. + /// Never matches, so it must not trigger any child statistics computation. + #[derive(Debug)] + struct NonMatchingProvider; + + impl StatisticsProvider for NonMatchingProvider { + fn matches(&self, _plan: &dyn ExecutionPlan) -> bool { + false + } + } + + /// A `ProjectionExec` (which requests `At`) over a child that errors when its + /// statistics are computed. + fn projection_over_erroring_child() -> Result> { let schema = make_schema(); let child: Arc = Arc::new(ErroringExec { input: make_source(1000), }); - let parent: Arc = Arc::new(ProjectionExec::try_new( + Ok(Arc::new(ProjectionExec::try_new( vec![(col("a", &schema)?, "a".to_string())], child, - )?); + )?)) + } + + #[test] + fn test_provider_override_skips_operator_child_walk() -> Result<()> { + // A matching provider requests Skip, so the erroring child is never + // resolved and its error never surfaces. + let parent = projection_over_erroring_child()?; let mut registry = StatisticsRegistry::new(); registry.register(Arc::new(SkipComputeProvider)); @@ -1479,6 +1558,21 @@ mod tests { Ok(()) } + #[test] + fn test_non_matching_provider_does_not_trigger_child_walk() -> Result<()> { + // A non-matching provider precedes the handling one; it must not resolve + // the erroring child, so the later provider still succeeds. + let parent = projection_over_erroring_child()?; + let registry = StatisticsRegistry::with_providers(vec![ + Arc::new(NonMatchingProvider), + Arc::new(SkipComputeProvider), + ]); + + let stats = compute(®istry, parent.as_ref())?; + assert!(matches!(stats.base.num_rows, Precision::Exact(7))); + Ok(()) + } + #[derive(Debug)] struct OverrideFilterProvider { fixed_selectivity: f64, diff --git a/datafusion/physical-plan/src/statistics.rs b/datafusion/physical-plan/src/statistics.rs index 10d9bfdf1a939..68b37db824d5f 100644 --- a/datafusion/physical-plan/src/statistics.rs +++ b/datafusion/physical-plan/src/statistics.rs @@ -288,6 +288,9 @@ impl StatisticsContext { } let partition = args.partition(); for provider in providers { + if !provider.matches(plan) { + continue; + } let requests = provider.child_stats_requests(plan, partition); let child_statistics = self.resolve_children(plan, children, &requests)?; let child_extended = From 4068c7b2080ef70a175e06af7026af6f991e0112 Mon Sep 17 00:00:00 2001 From: Alessandro Solimando Date: Tue, 28 Jul 2026 17:57:51 +0200 Subject: [PATCH 13/15] fix: drop redundant explicit rustdoc link targets --- .../physical-plan/src/operator_statistics/mod.rs | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/datafusion/physical-plan/src/operator_statistics/mod.rs b/datafusion/physical-plan/src/operator_statistics/mod.rs index c16a4a23703db..cbfe72e27860d 100644 --- a/datafusion/physical-plan/src/operator_statistics/mod.rs +++ b/datafusion/physical-plan/src/operator_statistics/mod.rs @@ -581,7 +581,7 @@ fn computed_with_row_count( Ok(StatisticsResult::Computed(ExtendedStatistics::new(base))) } -/// Statistics provider for [`FilterExec`](crate::filter::FilterExec) that uses +/// Statistics provider for [`FilterExec`] that uses /// pre-computed enhanced child statistics from the registry walk. /// /// Unlike the built-in fallback (which calls `statistics_from_inputs` and gets raw @@ -642,7 +642,7 @@ impl StatisticsProvider for FilterStatisticsProvider { } } -/// Statistics provider for [`ProjectionExec`](crate::projection::ProjectionExec) +/// Statistics provider for [`ProjectionExec`] /// that uses pre-computed enhanced child statistics from the registry walk. /// /// Maps enhanced child column statistics to output columns based on the @@ -680,7 +680,7 @@ impl StatisticsProvider for ProjectionStatisticsProvider { } /// Statistics provider for single-input operators with -/// [`CardinalityEffect::Equal`](crate::execution_plan::CardinalityEffect::Equal). +/// [`CardinalityEffect::Equal`]. /// /// These operators (Sort, Repartition, CoalescePartitions, etc.) don't /// transform statistics, so we pass through the enhanced child stats directly. @@ -719,7 +719,7 @@ impl StatisticsProvider for PassthroughStatisticsProvider { } } -/// Statistics provider for [`AggregateExec`](crate::aggregates::AggregateExec) +/// Statistics provider for [`AggregateExec`] /// that estimates output cardinality from the NDV of GROUP BY columns. /// /// For each GROUP BY column, looks up `distinct_count` from the enhanced @@ -950,8 +950,8 @@ impl StatisticsProvider for JoinStatisticsProvider { } } -/// Statistics provider for [`LocalLimitExec`](crate::limit::LocalLimitExec) and -/// [`GlobalLimitExec`](crate::limit::GlobalLimitExec). +/// Statistics provider for [`LocalLimitExec`] and +/// [`GlobalLimitExec`]. /// /// Caps output row count at the limit value, accounting for any leading skip offset /// in `GlobalLimitExec`. @@ -1003,7 +1003,7 @@ impl StatisticsProvider for LimitStatisticsProvider { } } -/// Statistics provider for [`UnionExec`](crate::union::UnionExec). +/// Statistics provider for [`UnionExec`]. /// /// Sums row counts across all inputs. #[derive(Debug, Default)] From f9355ca676aff4ab6d1dde4e1dced87f19591dd4 Mon Sep 17 00:00:00 2001 From: Alessandro Solimando Date: Thu, 30 Jul 2026 11:42:56 +0200 Subject: [PATCH 14/15] fix: skip statistics providers whose child walk fails --- .../src/operator_statistics/mod.rs | 64 +++++++++++++++++-- datafusion/physical-plan/src/statistics.rs | 35 ++++++++-- 2 files changed, 86 insertions(+), 13 deletions(-) diff --git a/datafusion/physical-plan/src/operator_statistics/mod.rs b/datafusion/physical-plan/src/operator_statistics/mod.rs index cbfe72e27860d..66303f6803959 100644 --- a/datafusion/physical-plan/src/operator_statistics/mod.rs +++ b/datafusion/physical-plan/src/operator_statistics/mod.rs @@ -1522,14 +1522,42 @@ mod tests { } } - /// Never matches, so it must not trigger any child statistics computation. + /// Never matches, but if consulted it resolves the child and returns its row + /// count, leaking an observable (non-error) value. #[derive(Debug)] - struct NonMatchingProvider; + struct NonMatchingLeakProvider; - impl StatisticsProvider for NonMatchingProvider { + impl StatisticsProvider for NonMatchingLeakProvider { fn matches(&self, _plan: &dyn ExecutionPlan) -> bool { false } + + fn compute_statistics( + &self, + plan: &dyn ExecutionPlan, + child_stats: &[ExtendedStatistics], + ) -> Result { + let mut stats = Statistics::new_unknown(plan.schema().as_ref()); + if let Some(child) = child_stats.first() { + stats.num_rows = child.base.num_rows; + } + Ok(StatisticsResult::Computed(ExtendedStatistics::from(stats))) + } + } + + /// Default `matches()` (always true) and `child_stats_requests` (`At(None)`), + /// but always delegates. + #[derive(Debug)] + struct DelegatingProvider; + + impl StatisticsProvider for DelegatingProvider { + fn compute_statistics( + &self, + _plan: &dyn ExecutionPlan, + _child_stats: &[ExtendedStatistics], + ) -> Result { + Ok(StatisticsResult::Delegate) + } } /// A `ProjectionExec` (which requests `At`) over a child that errors when its @@ -1559,12 +1587,34 @@ mod tests { } #[test] - fn test_non_matching_provider_does_not_trigger_child_walk() -> Result<()> { - // A non-matching provider precedes the handling one; it must not resolve - // the erroring child, so the later provider still succeeds. + fn test_matches_gate_skips_non_matching_provider() -> Result<()> { + // A non-matching provider precedes the handling one. The gate must skip it + // before it is consulted, otherwise it leaks the child row count (999) + // instead of the later provider's 7. + let schema = make_schema(); + let parent: Arc = Arc::new(ProjectionExec::try_new( + vec![(col("a", &schema)?, "a".to_string())], + make_source(999), + )?); + let registry = StatisticsRegistry::with_providers(vec![ + Arc::new(NonMatchingLeakProvider), + Arc::new(SkipComputeProvider), + ]); + + let stats = compute(®istry, parent.as_ref())?; + assert!(matches!(stats.base.num_rows, Precision::Exact(7))); + Ok(()) + } + + #[test] + fn test_delegating_provider_child_error_does_not_block_later_provider() -> Result<()> + { + // A delegating provider that keeps the default matcher resolves the + // erroring child speculatively; that failure must not stop the later + // child-skipping provider from succeeding. let parent = projection_over_erroring_child()?; let registry = StatisticsRegistry::with_providers(vec![ - Arc::new(NonMatchingProvider), + Arc::new(DelegatingProvider), Arc::new(SkipComputeProvider), ]); diff --git a/datafusion/physical-plan/src/statistics.rs b/datafusion/physical-plan/src/statistics.rs index 68b37db824d5f..0f50372d4bd6e 100644 --- a/datafusion/physical-plan/src/statistics.rs +++ b/datafusion/physical-plan/src/statistics.rs @@ -21,6 +21,7 @@ //! [`ExecutionPlan::statistics_from_inputs`]. use crate::ExecutionPlan; +use crate::displayable; use crate::operator_statistics::{ ExtendedStatistics, StatisticsRegistry, StatisticsResult, }; @@ -28,6 +29,7 @@ use datafusion_common::extensions::Extensions; use datafusion_common::{ Result, Statistics, assert_eq_or_internal_err, assert_or_internal_err, }; +use log::debug; use std::cell::RefCell; use std::collections::HashMap; use std::rc::Rc; @@ -256,11 +258,17 @@ impl StatisticsContext { children .iter() .zip(requests) - .map(|(child, directive)| match directive { - ChildStats::At(p) => self.compute_base( - child.as_ref(), - &StatisticsArgs::new().with_partition(*p), - ), + .enumerate() + .map(|(i, (child, directive))| match directive { + ChildStats::At(p) => self + .compute_base(child.as_ref(), &StatisticsArgs::new().with_partition(*p)) + .map_err(|e| { + e.context(format!( + "computing statistics for child {i} ({}) of {} at partition {p:?}", + child.name(), + plan.name() + )) + }), ChildStats::Skip => { Ok(Arc::new(Statistics::new_unknown(child.schema().as_ref()))) } @@ -292,7 +300,22 @@ impl StatisticsContext { continue; } let requests = provider.child_stats_requests(plan, partition); - let child_statistics = self.resolve_children(plan, children, &requests)?; + // A provider's child walk is speculative: on failure, skip the provider + // so a later one or the operator fallback can handle the node. Not + // error-swallowing, whoever genuinely needs the child resolves it again + // and the error resurfaces there; a matched provider's own `compute` + // error below stays fatal. + let child_statistics = match self.resolve_children(plan, children, &requests) + { + Ok(child_statistics) => child_statistics, + Err(e) => { + debug!( + "Statistics provider {provider:?} skipped for {}: child statistics resolution failed: {e}", + displayable(plan).one_line().to_string().trim_end() + ); + continue; + } + }; let child_extended = self.child_extended_stats(children, &requests, &child_statistics); if let StatisticsResult::Computed(computed) = From 08845ef8a34b4be5845c40e70a58b7743bb628fd Mon Sep 17 00:00:00 2001 From: Alessandro Solimando Date: Thu, 30 Jul 2026 19:47:39 +0200 Subject: [PATCH 15/15] fix: fail fast on invalid statistics provider child requests --- .../src/operator_statistics/mod.rs | 70 +++++++++++++++++-- datafusion/physical-plan/src/statistics.rs | 35 ++++++++-- 2 files changed, 95 insertions(+), 10 deletions(-) diff --git a/datafusion/physical-plan/src/operator_statistics/mod.rs b/datafusion/physical-plan/src/operator_statistics/mod.rs index 66303f6803959..2c9086280704e 100644 --- a/datafusion/physical-plan/src/operator_statistics/mod.rs +++ b/datafusion/physical-plan/src/operator_statistics/mod.rs @@ -1573,6 +1573,15 @@ mod tests { )?)) } + /// A `ProjectionExec` (which requests `At`) over a source with `num_rows` rows. + fn projection_over_source(num_rows: usize) -> Result> { + let schema = make_schema(); + Ok(Arc::new(ProjectionExec::try_new( + vec![(col("a", &schema)?, "a".to_string())], + make_source(num_rows), + )?)) + } + #[test] fn test_provider_override_skips_operator_child_walk() -> Result<()> { // A matching provider requests Skip, so the erroring child is never @@ -1591,11 +1600,7 @@ mod tests { // A non-matching provider precedes the handling one. The gate must skip it // before it is consulted, otherwise it leaks the child row count (999) // instead of the later provider's 7. - let schema = make_schema(); - let parent: Arc = Arc::new(ProjectionExec::try_new( - vec![(col("a", &schema)?, "a".to_string())], - make_source(999), - )?); + let parent = projection_over_source(999)?; let registry = StatisticsRegistry::with_providers(vec![ Arc::new(NonMatchingLeakProvider), Arc::new(SkipComputeProvider), @@ -1623,6 +1628,61 @@ mod tests { Ok(()) } + /// Returns the wrong number of child stat requests. + #[derive(Debug)] + struct MalformedRequestProvider; + + impl StatisticsProvider for MalformedRequestProvider { + fn child_stats_requests( + &self, + _plan: &dyn ExecutionPlan, + _partition: Option, + ) -> Vec { + Vec::new() + } + } + + /// Requests a non-existent partition of its child. + #[derive(Debug)] + struct InvalidPartitionProvider; + + impl StatisticsProvider for InvalidPartitionProvider { + fn child_stats_requests( + &self, + plan: &dyn ExecutionPlan, + _partition: Option, + ) -> Vec { + plan.children() + .iter() + .map(|_| ChildStats::At(Some(usize::MAX))) + .collect() + } + } + + #[test] + fn test_malformed_child_stats_requests_fails_fast() -> Result<()> { + // A wrong-length request set is a provider bug: it must error, not be + // skipped in favor of the later provider. + let parent = projection_over_source(10)?; + let registry = StatisticsRegistry::with_providers(vec![ + Arc::new(MalformedRequestProvider), + Arc::new(SkipComputeProvider), + ]); + assert!(compute(®istry, parent.as_ref()).is_err()); + Ok(()) + } + + #[test] + fn test_invalid_partition_request_fails_fast() -> Result<()> { + let parent = projection_over_source(10)?; + let registry = StatisticsRegistry::with_providers(vec![ + Arc::new(InvalidPartitionProvider), + Arc::new(SkipComputeProvider), + ]); + assert!(compute(®istry, parent.as_ref()).is_err()); + Ok(()) + } + #[derive(Debug)] struct OverrideFilterProvider { fixed_selectivity: f64, diff --git a/datafusion/physical-plan/src/statistics.rs b/datafusion/physical-plan/src/statistics.rs index 0f50372d4bd6e..efd8af0d726b0 100644 --- a/datafusion/physical-plan/src/statistics.rs +++ b/datafusion/physical-plan/src/statistics.rs @@ -229,6 +229,7 @@ impl StatisticsContext { Some(statistics) => statistics, None => { let requests = plan.child_stats_requests(partition); + self.validate_child_requests(plan, &children, &requests)?; let child_statistics = self.resolve_children(plan, &children, &requests)?; plan.statistics_from_inputs(&child_statistics, args)? @@ -238,15 +239,14 @@ impl StatisticsContext { Ok(statistics) } - /// Resolves each child's core statistics per `requests`: computes the child - /// at the requested partition (memoized), or supplies a - /// [`Statistics::new_unknown`] placeholder for [`ChildStats::Skip`]. - fn resolve_children( + /// Validates child stat `requests` against `plan`'s children: the count must + /// match, and each `At(Some(idx))` must be a valid partition of that child. + fn validate_child_requests( &self, plan: &dyn ExecutionPlan, children: &[&Arc], requests: &[ChildStats], - ) -> Result>> { + ) -> Result<()> { assert_eq_or_internal_err!( requests.len(), children.len(), @@ -255,6 +255,30 @@ impl StatisticsContext { requests.len(), children.len() ); + for (child, directive) in children.iter().zip(requests) { + if let ChildStats::At(Some(idx)) = directive { + let count = child.properties().partitioning.partition_count(); + assert_or_internal_err!( + *idx < count, + "{} requested invalid partition {idx} for child {} with {count} partitions", + plan.name(), + child.name() + ); + } + } + Ok(()) + } + + /// Resolves each child's core statistics per `requests`: computes the child + /// at the requested partition (memoized), or supplies a + /// [`Statistics::new_unknown`] placeholder for [`ChildStats::Skip`]. Callers + /// must validate `requests` via [`Self::validate_child_requests`] first. + fn resolve_children( + &self, + plan: &dyn ExecutionPlan, + children: &[&Arc], + requests: &[ChildStats], + ) -> Result>> { children .iter() .zip(requests) @@ -300,6 +324,7 @@ impl StatisticsContext { continue; } let requests = provider.child_stats_requests(plan, partition); + self.validate_child_requests(plan, children, &requests)?; // A provider's child walk is speculative: on failure, skip the provider // so a later one or the operator fallback can handle the node. Not // error-swallowing, whoever genuinely needs the child resolves it again