From 73895823a35dda0cf045ea05c13904356f1332f3 Mon Sep 17 00:00:00 2001 From: Qi Zhu Date: Wed, 9 Sep 2026 13:51:37 +0800 Subject: [PATCH 1/2] perf: reuse a registry-backed StatisticsContext across ensure_distribution get_repartition_requirement_status created a fresh StatisticsContext per child. StatisticsContext::compute recurses the child's whole subtree and carries a pointer-keyed memoization cache meant to be reused across a walk, so allocating a new context per child discards it every time and a single ensure_distribution pass recomputes shared subtree statistics O(depth) times. A fresh context also builds an empty StatisticsRegistry, so the rule never consulted registered statistics providers. Override EnsureRequirements::optimize_with_context and build one shared StatisticsContext from context.statistics_registry() per pass, threaded into ensure_distribution and get_repartition_requirement_status. This both shares the memoization cache and lets registered providers inform the distribution decision (no change by default: an empty registry behaves as before). optimize delegates to optimize_with_context via ConfigOnlyContext. The cache is keyed by raw node pointer, so it is reset after any node whose plan pointer actually changed (a rewrite can free a cached node and a later allocation could reuse its address); no-op nodes keep the cache warm. Tests: - ensure_distribution_shares_statistics_cache: counts a leaf's statistics computations under a deepening pass-through stack and asserts the shared cache saves progressively more with depth (a non-shared cache saves nothing). - ensure_distribution_uses_context_statistics_registry: a scan whose tiny real stats do not warrant round-robin is reported large by a registry provider, and the round-robin repartition then appears only via optimize_with_context. No plan changes by default; physical_optimizer and statistics suites pass. Thread the PhysicalOptimizerContext through ensure_distribution (replacing the bare ConfigOptions) so config, the statistics registry, and future CBO context are reachable without changing this public signature again. The shared StatisticsContext is built once per pass from context.statistics_registry() and threaded to get_repartition_requirement_status so the memoization cache is shared across the whole distribution walk; building it per call from the registry would lose that sharing. optimize delegates to optimize_with_context via ConfigOnlyContext. Cache is reset only when a node's plan pointer actually changed (a rewrite can free a cached node); robustness of the pointer key tracked separately in #25141. Tests: - ensure_distribution_shares_statistics_cache: the shared cache saves progressively more statistics recomputation as a pass-through stack deepens. - ensure_distribution_uses_context_statistics_registry: a scan reported large by a registry provider gets a round-robin repartition only via optimize_with_context, proving the registry is threaded through. No default behavior change; physical_optimizer and statistics suites pass. --- .../enforce_distribution.rs | 276 +++++++++++++++++- .../enforce_distribution.rs | 17 +- .../src/ensure_requirements/mod.rs | 33 ++- 3 files changed, 318 insertions(+), 8 deletions(-) diff --git a/datafusion/core/tests/physical_optimizer/enforce_distribution.rs b/datafusion/core/tests/physical_optimizer/enforce_distribution.rs index 4aeee422394ca..8513dae56a43e 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, @@ -55,15 +56,20 @@ 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; use datafusion_physical_optimizer::join_selection::JoinSelection; +use datafusion_physical_optimizer::optimizer::ConfigOnlyContext; use datafusion_physical_optimizer::output_requirements::OutputRequirements; 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, @@ -641,7 +647,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( + distribution_context, + &ConfigOnlyContext::new(&config), + &datafusion_physical_plan::statistics::StatisticsContext::new(), + ) + .map(|item| item.data.plan) } fn test_suite_default_config_options() -> ConfigOptions { @@ -764,7 +775,11 @@ impl TestConfig { // Then run ensure_distribution rule DistributionContext::new_default(adjusted) .transform_up(|distribution_context| { - ensure_distribution(distribution_context, &self.config) + ensure_distribution( + distribution_context, + &ConfigOnlyContext::new(&self.config), + &datafusion_physical_plan::statistics::StatisticsContext::new(), + ) }) .data() .and_then(check_integrity)?; @@ -4997,3 +5012,260 @@ 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 while + // `ensure_distribution` runs 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. + // + // `shared` uses one `StatisticsContext` for the whole pass (what + // `EnsureRequirements` does); `fresh` allocates a new context per node (the + // behavior before this change). Returns (shared_computes, fresh_computes). + fn run(depth: usize) -> Result<(usize, usize)> { + 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 + } + + 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; + + let shared_calls = Arc::new(AtomicUsize::new(0)); + let stats_ctx = datafusion_physical_plan::statistics::StatisticsContext::new(); + DistributionContext::new_default(deep_plan(depth, &shared_calls)).transform_up( + |ctx| { + // Reset only when the node's plan pointer actually changed, exactly + // as `EnsureRequirements` does (a rewrite can free a cached node). + let before = Arc::clone(&ctx.plan); + let result = ensure_distribution( + ctx, + &ConfigOnlyContext::new(&config), + &stats_ctx, + )?; + if !Arc::ptr_eq(&before, &result.data.plan) { + stats_ctx.reset_cache(); + } + Ok(result) + }, + )?; + let shared = shared_calls.load(Ordering::Relaxed); + + let fresh_calls = Arc::new(AtomicUsize::new(0)); + DistributionContext::new_default(deep_plan(depth, &fresh_calls)).transform_up( + |ctx| { + ensure_distribution( + ctx, + &ConfigOnlyContext::new(&config), + &datafusion_physical_plan::statistics::StatisticsContext::new(), + ) + }, + )?; + let fresh = fresh_calls.load(Ordering::Relaxed); + + Ok((shared, fresh)) + } + + let (shared_shallow, fresh_shallow) = run(4)?; + let (shared_deep, fresh_deep) = run(12)?; + + // Sharing strictly reduces statistics recomputation at any depth. A broken + // cache (e.g. reset 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 between `fresh` and `shared` widens as the plan gets deeper. That is the + // depth-scaling recomputation the shared cache removes; a cache that is not + // actually shared would save nothing and the gap would not grow. + 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 ae9774c9f8c2d..a250c0ec0caff 100644 --- a/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs +++ b/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs @@ -38,8 +38,8 @@ use crate::utils::{ is_sort_preserving_merge, }; +use crate::optimizer::PhysicalOptimizerContext; use arrow::compute::SortOptions; -use datafusion_common::config::ConfigOptions; use datafusion_common::error::Result; use datafusion_common::stats::Precision; use datafusion_common::tree_node::Transformed; @@ -1014,6 +1014,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(); @@ -1025,7 +1026,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 { @@ -1354,7 +1355,8 @@ fn enforce_distribution_relationships( )] pub fn ensure_distribution( dist_context: DistributionContext, - config: &ConfigOptions, + context: &dyn PhysicalOptimizerContext, + stats_ctx: &StatisticsContext, ) -> Result> { let dist_context = update_children(dist_context)?; @@ -1362,6 +1364,7 @@ pub fn ensure_distribution( return Ok(Transformed::no(dist_context)); } + let config = context.config_options(); let target_partitions = config.execution.target_partitions; // When `false`, round robin repartition will not be added to increase parallelism let enable_round_robin = config.optimizer.enable_round_robin_repartition; @@ -1450,8 +1453,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..dc45f8b00dac5 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 @@ -208,8 +219,28 @@ impl PhysicalOptimizerRule for EnsureRequirements { // 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(ctx, context, &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 From 31bdfad37e16a084cb6891b0213b037b00f4d6a2 Mon Sep 17 00:00:00 2001 From: Qi Zhu Date: Thu, 10 Sep 2026 14:09:37 +0800 Subject: [PATCH 2/2] test: update ensure_distribution callers after merging main The main merge pulled in #24809's preserve-fetch reoptimization tests, whose two ensure_distribution() calls still used the pre-refactor 2-arg (context, &config) form. Thread the PhysicalOptimizerContext and a StatisticsContext through them to match the current signature. --- .../physical_optimizer/enforce_distribution.rs | 16 ++++++++++++++-- 1 file changed, 14 insertions(+), 2 deletions(-) diff --git a/datafusion/core/tests/physical_optimizer/enforce_distribution.rs b/datafusion/core/tests/physical_optimizer/enforce_distribution.rs index 806cd656d07c7..898eb70b88848 100644 --- a/datafusion/core/tests/physical_optimizer/enforce_distribution.rs +++ b/datafusion/core/tests/physical_optimizer/enforce_distribution.rs @@ -5080,7 +5080,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( + context, + &ConfigOnlyContext::new(&config), + &datafusion_physical_plan::statistics::StatisticsContext::new(), + ) + })? .data; check_integrity(distribution)?; optimized = EnsureRequirements::new().optimize(optimized, &config)?; @@ -5155,7 +5161,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( + context, + &ConfigOnlyContext::new(&config), + &datafusion_physical_plan::statistics::StatisticsContext::new(), + ) + })? .data; check_integrity(distribution)?; plan = EnsureRequirements::new().optimize(plan, &config)?;