diff --git a/datafusion/core/src/physical_planner.rs b/datafusion/core/src/physical_planner.rs index b48bc19a5517c..ba2e05b279743 100644 --- a/datafusion/core/src/physical_planner.rs +++ b/datafusion/core/src/physical_planner.rs @@ -64,6 +64,7 @@ use arrow_schema::Field; use datafusion_catalog::ScanArgs; use datafusion_common::Column; use datafusion_common::HashMap as DFHashMap; +use datafusion_common::config::ConfigOptions; use datafusion_common::display::ToStringifiedPlan; use datafusion_common::format::ExplainAnalyzeCategories; use datafusion_common::tree_node::{Transformed, TreeNode, TreeNodeRecursion}; @@ -103,9 +104,11 @@ use datafusion_physical_expr::{ use datafusion_physical_plan::empty::EmptyExec; use datafusion_physical_plan::execution_plan::InvariantLevel; use datafusion_physical_plan::joins::PiecewiseMergeJoinExec; +use datafusion_physical_plan::operator_statistics::StatisticsRegistry; use datafusion_physical_plan::placeholder_row::PlaceholderRowExec; use datafusion_physical_plan::recursive_query::RecursiveQueryExec; use datafusion_physical_plan::scalar_subquery::{ScalarSubqueryExec, ScalarSubqueryLink}; +use datafusion_physical_plan::statistics::StatisticsContext; use datafusion_physical_plan::unnest::ListUnnest; use datafusion_session::{PhysicalOptimizerContext, PhysicalOptimizerRule, Session}; @@ -123,18 +126,34 @@ pub use datafusion_session::{ExtensionPlanner, PhysicalPlanner}; struct SessionOptimizerContext<'a> { session: &'a dyn Session, + statistics_context: StatisticsContext, +} + +impl<'a> SessionOptimizerContext<'a> { + fn new(session: &'a dyn Session) -> Self { + let statistics_context = match session.statistics_registry() { + Some(registry) => StatisticsContext::new_with_registry(registry.clone()), + None => StatisticsContext::new(), + }; + Self { + session, + statistics_context, + } + } } impl PhysicalOptimizerContext for SessionOptimizerContext<'_> { - fn config_options(&self) -> &datafusion_common::config::ConfigOptions { + fn config_options(&self) -> &ConfigOptions { self.session.config_options() } - fn statistics_registry( - &self, - ) -> Option<&datafusion_physical_plan::operator_statistics::StatisticsRegistry> { + fn statistics_registry(&self) -> Option<&StatisticsRegistry> { self.session.statistics_registry() } + + fn statistics_context(&self) -> Option<&StatisticsContext> { + Some(&self.statistics_context) + } } /// Default single node physical query planner that converts a @@ -3123,9 +3142,7 @@ impl DefaultPhysicalPlanner { InvariantChecker(InvariantLevel::Always).check(&plan)?; let mut new_plan = Arc::clone(&plan); - let optimizer_context = SessionOptimizerContext { - session: session_state, - }; + let optimizer_context = SessionOptimizerContext::new(session_state); for optimizer in optimizers { let before_schema = new_plan.schema(); new_plan = optimizer @@ -3536,11 +3553,14 @@ mod tests { use arrow::datatypes::{DataType, Field, Int32Type}; use arrow_schema::{FieldRef, SchemaRef}; use datafusion_catalog::CatalogProviderList; + use datafusion_common::Statistics; use datafusion_common::config::{ConfigOptions, TableOptions}; + use datafusion_common::stats::Precision; use datafusion_common::{ DFSchemaRef, ScalarValue, SplitPoint, TableReference, ToDFSchema as _, assert_batches_eq, assert_contains, }; + use datafusion_datasource::source::DataSourceExec; use datafusion_execution::TaskContext; use datafusion_execution::runtime_env::RuntimeEnv; use datafusion_expr::builder::subquery_alias; @@ -3558,12 +3578,18 @@ mod tests { use datafusion_functions_aggregate::expr_fn::sum; use datafusion_physical_expr::EquivalenceProperties; use datafusion_physical_plan::execution_plan::{Boundedness, EmissionType}; + use datafusion_physical_plan::operator_statistics::{ + ClosureStatisticsProvider, StatisticsRegistry, StatisticsResult, + }; + use datafusion_physical_plan::statistics::StatisticsArgs; use datafusion_physical_plan::{ChildrenPropertiesMode, ReplaceChildrenOptions}; use datafusion_session::QueryPlanner; + use parking_lot::Mutex as SyncMutex; + /// Records the root statistics computed with the shared statistics context #[derive(Debug)] struct ContextCheckingRule { - invoked: Arc, + recorded: Arc>>>, } impl PhysicalOptimizerRule for ContextCheckingRule { @@ -3581,7 +3607,12 @@ mod tests { context: &dyn PhysicalOptimizerContext, ) -> Result> { assert!(context.statistics_registry().is_some()); - self.invoked.store(true, AtomicOrdering::Relaxed); + let statistics_context = context + .statistics_context() + .expect("the planner shares a statistics context"); + let statistics = + statistics_context.compute_arc(&plan, &StatisticsArgs::new())?; + self.recorded.lock().push(statistics); Ok(plan) } @@ -3644,10 +3675,7 @@ mod tests { self.inner.physical_optimizers() } - fn statistics_registry( - &self, - ) -> Option<&datafusion_physical_plan::operator_statistics::StatisticsRegistry> - { + fn statistics_registry(&self) -> Option<&StatisticsRegistry> { self.inner.statistics_registry() } @@ -3816,15 +3844,13 @@ mod tests { #[tokio::test] async fn plans_with_non_session_state_implementation() -> Result<()> { - let invoked = Arc::new(AtomicBool::new(false)); + let recorded = Arc::new(SyncMutex::new(vec![])); let inner = SessionStateBuilder::new() .with_default_features() .with_physical_optimizer_rules(vec![Arc::new(ContextCheckingRule { - invoked: Arc::clone(&invoked), + recorded: Arc::clone(&recorded), })]) - .with_statistics_registry( - datafusion_physical_plan::operator_statistics::StatisticsRegistry::new(), - ) + .with_statistics_registry(StatisticsRegistry::new()) .build(); let query_planner_invoked = Arc::new(AtomicBool::new(false)); let session = TestSession { @@ -3839,7 +3865,75 @@ mod tests { let physical_plan = session.create_physical_plan(&logical_plan).await?; assert!(physical_plan.is::()); assert!(query_planner_invoked.load(AtomicOrdering::Relaxed)); - assert!(invoked.load(AtomicOrdering::Relaxed)); + assert_eq!(recorded.lock().len(), 1); + Ok(()) + } + + #[tokio::test] + async fn optimizer_rules_share_statistics_context() -> Result<()> { + let recorded = Arc::new(SyncMutex::new(vec![])); + let rule = || -> Arc { + Arc::new(ContextCheckingRule { + recorded: Arc::clone(&recorded), + }) + }; + let session_state = SessionStateBuilder::new() + .with_default_features() + .with_physical_optimizer_rules(vec![rule(), rule()]) + .with_statistics_registry(StatisticsRegistry::new()) + .build(); + + let logical_plan = LogicalPlanBuilder::empty(false).build()?; + session_state.create_physical_plan(&logical_plan).await?; + + // The second rule reads the statistics the first rule cached + let recorded = recorded.lock(); + assert_eq!(recorded.len(), 2); + assert!(Arc::ptr_eq(&recorded[0], &recorded[1])); + Ok(()) + } + + /// `AggregateStatistics` trusts an `Exact` row count from a registered + /// provider: it answers `COUNT(*)` without scanning + #[tokio::test] + async fn aggregate_statistics_consults_statistics_providers() -> Result<()> { + let provider = ClosureStatisticsProvider::with_matches( + |plan| plan.is::(), + |plan, child_stats| { + let child_stats = child_stats + .iter() + .map(|c| Arc::clone(c.base_arc())) + .collect::>(); + let mut stats = Arc::unwrap_or_clone( + plan.statistics_from_inputs(&child_stats, &StatisticsArgs::new())?, + ); + stats.num_rows = Precision::Exact(42); + Ok(StatisticsResult::Computed(stats.into())) + }, + ); + let state = SessionStateBuilder::new() + .with_default_features() + .with_statistics_registry(StatisticsRegistry::with_providers(vec![Arc::new( + provider, + )])) + .build(); + let ctx = SessionContext::new_with_state(state); + ctx.sql("CREATE TABLE t AS VALUES (1), (2), (3)") + .await? + .collect() + .await?; + + let batches = ctx.sql("SELECT COUNT(*) FROM t").await?.collect().await?; + assert_batches_eq!( + &[ + "+----------+", + "| count(*) |", + "+----------+", + "| 42 |", + "+----------+" + ], + &batches + ); Ok(()) } diff --git a/datafusion/physical-optimizer/src/aggregate_statistics.rs b/datafusion/physical-optimizer/src/aggregate_statistics.rs index 43b1abb4b68a9..a1260529431de 100644 --- a/datafusion/physical-optimizer/src/aggregate_statistics.rs +++ b/datafusion/physical-optimizer/src/aggregate_statistics.rs @@ -25,7 +25,7 @@ use datafusion_physical_plan::aggregates::{ }; use datafusion_physical_plan::placeholder_row::PlaceholderRowExec; use datafusion_physical_plan::projection::{ProjectionExec, ProjectionExpr}; -use datafusion_physical_plan::statistics::{StatisticsArgs, StatisticsContext}; +use datafusion_physical_plan::statistics::StatisticsArgs; use datafusion_physical_plan::udaf::{ AggregateFunctionExpr, StatisticsArgs as PlanStatisticsArgs, }; @@ -33,6 +33,7 @@ use datafusion_physical_plan::{ExecutionPlan, expressions}; use std::sync::Arc; use crate::PhysicalOptimizerRule; +use crate::optimizer::{ConfigOnlyContext, PhysicalOptimizerContext}; /// Optimizer that uses available statistics for aggregate functions #[derive(Default, Debug)] @@ -46,20 +47,26 @@ impl AggregateStatistics { } impl PhysicalOptimizerRule for AggregateStatistics { - #[cfg_attr(feature = "recursive_protection", recursive::recursive)] - #[expect(clippy::allow_attributes)] // See https://github.com/apache/datafusion/issues/18881#issuecomment-3621545670 - #[allow(clippy::only_used_in_recursion)] // See https://github.com/rust-lang/rust-clippy/issues/14566 fn optimize( &self, plan: Arc, config: &ConfigOptions, + ) -> Result> { + self.optimize_with_context(plan, &ConfigOnlyContext::new(config)) + } + + #[cfg_attr(feature = "recursive_protection", recursive::recursive)] + fn optimize_with_context( + &self, + plan: Arc, + context: &dyn PhysicalOptimizerContext, ) -> Result> { if let Some(partial_agg_exec) = take_optimizable(&plan) { let partial_agg_exec = partial_agg_exec .downcast_ref::() .expect("take_optimizable() ensures that this is a AggregateExec"); - let stats = StatisticsContext::new() - .compute(partial_agg_exec.input().as_ref(), &StatisticsArgs::new())?; + let stats = context + .compute_statistics(partial_agg_exec.input(), &StatisticsArgs::new())?; let mut projections = vec![]; for expr in partial_agg_exec.aggr_expr() { let field = expr.field(); @@ -92,13 +99,17 @@ impl PhysicalOptimizerRule for AggregateStatistics { )?)) } else { plan.map_children(|child| { - self.optimize(child, config).map(Transformed::yes) + self.optimize_with_context(child, context) + .map(Transformed::yes) }) .data() } } else { - plan.map_children(|child| self.optimize(child, config).map(Transformed::yes)) - .data() + plan.map_children(|child| { + self.optimize_with_context(child, context) + .map(Transformed::yes) + }) + .data() } } diff --git a/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs b/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs index 3c2426a5ef2b1..31f00fde90ee5 100644 --- a/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs +++ b/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs @@ -1063,10 +1063,8 @@ struct PlanSize { } impl PlanSize { - fn from_plan(plan: &dyn ExecutionPlan) -> Self { - let stats = StatisticsContext::new() - .compute(plan, &StatisticsArgs::new()) - .ok(); + fn from_plan(plan: &Arc, stats_ctx: &StatisticsContext) -> Self { + let stats = stats_ctx.compute_arc(plan, &StatisticsArgs::new()).ok(); Self { byte_size: stats .as_ref() @@ -1124,6 +1122,7 @@ fn enforce_distribution_relationships( input_distributions: &InputDistributionRequirements, children: &mut [DistributionChildState], target_partitions: usize, + stats_ctx: &StatisticsContext, ) -> Result<()> { let mut repartitioned_for_relationship = vec![false; children.len()]; @@ -1162,53 +1161,52 @@ fn enforce_distribution_relationships( } } - let best_satisfied_child: Option<(usize, Partitioning)> = match satisfied_children - .len() - { - 0 => None, - 1 => satisfied_children - .into_iter() - .next() - .map(|(i, p, _)| (i, p)), - _ => { - // Prefer native partitioned children over newly repartitioned ones - let native_children: Vec<_> = satisfied_children - .iter() - .filter(|(_, _, is_native)| *is_native) - .collect(); - if native_children.len() == 1 { - let (i, p, _) = native_children[0]; - Some((*i, p.clone())) - } else { - let pool = if !native_children.is_empty() { - native_children - } else { - satisfied_children.iter().collect() - }; - let candidates: Vec<_> = pool - .into_iter() - .map(|(idx, part, _)| { - let size = - PlanSize::from_plan(children[*idx].context.plan.as_ref()); - (size, *idx, part.clone()) - }) - .collect(); - - // Prefer a unique, strictly larger winner (`size_a > size_b`). - // Otherwise, fall back to standard distribution rather - // than choosing an arbitrary reference. - candidates + let best_satisfied_child: Option<(usize, Partitioning)> = + match satisfied_children.len() { + 0 => None, + 1 => satisfied_children + .into_iter() + .next() + .map(|(i, p, _)| (i, p)), + _ => { + // Prefer native partitioned children over newly repartitioned ones + let native_children: Vec<_> = satisfied_children .iter() - .find(|(size_a, idx_a, _)| { - size_a.is_known() - && candidates.iter().all(|(size_b, idx_b, _)| { - idx_a == idx_b || size_a > size_b - }) - }) - .map(|(_, idx, part)| (*idx, part.clone())) + .filter(|(_, _, is_native)| *is_native) + .collect(); + if native_children.len() == 1 { + let (i, p, _) = native_children[0]; + Some((*i, p.clone())) + } else { + let pool = if !native_children.is_empty() { + native_children + } else { + satisfied_children.iter().collect() + }; + let candidates: Vec<_> = pool + .into_iter() + .map(|(idx, part, _)| { + let plan = &children[*idx].context.plan; + let size = PlanSize::from_plan(plan, stats_ctx); + (size, *idx, part.clone()) + }) + .collect(); + + // Prefer a unique, strictly larger winner (`size_a > size_b`). + // Otherwise, fall back to standard distribution rather + // than choosing an arbitrary reference. + candidates + .iter() + .find(|(size_a, idx_a, _)| { + size_a.is_known() + && candidates.iter().all(|(size_b, idx_b, _)| { + idx_a == idx_b || size_a > size_b + }) + }) + .map(|(_, idx, part)| (*idx, part.clone())) + } } - } - }; + }; // Validate that best_satisfied_child can be adapted across all unsatisfied children in the group let best_satisfied_child = @@ -1666,6 +1664,7 @@ pub fn ensure_distribution_with_stats( &input_distributions, &mut children, target_partitions, + stats_ctx, )?; let children = children diff --git a/datafusion/physical-optimizer/src/ensure_requirements/mod.rs b/datafusion/physical-optimizer/src/ensure_requirements/mod.rs index 3909b56c17e58..56ca430d69973 100644 --- a/datafusion/physical-optimizer/src/ensure_requirements/mod.rs +++ b/datafusion/physical-optimizer/src/ensure_requirements/mod.rs @@ -150,13 +150,14 @@ pub mod enforce_sorting; use std::sync::Arc; use crate::PhysicalOptimizerRule; -use crate::optimizer::{ConfigOnlyContext, PhysicalOptimizerContext}; +use crate::optimizer::{ + ConfigOnlyContext, PhysicalOptimizerContext, with_statistics_context, +}; 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. /// @@ -221,17 +222,13 @@ 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). - 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_with_stats(ctx, config, &stats_ctx)) - .data()?; + let dist_ctx = with_statistics_context(context, |stats_ctx| { + dist_ctx + .transform_up(|ctx| { + ensure_distribution_with_stats(ctx, config, stats_ctx) + }) + .data() + })?; // Step 2b: Sorting enforcement (bottom-up) — runs on distribution-fixed plan let sort_ctx = PlanWithCorrespondingSort::new_default(dist_ctx.plan); diff --git a/datafusion/physical-optimizer/src/join_selection.rs b/datafusion/physical-optimizer/src/join_selection.rs index 2b342bece7040..c7f414d91bfa5 100644 --- a/datafusion/physical-optimizer/src/join_selection.rs +++ b/datafusion/physical-optimizer/src/join_selection.rs @@ -39,7 +39,7 @@ use datafusion_physical_plan::joins::{ CrossJoinExec, HashJoinExec, NestedLoopJoinExec, PartitionMode, StreamJoinPartitionMode, SymmetricHashJoinExec, }; -use datafusion_physical_plan::statistics::{StatisticsArgs, StatisticsContext}; +use datafusion_physical_plan::statistics::StatisticsArgs; use datafusion_physical_plan::{ExecutionPlan, ExecutionPlanProperties}; use std::sync::Arc; @@ -59,14 +59,10 @@ impl JoinSelection { /// Get statistics for a plan node, consulting the registry's providers if one /// is available. fn get_stats( - plan: &dyn ExecutionPlan, + plan: &Arc, context: &dyn PhysicalOptimizerContext, ) -> Result> { - let ctx = match context.statistics_registry() { - Some(reg) => StatisticsContext::new_with_registry(reg.clone()), - None => StatisticsContext::new(), - }; - ctx.compute(plan, &StatisticsArgs::new()) + context.compute_statistics(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. @@ -86,8 +82,8 @@ fn get_stats( /// Used configurations /// - `optimizer.join_reordering`: allows or forbids statistics-driven join swapping pub(crate) fn should_swap_join_order( - left: &dyn ExecutionPlan, - right: &dyn ExecutionPlan, + left: &Arc, + right: &Arc, context: &dyn PhysicalOptimizerContext, ) -> Result { if !context.config_options().optimizer.join_reordering { @@ -115,7 +111,7 @@ pub(crate) fn should_swap_join_order( } fn supports_collect_by_thresholds( - plan: &dyn ExecutionPlan, + plan: &Arc, threshold_byte_size: usize, threshold_num_rows: usize, context: &dyn PhysicalOptimizerContext, @@ -204,14 +200,14 @@ pub(crate) fn try_collect_left( let left_can_collect = ignore_threshold || supports_collect_by_thresholds( - &**left, + left, optimizer_config.hash_join_single_partition_threshold, optimizer_config.hash_join_single_partition_threshold_rows, context, ); let right_can_collect = ignore_threshold || supports_collect_by_thresholds( - &**right, + right, optimizer_config.hash_join_single_partition_threshold, optimizer_config.hash_join_single_partition_threshold_rows, context, @@ -221,7 +217,7 @@ pub(crate) fn try_collect_left( (true, true) => { // For null-aware joins, we only swap `LeftAnti` joins where the left side is > right side if can_swap_hash_join(hash_join) - && should_swap_join_order(&**left, &**right, context)? + && should_swap_join_order(left, right, context)? { Ok(Some(hash_join.swap_inputs(PartitionMode::CollectLeft)?)) } else { @@ -269,9 +265,7 @@ pub(crate) fn partitioned_hash_join( } else { PartitionMode::Partitioned }; - if can_swap_hash_join(hash_join) - && should_swap_join_order(&**left, &**right, context)? - { + if can_swap_hash_join(hash_join) && should_swap_join_order(left, right, context)? { hash_join.swap_inputs(partition_mode) } else { // Null-aware anti joins must use CollectLeft mode because they track probe-side state @@ -316,7 +310,7 @@ fn statistical_join_selection_subrule( let left = hash_join.left(); let right = hash_join.right(); if can_swap_hash_join(hash_join) - && should_swap_join_order(&**left, &**right, context)? + && should_swap_join_order(left, right, context)? { // Null-aware RightAnti only supports CollectLeft let partition_mode = if hash_join.null_aware { @@ -333,7 +327,7 @@ fn statistical_join_selection_subrule( } else if let Some(cross_join) = plan.downcast_ref::() { let left = cross_join.left(); let right = cross_join.right(); - if should_swap_join_order(&**left, &**right, context)? { + if should_swap_join_order(left, right, context)? { cross_join.swap_inputs().map(Some)? } else { None @@ -342,7 +336,7 @@ fn statistical_join_selection_subrule( let left = nl_join.left(); let right = nl_join.right(); if nl_join.join_type().supports_swap() - && should_swap_join_order(&**left, &**right, context)? + && should_swap_join_order(left, right, context)? { nl_join.swap_inputs().map(Some)? } else { diff --git a/datafusion/physical-optimizer/src/limit_pushdown.rs b/datafusion/physical-optimizer/src/limit_pushdown.rs index f88a2be14e984..226c41cf8d737 100644 --- a/datafusion/physical-optimizer/src/limit_pushdown.rs +++ b/datafusion/physical-optimizer/src/limit_pushdown.rs @@ -64,6 +64,9 @@ use std::fmt::Debug; use std::sync::Arc; use crate::PhysicalOptimizerRule; +use crate::optimizer::{ + ConfigOnlyContext, PhysicalOptimizerContext, with_statistics_context, +}; use datafusion_common::config::ConfigOptions; use datafusion_common::error::Result; @@ -111,7 +114,15 @@ impl PhysicalOptimizerRule for LimitPushdown { fn optimize( &self, plan: Arc, - _config: &ConfigOptions, + config: &ConfigOptions, + ) -> Result> { + self.optimize_with_context(plan, &ConfigOnlyContext::new(config)) + } + + fn optimize_with_context( + &self, + plan: Arc, + context: &dyn PhysicalOptimizerContext, ) -> Result> { let global_state = GlobalRequirements { fetch: None, @@ -119,7 +130,9 @@ impl PhysicalOptimizerRule for LimitPushdown { satisfied: false, preserve_order: false, }; - pushdown_limits(plan, global_state) + with_statistics_context(context, |stats_ctx| { + pushdown_limits(plan, global_state, stats_ctx) + }) } fn name(&self) -> &str { @@ -145,9 +158,33 @@ struct LimitInfo { /// /// If a limit is encountered, a [`TreeNodeRecursion::Stop`] is returned. Otherwise, /// return a [`TreeNodeRecursion::Continue`]. +/// +/// Computes statistics with a new [`StatisticsContext`] that has no statistics +/// providers. A context built from a statistics registry, as [`LimitPushdown`] +/// uses, also consults the registered providers, so switching to +/// [`pushdown_limit_helper_with_stats`] with such a context can change the +/// result when providers are registered. +#[deprecated( + since = "56.0.0", + note = "use `pushdown_limit_helper_with_stats` and share one `StatisticsContext` across calls" +)] pub fn pushdown_limit_helper( + pushdown_plan: Arc, + global_state: GlobalRequirements, +) -> Result<(Transformed>, GlobalRequirements)> { + pushdown_limit_helper_with_stats( + pushdown_plan, + global_state, + &StatisticsContext::new(), + ) +} + +/// Like [`pushdown_limit_helper`], but computes statistics with `stats_ctx`, +/// so its cache can be shared with other calls and other rules. +pub fn pushdown_limit_helper_with_stats( mut pushdown_plan: Arc, mut global_state: GlobalRequirements, + stats_ctx: &StatisticsContext, ) -> Result<(Transformed>, GlobalRequirements)> { // Extract limit, if exist, and return child inputs. if let Some(limit_info) = extract_limit(&pushdown_plan) { @@ -165,7 +202,7 @@ pub fn pushdown_limit_helper( global_state.satisfied = false; if let Some(fetch) = fetch - && limit_satisfied_by_input(&limit_info.input, skip, fetch)? + && limit_satisfied_by_input(&limit_info.input, skip, fetch, stats_ctx)? { // The input already produces at most `fetch` rows, so no new limit // node is needed. Mark satisfied so downstream won't re-add one, @@ -315,6 +352,7 @@ fn limit_satisfied_by_input( plan: &Arc, skip: usize, fetch: usize, + stats_ctx: &StatisticsContext, ) -> Result { if skip > 0 { return Ok(false); @@ -324,7 +362,7 @@ fn limit_satisfied_by_input( return Ok(false); } - let Some(num_rows) = limit_eliminable_exact_num_rows(plan)? else { + let Some(num_rows) = limit_eliminable_exact_num_rows(plan, stats_ctx)? else { return Ok(false); }; @@ -335,6 +373,7 @@ fn limit_satisfied_by_input( /// whose row-count guarantees are strong enough to remove a limit. fn limit_eliminable_exact_num_rows( plan: &Arc, + stats_ctx: &StatisticsContext, ) -> Result> { // Unwrap any wrapping ProjectionExec layers; projections preserve row count // but may derive statistics in ways that are not trustworthy, so we peek @@ -353,8 +392,8 @@ fn limit_eliminable_exact_num_rows( } if matches!( - StatisticsContext::new() - .compute(current.as_ref(), &StatisticsArgs::new())? + stats_ctx + .compute_arc(current, &StatisticsArgs::new())? .num_rows, Precision::Exact(0) ) { @@ -368,15 +407,17 @@ fn limit_eliminable_exact_num_rows( pub(crate) fn pushdown_limits( pushdown_plan: Arc, global_state: GlobalRequirements, + stats_ctx: &StatisticsContext, ) -> Result> { // Call pushdown_limit_helper. // This will either extract the limit node (returning the child), or apply the limit pushdown. let (mut new_node, mut global_state) = - pushdown_limit_helper(pushdown_plan, global_state)?; + pushdown_limit_helper_with_stats(pushdown_plan, global_state, stats_ctx)?; // While limits exist, continue combining the global_state. while new_node.tnr == TreeNodeRecursion::Stop { - (new_node, global_state) = pushdown_limit_helper(new_node.data, global_state)?; + (new_node, global_state) = + pushdown_limit_helper_with_stats(new_node.data, global_state, stats_ctx)?; } // Once a limit has been materialized above the current node, child @@ -396,6 +437,7 @@ pub(crate) fn pushdown_limits( let new_child = pushdown_limits( Arc::::clone(child), global_state.clone(), + stats_ctx, )?; // Tracking if any of the children changed changed |= !Arc::ptr_eq(child, &new_child); diff --git a/datafusion/physical-optimizer/src/optimizer.rs b/datafusion/physical-optimizer/src/optimizer.rs index 80c33e1c08b6a..81932594b85fe 100644 --- a/datafusion/physical-optimizer/src/optimizer.rs +++ b/datafusion/physical-optimizer/src/optimizer.rs @@ -40,23 +40,32 @@ use crate::limit_pushdown_past_window::LimitPushPastWindows; use crate::pushdown_sort::PushdownSort; use crate::window_topn::WindowTopN; use datafusion_common::config::ConfigOptions; +use datafusion_physical_plan::statistics::StatisticsContext; // Re-export from this module for backwards compatibility. pub use datafusion_session::{PhysicalOptimizerContext, PhysicalOptimizerRule}; +pub use datafusion_session::with_statistics_context; + /// Simple context wrapping [`ConfigOptions`] for backward compatibility. /// /// This struct provides a minimal implementation of [`PhysicalOptimizerContext`] -/// that only supplies configuration options. Used when no statistics registry -/// is available or needed. +/// that supplies configuration options and a [`StatisticsContext`] without a +/// statistics registry. Used when no statistics registry is available or +/// needed. pub struct ConfigOnlyContext<'a> { config: &'a ConfigOptions, + statistics_context: StatisticsContext, } impl<'a> ConfigOnlyContext<'a> { /// Create a new context wrapping the given config options. pub fn new(config: &'a ConfigOptions) -> Self { - Self { config } + Self { + config, + // No providers, matching `statistics_registry()`, which is `None` + statistics_context: StatisticsContext::new(), + } } } @@ -64,6 +73,10 @@ impl PhysicalOptimizerContext for ConfigOnlyContext<'_> { fn config_options(&self) -> &ConfigOptions { self.config } + + fn statistics_context(&self) -> Option<&StatisticsContext> { + Some(&self.statistics_context) + } } /// A rule-based physical optimizer. diff --git a/datafusion/physical-plan/src/operator_statistics/mod.rs b/datafusion/physical-plan/src/operator_statistics/mod.rs index 04783bd9f24cd..a0b65515c90ce 100644 --- a/datafusion/physical-plan/src/operator_statistics/mod.rs +++ b/datafusion/physical-plan/src/operator_statistics/mod.rs @@ -258,6 +258,11 @@ pub enum StatisticsResult { /// estimation logic. The chain of providers is traversed until one returns /// [`StatisticsResult::Computed`]. /// +/// As for operators, `Exact` statistics must be guaranteed values, because +/// optimizer rules use them for correctness, not just cost: they answer +/// `COUNT(*)`, `MIN` and `MAX` without scanning and remove limits. Return +/// `Inexact` for anything that is an estimate. +/// /// # Implementing a Custom Provider /// /// ```ignore diff --git a/datafusion/physical-plan/src/statistics.rs b/datafusion/physical-plan/src/statistics.rs index 7124c9107dd0a..c52636b21ff56 100644 --- a/datafusion/physical-plan/src/statistics.rs +++ b/datafusion/physical-plan/src/statistics.rs @@ -31,10 +31,9 @@ use datafusion_common::{ Result, Statistics, assert_eq_or_internal_err, assert_or_internal_err, }; use log::debug; -use std::cell::RefCell; +use parking_lot::Mutex; use std::collections::HashMap; use std::ptr::from_ref; -use std::rc::Rc; use std::sync::Arc; type CacheKey = (usize, Option); @@ -158,7 +157,14 @@ pub enum ChildStats { /// plan rewrites. Each entry holds a strong reference to the plan node it was /// computed for, so cached nodes (and their per-partition statistics) stay /// alive until [`Self::reset_cache`] is called or the context is dropped. Reset -/// a long-lived context at a lifecycle boundary to bound its memory. +/// a long-lived context at a lifecycle boundary to bound its memory (see +/// [`Self::reset_cache`]). +/// +/// Cached statistics are served for as long as the context lives, so a plan +/// node must not change its statistics in place (for example through interior +/// mutability) while a context holds it: the cache would return stale values. +/// Optimizer rules satisfy this because they replace nodes instead of changing +/// them. /// /// An optional [`StatisticsRegistry`] plugs providers into the walk: at each node /// they are consulted before the operator's built-in @@ -173,7 +179,7 @@ pub enum ChildStats { /// [`Self::compute_extended`] observes extensions; [`Self::compute`] returns core /// [`Statistics`] only. pub struct StatisticsContext { - cache: Rc>, + cache: Mutex, registry: StatisticsRegistry, } @@ -192,7 +198,7 @@ impl StatisticsContext { /// 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())), + cache: Mutex::new(StatsCache::default()), registry, } } @@ -200,10 +206,12 @@ impl StatisticsContext { /// Clears the memoization cache and releases its retained plan nodes. /// /// Resetting is optional for correctness: each cache entry retains the plan - /// node that supplied its pointer key. Use it to bound memory at a logical - /// lifecycle boundary, such as after an optimizer pass. + /// node that supplied its pointer key. Use it to bound memory at a + /// lifecycle boundary, such as the end of a query's physical optimization. + /// Do not reset a context that several optimizer rules share between those + /// rules: that discards the statistics that later rules would reuse. pub fn reset_cache(&self) { - let mut cache = self.cache.borrow_mut(); + let mut cache = self.cache.lock(); cache.statistics.clear(); cache.extensions.clear(); } @@ -637,7 +645,7 @@ impl StatisticsContext { partition: Option, ) -> Option> { self.cache - .borrow() + .lock() .statistics .get(&cache_key(plan, partition)) .map(|entry| Arc::clone(&entry.value)) @@ -667,7 +675,7 @@ impl StatisticsContext { statistics: Arc, ) { Self::store_cache_entry( - &mut self.cache.borrow_mut().statistics, + &mut self.cache.lock().statistics, owner, partition, statistics, @@ -680,7 +688,7 @@ impl StatisticsContext { partition: Option, ) -> Option { self.cache - .borrow() + .lock() .extensions .get(&cache_key(plan, partition)) .map(|entry| entry.value.clone()) @@ -693,7 +701,7 @@ impl StatisticsContext { extensions: Extensions, ) { Self::store_cache_entry( - &mut self.cache.borrow_mut().extensions, + &mut self.cache.lock().extensions, owner, partition, extensions, @@ -841,7 +849,7 @@ mod tests { let args = StatisticsArgs::new(); let s1 = ctx.compute_arc(&leaf, &args).unwrap(); - assert!(!ctx.cache.borrow().statistics.is_empty()); + assert!(!ctx.cache.lock().statistics.is_empty()); let s2 = ctx.compute_arc(&leaf, &args).unwrap(); assert!(Arc::ptr_eq(&s1, &s2)); @@ -915,7 +923,7 @@ mod tests { .compute(parent.as_ref(), &StatisticsArgs::new()) .unwrap(); assert!( - !ctx.cache.borrow().statistics.contains_key(&parent_key), + !ctx.cache.lock().statistics.contains_key(&parent_key), "borrowed roots must not be memoized" ); @@ -947,7 +955,7 @@ mod tests { let _ = ctx .compute_extended_arc(&leaf, &StatisticsArgs::new()) .unwrap(); - ctx.cache.borrow_mut().statistics.clear(); + ctx.cache.lock().statistics.clear(); drop(leaf); assert!(weak.upgrade().is_some()); @@ -960,9 +968,9 @@ mod tests { let leaf = make_stats_leaf(10); let ctx = StatisticsContext::new(); let _ = ctx.compute_arc(&leaf, &StatisticsArgs::new()).unwrap(); - assert!(!ctx.cache.borrow().statistics.is_empty()); + assert!(!ctx.cache.lock().statistics.is_empty()); ctx.reset_cache(); - assert!(ctx.cache.borrow().statistics.is_empty()); + assert!(ctx.cache.lock().statistics.is_empty()); } #[test] @@ -1037,7 +1045,7 @@ mod tests { assert_eq!(extensions.unwrap().get::(), Some(&Tag(7))); assert!( ctx.cache - .borrow() + .lock() .extensions .contains_key(&cache_key(leaf.as_ref(), None)) ); diff --git a/datafusion/session/src/lib.rs b/datafusion/session/src/lib.rs index 26c675c20f928..0fc99ad66eb79 100644 --- a/datafusion/session/src/lib.rs +++ b/datafusion/session/src/lib.rs @@ -56,7 +56,9 @@ pub use crate::catalog::{ CatalogProvider, CatalogProviderFactory, CatalogProviderList, EmptyCatalogProviderList, }; -pub use crate::physical_optimizer::{PhysicalOptimizerContext, PhysicalOptimizerRule}; +pub use crate::physical_optimizer::{ + PhysicalOptimizerContext, PhysicalOptimizerRule, with_statistics_context, +}; pub use crate::planner::{ ExtensionPlanner, PhysicalPlanner, QueryPlanner, UnsupportedQueryPlanner, }; diff --git a/datafusion/session/src/physical_optimizer.rs b/datafusion/session/src/physical_optimizer.rs index d59dce24a484f..e0b2b2d2ce0d7 100644 --- a/datafusion/session/src/physical_optimizer.rs +++ b/datafusion/session/src/physical_optimizer.rs @@ -20,10 +20,11 @@ use std::fmt::Debug; use std::sync::Arc; -use datafusion_common::Result; use datafusion_common::config::ConfigOptions; +use datafusion_common::{Result, Statistics}; use datafusion_physical_plan::ExecutionPlan; use datafusion_physical_plan::operator_statistics::StatisticsRegistry; +use datafusion_physical_plan::statistics::{StatisticsArgs, StatisticsContext}; /// Context available to physical optimizer rules. /// @@ -40,6 +41,80 @@ pub trait PhysicalOptimizerContext: Send + Sync { fn statistics_registry(&self) -> Option<&StatisticsRegistry> { None } + + /// Returns a [`StatisticsContext`] shared by every rule of one optimizer + /// run, so statistics computed by one rule are reused by later rules. + /// + /// The context must be built from [`Self::statistics_registry`]. Its cache + /// entries hold the plan nodes they were computed for, so it is safe to + /// share across plan rewrites (see [`StatisticsContext`]). + /// + /// Returns `None` if no shared context is available, in which case rules + /// create their own context. + fn statistics_context(&self) -> Option<&StatisticsContext> { + None + } + + /// Computes the statistics of `plan` with [`Self::statistics_context`], + /// or, if it is `None`, with a new context built from + /// [`Self::statistics_registry`]. + /// + /// Rules that compute statistics should use this, so they reuse the + /// statistics computed by earlier rules and consult the same statistics + /// providers as the built-in rules. + /// + /// # Example + /// + /// ``` + /// # use std::sync::Arc; + /// # use arrow_schema::Schema; + /// # use datafusion_common::config::ConfigOptions; + /// # use datafusion_physical_plan::ExecutionPlan; + /// # use datafusion_physical_plan::empty::EmptyExec; + /// # use datafusion_physical_plan::statistics::StatisticsArgs; + /// # use datafusion_session::PhysicalOptimizerContext; + /// # struct MyContext(ConfigOptions); + /// # impl PhysicalOptimizerContext for MyContext { + /// # fn config_options(&self) -> &ConfigOptions { + /// # &self.0 + /// # } + /// # } + /// # let context = MyContext(ConfigOptions::new()); + /// let plan: Arc = Arc::new(EmptyExec::new(Arc::new(Schema::empty()))); + /// let statistics = context.compute_statistics(&plan, &StatisticsArgs::new())?; + /// # Ok::<(), datafusion_common::DataFusionError>(()) + /// ``` + fn compute_statistics( + &self, + plan: &Arc, + args: &StatisticsArgs, + ) -> Result> { + with_statistics_context(self, |stats_ctx| stats_ctx.compute_arc(plan, args)) + } +} + +/// Calls `f` with [`PhysicalOptimizerContext::statistics_context`], or, if it +/// is `None`, with a new context built from +/// [`PhysicalOptimizerContext::statistics_registry`]. +/// +/// Use this to pass one [`StatisticsContext`] down a traversal. To compute +/// the statistics of a single plan, use +/// [`PhysicalOptimizerContext::compute_statistics`]. +pub fn with_statistics_context( + context: &C, + f: impl FnOnce(&StatisticsContext) -> R, +) -> R +where + C: PhysicalOptimizerContext + ?Sized, +{ + if let Some(shared) = context.statistics_context() { + return f(shared); + } + let local = match context.statistics_registry() { + Some(registry) => StatisticsContext::new_with_registry(registry.clone()), + None => StatisticsContext::new(), + }; + f(&local) } /// `PhysicalOptimizerRule` transforms one [`ExecutionPlan`] into another which 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 11a63a5aa5bca..7cba485bc311f 100644 --- a/docs/source/library-user-guide/upgrading/56.0.0.md +++ b/docs/source/library-user-guide/upgrading/56.0.0.md @@ -770,6 +770,40 @@ let plan = DistributionContext::new_default(plan) Callers that only need the previous behaviour can keep using the deprecated form, or pass a freshly constructed `StatisticsContext` per call. +### Physical optimizer rules share one `StatisticsContext` per query + +`PhysicalOptimizerContext` has a new `statistics_context` method. The physical +planner returns one `StatisticsContext` for all the rules of one +`optimize_physical_plan` call, so statistics computed by one rule are reused by +later rules. The method returns `None` by default, so existing implementations +are not affected. Custom rules that compute statistics should call the new +`PhysicalOptimizerContext::compute_statistics` method in `optimize_with_context`. +It uses the shared context, or, if there is none, a new context built from the +statistics registry: + +```rust,ignore +let statistics = context.compute_statistics(&plan, &StatisticsArgs::new())?; +``` + +To pass one `StatisticsContext` down a traversal, use +`datafusion_session::with_statistics_context`, which applies the same choice. + +`AggregateStatistics`, `LimitPushdown` and the co-partitioning choice in +`EnsureRequirements` now also consult the session's statistics providers, as +`JoinSelection` already did. Sessions without registered providers (the +default) are not affected. As for operators, `Exact` statistics from a provider +must be guaranteed values: `AggregateStatistics` answers `COUNT(*)`, `MIN` and +`MAX` from them, and `LimitPushdown` removes a limit over an `Exact(0)` input. + +For the same reason, +`datafusion_physical_optimizer::limit_pushdown::pushdown_limit_helper` is +deprecated in favour of `pushdown_limit_helper_with_stats`, which takes the +`StatisticsContext` from the caller. The deprecated form creates a new context +on every call, without statistics providers. When providers are registered, +switching to `pushdown_limit_helper_with_stats` with a context built from the +statistics registry can change the result, because the providers are then +consulted. + ### `MergeIntoOp` requires the SQL-visible target qualifier `MergeIntoOp` now stores the SQL-visible target qualifier separately from the