diff --git a/datafusion/core/tests/physical_optimizer/enforce_distribution.rs b/datafusion/core/tests/physical_optimizer/enforce_distribution.rs index 6b8e38d150c23..7c9d4984a6a68 100644 --- a/datafusion/core/tests/physical_optimizer/enforce_distribution.rs +++ b/datafusion/core/tests/physical_optimizer/enforce_distribution.rs @@ -18,6 +18,7 @@ use std::fmt::Debug; use std::ops::Deref; use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; use crate::physical_optimizer::test_utils::{ RequirementsTestExec, bounded_window_exec_with_can_repartition, check_integrity, @@ -56,6 +57,7 @@ use datafusion_physical_expr_common::physical_expr::PhysicalExpr; use datafusion_physical_expr_common::sort_expr::{ LexOrdering, OrderingRequirements, PhysicalSortExpr, }; +use datafusion_physical_optimizer::PhysicalOptimizerContext; use datafusion_physical_optimizer::PhysicalOptimizerRule; use datafusion_physical_optimizer::enforce_distribution::*; use datafusion_physical_optimizer::ensure_requirements::EnsureRequirements; @@ -65,6 +67,9 @@ use datafusion_physical_optimizer::sanity_checker::SanityCheckPlan; use datafusion_physical_plan::aggregates::{ AggregateExec, AggregateMode, PhysicalGroupBy, }; +use datafusion_physical_plan::operator_statistics::{ + ClosureStatisticsProvider, StatisticsRegistry, StatisticsResult, +}; use datafusion_physical_expr::{ Distribution, Partitioning, RangePartitioning, SplitPoint, @@ -643,7 +648,12 @@ fn ensure_distribution_helper( config.optimizer.repartition_file_scans = false; config.optimizer.repartition_file_min_size = 1024; config.optimizer.prefer_existing_sort = prefer_existing_sort; - ensure_distribution(distribution_context, &config).map(|item| item.data.plan) + ensure_distribution_with_stats( + distribution_context, + &config, + &datafusion_physical_plan::statistics::StatisticsContext::new(), + ) + .map(|item| item.data.plan) } fn test_suite_default_config_options() -> ConfigOptions { @@ -766,7 +776,11 @@ impl TestConfig { // Then run ensure_distribution rule DistributionContext::new_default(adjusted) .transform_up(|distribution_context| { - ensure_distribution(distribution_context, &self.config) + ensure_distribution_with_stats( + distribution_context, + &self.config, + &datafusion_physical_plan::statistics::StatisticsContext::new(), + ) }) .data() .and_then(check_integrity)?; @@ -5065,7 +5079,13 @@ async fn assert_reoptimized_fetch_values( if iteration > 0 { let distribution = DistributionContext::new_default(Arc::clone(&optimized)) - .transform_up(|context| ensure_distribution(context, &config))? + .transform_up(|context| { + ensure_distribution_with_stats( + context, + &config, + &datafusion_physical_plan::statistics::StatisticsContext::new(), + ) + })? .data; check_integrity(distribution)?; optimized = EnsureRequirements::new().optimize(optimized, &config)?; @@ -5140,7 +5160,13 @@ async fn check_fetch_below_filter( for iteration in 0..3 { if iteration > 0 { let distribution = DistributionContext::new_default(Arc::clone(&plan)) - .transform_up(|context| ensure_distribution(context, &config))? + .transform_up(|context| { + ensure_distribution_with_stats( + context, + &config, + &datafusion_physical_plan::statistics::StatisticsContext::new(), + ) + })? .data; check_integrity(distribution)?; plan = EnsureRequirements::new().optimize(plan, &config)?; @@ -5312,3 +5338,253 @@ fn ensure_distribution_reuses_plan_arc_when_no_redistribution_needed() -> Result ); Ok(()) } + +/// Single-child pass-through whose `statistics_from_inputs` increments a counter +/// every time it is actually computed (i.e. on a statistics-cache miss). Used to +/// observe how often `ensure_distribution` recomputes a node's statistics. +#[derive(Debug)] +struct CountingStatsExec { + input: Arc, + cache: Arc, + calls: Arc, +} + +impl CountingStatsExec { + fn new(input: Arc, calls: Arc) -> Self { + let cache = PlanProperties::new( + input.equivalence_properties().clone(), + input.output_partitioning().clone(), + input.pipeline_behavior(), + input.boundedness(), + ); + Self { + input, + cache: Arc::new(cache), + calls, + } + } +} + +impl DisplayAs for CountingStatsExec { + fn fmt_as( + &self, + _t: DisplayFormatType, + f: &mut std::fmt::Formatter, + ) -> std::fmt::Result { + write!(f, "CountingStatsExec") + } +} + +impl ExecutionPlan for CountingStatsExec { + fn name(&self) -> &'static str { + "CountingStatsExec" + } + + fn properties(&self) -> &Arc { + &self.cache + } + + fn children(&self) -> Vec<&Arc> { + vec![&self.input] + } + + fn replace_children( + self: Arc, + mut children: Vec>, + _: ReplaceChildrenOptions, + ) -> Result> { + assert_eq!(children.len(), 1); + Ok(Arc::new(Self::new( + children.pop().unwrap(), + Arc::clone(&self.calls), + ))) + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + Ok(TreeNodeRecursion::Continue) + } + + fn execute( + &self, + _partition: usize, + _context: Arc, + ) -> Result { + unreachable!(); + } + + fn statistics_from_inputs( + &self, + _input_stats: &[Arc], + _args: &datafusion_physical_plan::statistics::StatisticsArgs, + ) -> Result> { + self.calls.fetch_add(1, Ordering::Relaxed); + Ok(Arc::new(Statistics::new_unknown( + self.input.schema().as_ref(), + ))) + } +} + +/// Regression test for the shared statistics cache in `ensure_distribution`. +/// +/// A deep stack of pass-through operators sits over a counting leaf. Each +/// ancestor's distribution enforcement inspects its child's statistics, which +/// recurse to the leaf. With one `StatisticsContext` shared across the pass the +/// leaf is computed once; with a fresh context per node it is recomputed once +/// per ancestor. This directly detects a regression where the cache is not +/// actually shared (e.g. reset on every node), which no plan-output assertion +/// can catch because the optimized plan is identical either way. +#[test] +fn ensure_distribution_shares_statistics_cache() -> Result<()> { + // Count how many times a leaf's statistics are computed over a stack of + // `depth` pass-through operators sitting on top of it. Each ancestor's + // distribution enforcement inspects its child's statistics, which recurse to + // the leaf. + // + // The measured arm drives the real `EnsureRequirements` rule, so the sharing + // and the cache-reset condition under test are the ones the rule actually + // uses — reimplementing them here would keep passing even if the rule + // stopped sharing. The baseline arm allocates a fresh `StatisticsContext` + // per node, reproducing the behavior before this change. + fn deep_plan(depth: usize, calls: &Arc) -> Arc { + let mut plan: Arc = + Arc::new(CountingStatsExec::new(parquet_exec(), Arc::clone(calls))); + for _ in 0..depth { + plan = filter_exec(plan); + } + plan + } + + fn config() -> ConfigOptions { + let mut config = ConfigOptions::new(); + config.execution.target_partitions = 10; + // Keep the plan a fixpoint so no node is rebuilt and the shared cache is + // never reset; statistics are still computed for the round-robin decision. + config.optimizer.enable_round_robin_repartition = false; + config + } + + /// Leaf statistics computations performed by the real rule. + fn via_rule(depth: usize) -> Result { + let calls = Arc::new(AtomicUsize::new(0)); + EnsureRequirements::new().optimize(deep_plan(depth, &calls), &config())?; + Ok(calls.load(Ordering::Relaxed)) + } + + /// Leaf statistics computations with a fresh context per node. + fn per_node_context(depth: usize) -> Result { + let calls = Arc::new(AtomicUsize::new(0)); + let config = config(); + DistributionContext::new_default(deep_plan(depth, &calls)).transform_up( + |ctx| { + ensure_distribution_with_stats( + ctx, + &config, + &datafusion_physical_plan::statistics::StatisticsContext::new(), + ) + }, + )?; + Ok(calls.load(Ordering::Relaxed)) + } + + let (shared_shallow, fresh_shallow) = (via_rule(4)?, per_node_context(4)?); + let (shared_deep, fresh_deep) = (via_rule(12)?, per_node_context(12)?); + + // Sharing strictly reduces statistics recomputation at any depth. A rule that + // stopped sharing (or reset the cache on every node) would make these equal. + assert!( + shared_shallow < fresh_shallow && shared_deep < fresh_deep, + "shared cache must recompute less: shallow {shared_shallow} vs {fresh_shallow}, \ + deep {shared_deep} vs {fresh_deep}" + ); + + // Without sharing, each extra ancestor recomputes the leaf's subtree, so the + // gap widens as the plan gets deeper. That is the depth-scaling recomputation + // the shared cache removes. + let saved_shallow = fresh_shallow - shared_shallow; + let saved_deep = fresh_deep - shared_deep; + assert!( + saved_deep > saved_shallow, + "the shared cache should save more on deeper plans: \ + saved {saved_shallow} at depth 4, {saved_deep} at depth 12" + ); + + Ok(()) +} + +/// `EnsureRequirements::optimize_with_context` must thread the session's +/// statistics registry into the distribution pass, so registered providers can +/// influence cost-based decisions (here, whether a round-robin repartition is +/// worthwhile). A tiny single-partition scan does not warrant round-robin on its +/// real statistics; a provider that reports it as large flips that decision, but +/// only if the registry is actually threaded through. +#[test] +fn ensure_distribution_uses_context_statistics_registry() -> Result<()> { + let alias = vec![("a".to_string(), "a".to_string())]; + let plan = aggregate_exec_with_alias(parquet_exec_with_size(1, 100), alias); + + let mut config = ConfigOptions::new(); + config.execution.target_partitions = 10; + // Make the round-robin decision actually depend on the estimated row count. + config + .execution + .use_row_number_estimates_to_optimize_partitioning = true; + + // Default context: no registry, so the scan's real (tiny) statistics apply. + let plan_default = EnsureRequirements::new().optimize(plan.clone(), &config)?; + + // A provider that reports the scan as large. + let mut registry = StatisticsRegistry::new(); + registry.register(Arc::new(ClosureStatisticsProvider::with_matches( + |p| p.name() == "DataSourceExec", + |p, _child_stats| { + let mut stats = Statistics::new_unknown(&p.schema()); + stats.num_rows = Precision::Inexact(10_000_000); + Ok(StatisticsResult::Computed(stats.into())) + }, + ))); + + struct ContextWithRegistry { + config: ConfigOptions, + registry: StatisticsRegistry, + } + impl PhysicalOptimizerContext for ContextWithRegistry { + fn config_options(&self) -> &ConfigOptions { + &self.config + } + fn statistics_registry(&self) -> Option<&StatisticsRegistry> { + Some(&self.registry) + } + } + + let plan_registry = EnsureRequirements::new() + .optimize_with_context(plan, &ContextWithRegistry { config, registry })?; + + let s_default = displayable(plan_default.as_ref()).indent(true).to_string(); + let s_registry = displayable(plan_registry.as_ref()).indent(true).to_string(); + + // With the scan's tiny real stats, a round-robin repartition is not worth it. + assert!( + !s_default.contains("RoundRobinBatch"), + "default context (tiny stats) should not add a round-robin repartition:\n{s_default}" + ); + // The registry reports the scan as large, so the same rule now parallelizes + // it — proving the registry was threaded through `optimize_with_context`. + assert!( + s_registry.contains("RoundRobinBatch"), + "registry-reported large stats should add a round-robin repartition:\n{s_registry}" + ); + Ok(()) +} diff --git a/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs b/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs index 0d850d4113c74..776e3c3bf14bd 100644 --- a/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs +++ b/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs @@ -997,6 +997,7 @@ fn get_repartition_requirement_status( plan: &Arc, batch_size: usize, should_use_estimates: bool, + stats_ctx: &StatisticsContext, ) -> Result> { let mut needs_alignment = false; let children = plan.children(); @@ -1008,7 +1009,7 @@ fn get_repartition_requirement_status( { // Decide whether adding a round robin is beneficial depending on // the statistical information we have on the number of rows: - let roundrobin_beneficial_stats = match StatisticsContext::new() + let roundrobin_beneficial_stats = match stats_ctx .compute(child.as_ref(), &StatisticsArgs::new())? .num_rows { @@ -1323,6 +1324,23 @@ fn enforce_distribution_relationships( } } +/// Deprecated in favour of [`ensure_distribution_with_stats`]. +/// +/// This entry point allocates a fresh [`StatisticsContext`] per call, so a +/// bottom-up traversal recomputes each shared subtree's statistics once per +/// ancestor. Pass one context through the whole walk instead — see +/// [`ensure_distribution_with_stats`]. +#[deprecated( + since = "56.0.0", + note = "use `ensure_distribution_with_stats` and share one `StatisticsContext` across the traversal" +)] +pub fn ensure_distribution( + dist_context: DistributionContext, + config: &ConfigOptions, +) -> Result> { + ensure_distribution_with_stats(dist_context, config, &StatisticsContext::new()) +} + /// This function checks whether we need to add additional data exchange /// operators to satisfy distribution requirements. Since this function /// takes care of such requirements, we should avoid manually adding data @@ -1331,13 +1349,21 @@ fn enforce_distribution_relationships( /// This function is intended to be used in a bottom up traversal, as it /// can first repartition (or newly partition) at the datasources -- these /// source partitions may be later repartitioned with additional data exchange operators. +/// +/// `stats_ctx` carries the memoization cache used to answer the child +/// statistics queries behind the repartition decisions. Share one context +/// across the whole traversal so each subtree is computed once rather than +/// once per ancestor. The cache is keyed by raw plan-node pointers, so reset +/// it (via [`StatisticsContext::reset_cache`]) after any node whose plan +/// pointer actually changed. #[expect( deprecated, reason = "HashPartitioned is accepted during the KeyPartitioned migration" )] -pub fn ensure_distribution( +pub fn ensure_distribution_with_stats( dist_context: DistributionContext, config: &ConfigOptions, + stats_ctx: &StatisticsContext, ) -> Result> { let dist_context = update_children(dist_context)?; @@ -1429,8 +1455,12 @@ pub fn ensure_distribution( || plan.is::(); let input_distributions = plan.input_distribution_requirements(); - let repartition_status_flags = - get_repartition_requirement_status(&plan, batch_size, should_use_estimates)?; + let repartition_status_flags = get_repartition_requirement_status( + &plan, + batch_size, + should_use_estimates, + stats_ctx, + )?; // This loop iterates over all the children to: // - Increase parallelism for every child if it is beneficial. // - Satisfy the distribution requirements of every child, if it is not diff --git a/datafusion/physical-optimizer/src/ensure_requirements/mod.rs b/datafusion/physical-optimizer/src/ensure_requirements/mod.rs index 2bc57915b2315..08e708baa5271 100644 --- a/datafusion/physical-optimizer/src/ensure_requirements/mod.rs +++ b/datafusion/physical-optimizer/src/ensure_requirements/mod.rs @@ -150,11 +150,13 @@ pub mod enforce_sorting; use std::sync::Arc; use crate::PhysicalOptimizerRule; +use crate::optimizer::{ConfigOnlyContext, PhysicalOptimizerContext}; use datafusion_common::Result; use datafusion_common::config::ConfigOptions; use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode}; use datafusion_physical_plan::ExecutionPlan; +use datafusion_physical_plan::statistics::StatisticsContext; /// Optimizer rule that enforces both distribution and sorting requirements. /// @@ -180,6 +182,15 @@ impl PhysicalOptimizerRule for EnsureRequirements { plan: Arc, config: &ConfigOptions, ) -> Result> { + self.optimize_with_context(plan, &ConfigOnlyContext::new(config)) + } + + fn optimize_with_context( + &self, + plan: Arc, + context: &dyn PhysicalOptimizerContext, + ) -> Result> { + let config = context.config_options(); // Phase 0: Normalize `InterleaveExec` back to `UnionExec` (top-down). // Interleaves are distribution artifacts of Phase 2, which re-derives // them from the children's final partitioning. Keeping them would @@ -203,13 +214,35 @@ impl PhysicalOptimizerRule for EnsureRequirements { // Phase 2: Combined distribution + sorting enforcement (single bottom-up pass) // For each node: distribution first, then sorting. - use super::enforce_distribution::{DistributionContext, ensure_distribution}; + use super::enforce_distribution::{ + DistributionContext, ensure_distribution_with_stats, + }; use super::enforce_sorting::{PlanWithCorrespondingSort, ensure_sorting}; // Step 2a: Distribution enforcement (bottom-up) let dist_ctx = DistributionContext::new_default(plan); + // Share one statistics context across the whole distribution pass so each + // subtree's statistics are computed once instead of once per ancestor. + // Build it from the session's statistics registry so registered providers + // are consulted (an empty registry, the default, is unchanged behavior). + // `StatsCache` is keyed by raw node pointer, so reset it after any node + // whose plan pointer actually changed: a rewrite can free a cached node + // and a later allocation could reuse its address. A node that makes no + // change cannot free anything, so the cache safely persists across the + // no-op nodes that dominate a deep plan. + let stats_ctx = match context.statistics_registry() { + Some(registry) => StatisticsContext::new_with_registry(registry.clone()), + None => StatisticsContext::new(), + }; let dist_ctx = dist_ctx - .transform_up(|ctx| ensure_distribution(ctx, config)) + .transform_up(|ctx| { + let before = Arc::clone(&ctx.plan); + let result = ensure_distribution_with_stats(ctx, config, &stats_ctx)?; + if !Arc::ptr_eq(&before, &result.data.plan) { + stats_ctx.reset_cache(); + } + Ok(result) + }) .data()?; // Step 2b: Sorting enforcement (bottom-up) — runs on distribution-fixed plan diff --git a/docs/source/library-user-guide/upgrading/56.0.0.md b/docs/source/library-user-guide/upgrading/56.0.0.md index ef51e4c9a7ff6..b89a1d33daaba 100644 --- a/docs/source/library-user-guide/upgrading/56.0.0.md +++ b/docs/source/library-user-guide/upgrading/56.0.0.md @@ -357,3 +357,36 @@ let description = ChildFilterDescription::from_child_with_column_mapping( &child, )?; ``` + +### `ensure_distribution` is deprecated in favour of `ensure_distribution_with_stats` + +`datafusion_physical_optimizer::enforce_distribution::ensure_distribution` now +takes its `StatisticsContext` from the caller, so one context (and its +memoization cache) can be shared across a whole bottom-up traversal. The old +two-argument form is kept as a deprecated wrapper that allocates a fresh +context per call, which recomputes each shared subtree's statistics once per +ancestor. + +```rust,ignore +// Before +let plan = DistributionContext::new_default(plan) + .transform_up(|ctx| ensure_distribution(ctx, config)) + .data()?; + +// After: one context for the whole walk. `StatsCache` is keyed by raw plan-node +// pointers, so reset it whenever a node's plan pointer actually changed. +let stats_ctx = StatisticsContext::new(); +let plan = DistributionContext::new_default(plan) + .transform_up(|ctx| { + let before = Arc::clone(&ctx.plan); + let result = ensure_distribution_with_stats(ctx, config, &stats_ctx)?; + if !Arc::ptr_eq(&before, &result.data.plan) { + stats_ctx.reset_cache(); + } + Ok(result) + }) + .data()?; +``` + +Callers that only need the previous behaviour can keep using the deprecated +form, or pass a freshly constructed `StatisticsContext` per call.