You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/10/12 14:47:54 UTC
(datafusion) branch gh-readonly-queue/main/pr-26094-6663e4348049cc7b70109c9e0f8e39049794ea23 created (now a123121410)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a change to branch gh-readonly-queue/main/pr-26094-6663e4348049cc7b70109c9e0f8e39049794ea23
in repository https://gitbox.apache.org/repos/asf/datafusion.git
at a123121410 perf: share one StatisticsContext across all physical optimizer rules (#26094)
This branch includes the following new commits:
new a123121410 perf: share one StatisticsContext across all physical optimizer rules (#26094)
The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org
(datafusion) 01/01: perf: share one StatisticsContext across all physical optimizer rules (#26094)
Posted by gi...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-26094-6663e4348049cc7b70109c9e0f8e39049794ea23
in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit a123121410afd050797000bfcd45fa57fdb80812
Author: Alessandro Solimando <as...@apache.org>
AuthorDate: Mon Oct 12 14:47:32 2026 +0000
perf: share one StatisticsContext across all physical optimizer rules (#26094)
## Which issue does this PR close?
- Part of #8227.
- Related to #25141.
## Rationale for this change
Each physical optimizer rule that reads statistics creates its own
`StatisticsContext`, so the statistics of the same plan nodes are
computed again in every rule. `JoinSelection` creates a new context for
every `get_stats` call, so it has no caching at all, even within its own
pass.
#25098 showed the gain of sharing one context inside
`EnsureDistribution`. #25929 made cache entries keep the plan node they
were computed for, so a context can now be shared safely across plan
rewrites. This PR uses that to share one `StatisticsContext` across all
the rules of one `optimize_physical_plan` call.
On TPC-DS (98 queries) and TPC-H (21 queries), sf1 Parquet, planned with
a local harness (not part of this PR):
| | `main` | this PR |
|---|---|---|
| Statistics cache misses (statistics computed from scratch) | 46,022 |
25,428 (-45%) |
| Physical planning time, sum over all queries | 294.8 / 293.5 ms |
274.3 / 279.7 ms (about -6%) |
| Physical plans | | identical for all 119 queries |
The largest gain is TPC-DS q64: cache misses go from 3,785 to 531 and
planning time drops by about 21%.
## What changes are included in this PR?
- `StatisticsContext` stores its cache in a `parking_lot::Mutex` instead
of `Rc<RefCell>`, so it is `Send + Sync`. This is needed because
`PhysicalOptimizerContext: Send + Sync`. The lock is held only for
single map lookups and inserts, never across the recursive walk. The
`compute_statistics` benchmark shows no difference from `main` (all
cases within ±4%, in both directions).
- New `PhysicalOptimizerContext::statistics_context()`, which returns
`None` by default. `DefaultPhysicalPlanner` creates one context per
`optimize_physical_plan` call, built from the session's statistics
registry, and returns it to every rule.
- `ConfigOnlyContext` also owns a `StatisticsContext`, so a rule called
through `optimize()` shares one cache for its whole pass.
- `JoinSelection`, `EnsureRequirements` (including
`PlanSize::from_plan`), `AggregateStatistics` and `LimitPushdown` use
the shared context. When a context does not share one (for example, a
context received through FFI), the rules create a new context from the
statistics registry, as they did before.
- New default method `PhysicalOptimizerContext::compute_statistics`,
which custom rules can call on the context they already receive. Without
it, a custom rule would have to repeat the fallback itself (use the
shared context, otherwise build one from `statistics_registry()`), and
the obvious shortcut, `StatisticsContext::new()`, does not consult the
registered statistics providers. With it, custom rules behave exactly
like the built-in ones.
- The fallback itself is the public function `with_statistics_context`.
Rules that walk the plan (`EnsureRequirements`, `LimitPushdown`) use it
to get one `StatisticsContext` for the whole traversal, so a fallback
context still caches within the pass.
- `pushdown_limit_helper` is deprecated in favour of the new
`pushdown_limit_helper_with_stats`, which takes a `&StatisticsContext`.
This is the same pattern as `ensure_distribution_with_stats`.
- The shared cache keeps every plan node whose statistics were computed
alive until `optimize_physical_plan` returns, including nodes that a
later rule replaces. Only the rules that read statistics add entries,
and the memory is released at the end of physical planning.
## What is the testing strategy for this PR?
- New `optimizer_rules_share_statistics_context` test in
`physical_planner.rs`: two rules compute the root statistics through the
shared context, and the second one gets the `Arc` cached by the first.
- Existing tests cover the rewired rules. All sqllogictests pass with no
plan changes.
## Are there any user-facing changes?
No breaking changes.
- New default methods `PhysicalOptimizerContext::statistics_context()`
and `PhysicalOptimizerContext::compute_statistics()`.
- New public functions `datafusion_session::with_statistics_context` and
`pushdown_limit_helper_with_stats`; `pushdown_limit_helper` is
deprecated. All are described in the 56.0.0 upgrade guide.
- `StatisticsContext` is now `Send + Sync`.
- `AggregateStatistics`, `LimitPushdown` and `PlanSize::from_plan` now
consult the session's statistics providers, as `JoinSelection` and
`EnsureRequirements` already do. Without registered providers (the
default), their results are unchanged.
----
Disclaimer: I used AI to assist in the code generation, I have manually
reviewed the output and it matches my intention and understanding.
---------
Co-authored-by: Qi Zhu <zh...@gmail.com>
---
datafusion/core/src/physical_planner.rs | 84 ++++++++++++++-----
.../physical_optimizer/aggregate_statistics.rs | 49 +++++++++++
.../tests/physical_optimizer/limit_pushdown.rs | 48 +++++++++++
.../physical-optimizer/src/aggregate_statistics.rs | 29 +++++--
.../ensure_requirements/enforce_distribution.rs | 97 +++++++++++-----------
.../src/ensure_requirements/mod.rs | 23 +++--
.../physical-optimizer/src/join_selection.rs | 32 +++----
.../physical-optimizer/src/limit_pushdown.rs | 58 +++++++++++--
datafusion/physical-optimizer/src/optimizer.rs | 19 ++++-
.../physical-plan/src/operator_statistics/mod.rs | 5 ++
datafusion/physical-plan/src/statistics.rs | 44 ++++++----
datafusion/session/src/lib.rs | 4 +-
datafusion/session/src/physical_optimizer.rs | 77 ++++++++++++++++-
docs/source/library-user-guide/upgrading/56.0.0.md | 34 ++++++++
14 files changed, 463 insertions(+), 140 deletions(-)
diff --git a/datafusion/core/src/physical_planner.rs b/datafusion/core/src/physical_planner.rs
index b48bc19a55..1fc700a2ec 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,6 +3553,7 @@ 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::{
DFSchemaRef, ScalarValue, SplitPoint, TableReference, ToDFSchema as _,
@@ -3558,12 +3576,16 @@ 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::StatisticsRegistry;
+ 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<AtomicBool>,
+ recorded: Arc<SyncMutex<Vec<Arc<Statistics>>>>,
}
impl PhysicalOptimizerRule for ContextCheckingRule {
@@ -3581,7 +3603,12 @@ mod tests {
context: &dyn PhysicalOptimizerContext,
) -> Result<Arc<dyn ExecutionPlan>> {
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 +3671,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 +3840,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 +3861,31 @@ mod tests {
let physical_plan = session.create_physical_plan(&logical_plan).await?;
assert!(physical_plan.is::<EmptyExec>());
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<dyn PhysicalOptimizerRule + Send + Sync> {
+ 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(())
}
diff --git a/datafusion/core/tests/physical_optimizer/aggregate_statistics.rs b/datafusion/core/tests/physical_optimizer/aggregate_statistics.rs
index 2d22b60856..f346dddfd4 100644
--- a/datafusion/core/tests/physical_optimizer/aggregate_statistics.rs
+++ b/datafusion/core/tests/physical_optimizer/aggregate_statistics.rs
@@ -28,6 +28,7 @@ use datafusion::datasource::memory::MemTable;
use datafusion::datasource::memory::MemorySourceConfig;
use datafusion::datasource::physical_plan::ParquetSource;
use datafusion::datasource::source::DataSourceExec;
+use datafusion::execution::session_state::SessionStateBuilder;
use datafusion::prelude::{SessionConfig, SessionContext};
use datafusion_common::cast::as_int64_array;
use datafusion_common::config::ConfigOptions;
@@ -52,7 +53,11 @@ use datafusion_physical_plan::coalesce_partitions::CoalescePartitionsExec;
use datafusion_physical_plan::common;
use datafusion_physical_plan::displayable;
use datafusion_physical_plan::filter::FilterExec;
+use datafusion_physical_plan::operator_statistics::{
+ ClosureStatisticsProvider, StatisticsRegistry, StatisticsResult,
+};
use datafusion_physical_plan::projection::ProjectionExec;
+use datafusion_physical_plan::statistics::StatisticsArgs;
/// Mock data using a MemorySourceConfig which has an exact count statistic
fn mock_data() -> Result<Arc<DataSourceExec>> {
@@ -856,3 +861,47 @@ async fn test_sum_from_statistics() -> Result<()> {
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::<DataSourceExec>(),
+ |plan, child_stats| {
+ let child_stats = child_stats
+ .iter()
+ .map(|c| Arc::clone(c.base_arc()))
+ .collect::<Vec<_>>();
+ 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/core/tests/physical_optimizer/limit_pushdown.rs b/datafusion/core/tests/physical_optimizer/limit_pushdown.rs
index b8ebc80348..e9c04d1460 100644
--- a/datafusion/core/tests/physical_optimizer/limit_pushdown.rs
+++ b/datafusion/core/tests/physical_optimizer/limit_pushdown.rs
@@ -24,8 +24,11 @@ use crate::physical_optimizer::test_utils::{
use arrow::compute::SortOptions;
use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
+use datafusion::execution::session_state::SessionStateBuilder;
+use datafusion::prelude::{SessionConfig, SessionContext};
use datafusion_common::config::ConfigOptions;
use datafusion_common::error::Result;
+use datafusion_common::stats::Precision;
use datafusion_expr::{JoinType, Operator};
use datafusion_physical_expr::Partitioning;
use datafusion_physical_expr::expressions::{BinaryExpr, col, lit};
@@ -33,11 +36,16 @@ use datafusion_physical_expr_common::physical_expr::PhysicalExprRef;
use datafusion_physical_expr_common::sort_expr::{LexOrdering, PhysicalSortExpr};
use datafusion_physical_optimizer::PhysicalOptimizerRule;
use datafusion_physical_optimizer::limit_pushdown::LimitPushdown;
+use datafusion_physical_plan::aggregates::AggregateExec;
use datafusion_physical_plan::empty::EmptyExec;
use datafusion_physical_plan::filter::FilterExec;
use datafusion_physical_plan::joins::NestedLoopJoinExec;
+use datafusion_physical_plan::operator_statistics::{
+ ClosureStatisticsProvider, StatisticsRegistry, StatisticsResult,
+};
use datafusion_physical_plan::projection::ProjectionExec;
use datafusion_physical_plan::repartition::RepartitionExec;
+use datafusion_physical_plan::statistics::StatisticsArgs;
use datafusion_physical_plan::{ExecutionPlan, get_plan_string};
fn create_schema() -> SchemaRef {
@@ -834,3 +842,43 @@ fn outer_offset_with_same_sort_key_still_pushes_limit() -> Result<()> {
Ok(())
}
+
+/// `LimitPushdown` trusts an `Exact(0)` row count from a registered
+/// provider: it removes a limit over that input
+#[tokio::test]
+async fn limit_pushdown_consults_statistics_providers() -> Result<()> {
+ let provider = ClosureStatisticsProvider::with_matches(
+ |plan| plan.is::<AggregateExec>(),
+ |plan, child_stats| {
+ let child_stats = child_stats
+ .iter()
+ .map(|c| Arc::clone(c.base_arc()))
+ .collect::<Vec<_>>();
+ let mut stats = Arc::unwrap_or_clone(
+ plan.statistics_from_inputs(&child_stats, &StatisticsArgs::new())?,
+ );
+ stats.num_rows = Precision::Exact(0);
+ Ok(StatisticsResult::Computed(stats.into()))
+ },
+ );
+ let state = SessionStateBuilder::new()
+ .with_default_features()
+ .with_config(SessionConfig::new().with_target_partitions(1))
+ .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 column1, count(*) FROM t GROUP BY column1 LIMIT 1")
+ .await?
+ .collect()
+ .await?;
+ assert_eq!(batches.iter().map(|b| b.num_rows()).sum::<usize>(), 3);
+ Ok(())
+}
diff --git a/datafusion/physical-optimizer/src/aggregate_statistics.rs b/datafusion/physical-optimizer/src/aggregate_statistics.rs
index 43b1abb4b6..a126052943 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<dyn ExecutionPlan>,
config: &ConfigOptions,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ self.optimize_with_context(plan, &ConfigOnlyContext::new(config))
+ }
+
+ #[cfg_attr(feature = "recursive_protection", recursive::recursive)]
+ fn optimize_with_context(
+ &self,
+ plan: Arc<dyn ExecutionPlan>,
+ context: &dyn PhysicalOptimizerContext,
) -> Result<Arc<dyn ExecutionPlan>> {
if let Some(partial_agg_exec) = take_optimizable(&plan) {
let partial_agg_exec = partial_agg_exec
.downcast_ref::<AggregateExec>()
.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 3c2426a5ef..31f00fde90 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<dyn ExecutionPlan>, 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 3909b56c17..56ca430d69 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 2b342bece7..c7f414d91b 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<dyn ExecutionPlan>,
context: &dyn PhysicalOptimizerContext,
) -> Result<Arc<Statistics>> {
- 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<dyn ExecutionPlan>,
+ right: &Arc<dyn ExecutionPlan>,
context: &dyn PhysicalOptimizerContext,
) -> Result<bool> {
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<dyn ExecutionPlan>,
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::<CrossJoinExec>() {
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 f88a2be14e..226c41cf8d 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<dyn ExecutionPlan>,
- _config: &ConfigOptions,
+ config: &ConfigOptions,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ self.optimize_with_context(plan, &ConfigOnlyContext::new(config))
+ }
+
+ fn optimize_with_context(
+ &self,
+ plan: Arc<dyn ExecutionPlan>,
+ context: &dyn PhysicalOptimizerContext,
) -> Result<Arc<dyn ExecutionPlan>> {
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<dyn ExecutionPlan>,
+ global_state: GlobalRequirements,
+) -> Result<(Transformed<Arc<dyn ExecutionPlan>>, 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<dyn ExecutionPlan>,
mut global_state: GlobalRequirements,
+ stats_ctx: &StatisticsContext,
) -> Result<(Transformed<Arc<dyn ExecutionPlan>>, 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<dyn ExecutionPlan>,
skip: usize,
fetch: usize,
+ stats_ctx: &StatisticsContext,
) -> Result<bool> {
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<dyn ExecutionPlan>,
+ stats_ctx: &StatisticsContext,
) -> Result<Option<usize>> {
// 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<dyn ExecutionPlan>,
global_state: GlobalRequirements,
+ stats_ctx: &StatisticsContext,
) -> Result<Arc<dyn ExecutionPlan>> {
// 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::<dyn ExecutionPlan>::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 80c33e1c08..81932594b8 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 42cff12903..9696bdf721 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 7124c9107d..c52636b21f 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<usize>);
@@ -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<RefCell<StatsCache>>,
+ cache: Mutex<StatsCache>,
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<usize>,
) -> Option<Arc<Statistics>> {
self.cache
- .borrow()
+ .lock()
.statistics
.get(&cache_key(plan, partition))
.map(|entry| Arc::clone(&entry.value))
@@ -667,7 +675,7 @@ impl StatisticsContext {
statistics: Arc<Statistics>,
) {
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<usize>,
) -> Option<Extensions> {
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::<Tag>(), 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 26c675c20f..0fc99ad66e 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 d59dce24a4..e0b2b2d2ce 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<dyn ExecutionPlan> = 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<dyn ExecutionPlan>,
+ args: &StatisticsArgs,
+ ) -> Result<Arc<Statistics>> {
+ 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<C, R>(
+ 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 4bc75d462d..20c28596ec 100644
--- a/docs/source/library-user-guide/upgrading/56.0.0.md
+++ b/docs/source/library-user-guide/upgrading/56.0.0.md
@@ -784,6 +784,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
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org