diff --git a/datafusion/core/tests/physical_optimizer/enforce_distribution.rs b/datafusion/core/tests/physical_optimizer/enforce_distribution.rs index c805be44552cc..295e3c369068f 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, @@ -38,6 +39,7 @@ use datafusion::datasource::physical_plan::{CsvSource, ParquetSource}; use datafusion::datasource::source::DataSourceExec; use datafusion::prelude::{SessionConfig, SessionContext}; use datafusion_common::ScalarValue; +use datafusion_common::Statistics; use datafusion_common::config::CsvOptions; use datafusion_common::error::Result; use datafusion_common::tree_node::{ @@ -636,7 +638,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 { @@ -753,7 +760,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)?; @@ -4644,3 +4655,188 @@ 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(()) +} diff --git a/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs b/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs index 5d8d5bac4a385..3b07570506c07 100644 --- a/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs +++ b/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs @@ -1002,6 +1002,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(); @@ -1013,7 +1014,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 { @@ -1142,6 +1143,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 @@ -1150,13 +1168,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)?; @@ -1253,8 +1279,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 41a03bb031629..19403052b82aa 100644 --- a/datafusion/physical-optimizer/src/ensure_requirements/mod.rs +++ b/datafusion/physical-optimizer/src/ensure_requirements/mod.rs @@ -153,6 +153,7 @@ 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. /// @@ -194,13 +195,34 @@ 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. + // `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. + // Upstream builds this from the session's statistics registry; this + // branch has no registry plumbing inside StatisticsContext yet, so the + // context is created empty. The memoization, which is what the change + // is for, is unaffected. + let stats_ctx = 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