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/06/18 19:38:33 UTC
(datafusion) branch gh-readonly-queue/main/pr-21815-54b7dd99bcda9a3015485388713f632cc089a637 created (now 0838a4ddb9)
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-21815-54b7dd99bcda9a3015485388713f632cc089a637
in repository https://gitbox.apache.org/repos/asf/datafusion.git
at 0838a4ddb9 Add StatisticsContext parameter to partition_statistics (#21815)
This branch includes the following new commits:
new 0838a4ddb9 Add StatisticsContext parameter to partition_statistics (#21815)
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: Add StatisticsContext parameter to partition_statistics (#21815)
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-21815-54b7dd99bcda9a3015485388713f632cc089a637
in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit 0838a4ddb902535b0e95a1c5a254be7e9c7fe9bf
Author: Alessandro Solimando <as...@apache.org>
AuthorDate: Thu Jun 18 21:38:23 2026 +0200
Add StatisticsContext parameter to partition_statistics (#21815)
## Which issue does this PR close?
Closes #20184
## Rationale for this change
`ExecutionPlan::partition_statistics` forces each operator to re-fetch
child statistics internally, causing redundant subtree walks in deep
plans.
## What changes are included in this PR?
- Deprecate `partition_statistics` in favor of
`statistics_with_args(&self, args: &StatisticsArgs)`, an extensible
signature that won't require downstream churn when new parameters are
added
- `StatisticsArgs` carries the partition index and a shared per-call
`StatsCache`, eliminating redundant subtree walks within a single
`compute_statistics` call
- Child stats are pre-computed with `partition=None` and cached;
operators look them up via `args.child_stats_of(child)` (overall) or
`args.child_stats_for(child)` (partition-aware)
- Criterion micro-benchmark on three plan shapes from #19795
## Tests
Existing tests pass unchanged. New unit test verifies the caching
contract.
## Test plan
- [x] `cargo fmt --all`
- [x] `cargo clippy --all-targets --all-features -- -D warnings`
- [x] `cargo test --profile ci --all-features` on affected crates
- [x] Criterion benchmark: ~26x (coalesce chain), ~5x (cross-join tree),
~25x (filter chain) speedup
----
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: Andrew Lamb <an...@nerdnetworks.org>
Co-authored-by: Claude Opus 4.8 (1M context) <no...@anthropic.com>
---
.../examples/relation_planner/table_sample.rs | 8 +-
datafusion/core/src/datasource/file_format/csv.rs | 9 +-
datafusion/core/src/datasource/file_format/json.rs | 9 +-
.../core/src/datasource/file_format/parquet.rs | 11 +-
datafusion/core/src/datasource/listing/table.rs | 30 +-
datafusion/core/tests/custom_sources_cases/mod.rs | 5 +-
.../core/tests/custom_sources_cases/statistics.rs | 21 +-
datafusion/core/tests/parquet/file_statistics.rs | 40 ++-
.../tests/physical_optimizer/join_selection.rs | 34 +--
.../physical_optimizer/partition_statistics.rs | 186 +++++++++---
.../core/tests/physical_optimizer/test_utils.rs | 4 +-
datafusion/core/tests/sql/path_partition.rs | 5 +-
datafusion/datasource/src/file_scan_config/mod.rs | 6 +-
datafusion/datasource/src/memory.rs | 3 +-
datafusion/datasource/src/source.rs | 5 +-
datafusion/ffi/src/execution_plan.rs | 112 ++++++--
datafusion/ffi/tests/ffi_execution_plan.rs | 1 +
.../physical-optimizer/src/aggregate_statistics.rs | 13 +-
.../ensure_requirements/enforce_distribution.rs | 5 +-
.../physical-optimizer/src/join_selection.rs | 3 +-
.../physical-optimizer/src/limit_pushdown.rs | 5 +-
.../physical-optimizer/src/output_requirements.rs | 5 +-
datafusion/physical-plan/Cargo.toml | 8 +-
.../physical-plan/benches/compute_statistics.rs | 312 +++++++++++++++++++++
datafusion/physical-plan/src/aggregates/mod.rs | 47 ++--
datafusion/physical-plan/src/buffer.rs | 5 +-
datafusion/physical-plan/src/coalesce_batches.rs | 7 +-
.../physical-plan/src/coalesce_partitions.rs | 6 +-
datafusion/physical-plan/src/coop.rs | 5 +-
datafusion/physical-plan/src/display.rs | 18 +-
datafusion/physical-plan/src/empty.rs | 5 +-
datafusion/physical-plan/src/execution_plan.rs | 32 ++-
datafusion/physical-plan/src/filter.rs | 79 +++---
datafusion/physical-plan/src/joins/cross_join.rs | 14 +-
.../physical-plan/src/joins/hash_join/exec.rs | 56 ++--
.../physical-plan/src/joins/nested_loop_join.rs | 22 +-
.../src/joins/sort_merge_join/exec.rs | 16 +-
.../src/joins/sort_merge_join/tests.rs | 7 +-
datafusion/physical-plan/src/lib.rs | 2 +
datafusion/physical-plan/src/limit.rs | 26 +-
.../physical-plan/src/operator_statistics/mod.rs | 15 +-
datafusion/physical-plan/src/placeholder_row.rs | 5 +-
datafusion/physical-plan/src/projection.rs | 13 +-
datafusion/physical-plan/src/repartition/mod.rs | 10 +-
datafusion/physical-plan/src/scalar_subquery.rs | 5 +-
datafusion/physical-plan/src/sorts/partial_sort.rs | 5 +-
datafusion/physical-plan/src/sorts/sort.rs | 12 +-
.../src/sorts/sort_preserving_merge.rs | 5 +-
datafusion/physical-plan/src/statistics.rs | 229 +++++++++++++++
datafusion/physical-plan/src/test.rs | 5 +-
datafusion/physical-plan/src/test/exec.rs | 14 +-
datafusion/physical-plan/src/union.rs | 77 ++---
.../src/windows/bounded_window_agg_exec.rs | 8 +-
.../physical-plan/src/windows/window_agg_exec.rs | 8 +-
datafusion/physical-plan/src/work_table.rs | 3 +-
docs/source/library-user-guide/upgrading/55.0.0.md | 71 +++++
56 files changed, 1318 insertions(+), 354 deletions(-)
diff --git a/datafusion-examples/examples/relation_planner/table_sample.rs b/datafusion-examples/examples/relation_planner/table_sample.rs
index 46826216e2..6df1113e47 100644
--- a/datafusion-examples/examples/relation_planner/table_sample.rs
+++ b/datafusion-examples/examples/relation_planner/table_sample.rs
@@ -108,7 +108,7 @@ use datafusion::{
},
physical_expr::EquivalenceProperties,
physical_plan::{
- DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties,
+ DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, StatisticsArgs,
metrics::{BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet, RecordOutput},
},
physical_planner::{DefaultPhysicalPlanner, ExtensionPlanner, PhysicalPlanner},
@@ -722,8 +722,10 @@ impl ExecutionPlan for SampleExec {
Some(self.metrics.clone_inner())
}
- fn partition_statistics(&self, partition: Option<usize>) -> Result<Arc<Statistics>> {
- let mut stats = Arc::unwrap_or_clone(self.input.partition_statistics(partition)?);
+ fn statistics_with_args(&self, args: &StatisticsArgs) -> Result<Arc<Statistics>> {
+ let mut stats = Arc::unwrap_or_clone(
+ args.compute_child_statistics(&self.input, args.partition())?,
+ );
let ratio = self.upper_bound - self.lower_bound;
// Scale statistics by sampling ratio (inexact due to randomness)
diff --git a/datafusion/core/src/datasource/file_format/csv.rs b/datafusion/core/src/datasource/file_format/csv.rs
index a068b4f5c0..9392d6daec 100644
--- a/datafusion/core/src/datasource/file_format/csv.rs
+++ b/datafusion/core/src/datasource/file_format/csv.rs
@@ -45,6 +45,7 @@ mod tests {
use datafusion_datasource::file_format::FileFormat;
use datafusion_datasource::write::BatchSerializer;
use datafusion_expr::{col, lit};
+ use datafusion_physical_plan::statistics::StatisticsArgs;
use datafusion_physical_plan::{ExecutionPlan, collect};
use arrow::array::{
@@ -215,9 +216,13 @@ mod tests {
assert_eq!(tt_batches, 50 /* 100/2 */);
// test metadata
- assert_eq!(exec.partition_statistics(None)?.num_rows, Precision::Absent);
assert_eq!(
- exec.partition_statistics(None)?.total_byte_size,
+ exec.statistics_with_args(&StatisticsArgs::new())?.num_rows,
+ Precision::Absent
+ );
+ assert_eq!(
+ exec.statistics_with_args(&StatisticsArgs::new())?
+ .total_byte_size,
Precision::Absent
);
diff --git a/datafusion/core/src/datasource/file_format/json.rs b/datafusion/core/src/datasource/file_format/json.rs
index 5b3e227056..5dd3817829 100644
--- a/datafusion/core/src/datasource/file_format/json.rs
+++ b/datafusion/core/src/datasource/file_format/json.rs
@@ -36,6 +36,7 @@ mod tests {
BatchDeserializer, DecoderDeserializer, DeserializerOutput,
};
use datafusion_datasource::file_format::FileFormat;
+ use datafusion_physical_plan::statistics::StatisticsArgs;
use datafusion_physical_plan::{ExecutionPlan, collect};
use arrow::compute::concat_batches;
@@ -117,9 +118,13 @@ mod tests {
assert_eq!(tt_batches, 6 /* 12/2 */);
// test metadata
- assert_eq!(exec.partition_statistics(None)?.num_rows, Precision::Absent);
assert_eq!(
- exec.partition_statistics(None)?.total_byte_size,
+ exec.statistics_with_args(&StatisticsArgs::new())?.num_rows,
+ Precision::Absent
+ );
+ assert_eq!(
+ exec.statistics_with_args(&StatisticsArgs::new())?
+ .total_byte_size,
Precision::Absent
);
diff --git a/datafusion/core/src/datasource/file_format/parquet.rs b/datafusion/core/src/datasource/file_format/parquet.rs
index c977deab32..5f7fc2eebf 100644
--- a/datafusion/core/src/datasource/file_format/parquet.rs
+++ b/datafusion/core/src/datasource/file_format/parquet.rs
@@ -141,6 +141,7 @@ mod tests {
use datafusion_execution::object_store::ObjectStoreUrl;
use datafusion_execution::runtime_env::RuntimeEnv;
use datafusion_expr::dml::InsertOp;
+ use datafusion_physical_plan::statistics::StatisticsArgs;
use datafusion_physical_plan::stream::RecordBatchStreamAdapter;
use datafusion_physical_plan::{ExecutionPlan, collect};
@@ -715,12 +716,13 @@ mod tests {
// test metadata
assert_eq!(
- exec.partition_statistics(None)?.num_rows,
+ exec.statistics_with_args(&StatisticsArgs::new())?.num_rows,
Precision::Exact(8)
);
// TODO correct byte size: https://github.com/apache/datafusion/issues/14936
assert_eq!(
- exec.partition_statistics(None)?.total_byte_size,
+ exec.statistics_with_args(&StatisticsArgs::new())?
+ .total_byte_size,
Precision::Absent,
);
@@ -764,11 +766,12 @@ mod tests {
// note: even if the limit is set, the executor rounds up to the batch size
assert_eq!(
- exec.partition_statistics(None)?.num_rows,
+ exec.statistics_with_args(&StatisticsArgs::new())?.num_rows,
Precision::Exact(8)
);
assert_eq!(
- exec.partition_statistics(None)?.total_byte_size,
+ exec.statistics_with_args(&StatisticsArgs::new())?
+ .total_byte_size,
Precision::Absent,
);
let batches = collect(exec, task_ctx).await?;
diff --git a/datafusion/core/src/datasource/listing/table.rs b/datafusion/core/src/datasource/listing/table.rs
index 8c543fdd0a..50b3855a0a 100644
--- a/datafusion/core/src/datasource/listing/table.rs
+++ b/datafusion/core/src/datasource/listing/table.rs
@@ -142,6 +142,7 @@ mod tests {
use datafusion_physical_expr::expressions::binary;
use datafusion_physical_expr_common::sort_expr::LexOrdering;
use datafusion_physical_plan::empty::EmptyExec;
+ use datafusion_physical_plan::statistics::StatisticsArgs;
use datafusion_physical_plan::{ExecutionPlanProperties, collect};
use std::collections::HashMap;
use std::io::Write;
@@ -245,11 +246,12 @@ mod tests {
// test metadata
assert_eq!(
- exec.partition_statistics(None)?.num_rows,
+ exec.statistics_with_args(&StatisticsArgs::new())?.num_rows,
Precision::Exact(8)
);
assert_eq!(
- exec.partition_statistics(None)?.total_byte_size,
+ exec.statistics_with_args(&StatisticsArgs::new())?
+ .total_byte_size,
Precision::Absent,
);
@@ -1360,13 +1362,17 @@ mod tests {
let exec_default = table_default.scan(&state, None, &[], None).await?;
assert_eq!(
- exec_default.partition_statistics(None)?.num_rows,
+ exec_default
+ .statistics_with_args(&StatisticsArgs::new())?
+ .num_rows,
Precision::Exact(8)
);
// TODO correct byte size: https://github.com/apache/datafusion/issues/14936
assert_eq!(
- exec_default.partition_statistics(None)?.total_byte_size,
+ exec_default
+ .statistics_with_args(&StatisticsArgs::new())?
+ .total_byte_size,
Precision::Absent
);
@@ -1382,11 +1388,15 @@ mod tests {
let exec_disabled = table_disabled.scan(&state, None, &[], None).await?;
assert_eq!(
- exec_disabled.partition_statistics(None)?.num_rows,
+ exec_disabled
+ .statistics_with_args(&StatisticsArgs::new())?
+ .num_rows,
Precision::Absent
);
assert_eq!(
- exec_disabled.partition_statistics(None)?.total_byte_size,
+ exec_disabled
+ .statistics_with_args(&StatisticsArgs::new())?
+ .total_byte_size,
Precision::Absent
);
@@ -1402,12 +1412,16 @@ mod tests {
let exec_enabled = table_enabled.scan(&state, None, &[], None).await?;
assert_eq!(
- exec_enabled.partition_statistics(None)?.num_rows,
+ exec_enabled
+ .statistics_with_args(&StatisticsArgs::new())?
+ .num_rows,
Precision::Exact(8)
);
// TODO correct byte size: https://github.com/apache/datafusion/issues/14936
assert_eq!(
- exec_enabled.partition_statistics(None)?.total_byte_size,
+ exec_enabled
+ .statistics_with_args(&StatisticsArgs::new())?
+ .total_byte_size,
Precision::Absent,
);
diff --git a/datafusion/core/tests/custom_sources_cases/mod.rs b/datafusion/core/tests/custom_sources_cases/mod.rs
index 06b3701cbe..0b0df57e5a 100644
--- a/datafusion/core/tests/custom_sources_cases/mod.rs
+++ b/datafusion/core/tests/custom_sources_cases/mod.rs
@@ -40,6 +40,7 @@ use datafusion_common::project_schema;
use datafusion_common::stats::Precision;
use datafusion_physical_expr::EquivalenceProperties;
use datafusion_physical_plan::PlanProperties;
+use datafusion_physical_plan::StatisticsArgs;
use datafusion_physical_plan::execution_plan::{Boundedness, EmissionType};
use datafusion_physical_plan::placeholder_row::PlaceholderRowExec;
@@ -178,8 +179,8 @@ impl ExecutionPlan for CustomExecutionPlan {
Ok(Box::pin(TestCustomRecordBatchStream { nb_batch: 1 }))
}
- fn partition_statistics(&self, partition: Option<usize>) -> Result<Arc<Statistics>> {
- if partition.is_some() {
+ fn statistics_with_args(&self, args: &StatisticsArgs) -> Result<Arc<Statistics>> {
+ if args.partition().is_some() {
return Ok(Arc::new(Statistics::new_unknown(&self.schema())));
}
let batch = TEST_CUSTOM_RECORD_BATCH!().unwrap();
diff --git a/datafusion/core/tests/custom_sources_cases/statistics.rs b/datafusion/core/tests/custom_sources_cases/statistics.rs
index 14406c2316..1ea2b202b1 100644
--- a/datafusion/core/tests/custom_sources_cases/statistics.rs
+++ b/datafusion/core/tests/custom_sources_cases/statistics.rs
@@ -35,6 +35,7 @@ use datafusion::{
use datafusion_catalog::Session;
use datafusion_common::{project_schema, stats::Precision};
use datafusion_physical_expr::EquivalenceProperties;
+use datafusion_physical_plan::StatisticsArgs;
use datafusion_physical_plan::execution_plan::{Boundedness, EmissionType};
use async_trait::async_trait;
@@ -173,8 +174,8 @@ impl ExecutionPlan for StatisticsValidation {
unimplemented!("This plan only serves for testing statistics")
}
- fn partition_statistics(&self, partition: Option<usize>) -> Result<Arc<Statistics>> {
- if partition.is_some() {
+ fn statistics_with_args(&self, args: &StatisticsArgs) -> Result<Arc<Statistics>> {
+ if args.partition().is_some() {
Ok(Arc::new(Statistics::new_unknown(&self.schema)))
} else {
Ok(Arc::new(self.stats.clone()))
@@ -230,7 +231,10 @@ async fn sql_basic() -> Result<()> {
let physical_plan = df.create_physical_plan().await.unwrap();
// the statistics should be those of the source
- assert_eq!(stats, *physical_plan.partition_statistics(None)?);
+ assert_eq!(
+ stats,
+ *physical_plan.statistics_with_args(&StatisticsArgs::new())?
+ );
Ok(())
}
@@ -246,7 +250,7 @@ async fn sql_filter() -> Result<()> {
.unwrap();
let physical_plan = df.create_physical_plan().await.unwrap();
- let stats = physical_plan.partition_statistics(None)?;
+ let stats = physical_plan.statistics_with_args(&StatisticsArgs::new())?;
assert_eq!(stats.num_rows, Precision::Inexact(7));
Ok(())
@@ -261,7 +265,7 @@ async fn sql_limit() -> Result<()> {
let physical_plan = df.create_physical_plan().await.unwrap();
// when the limit is smaller than the original number of lines we mark the statistics as inexact
// and cap NDV at the new row count
- let limit_stats = physical_plan.partition_statistics(None)?;
+ let limit_stats = physical_plan.statistics_with_args(&StatisticsArgs::new())?;
assert_eq!(limit_stats.num_rows, Precision::Exact(5));
// c1: NDV=2 stays at 2 (already below limit of 5)
assert_eq!(
@@ -280,7 +284,10 @@ async fn sql_limit() -> Result<()> {
.unwrap();
let physical_plan = df.create_physical_plan().await.unwrap();
// when the limit is larger than the original number of lines, statistics remain unchanged
- assert_eq!(stats, *physical_plan.partition_statistics(None)?);
+ assert_eq!(
+ stats,
+ *physical_plan.statistics_with_args(&StatisticsArgs::new())?
+ );
Ok(())
}
@@ -297,7 +304,7 @@ async fn sql_window() -> Result<()> {
let physical_plan = df.create_physical_plan().await.unwrap();
- let result = physical_plan.partition_statistics(None)?;
+ let result = physical_plan.statistics_with_args(&StatisticsArgs::new())?;
assert_eq!(stats.num_rows, result.num_rows);
let col_stats = &result.column_statistics;
diff --git a/datafusion/core/tests/parquet/file_statistics.rs b/datafusion/core/tests/parquet/file_statistics.rs
index 900bb4d239..45c0b66a6c 100644
--- a/datafusion/core/tests/parquet/file_statistics.rs
+++ b/datafusion/core/tests/parquet/file_statistics.rs
@@ -44,6 +44,7 @@ use datafusion_physical_optimizer::PhysicalOptimizerRule;
use datafusion_physical_optimizer::filter_pushdown::FilterPushdown;
use datafusion_physical_plan::ExecutionPlan;
use datafusion_physical_plan::filter::FilterExec;
+use datafusion_physical_plan::statistics::StatisticsArgs;
use tempfile::tempdir;
#[tokio::test]
@@ -63,7 +64,9 @@ async fn check_stats_precision_with_filter_pushdown() {
// Scan without filter, stats are exact
let exec = table.scan(&state, None, &[], None).await.unwrap();
assert_eq!(
- exec.partition_statistics(None).unwrap().num_rows,
+ exec.statistics_with_args(&StatisticsArgs::new())
+ .unwrap()
+ .num_rows,
Precision::Exact(8),
"Stats without filter should be exact"
);
@@ -95,7 +98,10 @@ async fn check_stats_precision_with_filter_pushdown() {
);
// Scan with filter pushdown, stats are inexact
assert_eq!(
- optimized_exec.partition_statistics(None).unwrap().num_rows,
+ optimized_exec
+ .statistics_with_args(&StatisticsArgs::new())
+ .unwrap()
+ .num_rows,
Precision::Inexact(8),
"Stats after filter pushdown should be inexact"
);
@@ -126,11 +132,17 @@ async fn load_table_stats_with_session_level_cache() {
let exec1 = table1.scan(&state1, None, &[], None).await.unwrap();
assert_eq!(
- exec1.partition_statistics(None).unwrap().num_rows,
+ exec1
+ .statistics_with_args(&StatisticsArgs::new())
+ .unwrap()
+ .num_rows,
Precision::Exact(8)
);
assert_eq!(
- exec1.partition_statistics(None).unwrap().total_byte_size,
+ exec1
+ .statistics_with_args(&StatisticsArgs::new())
+ .unwrap()
+ .total_byte_size,
// Byte size is absent because we cannot estimate the output size
// of the Arrow data since there are variable length columns.
Precision::Absent,
@@ -142,11 +154,17 @@ async fn load_table_stats_with_session_level_cache() {
assert_eq!(get_static_cache_size(&state2), 0);
let exec2 = table2.scan(&state2, None, &[], None).await.unwrap();
assert_eq!(
- exec2.partition_statistics(None).unwrap().num_rows,
+ exec2
+ .statistics_with_args(&StatisticsArgs::new())
+ .unwrap()
+ .num_rows,
Precision::Exact(8)
);
assert_eq!(
- exec2.partition_statistics(None).unwrap().total_byte_size,
+ exec2
+ .statistics_with_args(&StatisticsArgs::new())
+ .unwrap()
+ .total_byte_size,
// Absent because the data contains variable length columns
Precision::Absent,
);
@@ -157,11 +175,17 @@ async fn load_table_stats_with_session_level_cache() {
assert_eq!(get_static_cache_size(&state1), 1);
let exec3 = table1.scan(&state1, None, &[], None).await.unwrap();
assert_eq!(
- exec3.partition_statistics(None).unwrap().num_rows,
+ exec3
+ .statistics_with_args(&StatisticsArgs::new())
+ .unwrap()
+ .num_rows,
Precision::Exact(8)
);
assert_eq!(
- exec3.partition_statistics(None).unwrap().total_byte_size,
+ exec3
+ .statistics_with_args(&StatisticsArgs::new())
+ .unwrap()
+ .total_byte_size,
// Absent because the data contains variable length columns
Precision::Absent,
);
diff --git a/datafusion/core/tests/physical_optimizer/join_selection.rs b/datafusion/core/tests/physical_optimizer/join_selection.rs
index 29a2b59e57..80e0a3f23e 100644
--- a/datafusion/core/tests/physical_optimizer/join_selection.rs
+++ b/datafusion/core/tests/physical_optimizer/join_selection.rs
@@ -44,7 +44,7 @@ use datafusion_physical_plan::joins::utils::JoinFilter;
use datafusion_physical_plan::joins::{HashJoinExec, NestedLoopJoinExec, PartitionMode};
use datafusion_physical_plan::projection::ProjectionExec;
use datafusion_physical_plan::{
- DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties,
+ DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, StatisticsArgs,
execution_plan::{Boundedness, EmissionType},
};
@@ -250,7 +250,7 @@ async fn test_join_with_swap() {
assert_eq!(
swapped_join
.left()
- .partition_statistics(None)
+ .statistics_with_args(&StatisticsArgs::new())
.unwrap()
.total_byte_size,
Precision::Inexact(8192)
@@ -258,7 +258,7 @@ async fn test_join_with_swap() {
assert_eq!(
swapped_join
.right()
- .partition_statistics(None)
+ .statistics_with_args(&StatisticsArgs::new())
.unwrap()
.total_byte_size,
Precision::Inexact(2097152)
@@ -298,7 +298,7 @@ async fn test_left_join_no_swap() {
assert_eq!(
swapped_join
.left()
- .partition_statistics(None)
+ .statistics_with_args(&StatisticsArgs::new())
.unwrap()
.total_byte_size,
Precision::Inexact(8192)
@@ -306,7 +306,7 @@ async fn test_left_join_no_swap() {
assert_eq!(
swapped_join
.right()
- .partition_statistics(None)
+ .statistics_with_args(&StatisticsArgs::new())
.unwrap()
.total_byte_size,
Precision::Inexact(2097152)
@@ -349,7 +349,7 @@ async fn test_join_with_swap_semi() {
assert_eq!(
swapped_join
.left()
- .partition_statistics(None)
+ .statistics_with_args(&StatisticsArgs::new())
.unwrap()
.total_byte_size,
Precision::Inexact(8192)
@@ -357,7 +357,7 @@ async fn test_join_with_swap_semi() {
assert_eq!(
swapped_join
.right()
- .partition_statistics(None)
+ .statistics_with_args(&StatisticsArgs::new())
.unwrap()
.total_byte_size,
Precision::Inexact(2097152)
@@ -402,7 +402,7 @@ async fn test_join_with_swap_mark() {
assert_eq!(
swapped_join
.left()
- .partition_statistics(None)
+ .statistics_with_args(&StatisticsArgs::new())
.unwrap()
.total_byte_size,
Precision::Inexact(8192)
@@ -410,7 +410,7 @@ async fn test_join_with_swap_mark() {
assert_eq!(
swapped_join
.right()
- .partition_statistics(None)
+ .statistics_with_args(&StatisticsArgs::new())
.unwrap()
.total_byte_size,
Precision::Inexact(2097152)
@@ -530,7 +530,7 @@ async fn test_join_no_swap() {
assert_eq!(
swapped_join
.left()
- .partition_statistics(None)
+ .statistics_with_args(&StatisticsArgs::new())
.unwrap()
.total_byte_size,
Precision::Inexact(8192)
@@ -538,7 +538,7 @@ async fn test_join_no_swap() {
assert_eq!(
swapped_join
.right()
- .partition_statistics(None)
+ .statistics_with_args(&StatisticsArgs::new())
.unwrap()
.total_byte_size,
Precision::Inexact(2097152)
@@ -605,7 +605,7 @@ async fn test_nl_join_with_swap(join_type: JoinType) {
assert_eq!(
swapped_join
.left()
- .partition_statistics(None)
+ .statistics_with_args(&StatisticsArgs::new())
.unwrap()
.total_byte_size,
Precision::Inexact(8192)
@@ -613,7 +613,7 @@ async fn test_nl_join_with_swap(join_type: JoinType) {
assert_eq!(
swapped_join
.right()
- .partition_statistics(None)
+ .statistics_with_args(&StatisticsArgs::new())
.unwrap()
.total_byte_size,
Precision::Inexact(2097152)
@@ -678,7 +678,7 @@ async fn test_nl_join_with_swap_no_proj(join_type: JoinType) {
assert_eq!(
swapped_join
.left()
- .partition_statistics(None)
+ .statistics_with_args(&StatisticsArgs::new())
.unwrap()
.total_byte_size,
Precision::Inexact(8192)
@@ -686,7 +686,7 @@ async fn test_nl_join_with_swap_no_proj(join_type: JoinType) {
assert_eq!(
swapped_join
.right()
- .partition_statistics(None)
+ .statistics_with_args(&StatisticsArgs::new())
.unwrap()
.total_byte_size,
Precision::Inexact(2097152)
@@ -1152,8 +1152,8 @@ impl ExecutionPlan for StatisticsExec {
unimplemented!("This plan only serves for testing statistics")
}
- fn partition_statistics(&self, partition: Option<usize>) -> Result<Arc<Statistics>> {
- Ok(Arc::new(if partition.is_some() {
+ fn statistics_with_args(&self, args: &StatisticsArgs) -> Result<Arc<Statistics>> {
+ Ok(Arc::new(if args.partition().is_some() {
Statistics::new_unknown(&self.schema)
} else {
self.stats.clone()
diff --git a/datafusion/core/tests/physical_optimizer/partition_statistics.rs b/datafusion/core/tests/physical_optimizer/partition_statistics.rs
index e1bf22201d..6a79c668bd 100644
--- a/datafusion/core/tests/physical_optimizer/partition_statistics.rs
+++ b/datafusion/core/tests/physical_optimizer/partition_statistics.rs
@@ -55,6 +55,7 @@ mod test {
use datafusion_physical_plan::projection::{ProjectionExec, ProjectionExpr};
use datafusion_physical_plan::repartition::RepartitionExec;
use datafusion_physical_plan::sorts::sort::SortExec;
+ use datafusion_physical_plan::statistics::StatisticsArgs;
use datafusion_physical_plan::union::{InterleaveExec, UnionExec};
use datafusion_physical_plan::windows::{WindowAggExec, create_window_expr};
use datafusion_physical_plan::{
@@ -238,7 +239,11 @@ mod test {
async fn test_statistics_by_partition_of_data_source() -> Result<()> {
let scan = create_scan_exec_with_statistics(None, Some(2)).await;
let statistics = (0..scan.output_partitioning().partition_count())
- .map(|idx| scan.partition_statistics(Some(idx)))
+ .map(|idx| {
+ scan.statistics_with_args(
+ &StatisticsArgs::new().with_partition(Some(idx)),
+ )
+ })
.collect::<Result<Vec<_>>>()?;
// Partition 1: ids [3,4], dates [2025-03-01, 2025-03-02]
let expected_statistic_partition_1 = create_partition_statistics(
@@ -282,7 +287,11 @@ mod test {
let projection: Arc<dyn ExecutionPlan> =
Arc::new(ProjectionExec::try_new(exprs, scan)?);
let statistics = (0..projection.output_partitioning().partition_count())
- .map(|idx| projection.partition_statistics(Some(idx)))
+ .map(|idx| {
+ projection.statistics_with_args(
+ &StatisticsArgs::new().with_partition(Some(idx)),
+ )
+ })
.collect::<Result<Vec<_>>>()?;
// Projection only includes id column, not the date partition column
let expected_statistic_partition_1 =
@@ -314,7 +323,11 @@ mod test {
let sort = SortExec::new(ordering.clone().into(), scan_1);
let sort_exec: Arc<dyn ExecutionPlan> = Arc::new(sort);
let statistics = (0..sort_exec.output_partitioning().partition_count())
- .map(|idx| sort_exec.partition_statistics(Some(idx)))
+ .map(|idx| {
+ sort_exec.statistics_with_args(
+ &StatisticsArgs::new().with_partition(Some(idx)),
+ )
+ })
.collect::<Result<Vec<_>>>()?;
// All 4 files merged: ids [1-4], dates [2025-03-01, 2025-03-04]
let expected_statistic_partition = create_partition_statistics(
@@ -353,7 +366,11 @@ mod test {
Some((DATE_2025_03_03, DATE_2025_03_04)),
);
let statistics = (0..sort_exec.output_partitioning().partition_count())
- .map(|idx| sort_exec.partition_statistics(Some(idx)))
+ .map(|idx| {
+ sort_exec.statistics_with_args(
+ &StatisticsArgs::new().with_partition(Some(idx)),
+ )
+ })
.collect::<Result<Vec<_>>>()?;
assert_eq!(statistics.len(), 2);
assert_eq!(*statistics[0], expected_statistic_partition_1);
@@ -380,7 +397,7 @@ mod test {
)?;
let filter: Arc<dyn ExecutionPlan> =
Arc::new(FilterExec::try_new(predicate, scan)?);
- let full_statistics = filter.partition_statistics(None)?;
+ let full_statistics = filter.statistics_with_args(&StatisticsArgs::new())?;
let expected_full_statistic = Statistics {
num_rows: Precision::Inexact(0),
total_byte_size: Precision::Inexact(0),
@@ -406,7 +423,11 @@ mod test {
assert_eq!(*full_statistics, expected_full_statistic);
let statistics = (0..filter.output_partitioning().partition_count())
- .map(|idx| filter.partition_statistics(Some(idx)))
+ .map(|idx| {
+ filter.statistics_with_args(
+ &StatisticsArgs::new().with_partition(Some(idx)),
+ )
+ })
.collect::<Result<Vec<_>>>()?;
assert_eq!(statistics.len(), 2);
let expected_partition_statistic = Statistics {
@@ -442,7 +463,11 @@ mod test {
let union_exec: Arc<dyn ExecutionPlan> =
UnionExec::try_new(vec![scan.clone(), scan])?;
let statistics = (0..union_exec.output_partitioning().partition_count())
- .map(|idx| union_exec.partition_statistics(Some(idx)))
+ .map(|idx| {
+ union_exec.statistics_with_args(
+ &StatisticsArgs::new().with_partition(Some(idx)),
+ )
+ })
.collect::<Result<Vec<_>>>()?;
// Check that we have 4 partitions (2 from each scan)
assert_eq!(statistics.len(), 4);
@@ -505,7 +530,11 @@ mod test {
// Verify the result of partition statistics
let stats = (0..interleave.output_partitioning().partition_count())
- .map(|idx| interleave.partition_statistics(Some(idx)))
+ .map(|idx| {
+ interleave.statistics_with_args(
+ &StatisticsArgs::new().with_partition(Some(idx)),
+ )
+ })
.collect::<Result<Vec<_>>>()?;
assert_eq!(stats.len(), 2);
@@ -551,7 +580,11 @@ mod test {
let cross_join: Arc<dyn ExecutionPlan> =
Arc::new(CrossJoinExec::new(left_scan, right_scan));
let statistics = (0..cross_join.output_partitioning().partition_count())
- .map(|idx| cross_join.partition_statistics(Some(idx)))
+ .map(|idx| {
+ cross_join.statistics_with_args(
+ &StatisticsArgs::new().with_partition(Some(idx)),
+ )
+ })
.collect::<Result<Vec<_>>>()?;
// Check that we have 2 partitions
assert_eq!(statistics.len(), 2);
@@ -658,7 +691,8 @@ mod test {
// Test partition_statistics(None) - returns overall statistics
// For RightSemi join, output columns come from right side only
- let full_statistics = nested_loop_join.partition_statistics(None)?;
+ let full_statistics =
+ nested_loop_join.statistics_with_args(&StatisticsArgs::new())?;
// With empty join columns, estimate_join_statistics returns Inexact row count
// based on the outer side (right side for RightSemi)
let expected_full_statistics = create_partition_statistics(
@@ -693,7 +727,11 @@ mod test {
.to_inexact();
let statistics = (0..nested_loop_join.output_partitioning().partition_count())
- .map(|idx| nested_loop_join.partition_statistics(Some(idx)))
+ .map(|idx| {
+ nested_loop_join.statistics_with_args(
+ &StatisticsArgs::new().with_partition(Some(idx)),
+ )
+ })
.collect::<Result<Vec<_>>>()?;
assert_eq!(statistics.len(), 2);
assert_eq!(*statistics[0], expected_statistic_partition_1);
@@ -723,7 +761,11 @@ mod test {
Some((DATE_2025_03_01, DATE_2025_03_04)),
);
let statistics = (0..coalesce_partitions.output_partitioning().partition_count())
- .map(|idx| coalesce_partitions.partition_statistics(Some(idx)))
+ .map(|idx| {
+ coalesce_partitions.statistics_with_args(
+ &StatisticsArgs::new().with_partition(Some(idx)),
+ )
+ })
.collect::<Result<Vec<_>>>()?;
assert_eq!(statistics.len(), 1);
assert_eq!(*statistics[0], expected_statistic_partition);
@@ -740,7 +782,11 @@ mod test {
let local_limit: Arc<dyn ExecutionPlan> =
Arc::new(LocalLimitExec::new(scan.clone(), 1));
let statistics = (0..local_limit.output_partitioning().partition_count())
- .map(|idx| local_limit.partition_statistics(Some(idx)))
+ .map(|idx| {
+ local_limit.statistics_with_args(
+ &StatisticsArgs::new().with_partition(Some(idx)),
+ )
+ })
.collect::<Result<Vec<_>>>()?;
assert_eq!(statistics.len(), 2);
let mut expected_0 = Statistics::clone(&statistics[0]);
@@ -767,7 +813,11 @@ mod test {
let global_limit: Arc<dyn ExecutionPlan> =
Arc::new(GlobalLimitExec::new(scan.clone(), 0, Some(2)));
let statistics = (0..global_limit.output_partitioning().partition_count())
- .map(|idx| global_limit.partition_statistics(Some(idx)))
+ .map(|idx| {
+ global_limit.statistics_with_args(
+ &StatisticsArgs::new().with_partition(Some(idx)),
+ )
+ })
.collect::<Result<Vec<_>>>()?;
assert_eq!(statistics.len(), 1);
// GlobalLimit takes from first partition: ids [3,4], dates [2025-03-01, 2025-03-02]
@@ -826,7 +876,8 @@ mod test {
@"AggregateExec: mode=Partial, gby=[id@0 as id, 1 + id@0 as expr], aggr=[COUNT(c)]"
);
- let p0_statistics = aggregate_exec_partial.partition_statistics(Some(0))?;
+ let p0_statistics = aggregate_exec_partial
+ .statistics_with_args(&StatisticsArgs::new().with_partition(Some(0)))?;
// Aggregate doesn't propagate num_rows and ColumnStatistics byte_size from input
let expected_p0_statistics = Statistics {
@@ -865,7 +916,8 @@ mod test {
],
};
- let p1_statistics = aggregate_exec_partial.partition_statistics(Some(1))?;
+ let p1_statistics = aggregate_exec_partial
+ .statistics_with_args(&StatisticsArgs::new().with_partition(Some(1)))?;
assert_eq!(*p1_statistics, expected_p1_statistics);
validate_statistics_with_data(
@@ -887,10 +939,12 @@ mod test {
aggregate_exec_partial.schema(),
)?);
- let p0_statistics = agg_final.partition_statistics(Some(0))?;
+ let p0_statistics = agg_final
+ .statistics_with_args(&StatisticsArgs::new().with_partition(Some(0)))?;
assert_eq!(*p0_statistics, expected_p0_statistics);
- let p1_statistics = agg_final.partition_statistics(Some(1))?;
+ let p1_statistics = agg_final
+ .statistics_with_args(&StatisticsArgs::new().with_partition(Some(1)))?;
assert_eq!(*p1_statistics, expected_p1_statistics);
validate_statistics_with_data(
@@ -935,8 +989,16 @@ mod test {
],
};
- assert_eq!(empty_stat, *agg_partial.partition_statistics(Some(0))?);
- assert_eq!(empty_stat, *agg_partial.partition_statistics(Some(1))?);
+ assert_eq!(
+ empty_stat,
+ *agg_partial
+ .statistics_with_args(&StatisticsArgs::new().with_partition(Some(0)))?
+ );
+ assert_eq!(
+ empty_stat,
+ *agg_partial
+ .statistics_with_args(&StatisticsArgs::new().with_partition(Some(1)))?
+ );
validate_statistics_with_data(
agg_partial.clone(),
vec![ExpectedStatistics::Empty, ExpectedStatistics::Empty],
@@ -962,8 +1024,16 @@ mod test {
agg_partial.schema(),
)?);
- assert_eq!(empty_stat, *agg_final.partition_statistics(Some(0))?);
- assert_eq!(empty_stat, *agg_final.partition_statistics(Some(1))?);
+ assert_eq!(
+ empty_stat,
+ *agg_final
+ .statistics_with_args(&StatisticsArgs::new().with_partition(Some(0)))?
+ );
+ assert_eq!(
+ empty_stat,
+ *agg_final
+ .statistics_with_args(&StatisticsArgs::new().with_partition(Some(1)))?
+ );
validate_statistics_with_data(
agg_final,
@@ -989,11 +1059,13 @@ mod test {
};
assert_eq!(
expect_partial_stat,
- *agg_partial.partition_statistics(Some(0))?
+ *agg_partial
+ .statistics_with_args(&StatisticsArgs::new().with_partition(Some(0)))?
);
assert_eq!(
expect_partial_stat,
- *agg_partial.partition_statistics(Some(1))?
+ *agg_partial
+ .statistics_with_args(&StatisticsArgs::new().with_partition(Some(1)))?
);
let expect_partial_overall_stat = Statistics {
@@ -1003,7 +1075,7 @@ mod test {
};
assert_eq!(
expect_partial_overall_stat,
- *agg_partial.partition_statistics(None)?
+ *agg_partial.statistics_with_args(&StatisticsArgs::new())?
);
// Verify that the partial aggregate emits one accumulator-state row per
@@ -1036,7 +1108,11 @@ mod test {
column_statistics: vec![ColumnStatistics::new_unknown()],
};
- assert_eq!(expect_stat, *agg_final.partition_statistics(Some(0))?);
+ assert_eq!(
+ expect_stat,
+ *agg_final
+ .statistics_with_args(&StatisticsArgs::new().with_partition(Some(0)))?
+ );
// Verify that the aggregate final result has exactly one partition with one row
let mut partitions = execute_stream_partitioned(
@@ -1064,7 +1140,8 @@ mod test {
let mut all_batches = vec![];
for (i, partition_stream) in partitions.into_iter().enumerate() {
let batches: Vec<RecordBatch> = partition_stream.try_collect().await?;
- let actual = plan.partition_statistics(Some(i))?;
+ let actual = plan
+ .statistics_with_args(&StatisticsArgs::new().with_partition(Some(i)))?;
let expected = compute_record_batch_statistics(
std::slice::from_ref(&batches),
&schema,
@@ -1074,7 +1151,7 @@ mod test {
all_batches.push(batches);
}
- let actual = plan.partition_statistics(None)?;
+ let actual = plan.statistics_with_args(&StatisticsArgs::new())?;
let expected = compute_record_batch_statistics(&all_batches, &schema, None);
assert_eq!(*actual, expected);
@@ -1091,7 +1168,11 @@ mod test {
)?);
let statistics = (0..repartition.partitioning().partition_count())
- .map(|idx| repartition.partition_statistics(Some(idx)))
+ .map(|idx| {
+ repartition.statistics_with_args(
+ &StatisticsArgs::new().with_partition(Some(idx)),
+ )
+ })
.collect::<Result<Vec<_>>>()?;
assert_eq!(statistics.len(), 3);
@@ -1142,7 +1223,8 @@ mod test {
Partitioning::RoundRobinBatch(2),
)?);
- let result = repartition.partition_statistics(Some(2));
+ let result = repartition
+ .statistics_with_args(&StatisticsArgs::new().with_partition(Some(2)));
assert!(result.is_err());
let error = result.unwrap_err();
assert!(
@@ -1171,7 +1253,8 @@ mod test {
Partitioning::RoundRobinBatch(0),
)?);
- let result = repartition.partition_statistics(Some(0))?;
+ let result = repartition
+ .statistics_with_args(&StatisticsArgs::new().with_partition(Some(0)))?;
assert_eq!(*result, Statistics::new_unknown(&scan_schema));
// Verify that the result has exactly 0 partitions
@@ -1198,7 +1281,11 @@ mod test {
// Verify the result of partition statistics of repartition
let stats = (0..repartition.partitioning().partition_count())
- .map(|idx| repartition.partition_statistics(Some(idx)))
+ .map(|idx| {
+ repartition.statistics_with_args(
+ &StatisticsArgs::new().with_partition(Some(idx)),
+ )
+ })
.collect::<Result<Vec<_>>>()?;
assert_eq!(stats.len(), 2);
@@ -1256,7 +1343,11 @@ mod test {
// Verify partition statistics are properly propagated (not unknown)
let statistics = (0..window_agg.output_partitioning().partition_count())
- .map(|idx| window_agg.partition_statistics(Some(idx)))
+ .map(|idx| {
+ window_agg.statistics_with_args(
+ &StatisticsArgs::new().with_partition(Some(idx)),
+ )
+ })
.collect::<Result<Vec<_>>>()?;
assert_eq!(statistics.len(), 2);
@@ -1342,7 +1433,8 @@ mod test {
// Try to test with single partition
let empty_single = Arc::new(EmptyExec::new(Arc::clone(&schema)));
- let stats = empty_single.partition_statistics(Some(0))?;
+ let stats = empty_single
+ .statistics_with_args(&StatisticsArgs::new().with_partition(Some(0)))?;
assert_eq!(stats.num_rows, Precision::Exact(0));
assert_eq!(stats.total_byte_size, Precision::Exact(0));
assert_eq!(stats.column_statistics.len(), 2);
@@ -1357,7 +1449,7 @@ mod test {
assert_eq!(col_stat.byte_size, Precision::Exact(0));
}
- let overall_stats = empty_single.partition_statistics(None)?;
+ let overall_stats = empty_single.statistics_with_args(&StatisticsArgs::new())?;
assert_eq!(stats, overall_stats);
validate_statistics_with_data(empty_single, vec![ExpectedStatistics::Empty], 0)
@@ -1368,7 +1460,11 @@ mod test {
Arc::new(EmptyExec::new(Arc::clone(&schema)).with_partitions(3));
let statistics = (0..empty_multi.output_partitioning().partition_count())
- .map(|idx| empty_multi.partition_statistics(Some(idx)))
+ .map(|idx| {
+ empty_multi.statistics_with_args(
+ &StatisticsArgs::new().with_partition(Some(idx)),
+ )
+ })
.collect::<Result<Vec<_>>>()?;
assert_eq!(statistics.len(), 3);
@@ -1428,7 +1524,11 @@ mod test {
// Test partition statistics for CollectLeft mode
let statistics = (0..collect_left_join.output_partitioning().partition_count())
- .map(|idx| collect_left_join.partition_statistics(Some(idx)))
+ .map(|idx| {
+ collect_left_join.statistics_with_args(
+ &StatisticsArgs::new().with_partition(Some(idx)),
+ )
+ })
.collect::<Result<Vec<_>>>()?;
// Check that we have the expected number of partitions
@@ -1504,7 +1604,11 @@ mod test {
// Test partition statistics for Partitioned mode
let statistics = (0..partitioned_join.output_partitioning().partition_count())
- .map(|idx| partitioned_join.partition_statistics(Some(idx)))
+ .map(|idx| {
+ partitioned_join.statistics_with_args(
+ &StatisticsArgs::new().with_partition(Some(idx)),
+ )
+ })
.collect::<Result<Vec<_>>>()?;
// Check that we have the expected number of partitions
@@ -1578,7 +1682,11 @@ mod test {
// Test partition statistics for Auto mode
let statistics = (0..auto_join.output_partitioning().partition_count())
- .map(|idx| auto_join.partition_statistics(Some(idx)))
+ .map(|idx| {
+ auto_join.statistics_with_args(
+ &StatisticsArgs::new().with_partition(Some(idx)),
+ )
+ })
.collect::<Result<Vec<_>>>()?;
// Check that we have the expected number of partitions
diff --git a/datafusion/core/tests/physical_optimizer/test_utils.rs b/datafusion/core/tests/physical_optimizer/test_utils.rs
index 09225cb038..d71f9be5a2 100644
--- a/datafusion/core/tests/physical_optimizer/test_utils.rs
+++ b/datafusion/core/tests/physical_optimizer/test_utils.rs
@@ -68,7 +68,7 @@ use datafusion_physical_plan::union::UnionExec;
use datafusion_physical_plan::windows::{BoundedWindowAggExec, create_window_expr};
use datafusion_physical_plan::{
DisplayAs, DisplayFormatType, ExecutionPlan, InputOrderMode, Partitioning,
- PlanProperties, SortOrderPushdownResult, displayable,
+ PlanProperties, SortOrderPushdownResult, StatisticsArgs, displayable,
};
/// Create a non sorted parquet exec
@@ -967,7 +967,7 @@ impl ExecutionPlan for TestScan {
internal_err!("TestScan is for testing optimizer only, not for execution")
}
- fn partition_statistics(&self, _partition: Option<usize>) -> Result<Arc<Statistics>> {
+ fn statistics_with_args(&self, _args: &StatisticsArgs) -> Result<Arc<Statistics>> {
Ok(Arc::new(Statistics::new_unknown(&self.schema)))
}
diff --git a/datafusion/core/tests/sql/path_partition.rs b/datafusion/core/tests/sql/path_partition.rs
index 20f49de8aa..de6349d129 100644
--- a/datafusion/core/tests/sql/path_partition.rs
+++ b/datafusion/core/tests/sql/path_partition.rs
@@ -38,6 +38,7 @@ use datafusion_common::ScalarValue;
use datafusion_common::stats::Precision;
use datafusion_common::test_util::batches_to_sort_string;
use datafusion_execution::config::SessionConfig;
+use datafusion_physical_plan::statistics::StatisticsArgs;
use async_trait::async_trait;
use bytes::Bytes;
@@ -462,7 +463,7 @@ async fn parquet_statistics() -> Result<()> {
assert_eq!(schema.fields().len(), 4);
let stat_cols = physical_plan
- .partition_statistics(None)?
+ .statistics_with_args(&StatisticsArgs::new())?
.column_statistics
.clone();
assert_eq!(stat_cols.len(), 4);
@@ -489,7 +490,7 @@ async fn parquet_statistics() -> Result<()> {
assert_eq!(schema.fields().len(), 2);
let stat_cols = physical_plan
- .partition_statistics(None)?
+ .statistics_with_args(&StatisticsArgs::new())?
.column_statistics
.clone();
assert_eq!(stat_cols.len(), 2);
diff --git a/datafusion/datasource/src/file_scan_config/mod.rs b/datafusion/datasource/src/file_scan_config/mod.rs
index 3ebd588a07..b1ba0584c9 100644
--- a/datafusion/datasource/src/file_scan_config/mod.rs
+++ b/datafusion/datasource/src/file_scan_config/mod.rs
@@ -2276,7 +2276,7 @@ mod tests {
// of just the projected ones.
use crate::source::DataSourceExec;
- use datafusion_physical_plan::ExecutionPlan;
+ use datafusion_physical_plan::statistics::StatisticsArgs;
// Create a schema with 4 columns
let schema = Arc::new(Schema::new(vec![
@@ -2330,7 +2330,9 @@ mod tests {
let exec = DataSourceExec::from_data_source(config);
// Get statistics for partition 0
- let partition_stats = exec.partition_statistics(Some(0)).unwrap();
+ let partition_stats = exec
+ .statistics_with_args(&StatisticsArgs::new().with_partition(Some(0)))
+ .unwrap();
// Verify that only 2 columns are in the statistics (the projected ones)
assert_eq!(
diff --git a/datafusion/datasource/src/memory.rs b/datafusion/datasource/src/memory.rs
index f073b09c54..a4e30d7f0b 100644
--- a/datafusion/datasource/src/memory.rs
+++ b/datafusion/datasource/src/memory.rs
@@ -853,6 +853,7 @@ mod tests {
use datafusion_common::stats::{ColumnStatistics, Precision};
use datafusion_physical_expr::PhysicalSortExpr;
use datafusion_physical_plan::expressions::lit;
+ use datafusion_physical_plan::statistics::StatisticsArgs;
use datafusion_physical_plan::ExecutionPlan;
@@ -985,7 +986,7 @@ mod tests {
let values = MemorySourceConfig::try_new_as_values(schema, data)?;
assert_eq!(
- *values.partition_statistics(None)?,
+ *values.statistics_with_args(&StatisticsArgs::new())?,
Statistics {
num_rows: Precision::Exact(rows),
total_byte_size: Precision::Exact(8), // not important
diff --git a/datafusion/datasource/src/source.rs b/datafusion/datasource/src/source.rs
index af4bc09504..b7e920f53f 100644
--- a/datafusion/datasource/src/source.rs
+++ b/datafusion/datasource/src/source.rs
@@ -45,6 +45,7 @@ use datafusion_execution::{SendableRecordBatchStream, TaskContext};
use datafusion_physical_expr::{EquivalenceProperties, Partitioning, PhysicalExpr};
use datafusion_physical_expr_common::sort_expr::{LexOrdering, PhysicalSortExpr};
use datafusion_physical_plan::SortOrderPushdownResult;
+use datafusion_physical_plan::StatisticsArgs;
use datafusion_physical_plan::filter_pushdown::{
ChildPushdownResult, FilterPushdownPhase, FilterPushdownPropagation, PushedDown,
};
@@ -426,8 +427,8 @@ impl ExecutionPlan for DataSourceExec {
Some(metrics)
}
- fn partition_statistics(&self, partition: Option<usize>) -> Result<Arc<Statistics>> {
- self.data_source.partition_statistics(partition)
+ fn statistics_with_args(&self, args: &StatisticsArgs) -> Result<Arc<Statistics>> {
+ self.data_source.partition_statistics(args.partition())
}
fn with_fetch(&self, limit: Option<usize>) -> Option<Arc<dyn ExecutionPlan>> {
diff --git a/datafusion/ffi/src/execution_plan.rs b/datafusion/ffi/src/execution_plan.rs
index f942916ea1..738f87fd61 100644
--- a/datafusion/ffi/src/execution_plan.rs
+++ b/datafusion/ffi/src/execution_plan.rs
@@ -24,7 +24,7 @@ use datafusion_common::{DataFusionError, Result, Statistics};
use datafusion_execution::{SendableRecordBatchStream, TaskContext};
use datafusion_physical_expr_common::metrics::MetricsSet;
use datafusion_physical_plan::{
- DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties,
+ DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, StatisticsArgs,
};
use stabby::string::String as SString;
use stabby::vec::Vec as SVec;
@@ -210,7 +210,7 @@ unsafe extern "C" fn partition_statistics_fn_wrapper(
) -> FFI_Result<SVec<u8>> {
let partition: Option<usize> = partition.into();
plan.inner()
- .partition_statistics(partition)
+ .statistics_with_args(&StatisticsArgs::new().with_partition(partition))
.map(|stats| SVec::from(serialize_statistics(stats.as_ref()).as_slice()))
.into()
}
@@ -556,9 +556,9 @@ pub mod tests {
self.metrics.clone()
}
- fn partition_statistics(
+ fn statistics_with_args(
&self,
- _partition: Option<usize>,
+ _args: &StatisticsArgs,
) -> Result<Arc<Statistics>> {
Ok(Arc::new(self.statistics.clone().unwrap_or_else(|| {
Statistics::new_unknown(self.props.eq_properties.schema())
@@ -672,27 +672,34 @@ pub mod tests {
Ok(())
}
- #[test]
- fn test_ffi_execution_plan_partition_statistics_round_trip() -> Result<()> {
+ /// Build an `EmptyExec` carrying `statistics`, then export it across the
+ /// (mock) FFI boundary and return the resulting foreign plan.
+ #[cfg(test)]
+ fn export_empty_exec_over_ffi(
+ schema: &arrow::datatypes::SchemaRef,
+ statistics: Option<Statistics>,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ let mut plan = EmptyExec::new(Arc::clone(schema));
+ if let Some(statistics) = statistics {
+ plan = plan.with_statistics(statistics);
+ }
+ let mut local = FFI_ExecutionPlan::new(Arc::new(plan), None);
+ local.library_marker_id = crate::mock_foreign_marker_id;
+ let foreign: Arc<dyn ExecutionPlan> = (&local).try_into()?;
+ Ok(foreign)
+ }
+
+ /// Schema and a fully-populated `Statistics` (including `ScalarValue`-typed
+ /// min/max) shared by the FFI statistics round-trip tests.
+ #[cfg(test)]
+ fn stats_round_trip_fixture() -> (arrow::datatypes::SchemaRef, Statistics) {
use datafusion_common::stats::Precision;
use datafusion_common::{ColumnStatistics, ScalarValue};
let schema = Arc::new(arrow::datatypes::Schema::new(vec![
arrow::datatypes::Field::new("a", arrow::datatypes::DataType::Int32, true),
]));
-
- // Plans without explicit statistics return Statistics::new_unknown across
- // the boundary.
- let bare_plan = Arc::new(EmptyExec::new(Arc::clone(&schema)));
- let mut bare_local = FFI_ExecutionPlan::new(bare_plan, None);
- bare_local.library_marker_id = crate::mock_foreign_marker_id;
- let bare_foreign: Arc<dyn ExecutionPlan> = (&bare_local).try_into()?;
- let bare_stats = bare_foreign.partition_statistics(None)?;
- assert_eq!(bare_stats.as_ref(), &Statistics::new_unknown(&schema));
-
- // Plans with statistics round-trip them faithfully, including
- // ScalarValue-typed min/max.
- let original_stats = Statistics {
+ let statistics = Statistics {
num_rows: Precision::Exact(7),
total_byte_size: Precision::Inexact(128),
column_statistics: vec![ColumnStatistics {
@@ -704,18 +711,67 @@ pub mod tests {
byte_size: Precision::Exact(28),
}],
};
- let stats_plan = Arc::new(
- EmptyExec::new(Arc::clone(&schema)).with_statistics(original_stats.clone()),
+ (schema, statistics)
+ }
+
+ /// Statistics survive an FFI round trip when queried through the
+ /// **deprecated** `partition_statistics` entry point on the foreign plan.
+ #[test]
+ #[expect(deprecated)]
+ fn test_ffi_execution_plan_partition_statistics_round_trip() -> Result<()> {
+ let (schema, original_stats) = stats_round_trip_fixture();
+
+ // A plan without explicit statistics reports new_unknown.
+ let bare = export_empty_exec_over_ffi(&schema, None)?;
+ assert_eq!(
+ bare.partition_statistics(None)?.as_ref(),
+ &Statistics::new_unknown(&schema)
);
- let mut stats_local = FFI_ExecutionPlan::new(stats_plan, None);
- stats_local.library_marker_id = crate::mock_foreign_marker_id;
- let stats_foreign: Arc<dyn ExecutionPlan> = (&stats_local).try_into()?;
- let observed = stats_foreign.partition_statistics(None)?;
- assert_eq!(observed.as_ref(), &original_stats);
+ // A plan with statistics round-trips them for overall and per-partition queries.
+ let with_stats =
+ export_empty_exec_over_ffi(&schema, Some(original_stats.clone()))?;
+ assert_eq!(
+ with_stats.partition_statistics(None)?.as_ref(),
+ &original_stats
+ );
+ assert_eq!(
+ with_stats.partition_statistics(Some(1))?.as_ref(),
+ &original_stats
+ );
- let observed_partition = stats_foreign.partition_statistics(Some(1))?;
- assert_eq!(observed_partition.as_ref(), &original_stats);
+ Ok(())
+ }
+
+ /// Same round trip as
+ /// [`test_ffi_execution_plan_partition_statistics_round_trip`], but queried
+ /// through the **new** `statistics_with_args` entry point.
+ #[test]
+ fn test_ffi_execution_plan_statistics_with_args_round_trip() -> Result<()> {
+ let (schema, original_stats) = stats_round_trip_fixture();
+
+ // A plan without explicit statistics reports new_unknown.
+ let bare = export_empty_exec_over_ffi(&schema, None)?;
+ assert_eq!(
+ bare.statistics_with_args(&StatisticsArgs::new())?.as_ref(),
+ &Statistics::new_unknown(&schema)
+ );
+
+ // A plan with statistics round-trips them for overall and per-partition queries.
+ let with_stats =
+ export_empty_exec_over_ffi(&schema, Some(original_stats.clone()))?;
+ assert_eq!(
+ with_stats
+ .statistics_with_args(&StatisticsArgs::new())?
+ .as_ref(),
+ &original_stats
+ );
+ assert_eq!(
+ with_stats
+ .statistics_with_args(&StatisticsArgs::new().with_partition(Some(1)))?
+ .as_ref(),
+ &original_stats
+ );
Ok(())
}
diff --git a/datafusion/ffi/tests/ffi_execution_plan.rs b/datafusion/ffi/tests/ffi_execution_plan.rs
index bd84e064de..7d04e828bd 100644
--- a/datafusion/ffi/tests/ffi_execution_plan.rs
+++ b/datafusion/ffi/tests/ffi_execution_plan.rs
@@ -29,6 +29,7 @@ mod tests {
use std::sync::Arc;
#[test]
+ #[expect(deprecated)]
fn test_ffi_execution_plan_partition_statistics_cross_library()
-> Result<(), DataFusionError> {
let module = get_module()?;
diff --git a/datafusion/physical-optimizer/src/aggregate_statistics.rs b/datafusion/physical-optimizer/src/aggregate_statistics.rs
index d0be53d59b..b83f4ed730 100644
--- a/datafusion/physical-optimizer/src/aggregate_statistics.rs
+++ b/datafusion/physical-optimizer/src/aggregate_statistics.rs
@@ -25,7 +25,10 @@ use datafusion_physical_plan::aggregates::{
};
use datafusion_physical_plan::placeholder_row::PlaceholderRowExec;
use datafusion_physical_plan::projection::{ProjectionExec, ProjectionExpr};
-use datafusion_physical_plan::udaf::{AggregateFunctionExpr, StatisticsArgs};
+use datafusion_physical_plan::statistics::StatisticsArgs;
+use datafusion_physical_plan::udaf::{
+ AggregateFunctionExpr, StatisticsArgs as PlanStatisticsArgs,
+};
use datafusion_physical_plan::{ExecutionPlan, expressions};
use std::sync::Arc;
@@ -55,12 +58,14 @@ impl PhysicalOptimizerRule for AggregateStatistics {
let partial_agg_exec = partial_agg_exec
.downcast_ref::<AggregateExec>()
.expect("take_optimizable() ensures that this is a AggregateExec");
- let stats = partial_agg_exec.input().partition_statistics(None)?;
+ let stats = partial_agg_exec
+ .input()
+ .statistics_with_args(&StatisticsArgs::new())?;
let mut projections = vec![];
for expr in partial_agg_exec.aggr_expr() {
let field = expr.field();
let args = expr.expressions();
- let statistics_args = StatisticsArgs {
+ let statistics_args = PlanStatisticsArgs {
statistics: &stats,
return_type: field.data_type(),
is_distinct: expr.is_distinct(),
@@ -148,7 +153,7 @@ fn take_optimizable(plan: &Arc<dyn ExecutionPlan>) -> Option<Arc<dyn ExecutionPl
/// If this agg_expr is a max that is exactly defined in the statistics, return it.
fn take_optimizable_value_from_statistics(
- statistics_args: &StatisticsArgs,
+ statistics_args: &PlanStatisticsArgs,
agg_expr: &AggregateFunctionExpr,
) -> Option<(ScalarValue, String)> {
let value = agg_expr.fun().value_from_stats(statistics_args);
diff --git a/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs b/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs
index ada7b6d741..6d9550fa50 100644
--- a/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs
+++ b/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs
@@ -61,6 +61,7 @@ use datafusion_physical_plan::joins::{
use datafusion_physical_plan::projection::{ProjectionExec, ProjectionExpr};
use datafusion_physical_plan::repartition::RepartitionExec;
use datafusion_physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec;
+use datafusion_physical_plan::statistics::StatisticsArgs;
use datafusion_physical_plan::tree_node::PlanContext;
use datafusion_physical_plan::union::{InterleaveExec, UnionExec, can_interleave};
use datafusion_physical_plan::windows::WindowAggExec;
@@ -1015,7 +1016,9 @@ fn get_repartition_requirement_status(
{
// Decide whether adding a round robin is beneficial depending on
// the statistical information we have on the number of rows:
- let roundrobin_beneficial_stats = match child.partition_statistics(None)?.num_rows
+ let roundrobin_beneficial_stats = match child
+ .statistics_with_args(&StatisticsArgs::new())?
+ .num_rows
{
Precision::Exact(n_rows) => n_rows > batch_size,
Precision::Inexact(n_rows) => !should_use_estimates || (n_rows > batch_size),
diff --git a/datafusion/physical-optimizer/src/join_selection.rs b/datafusion/physical-optimizer/src/join_selection.rs
index 74c6cbb19a..82294825b6 100644
--- a/datafusion/physical-optimizer/src/join_selection.rs
+++ b/datafusion/physical-optimizer/src/join_selection.rs
@@ -40,6 +40,7 @@ use datafusion_physical_plan::joins::{
StreamJoinPartitionMode, SymmetricHashJoinExec,
};
use datafusion_physical_plan::operator_statistics::StatisticsRegistry;
+use datafusion_physical_plan::statistics::StatisticsArgs;
use datafusion_physical_plan::{ExecutionPlan, ExecutionPlanProperties};
use std::sync::Arc;
@@ -65,7 +66,7 @@ fn get_stats(
reg.compute(plan)
.map(|s| Arc::<Statistics>::clone(s.base_arc()))
} else {
- plan.partition_statistics(None)
+ plan.statistics_with_args(&StatisticsArgs::new())
}
}
diff --git a/datafusion/physical-optimizer/src/limit_pushdown.rs b/datafusion/physical-optimizer/src/limit_pushdown.rs
index 63c4f21bd9..224084d576 100644
--- a/datafusion/physical-optimizer/src/limit_pushdown.rs
+++ b/datafusion/physical-optimizer/src/limit_pushdown.rs
@@ -76,6 +76,7 @@ use datafusion_physical_plan::limit::{GlobalLimitExec, LocalLimitExec};
use datafusion_physical_plan::placeholder_row::PlaceholderRowExec;
use datafusion_physical_plan::projection::ProjectionExec;
use datafusion_physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec;
+use datafusion_physical_plan::statistics::StatisticsArgs;
use datafusion_physical_plan::{ExecutionPlan, ExecutionPlanProperties};
/// This rule inspects [`ExecutionPlan`]'s and pushes down the fetch limit from
/// the parent to the child if applicable.
@@ -351,7 +352,9 @@ fn limit_eliminable_exact_num_rows(
}
if matches!(
- current.partition_statistics(None)?.num_rows,
+ current
+ .statistics_with_args(&StatisticsArgs::new())?
+ .num_rows,
Precision::Exact(0)
) {
return Ok(Some(0));
diff --git a/datafusion/physical-optimizer/src/output_requirements.rs b/datafusion/physical-optimizer/src/output_requirements.rs
index 899abcc88b..fb91ae46a2 100644
--- a/datafusion/physical-optimizer/src/output_requirements.rs
+++ b/datafusion/physical-optimizer/src/output_requirements.rs
@@ -32,6 +32,7 @@ use datafusion_common::{Result, Statistics};
use datafusion_execution::TaskContext;
use datafusion_physical_expr::Distribution;
use datafusion_physical_expr_common::sort_expr::OrderingRequirements;
+use datafusion_physical_plan::StatisticsArgs;
use datafusion_physical_plan::execution_plan::Boundedness;
use datafusion_physical_plan::projection::{
ProjectionExec, make_with_child, update_expr, update_ordering_requirement,
@@ -242,8 +243,8 @@ impl ExecutionPlan for OutputRequirementExec {
unreachable!();
}
- fn partition_statistics(&self, partition: Option<usize>) -> Result<Arc<Statistics>> {
- self.input.partition_statistics(partition)
+ fn statistics_with_args(&self, args: &StatisticsArgs) -> Result<Arc<Statistics>> {
+ args.compute_child_statistics(&self.input, args.partition())
}
fn try_swapping_with_projection(
diff --git a/datafusion/physical-plan/Cargo.toml b/datafusion/physical-plan/Cargo.toml
index 4b2b31febe..b4fc8f9d01 100644
--- a/datafusion/physical-plan/Cargo.toml
+++ b/datafusion/physical-plan/Cargo.toml
@@ -126,13 +126,17 @@ name = "aggregate_vectorized"
[[bench]]
harness = false
-name = "hash_join_semi_anti"
-required-features = ["test_utils"]
+name = "compute_statistics"
[[bench]]
harness = false
name = "dictionary_group_values"
+[[bench]]
+harness = false
+name = "hash_join_semi_anti"
+required-features = ["test_utils"]
+
[[bench]]
harness = false
name = "multi_group_by"
diff --git a/datafusion/physical-plan/benches/compute_statistics.rs b/datafusion/physical-plan/benches/compute_statistics.rs
new file mode 100644
index 0000000000..04b5612563
--- /dev/null
+++ b/datafusion/physical-plan/benches/compute_statistics.rs
@@ -0,0 +1,312 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Benchmarks for `compute_statistics` with `StatsCache`.
+//!
+//! Demonstrates that caching eliminates redundant subtree walks in plans
+//! containing partition-merging operators (CoalescePartitionsExec) and
+//! binary join trees (CrossJoinExec).
+//!
+//! The plan shapes here mirror the reproducers from the planning-speed
+//! EPIC (<https://github.com/apache/datafusion/issues/19795>):
+//! - Coalesce chain: deep linear plans (e.g. deeply nested subqueries)
+//! - Cross-join tree: balanced binary trees from multi-way joins
+//! (mirrors the `physical_many_self_joins` sql_planner benchmark)
+
+use std::fmt;
+use std::sync::Arc;
+
+use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
+use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main};
+use datafusion_common::ScalarValue;
+use datafusion_common::{Result, Statistics};
+use datafusion_execution::TaskContext;
+use datafusion_physical_expr::EquivalenceProperties;
+use datafusion_physical_expr::PhysicalExpr;
+use datafusion_physical_expr::expressions::Literal;
+use datafusion_physical_plan::coalesce_partitions::CoalescePartitionsExec;
+use datafusion_physical_plan::execution_plan::{
+ Boundedness, EmissionType, ExecutionPlan, PlanProperties,
+};
+use datafusion_physical_plan::filter::FilterExec;
+use datafusion_physical_plan::joins::CrossJoinExec;
+use datafusion_physical_plan::statistics::StatisticsArgs;
+use datafusion_physical_plan::{
+ DisplayAs, DisplayFormatType, Partitioning, SendableRecordBatchStream,
+};
+
+/// Minimal leaf node for benchmarking
+#[derive(Debug)]
+struct BenchLeaf {
+ schema: SchemaRef,
+ cache: Arc<PlanProperties>,
+}
+
+impl BenchLeaf {
+ fn new(col_name: &str) -> Self {
+ let schema = Arc::new(Schema::new(vec![Field::new(
+ col_name,
+ DataType::Int32,
+ false,
+ )]));
+ let cache = Arc::new(PlanProperties::new(
+ EquivalenceProperties::new(Arc::clone(&schema)),
+ Partitioning::UnknownPartitioning(2),
+ EmissionType::Incremental,
+ Boundedness::Bounded,
+ ));
+ Self { schema, cache }
+ }
+}
+
+impl DisplayAs for BenchLeaf {
+ fn fmt_as(&self, _t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result {
+ write!(f, "BenchLeaf")
+ }
+}
+
+impl ExecutionPlan for BenchLeaf {
+ fn name(&self) -> &str {
+ "BenchLeaf"
+ }
+
+ fn schema(&self) -> SchemaRef {
+ Arc::clone(&self.schema)
+ }
+
+ fn properties(&self) -> &Arc<PlanProperties> {
+ &self.cache
+ }
+
+ fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
+ vec![]
+ }
+
+ fn with_new_children(
+ self: Arc<Self>,
+ _children: Vec<Arc<dyn ExecutionPlan>>,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ Ok(self)
+ }
+
+ fn execute(
+ &self,
+ _partition: usize,
+ _context: Arc<TaskContext>,
+ ) -> Result<SendableRecordBatchStream> {
+ unimplemented!()
+ }
+
+ fn statistics_with_args(&self, _args: &StatisticsArgs) -> Result<Arc<Statistics>> {
+ Ok(Arc::new(Statistics::new_unknown(&self.schema)))
+ }
+}
+
+/// Build: CoalescePartitions^depth -> BenchLeaf
+fn build_coalesce_chain(depth: usize) -> Arc<dyn ExecutionPlan> {
+ let mut plan: Arc<dyn ExecutionPlan> = Arc::new(BenchLeaf::new("a"));
+ for _ in 0..depth {
+ plan = Arc::new(CoalescePartitionsExec::new(plan));
+ }
+ plan
+}
+
+/// Build a balanced binary tree of CrossJoinExec with 2^depth leaves.
+/// Mirrors the plan shape produced by multi-way self-joins like the
+/// `physical_many_self_joins` benchmark in sql_planner.rs (#19795).
+fn build_cross_join_tree(depth: usize, next_col: &mut usize) -> Arc<dyn ExecutionPlan> {
+ if depth == 0 {
+ let col_name = format!("c{next_col}");
+ *next_col += 1;
+ return Arc::new(BenchLeaf::new(&col_name));
+ }
+ let left = build_cross_join_tree(depth - 1, next_col);
+ let right = build_cross_join_tree(depth - 1, next_col);
+ Arc::new(CrossJoinExec::new(left, right))
+}
+
+/// Build: Filter^depth -> BenchLeaf (always-true predicate).
+fn build_filter_chain(depth: usize) -> Arc<dyn ExecutionPlan> {
+ let mut plan: Arc<dyn ExecutionPlan> = Arc::new(BenchLeaf::new("a"));
+ let predicate: Arc<dyn PhysicalExpr> =
+ Arc::new(Literal::new(ScalarValue::Boolean(Some(true))));
+ for _ in 0..depth {
+ plan = Arc::new(
+ FilterExec::try_new(Arc::clone(&predicate), plan)
+ .expect("FilterExec::try_new failed"),
+ );
+ }
+ plan
+}
+
+/// Build a mixed chain alternating partition-merging and partition-preserving
+/// operators: (Coalesce -> Filter -> Filter) repeated `groups` times -> BenchLeaf.
+/// Exercises the cache with both None and Some(p) lookups in the same walk.
+fn build_mixed_chain(groups: usize) -> Arc<dyn ExecutionPlan> {
+ let mut plan: Arc<dyn ExecutionPlan> = Arc::new(BenchLeaf::new("a"));
+ let predicate: Arc<dyn PhysicalExpr> =
+ Arc::new(Literal::new(ScalarValue::Boolean(Some(true))));
+ for _ in 0..groups {
+ // Two partition-preserving filters
+ for _ in 0..2 {
+ plan = Arc::new(
+ FilterExec::try_new(Arc::clone(&predicate), plan)
+ .expect("FilterExec::try_new failed"),
+ );
+ }
+ // One partition-merging coalesce
+ plan = Arc::new(CoalescePartitionsExec::new(plan));
+ }
+ plan
+}
+
+/// Recursive walk without a shared cross-node cache, simulating pre-cache behavior.
+/// Each operator's internal `compute_child_statistics` call triggers a fresh
+/// subtree walk, resulting in O(n^2) total node visits for a chain of depth n.
+///
+/// Note: each `compute_child_statistics` re-walk still benefits from its own
+/// ephemeral cache; only the cross-node sharing is removed.
+fn compute_statistics_without_shared_cache(
+ plan: &dyn ExecutionPlan,
+ partition: Option<usize>,
+) -> Result<Arc<Statistics>> {
+ for child in plan.children() {
+ compute_statistics_without_shared_cache(child.as_ref(), None)?;
+ }
+ let args = StatisticsArgs::new().with_partition(partition);
+ plan.statistics_with_args(&args)
+}
+
+fn bench_compute_statistics(c: &mut Criterion) {
+ // --- Coalesce chain (linear plan) ---
+ // Deep linear plans arise from deeply nested subqueries, CTEs, etc.
+ let mut group = c.benchmark_group("compute_statistics_coalesce_chain");
+ for depth in [10, 20, 50] {
+ let plan = build_coalesce_chain(depth);
+ group.bench_with_input(BenchmarkId::new("cached", depth), &plan, |b, plan| {
+ b.iter(|| plan.statistics_with_args(&StatisticsArgs::new()).unwrap());
+ });
+ group.bench_with_input(
+ BenchmarkId::new("no_shared_cache", depth),
+ &plan,
+ |b, plan| {
+ b.iter(|| {
+ compute_statistics_without_shared_cache(plan.as_ref(), None).unwrap()
+ });
+ },
+ );
+ }
+ group.finish();
+
+ // --- Cross-join tree (balanced binary plan) ---
+ // Binary trees arise from multi-way joins (e.g. physical_many_self_joins
+ // in sql_planner.rs, see #19795). CrossJoinExec calls
+ // compute_child_statistics for per-partition stats, re-walking the left
+ // subtree at each node. The gap between cached/uncached is smaller than
+ // the linear chain because only the left child triggers a re-walk.
+ let mut group = c.benchmark_group("compute_statistics_cross_join_tree");
+ for depth in [3, 5, 7] {
+ let mut next_col = 0;
+ let plan = build_cross_join_tree(depth, &mut next_col);
+ let label = format!("depth={depth}_leaves={}", 1usize << depth);
+ group.bench_with_input(BenchmarkId::new("cached", &label), &plan, |b, plan| {
+ b.iter(|| {
+ plan.statistics_with_args(&StatisticsArgs::new().with_partition(Some(0)))
+ .unwrap()
+ });
+ });
+ group.bench_with_input(
+ BenchmarkId::new("no_shared_cache", &label),
+ &plan,
+ |b, plan| {
+ b.iter(|| {
+ compute_statistics_without_shared_cache(plan.as_ref(), Some(0))
+ .unwrap()
+ });
+ },
+ );
+ }
+ group.finish();
+
+ // --- Filter chain (partition-preserving linear plan) ---
+ // When called with Some(0), the framework first walks the entire tree
+ // computing None stats, then each filter requests Some(0) on demand.
+ // Both walks are cached, so the total cost is ~2n vs n node visits for None.
+ let mut group = c.benchmark_group("compute_statistics_filter_chain");
+ for depth in [10, 20, 50] {
+ let plan = build_filter_chain(depth);
+ group.bench_with_input(
+ BenchmarkId::new("cached_partition", depth),
+ &plan,
+ |b, plan| {
+ b.iter(|| {
+ plan.statistics_with_args(
+ &StatisticsArgs::new().with_partition(Some(0)),
+ )
+ .unwrap()
+ });
+ },
+ );
+ group.bench_with_input(
+ BenchmarkId::new("cached_overall", depth),
+ &plan,
+ |b, plan| {
+ b.iter(|| plan.statistics_with_args(&StatisticsArgs::new()).unwrap());
+ },
+ );
+ group.bench_with_input(
+ BenchmarkId::new("no_shared_cache", depth),
+ &plan,
+ |b, plan| {
+ b.iter(|| {
+ compute_statistics_without_shared_cache(plan.as_ref(), Some(0))
+ .unwrap()
+ });
+ },
+ );
+ }
+ group.finish();
+
+ // --- Mixed chain (partition-preserving + partition-merging) ---
+ // Alternates Filter (preserving) and CoalescePartitions (merging) to
+ // exercise the cache with both None and Some(p) lookups in a single walk.
+ let mut group = c.benchmark_group("compute_statistics_mixed_chain");
+ for groups in [3, 5, 10] {
+ let plan = build_mixed_chain(groups);
+ let depth = groups * 3; // 2 filters + 1 coalesce per group
+ group.bench_with_input(BenchmarkId::new("cached", depth), &plan, |b, plan| {
+ b.iter(|| {
+ plan.statistics_with_args(&StatisticsArgs::new().with_partition(Some(0)))
+ .unwrap()
+ });
+ });
+ group.bench_with_input(
+ BenchmarkId::new("no_shared_cache", depth),
+ &plan,
+ |b, plan| {
+ b.iter(|| {
+ compute_statistics_without_shared_cache(plan.as_ref(), Some(0))
+ .unwrap()
+ });
+ },
+ );
+ }
+ group.finish();
+}
+
+criterion_group!(benches, bench_compute_statistics);
+criterion_main!(benches);
diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs
index 54e44aa86d..e1c598e02d 100644
--- a/datafusion/physical-plan/src/aggregates/mod.rs
+++ b/datafusion/physical-plan/src/aggregates/mod.rs
@@ -33,6 +33,7 @@ use crate::filter_pushdown::{
FilterPushdownPropagation, PushedDownPredicate,
};
use crate::metrics::{ExecutionPlanMetricsSet, MetricsSet};
+use crate::statistics::StatisticsArgs;
use crate::{
DisplayFormatType, Distribution, ExecutionPlan, InputOrderMode,
SendableRecordBatchStream, Statistics, check_if_same_properties,
@@ -1764,10 +1765,11 @@ impl ExecutionPlan for AggregateExec {
Some(self.metrics.clone_inner())
}
- fn partition_statistics(&self, partition: Option<usize>) -> Result<Arc<Statistics>> {
- let child_statistics = self.input().partition_statistics(partition)?;
+ fn statistics_with_args(&self, args: &StatisticsArgs) -> Result<Arc<Statistics>> {
+ let child_statistics =
+ args.compute_child_statistics(&self.input, args.partition())?;
Ok(Arc::new(
- self.statistics_inner(&child_statistics, partition)?,
+ self.statistics_inner(&child_statistics, args.partition())?,
))
}
@@ -2417,6 +2419,7 @@ mod tests {
use crate::execution_plan::Boundedness;
use crate::expressions::col;
use crate::metrics::MetricValue;
+ use crate::statistics::StatisticsArgs;
use crate::test::TestMemoryExec;
use crate::test::assert_is_pending;
use crate::test::exec::{
@@ -2802,7 +2805,7 @@ mod tests {
)?);
// Verify statistics are preserved proportionally through aggregation
- let final_stats = merged_aggregate.partition_statistics(None)?;
+ let final_stats = merged_aggregate.statistics_with_args(&StatisticsArgs::new())?;
assert!(final_stats.total_byte_size.get_value().is_some());
let task_ctx = if spill {
@@ -2937,11 +2940,8 @@ mod tests {
Ok(Box::pin(stream))
}
- fn partition_statistics(
- &self,
- partition: Option<usize>,
- ) -> Result<Arc<Statistics>> {
- if partition.is_some() {
+ fn statistics_with_args(&self, args: &StatisticsArgs) -> Result<Arc<Statistics>> {
+ if args.partition().is_some() {
return Ok(Arc::new(Statistics::new_unknown(self.schema().as_ref())));
}
let (_, batches) = some_data();
@@ -4741,7 +4741,7 @@ mod tests {
PhysicalGroupBy::default(),
None,
)?;
- let stats = agg.partition_statistics(None)?;
+ let stats = agg.statistics_with_args(&StatisticsArgs::new())?;
assert_eq!(stats.total_byte_size, Precision::Absent);
let zero_row_stats = Statistics {
@@ -4758,7 +4758,7 @@ mod tests {
PhysicalGroupBy::default(),
None,
)?;
- let stats_zero = agg_zero.partition_statistics(None)?;
+ let stats_zero = agg_zero.statistics_with_args(&StatisticsArgs::new())?;
assert_eq!(stats_zero.total_byte_size, Precision::Absent);
let single_input =
@@ -4778,7 +4778,8 @@ mod tests {
.partition_count(),
1
);
- let single_stats_zero = single_agg_zero.partition_statistics(None)?;
+ let single_stats_zero =
+ single_agg_zero.statistics_with_args(&StatisticsArgs::new())?;
assert_eq!(single_stats_zero.num_rows, Precision::Exact(1));
Ok(())
@@ -5151,7 +5152,7 @@ mod tests {
let agg =
build_test_aggregate(&schema, input_stats, group_by, case.limit_options)?;
- let stats = agg.partition_statistics(None)?;
+ let stats = agg.statistics_with_args(&StatisticsArgs::new())?;
assert_eq!(
stats.num_rows, case.expected_num_rows,
"FAILED: '{}' — expected {:?}, got {:?}",
@@ -5190,7 +5191,7 @@ mod tests {
None,
)?;
- let stats = agg.partition_statistics(None)?;
+ let stats = agg.statistics_with_args(&StatisticsArgs::new())?;
assert_eq!(
stats.column_statistics[0].distinct_count,
Precision::Exact(100),
@@ -5244,7 +5245,7 @@ mod tests {
let agg = build_test_aggregate(&schema, input_stats, grouping_set, None)?;
- let stats = agg.partition_statistics(None)?;
+ let stats = agg.statistics_with_args(&StatisticsArgs::new())?;
// Per-set NDV: (a,NULL)=100, (NULL,b)=50, (a,b)=100*50=5000
// Total = 100 + 50 + 5000 = 5150
assert_eq!(
@@ -5274,7 +5275,9 @@ mod tests {
Arc::clone(&schema),
)?;
assert_eq!(
- single_agg.partition_statistics(None)?.num_rows,
+ single_agg
+ .statistics_with_args(&StatisticsArgs::new())?
+ .num_rows,
Precision::Exact(2)
);
@@ -5300,7 +5303,11 @@ mod tests {
let task_ctx = Arc::new(TaskContext::default());
for partition in 0..2 {
assert_eq!(
- partial_agg.partition_statistics(Some(partition))?.num_rows,
+ partial_agg
+ .statistics_with_args(
+ &StatisticsArgs::new().with_partition(Some(partition))
+ )?
+ .num_rows,
Precision::Exact(2)
);
let result =
@@ -5309,7 +5316,9 @@ mod tests {
}
assert_eq!(
- partial_agg.partition_statistics(None)?.num_rows,
+ partial_agg
+ .statistics_with_args(&StatisticsArgs::new())?
+ .num_rows,
Precision::Exact(4)
);
@@ -5353,7 +5362,7 @@ mod tests {
PhysicalGroupBy::new_single(vec![(expr_a_plus_b, "a+b".to_string())]);
let agg = build_test_aggregate(&schema, input_stats, group_by, None)?;
- let stats = agg.partition_statistics(None)?;
+ let stats = agg.statistics_with_args(&StatisticsArgs::new())?;
assert_eq!(
stats.num_rows,
Precision::Inexact(1_000_000),
diff --git a/datafusion/physical-plan/src/buffer.rs b/datafusion/physical-plan/src/buffer.rs
index 2985dc5766..871d3c4d3f 100644
--- a/datafusion/physical-plan/src/buffer.rs
+++ b/datafusion/physical-plan/src/buffer.rs
@@ -24,6 +24,7 @@ use crate::filter_pushdown::{
FilterPushdownPropagation,
};
use crate::projection::ProjectionExec;
+use crate::statistics::StatisticsArgs;
use crate::stream::RecordBatchStreamAdapter;
use crate::{
DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, SortOrderPushdownResult,
@@ -237,8 +238,8 @@ impl ExecutionPlan for BufferExec {
Some(self.metrics.clone_inner())
}
- fn partition_statistics(&self, partition: Option<usize>) -> Result<Arc<Statistics>> {
- self.input.partition_statistics(partition)
+ fn statistics_with_args(&self, args: &StatisticsArgs) -> Result<Arc<Statistics>> {
+ args.compute_child_statistics(&self.input, args.partition())
}
fn supports_limit_pushdown(&self) -> bool {
diff --git a/datafusion/physical-plan/src/coalesce_batches.rs b/datafusion/physical-plan/src/coalesce_batches.rs
index 76b2f63798..59b3138b55 100644
--- a/datafusion/physical-plan/src/coalesce_batches.rs
+++ b/datafusion/physical-plan/src/coalesce_batches.rs
@@ -24,6 +24,7 @@ use std::task::{Context, Poll};
use super::metrics::{BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet};
use super::{DisplayAs, ExecutionPlanProperties, PlanProperties, Statistics};
use crate::projection::ProjectionExec;
+use crate::statistics::StatisticsArgs;
use crate::stream::EmptyRecordBatchStream;
use crate::{
DisplayFormatType, ExecutionPlan, RecordBatchStream, SendableRecordBatchStream,
@@ -215,8 +216,10 @@ impl ExecutionPlan for CoalesceBatchesExec {
Some(self.metrics.clone_inner())
}
- fn partition_statistics(&self, partition: Option<usize>) -> Result<Arc<Statistics>> {
- let stats = Arc::unwrap_or_clone(self.input.partition_statistics(partition)?);
+ fn statistics_with_args(&self, args: &StatisticsArgs) -> Result<Arc<Statistics>> {
+ let stats = Arc::unwrap_or_clone(
+ args.compute_child_statistics(&self.input, args.partition())?,
+ );
Ok(Arc::new(stats.with_fetch(self.fetch, 0, 1)?))
}
diff --git a/datafusion/physical-plan/src/coalesce_partitions.rs b/datafusion/physical-plan/src/coalesce_partitions.rs
index fa200ef845..0a8c5f7888 100644
--- a/datafusion/physical-plan/src/coalesce_partitions.rs
+++ b/datafusion/physical-plan/src/coalesce_partitions.rs
@@ -30,6 +30,7 @@ use crate::execution_plan::{CardinalityEffect, EvaluationType, SchedulingType};
use crate::filter_pushdown::{FilterDescription, FilterPushdownPhase};
use crate::projection::{ProjectionExec, make_with_child};
use crate::sort_pushdown::SortOrderPushdownResult;
+use crate::statistics::StatisticsArgs;
use crate::{DisplayFormatType, ExecutionPlan, Partitioning, check_if_same_properties};
use datafusion_physical_expr_common::sort_expr::PhysicalSortExpr;
@@ -231,8 +232,9 @@ impl ExecutionPlan for CoalescePartitionsExec {
Some(self.metrics.clone_inner())
}
- fn partition_statistics(&self, _partition: Option<usize>) -> Result<Arc<Statistics>> {
- let stats = Arc::unwrap_or_clone(self.input.partition_statistics(None)?);
+ fn statistics_with_args(&self, args: &StatisticsArgs) -> Result<Arc<Statistics>> {
+ let stats =
+ Arc::unwrap_or_clone(args.compute_child_statistics(&self.input, None)?);
Ok(Arc::new(stats.with_fetch(self.fetch, 0, 1)?))
}
diff --git a/datafusion/physical-plan/src/coop.rs b/datafusion/physical-plan/src/coop.rs
index 111999b71c..7bd84a3a6b 100644
--- a/datafusion/physical-plan/src/coop.rs
+++ b/datafusion/physical-plan/src/coop.rs
@@ -84,6 +84,7 @@ use crate::filter_pushdown::{
FilterPushdownPropagation,
};
use crate::projection::ProjectionExec;
+use crate::statistics::StatisticsArgs;
use crate::{
DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, RecordBatchStream,
SendableRecordBatchStream, SortOrderPushdownResult, check_if_same_properties,
@@ -298,8 +299,8 @@ impl ExecutionPlan for CooperativeExec {
Ok(make_cooperative(child_stream))
}
- fn partition_statistics(&self, partition: Option<usize>) -> Result<Arc<Statistics>> {
- self.input.partition_statistics(partition)
+ fn statistics_with_args(&self, args: &StatisticsArgs) -> Result<Arc<Statistics>> {
+ args.compute_child_statistics(&self.input, args.partition())
}
fn supports_limit_pushdown(&self) -> bool {
diff --git a/datafusion/physical-plan/src/display.rs b/datafusion/physical-plan/src/display.rs
index 4642a9a4b1..164637f760 100644
--- a/datafusion/physical-plan/src/display.rs
+++ b/datafusion/physical-plan/src/display.rs
@@ -32,6 +32,8 @@ use datafusion_physical_expr::LexOrdering;
use crate::metrics::{MetricCategory, MetricType, MetricValue};
use crate::render_tree::RenderTree;
+use crate::statistics::StatisticsArgs;
+
use super::{ExecutionPlan, ExecutionPlanVisitor, accept};
/// Options for controlling how each [`ExecutionPlan`] should format itself
@@ -579,7 +581,9 @@ impl ExecutionPlanVisitor for IndentVisitor<'_, '_> {
}
}
if self.show_statistics {
- let stats = plan.partition_statistics(None).map_err(|_e| fmt::Error)?;
+ let stats = plan
+ .statistics_with_args(&StatisticsArgs::default())
+ .map_err(|_e| fmt::Error)?;
write!(self.f, ", statistics=[{stats}]")?;
}
if self.show_schema {
@@ -675,7 +679,9 @@ impl ExecutionPlanVisitor for GraphvizVisitor<'_, '_> {
};
let statistics = if self.show_statistics {
- let stats = plan.partition_statistics(None).map_err(|_e| fmt::Error)?;
+ let stats = plan
+ .statistics_with_args(&StatisticsArgs::new())
+ .map_err(|_e| fmt::Error)?;
format!("statistics=[{stats}]")
} else {
"".to_string()
@@ -1445,6 +1451,7 @@ mod tests {
use datafusion_common::{Result, Statistics, internal_datafusion_err};
use datafusion_execution::{SendableRecordBatchStream, TaskContext};
+ use crate::statistics::StatisticsArgs;
use crate::{DisplayAs, ExecutionPlan, PlanProperties};
use super::DisplayableExecutionPlan;
@@ -1494,11 +1501,8 @@ mod tests {
todo!()
}
- fn partition_statistics(
- &self,
- partition: Option<usize>,
- ) -> Result<Arc<Statistics>> {
- if partition.is_some() {
+ fn statistics_with_args(&self, args: &StatisticsArgs) -> Result<Arc<Statistics>> {
+ if args.partition().is_some() {
return Ok(Arc::new(Statistics::new_unknown(self.schema().as_ref())));
}
match self {
diff --git a/datafusion/physical-plan/src/empty.rs b/datafusion/physical-plan/src/empty.rs
index 2e7f982a51..a8f4af5b3d 100644
--- a/datafusion/physical-plan/src/empty.rs
+++ b/datafusion/physical-plan/src/empty.rs
@@ -34,6 +34,7 @@ use datafusion_execution::TaskContext;
use datafusion_physical_expr::EquivalenceProperties;
use crate::execution_plan::SchedulingType;
+use crate::statistics::StatisticsArgs;
use log::trace;
/// Execution plan for empty relation with produce_one_row=false
@@ -151,8 +152,8 @@ impl ExecutionPlan for EmptyExec {
)?))
}
- fn partition_statistics(&self, partition: Option<usize>) -> Result<Arc<Statistics>> {
- if let Some(partition) = partition {
+ fn statistics_with_args(&self, args: &StatisticsArgs) -> Result<Arc<Statistics>> {
+ if let Some(partition) = args.partition() {
assert_or_internal_err!(
partition < self.partitions,
"EmptyExec invalid partition {} (expected less than {})",
diff --git a/datafusion/physical-plan/src/execution_plan.rs b/datafusion/physical-plan/src/execution_plan.rs
index 8577e86f00..76abf73e0e 100644
--- a/datafusion/physical-plan/src/execution_plan.rs
+++ b/datafusion/physical-plan/src/execution_plan.rs
@@ -47,6 +47,7 @@ use crate::metrics::MetricsSet;
use crate::projection::ProjectionExec;
use crate::repartition::RepartitionExec;
use crate::sorts::sort_preserving_merge::SortPreservingMergeExec;
+use crate::statistics::StatisticsArgs;
use crate::stream::RecordBatchStreamAdapter;
use arrow::array::{Array, RecordBatch};
@@ -496,9 +497,11 @@ pub trait ExecutionPlan: Any + Debug + DisplayAs + Send + Sync {
}
/// Returns statistics for a specific partition of this `ExecutionPlan` node.
- /// If statistics are not available, should return [`Statistics::new_unknown`]
- /// (the default), not an error.
- /// If `partition` is `None`, it returns statistics for the entire plan.
+ ///
+ /// Deprecated: use [`Self::statistics_with_args`] instead,
+ /// which accepts a [`StatisticsArgs`] carrying pre-computed child
+ /// statistics.
+ #[deprecated(since = "55.0.0", note = "Use statistics_with_args instead")]
fn partition_statistics(&self, partition: Option<usize>) -> Result<Arc<Statistics>> {
if let Some(idx) = partition {
// Validate partition index
@@ -513,6 +516,21 @@ pub trait ExecutionPlan: Any + Debug + DisplayAs + Send + Sync {
Ok(Arc::new(Statistics::new_unknown(&self.schema())))
}
+ /// Returns statistics for a specific partition of this `ExecutionPlan` node.
+ /// If statistics are not available, should return [`Statistics::new_unknown`]
+ /// (the default), not an error.
+ /// If `partition` is `None`, it returns statistics for all partitions.
+ ///
+ /// [`StatisticsArgs`] carries the partition index and a shared cache.
+ /// Create one with [`StatisticsArgs::new`] and pass it to this method.
+ ///
+ /// [`StatisticsArgs`]: crate::statistics::StatisticsArgs
+ /// [`StatisticsArgs::new`]: crate::statistics::StatisticsArgs::new
+ fn statistics_with_args(&self, args: &StatisticsArgs) -> Result<Arc<Statistics>> {
+ #[expect(deprecated)]
+ self.partition_statistics(args.partition())
+ }
+
/// Returns `true` if a limit can be safely pushed down through this
/// `ExecutionPlan` node.
///
@@ -1639,9 +1657,9 @@ mod tests {
unimplemented!()
}
- fn partition_statistics(
+ fn statistics_with_args(
&self,
- _partition: Option<usize>,
+ _args: &StatisticsArgs,
) -> Result<Arc<Statistics>> {
unimplemented!()
}
@@ -1701,9 +1719,9 @@ mod tests {
unimplemented!()
}
- fn partition_statistics(
+ fn statistics_with_args(
&self,
- _partition: Option<usize>,
+ _args: &StatisticsArgs,
) -> Result<Arc<Statistics>> {
unimplemented!()
}
diff --git a/datafusion/physical-plan/src/filter.rs b/datafusion/physical-plan/src/filter.rs
index 11d36192f3..d23dd38042 100644
--- a/datafusion/physical-plan/src/filter.rs
+++ b/datafusion/physical-plan/src/filter.rs
@@ -42,6 +42,7 @@ use crate::projection::{
EmbeddedProjection, ProjectionExec, ProjectionExpr, make_with_child,
try_embed_projection, update_expr,
};
+use crate::statistics::StatisticsArgs;
use crate::stream::EmptyRecordBatchStream;
use crate::{
DisplayFormatType, ExecutionPlan,
@@ -421,7 +422,7 @@ impl FilterExec {
let schema = input.schema();
let stats = Self::statistics_helper(
&schema,
- Arc::unwrap_or_clone(input.partition_statistics(None)?),
+ Arc::unwrap_or_clone(input.statistics_with_args(&StatisticsArgs::new())?),
predicate,
default_selectivity,
)?;
@@ -590,9 +591,10 @@ impl ExecutionPlan for FilterExec {
/// The output statistics of a filtering operation can be estimated if the
/// predicate's selectivity value can be determined for the incoming data.
- fn partition_statistics(&self, partition: Option<usize>) -> Result<Arc<Statistics>> {
- let input_stats =
- Arc::unwrap_or_clone(self.input.partition_statistics(partition)?);
+ fn statistics_with_args(&self, args: &StatisticsArgs) -> Result<Arc<Statistics>> {
+ let input_stats = Arc::unwrap_or_clone(
+ args.compute_child_statistics(&self.input, args.partition())?,
+ );
let stats = Self::statistics_helper(
&self.input.schema(),
input_stats,
@@ -1266,6 +1268,7 @@ mod tests {
use super::*;
use crate::empty::EmptyExec;
use crate::expressions::*;
+ use crate::statistics::StatisticsArgs;
use crate::test;
use crate::test::exec::StatisticsExec;
use arrow::datatypes::{Field, Schema, UnionFields, UnionMode};
@@ -1342,7 +1345,7 @@ mod tests {
let filter: Arc<dyn ExecutionPlan> =
Arc::new(FilterExec::try_new(predicate, input)?);
- let statistics = filter.partition_statistics(None)?;
+ let statistics = filter.statistics_with_args(&StatisticsArgs::new())?;
assert_eq!(statistics.num_rows, Precision::Inexact(25));
assert_eq!(
statistics.total_byte_size,
@@ -1394,7 +1397,7 @@ mod tests {
sub_filter,
)?);
- let statistics = filter.partition_statistics(None)?;
+ let statistics = filter.statistics_with_args(&StatisticsArgs::new())?;
assert_eq!(statistics.num_rows, Precision::Inexact(16));
assert_eq!(
statistics.column_statistics,
@@ -1456,7 +1459,7 @@ mod tests {
binary(col("a", &schema)?, Operator::GtEq, lit(10i32), &schema)?,
b_gt_5,
)?);
- let statistics = filter.partition_statistics(None)?;
+ let statistics = filter.statistics_with_args(&StatisticsArgs::new())?;
// On a uniform distribution, only fifteen rows will satisfy the
// filter that 'a' proposed (a >= 10 AND a <= 25) (15/100) and only
// 5 rows will satisfy the filter that 'b' proposed (b > 45) (5/50).
@@ -1506,7 +1509,7 @@ mod tests {
let filter: Arc<dyn ExecutionPlan> =
Arc::new(FilterExec::try_new(predicate, input)?);
- let statistics = filter.partition_statistics(None)?;
+ let statistics = filter.statistics_with_args(&StatisticsArgs::new())?;
assert_eq!(statistics.num_rows, Precision::Absent);
Ok(())
@@ -1579,7 +1582,7 @@ mod tests {
));
let filter: Arc<dyn ExecutionPlan> =
Arc::new(FilterExec::try_new(predicate, input)?);
- let statistics = filter.partition_statistics(None)?;
+ let statistics = filter.statistics_with_args(&StatisticsArgs::new())?;
// 0.5 (from a) * 0.333333... (from b) * 0.798387... (from c) ≈ 0.1330...
// num_rows after ceil => 133.0... => 134
// total_byte_size after ceil => 532.0... => 533
@@ -1677,13 +1680,16 @@ mod tests {
// The filter predicate passes all (non-null) entries, so min/max/NDV
// are unchanged. `a < 200` and `1 <= b` are null-rejecting, though, so
// both columns lose any nulls regardless of selectivity.
- let mut expected = input.partition_statistics(None)?.column_statistics.clone();
+ let mut expected = input
+ .statistics_with_args(&StatisticsArgs::new())?
+ .column_statistics
+ .clone();
for col in &mut expected {
col.null_count = Precision::Exact(0);
}
let filter: Arc<dyn ExecutionPlan> =
Arc::new(FilterExec::try_new(predicate, input)?);
- let statistics = filter.partition_statistics(None)?;
+ let statistics = filter.statistics_with_args(&StatisticsArgs::new())?;
assert_eq!(statistics.num_rows, Precision::Inexact(1000));
assert_eq!(statistics.total_byte_size, Precision::Inexact(4000));
@@ -1736,7 +1742,7 @@ mod tests {
));
let filter: Arc<dyn ExecutionPlan> =
Arc::new(FilterExec::try_new(predicate, input)?);
- let statistics = filter.partition_statistics(None)?;
+ let statistics = filter.statistics_with_args(&StatisticsArgs::new())?;
assert_eq!(statistics.num_rows, Precision::Inexact(0));
assert_eq!(statistics.total_byte_size, Precision::Inexact(0));
@@ -1823,7 +1829,7 @@ mod tests {
Arc::new(FilterExec::try_new(outer_predicate, inner_filter)?);
// Should succeed without error
- let statistics = outer_filter.partition_statistics(None)?;
+ let statistics = outer_filter.statistics_with_args(&StatisticsArgs::new())?;
assert_eq!(statistics.num_rows, Precision::Inexact(0));
Ok(())
@@ -1862,7 +1868,7 @@ mod tests {
));
let filter: Arc<dyn ExecutionPlan> =
Arc::new(FilterExec::try_new(predicate, input)?);
- let statistics = filter.partition_statistics(None)?;
+ let statistics = filter.statistics_with_args(&StatisticsArgs::new())?;
assert_eq!(statistics.num_rows, Precision::Inexact(490));
assert_eq!(statistics.total_byte_size, Precision::Inexact(1960));
@@ -1916,7 +1922,7 @@ mod tests {
));
let filter: Arc<dyn ExecutionPlan> =
Arc::new(FilterExec::try_new(predicate, input)?);
- let filter_statistics = filter.partition_statistics(None)?;
+ let filter_statistics = filter.statistics_with_args(&StatisticsArgs::new())?;
let expected_filter_statistics = Statistics {
num_rows: Precision::Absent,
@@ -1953,7 +1959,7 @@ mod tests {
));
let filter: Arc<dyn ExecutionPlan> =
Arc::new(FilterExec::try_new(predicate, input)?);
- let filter_statistics = filter.partition_statistics(None)?;
+ let filter_statistics = filter.statistics_with_args(&StatisticsArgs::new())?;
// First column is "a", and it is a column with only one value after the filter.
assert!(filter_statistics.column_statistics[0].is_singleton());
@@ -2000,11 +2006,11 @@ mod tests {
Arc::new(Literal::new(ScalarValue::Decimal128(Some(10), 10, 10))),
));
let filter = FilterExec::try_new(predicate, input)?;
- let statistics = filter.partition_statistics(None)?;
+ let statistics = filter.statistics_with_args(&StatisticsArgs::new())?;
assert_eq!(statistics.num_rows, Precision::Inexact(200));
assert_eq!(statistics.total_byte_size, Precision::Inexact(800));
let filter = filter.with_default_selectivity(40)?;
- let statistics = filter.partition_statistics(None)?;
+ let statistics = filter.statistics_with_args(&StatisticsArgs::new())?;
assert_eq!(statistics.num_rows, Precision::Inexact(400));
assert_eq!(statistics.total_byte_size, Precision::Inexact(1600));
Ok(())
@@ -2039,7 +2045,7 @@ mod tests {
Arc::new(EmptyExec::new(Arc::clone(&schema))),
)?;
- exec.partition_statistics(None).unwrap();
+ exec.statistics_with_args(&StatisticsArgs::new()).unwrap();
Ok(())
}
@@ -2195,8 +2201,8 @@ mod tests {
assert_eq!(filter1.projection(), filter2.projection());
// Verify statistics are the same
- let stats1 = filter1.partition_statistics(None)?;
- let stats2 = filter2.partition_statistics(None)?;
+ let stats1 = filter1.statistics_with_args(&StatisticsArgs::new())?;
+ let stats2 = filter2.statistics_with_args(&StatisticsArgs::new())?;
assert_eq!(stats1.num_rows, stats2.num_rows);
assert_eq!(stats1.total_byte_size, stats2.total_byte_size);
@@ -2249,7 +2255,7 @@ mod tests {
.unwrap()
.build()?;
- let statistics = filter.partition_statistics(None)?;
+ let statistics = filter.statistics_with_args(&StatisticsArgs::new())?;
// Verify statistics reflect both filtering and projection
assert!(matches!(statistics.num_rows, Precision::Inexact(_)));
@@ -2480,7 +2486,7 @@ mod tests {
let filter: Arc<dyn ExecutionPlan> =
Arc::new(FilterExec::try_new(predicate, input)?);
- let statistics = filter.partition_statistics(None)?;
+ let statistics = filter.statistics_with_args(&StatisticsArgs::new())?;
let col_b_stats = &statistics.column_statistics[1];
assert_eq!(col_b_stats.min_value, Precision::Absent);
assert_eq!(col_b_stats.max_value, Precision::Absent);
@@ -2767,7 +2773,7 @@ mod tests {
));
let filter: Arc<dyn ExecutionPlan> =
Arc::new(FilterExec::try_new(predicate, input)?);
- let statistics = filter.partition_statistics(None)?;
+ let statistics = filter.statistics_with_args(&StatisticsArgs::new())?;
for (i, expected) in expected_ndvs.iter().enumerate() {
assert_eq!(
@@ -2828,7 +2834,7 @@ mod tests {
let input = Arc::new(StatisticsExec::new(input_stats, schema));
let filter: Arc<dyn ExecutionPlan> =
Arc::new(FilterExec::try_new(predicate, input)?);
- let statistics = filter.partition_statistics(None)?;
+ let statistics = filter.statistics_with_args(&StatisticsArgs::new())?;
assert_eq!(
statistics.num_rows,
@@ -2905,7 +2911,7 @@ mod tests {
));
let filter: Arc<dyn ExecutionPlan> =
Arc::new(FilterExec::try_new(predicate, input)?);
- let statistics = filter.partition_statistics(None)?;
+ let statistics = filter.statistics_with_args(&StatisticsArgs::new())?;
// Equality predicates collapse NDV and reject nulls for their columns.
assert_eq!(
statistics.column_statistics[0].distinct_count,
@@ -2958,7 +2964,7 @@ mod tests {
));
let filter: Arc<dyn ExecutionPlan> =
Arc::new(FilterExec::try_new(predicate, input)?);
- let statistics = filter.partition_statistics(None)?;
+ let statistics = filter.statistics_with_args(&StatisticsArgs::new())?;
assert_eq!(
statistics.column_statistics[0].distinct_count,
Precision::Exact(1)
@@ -2991,7 +2997,7 @@ mod tests {
));
let filter: Arc<dyn ExecutionPlan> =
Arc::new(FilterExec::try_new(predicate, input)?);
- let statistics = filter.partition_statistics(None)?;
+ let statistics = filter.statistics_with_args(&StatisticsArgs::new())?;
assert_eq!(
statistics.column_statistics[0].distinct_count,
Precision::Exact(1)
@@ -3024,7 +3030,7 @@ mod tests {
));
let filter: Arc<dyn ExecutionPlan> =
Arc::new(FilterExec::try_new(predicate, input)?);
- let statistics = filter.partition_statistics(None)?;
+ let statistics = filter.statistics_with_args(&StatisticsArgs::new())?;
assert_eq!(
statistics.column_statistics[0].distinct_count,
Precision::Exact(1)
@@ -3057,7 +3063,7 @@ mod tests {
));
let filter: Arc<dyn ExecutionPlan> =
Arc::new(FilterExec::try_new(predicate, input)?);
- let statistics = filter.partition_statistics(None)?;
+ let statistics = filter.statistics_with_args(&StatisticsArgs::new())?;
assert_eq!(
statistics.column_statistics[0].distinct_count,
Precision::Exact(1)
@@ -3091,7 +3097,7 @@ mod tests {
));
let filter: Arc<dyn ExecutionPlan> =
Arc::new(FilterExec::try_new(predicate, input)?);
- let statistics = filter.partition_statistics(None)?;
+ let statistics = filter.statistics_with_args(&StatisticsArgs::new())?;
assert_eq!(
statistics.column_statistics[0].distinct_count,
Precision::Exact(1)
@@ -3137,7 +3143,7 @@ mod tests {
));
let filter: Arc<dyn ExecutionPlan> =
Arc::new(FilterExec::try_new(predicate, input)?);
- let statistics = filter.partition_statistics(None)?;
+ let statistics = filter.statistics_with_args(&StatisticsArgs::new())?;
assert_eq!(
statistics.column_statistics[0].distinct_count,
Precision::Exact(1)
@@ -3439,7 +3445,8 @@ mod tests {
let filter: Arc<dyn ExecutionPlan> =
Arc::new(FilterExec::try_new(predicate, input)?);
- let statistics = filter.partition_statistics(None)?;
+ let statistics = filter.statistics_with_args(&StatisticsArgs::new())?;
+ // Filter estimates ~10 rows (selectivity = 10/100)
assert_eq!(statistics.num_rows, Precision::Inexact(10));
let ndv = &statistics.column_statistics[0].distinct_count;
assert!(
@@ -3484,7 +3491,7 @@ mod tests {
let filter: Arc<dyn ExecutionPlan> =
Arc::new(FilterExec::try_new(predicate, input)?);
- let statistics = filter.partition_statistics(None)?;
+ let statistics = filter.statistics_with_args(&StatisticsArgs::new())?;
assert_eq!(statistics.num_rows, Precision::Inexact(20));
assert_eq!(
statistics.column_statistics[0].null_count,
@@ -3527,7 +3534,7 @@ mod tests {
let filter: Arc<dyn ExecutionPlan> =
Arc::new(FilterExec::try_new(predicate, input)?);
- let statistics = filter.partition_statistics(None)?;
+ let statistics = filter.statistics_with_args(&StatisticsArgs::new())?;
assert_eq!(statistics.num_rows, Precision::Inexact(20));
assert_eq!(
statistics.column_statistics[0].null_count,
@@ -3568,7 +3575,7 @@ mod tests {
let filter: Arc<dyn ExecutionPlan> =
Arc::new(FilterExec::try_new(predicate, input)?);
- let statistics = filter.partition_statistics(None)?;
+ let statistics = filter.statistics_with_args(&StatisticsArgs::new())?;
assert_eq!(statistics.num_rows, Precision::Inexact(20));
assert_eq!(
statistics.column_statistics[0].null_count,
diff --git a/datafusion/physical-plan/src/joins/cross_join.rs b/datafusion/physical-plan/src/joins/cross_join.rs
index 45b34692ab..79295ba2fb 100644
--- a/datafusion/physical-plan/src/joins/cross_join.rs
+++ b/datafusion/physical-plan/src/joins/cross_join.rs
@@ -31,6 +31,7 @@ use crate::projection::{
ProjectionExec, join_allows_pushdown, join_table_borders, new_join_children,
physical_to_column_exprs,
};
+use crate::statistics::StatisticsArgs;
use crate::stream::EmptyRecordBatchStream;
use crate::{
ColumnStatistics, DisplayAs, DisplayFormatType, Distribution, ExecutionPlan,
@@ -371,11 +372,14 @@ impl ExecutionPlan for CrossJoinExec {
}
}
- fn partition_statistics(&self, partition: Option<usize>) -> Result<Arc<Statistics>> {
- // Get the all partitions statistics of the left
- let left_stats = Arc::unwrap_or_clone(self.left.partition_statistics(None)?);
- let right_stats =
- Arc::unwrap_or_clone(self.right.partition_statistics(partition)?);
+ fn statistics_with_args(&self, args: &StatisticsArgs) -> Result<Arc<Statistics>> {
+ // Left side is always broadcast, so it always needs overall stats
+ let left_stats =
+ Arc::unwrap_or_clone(args.compute_child_statistics(&self.left, None)?);
+ // Right side is partitioned, so it needs per-partition stats
+ let right_stats = Arc::unwrap_or_clone(
+ args.compute_child_statistics(&self.right, args.partition())?,
+ );
Ok(Arc::new(stats_cartesian_product(left_stats, right_stats)))
}
diff --git a/datafusion/physical-plan/src/joins/hash_join/exec.rs b/datafusion/physical-plan/src/joins/hash_join/exec.rs
index 5ef767ebdd..bb9ebcd4e6 100644
--- a/datafusion/physical-plan/src/joins/hash_join/exec.rs
+++ b/datafusion/physical-plan/src/joins/hash_join/exec.rs
@@ -52,6 +52,7 @@ use crate::projection::{
try_pushdown_through_join,
};
use crate::repartition::REPARTITION_RANDOM_STATE;
+use crate::statistics::StatisticsArgs;
use crate::{
DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, Partitioning,
PlanProperties, SendableRecordBatchStream, Statistics,
@@ -1449,14 +1450,14 @@ impl ExecutionPlan for HashJoinExec {
Some(self.metrics.clone_inner())
}
- fn partition_statistics(&self, partition: Option<usize>) -> Result<Arc<Statistics>> {
- let stats = match (partition, self.mode) {
- // For CollectLeft mode, the left side is collected into a single partition,
- // so all left partitions are available to each output partition.
- // For the right side, we need the specific partition statistics.
- (Some(partition), PartitionMode::CollectLeft) => {
- let left_stats = self.left.partition_statistics(None)?;
- let right_stats = self.right.partition_statistics(Some(partition))?;
+ fn statistics_with_args(&self, args: &StatisticsArgs) -> Result<Arc<Statistics>> {
+ let stats = match (args.partition(), self.mode) {
+ // Left side is broadcast, so it always needs overall stats
+ // Right side is partitioned, so it needs per-partition stats
+ (Some(_), PartitionMode::CollectLeft) => {
+ let left_stats = args.compute_child_statistics(&self.left, None)?;
+ let right_stats =
+ args.compute_child_statistics(&self.right, args.partition())?;
estimate_join_statistics(
Arc::unwrap_or_clone(left_stats),
@@ -1468,11 +1469,13 @@ impl ExecutionPlan for HashJoinExec {
)?
}
- // For Partitioned mode, both sides are partitioned, so each output partition
- // only has access to the corresponding partition from both sides.
- (Some(partition), PartitionMode::Partitioned) => {
- let left_stats = self.left.partition_statistics(Some(partition))?;
- let right_stats = self.right.partition_statistics(Some(partition))?;
+ // For Partitioned mode, both sides are hash-partitioned symmetrically,
+ // so each output partition uses the matching partition from both sides.
+ (Some(_), PartitionMode::Partitioned) => {
+ let left_stats =
+ args.compute_child_statistics(&self.left, args.partition())?;
+ let right_stats =
+ args.compute_child_statistics(&self.right, args.partition())?;
estimate_join_statistics(
Arc::unwrap_or_clone(left_stats),
@@ -1484,14 +1487,25 @@ impl ExecutionPlan for HashJoinExec {
)?
}
- // For Auto mode or when no specific partition is requested, fall back to
- // the current behavior of getting all partition statistics.
- (None, _) | (Some(_), PartitionMode::Auto) => {
- // TODO stats: it is not possible in general to know the output size of joins
- // There are some special cases though, for example:
- // - `A LEFT JOIN B ON A.col=B.col` with `COUNT_DISTINCT(B.col)=COUNT(B.col)`
- let left_stats = self.left.partition_statistics(None)?;
- let right_stats = self.right.partition_statistics(None)?;
+ // Overall stats requested, look up overall child stats.
+ (None, _) => {
+ let left_stats = args.compute_child_statistics(&self.left, None)?;
+ let right_stats = args.compute_child_statistics(&self.right, None)?;
+ estimate_join_statistics(
+ Arc::unwrap_or_clone(left_stats),
+ Arc::unwrap_or_clone(right_stats),
+ &self.on,
+ self.null_equality,
+ &self.join_type,
+ &self.join_schema,
+ )?
+ }
+
+ // Auto mode hasn't decided partitioning yet, so it needs
+ // overall stats from both sides.
+ (Some(_), PartitionMode::Auto) => {
+ let left_stats = args.compute_child_statistics(&self.left, None)?;
+ let right_stats = args.compute_child_statistics(&self.right, None)?;
estimate_join_statistics(
Arc::unwrap_or_clone(left_stats),
Arc::unwrap_or_clone(right_stats),
diff --git a/datafusion/physical-plan/src/joins/nested_loop_join.rs b/datafusion/physical-plan/src/joins/nested_loop_join.rs
index d13e172352..db552fed96 100644
--- a/datafusion/physical-plan/src/joins/nested_loop_join.rs
+++ b/datafusion/physical-plan/src/joins/nested_loop_join.rs
@@ -42,6 +42,7 @@ use crate::projection::{
EmbeddedProjection, JoinData, ProjectionExec, try_embed_projection,
try_pushdown_through_join,
};
+use crate::statistics::StatisticsArgs;
use crate::{
DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, ExecutionPlanProperties,
PlanProperties, RecordBatchStream, SendableRecordBatchStream,
@@ -690,7 +691,7 @@ impl ExecutionPlan for NestedLoopJoinExec {
Some(self.metrics.clone_inner())
}
- fn partition_statistics(&self, partition: Option<usize>) -> Result<Arc<Statistics>> {
+ fn statistics_with_args(&self, args: &StatisticsArgs) -> Result<Arc<Statistics>> {
// NestedLoopJoinExec is designed for joins without equijoin keys in the
// ON clause (e.g., `t1 JOIN t2 ON (t1.v1 + t2.v1) % 2 = 0`). Any join
// predicates are stored in `self.filter`, but `estimate_join_statistics`
@@ -700,15 +701,13 @@ impl ExecutionPlan for NestedLoopJoinExec {
// unknown row counts.
let join_columns = Vec::new();
- // Left side is always a single partition (Distribution::SinglePartition),
- // so we always request overall stats with `None`. Right side can have
- // multiple partitions, so we forward the partition parameter to get
- // partition-specific statistics when requested.
- let left_stats = Arc::unwrap_or_clone(self.left.partition_statistics(None)?);
- let right_stats = Arc::unwrap_or_clone(match partition {
- Some(partition) => self.right.partition_statistics(Some(partition))?,
- None => self.right.partition_statistics(None)?,
- });
+ // Left side is always broadcast, so it always needs overall stats
+ let left_stats =
+ Arc::unwrap_or_clone(args.compute_child_statistics(&self.left, None)?);
+ // Right side is partitioned, so it needs per-partition stats
+ let right_stats = Arc::unwrap_or_clone(
+ args.compute_child_statistics(&self.right, args.partition())?,
+ );
let stats = estimate_join_statistics(
left_stats,
@@ -3064,6 +3063,7 @@ fn build_unmatched_batch(
#[cfg(test)]
pub(crate) mod tests {
use super::*;
+ use crate::statistics::StatisticsArgs;
use crate::test::{TestMemoryExec, assert_join_metrics};
use crate::{
common, expressions::Column, repartition::RepartitionExec, test::build_table_i32,
@@ -3443,7 +3443,7 @@ pub(crate) mod tests {
&JoinType::Left,
Some(vec![1, 2]),
)?;
- let stats = nested_loop_join.partition_statistics(None)?;
+ let stats = nested_loop_join.statistics_with_args(&StatisticsArgs::new())?;
assert_eq!(
nested_loop_join.schema().fields().len(),
stats.column_statistics.len(),
diff --git a/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs b/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs
index a86cb647e4..a8d25fd002 100644
--- a/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs
+++ b/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs
@@ -38,6 +38,7 @@ use crate::projection::{
physical_to_column_exprs, update_join_on,
};
use crate::spill::spill_manager::SpillManager;
+use crate::statistics::StatisticsArgs;
use crate::{
DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, ExecutionPlanProperties,
PlanProperties, SendableRecordBatchStream, Statistics, check_if_same_properties,
@@ -563,21 +564,20 @@ impl ExecutionPlan for SortMergeJoinExec {
Some(self.metrics.clone_inner())
}
- fn partition_statistics(&self, partition: Option<usize>) -> Result<Arc<Statistics>> {
+ fn statistics_with_args(&self, args: &StatisticsArgs) -> Result<Arc<Statistics>> {
// SortMergeJoinExec uses symmetric hash partitioning where both left and right
// inputs are hash-partitioned on the join keys. This means partition `i` of the
// left input is joined with partition `i` of the right input.
//
- // Therefore, partition-specific statistics can be computed by getting the
- // partition-specific statistics from both children and combining them via
- // `estimate_join_statistics`.
- //
// TODO stats: it is not possible in general to know the output size of joins
// There are some special cases though, for example:
// - `A LEFT JOIN B ON A.col=B.col` with `COUNT_DISTINCT(B.col)=COUNT(B.col)`
- let left_stats = Arc::unwrap_or_clone(self.left.partition_statistics(partition)?);
- let right_stats =
- Arc::unwrap_or_clone(self.right.partition_statistics(partition)?);
+ let left_stats = Arc::unwrap_or_clone(
+ args.compute_child_statistics(&self.left, args.partition())?,
+ );
+ let right_stats = Arc::unwrap_or_clone(
+ args.compute_child_statistics(&self.right, args.partition())?,
+ );
Ok(Arc::new(estimate_join_statistics(
left_stats,
right_stats,
diff --git a/datafusion/physical-plan/src/joins/sort_merge_join/tests.rs b/datafusion/physical-plan/src/joins/sort_merge_join/tests.rs
index 0347299dd0..338c5111d2 100644
--- a/datafusion/physical-plan/src/joins/sort_merge_join/tests.rs
+++ b/datafusion/physical-plan/src/joins/sort_merge_join/tests.rs
@@ -3383,7 +3383,7 @@ async fn test_left_outer_join_filtered_mask() -> Result<()> {
#[test]
fn test_partition_statistics() -> Result<()> {
- use crate::ExecutionPlan;
+ use crate::statistics::StatisticsArgs;
use datafusion_common::stats::Precision;
let left = build_table(
@@ -3420,7 +3420,7 @@ fn test_partition_statistics() -> Result<()> {
// Test aggregate statistics (partition = None)
// Should return meaningful statistics computed from both inputs
- let stats = join_exec.partition_statistics(None)?;
+ let stats = join_exec.statistics_with_args(&StatisticsArgs::new())?;
assert_eq!(
stats.column_statistics.len(),
expected_cols,
@@ -3438,7 +3438,8 @@ fn test_partition_statistics() -> Result<()> {
// Since the child TestMemoryExec returns unknown stats for specific partitions,
// the join output will also have Absent num_rows. This is expected behavior
// as the statistics depend on what the children can provide.
- let partition_stats = join_exec.partition_statistics(Some(0))?;
+ let partition_stats = join_exec
+ .statistics_with_args(&StatisticsArgs::new().with_partition(Some(0)))?;
assert_eq!(
partition_stats.column_statistics.len(),
expected_cols,
diff --git a/datafusion/physical-plan/src/lib.rs b/datafusion/physical-plan/src/lib.rs
index c7b1d4729e..6cc6e44c32 100644
--- a/datafusion/physical-plan/src/lib.rs
+++ b/datafusion/physical-plan/src/lib.rs
@@ -49,6 +49,7 @@ pub use crate::execution_plan::{
pub use crate::metrics::Metric;
pub use crate::ordering::InputOrderMode;
pub use crate::sort_pushdown::SortOrderPushdownResult;
+pub use crate::statistics::StatisticsArgs;
pub use crate::stream::EmptyRecordBatchStream;
pub use crate::topk::TopK;
pub use crate::visitor::{ExecutionPlanVisitor, accept, visit_execution_plan};
@@ -89,6 +90,7 @@ pub mod scalar_subquery;
pub mod sort_pushdown;
pub mod sorts;
pub mod spill;
+pub mod statistics;
pub mod stream;
pub mod streaming;
pub mod tree_node;
diff --git a/datafusion/physical-plan/src/limit.rs b/datafusion/physical-plan/src/limit.rs
index 7f42c33a79..1e4b5e5bb6 100644
--- a/datafusion/physical-plan/src/limit.rs
+++ b/datafusion/physical-plan/src/limit.rs
@@ -27,6 +27,7 @@ use super::{
SendableRecordBatchStream, Statistics,
};
use crate::execution_plan::{Boundedness, CardinalityEffect};
+use crate::statistics::StatisticsArgs;
use crate::{
DisplayFormatType, Distribution, ExecutionPlan, Partitioning,
check_if_same_properties,
@@ -219,8 +220,10 @@ impl ExecutionPlan for GlobalLimitExec {
Some(self.metrics.clone_inner())
}
- fn partition_statistics(&self, partition: Option<usize>) -> Result<Arc<Statistics>> {
- let stats = Arc::unwrap_or_clone(self.input.partition_statistics(partition)?);
+ fn statistics_with_args(&self, args: &StatisticsArgs) -> Result<Arc<Statistics>> {
+ let stats = Arc::unwrap_or_clone(
+ args.compute_child_statistics(&self.input, args.partition())?,
+ );
Ok(Arc::new(stats.with_fetch(self.fetch, self.skip, 1)?))
}
@@ -382,8 +385,10 @@ impl ExecutionPlan for LocalLimitExec {
Some(self.metrics.clone_inner())
}
- fn partition_statistics(&self, partition: Option<usize>) -> Result<Arc<Statistics>> {
- let stats = Arc::unwrap_or_clone(self.input.partition_statistics(partition)?);
+ fn statistics_with_args(&self, args: &StatisticsArgs) -> Result<Arc<Statistics>> {
+ let stats = Arc::unwrap_or_clone(
+ args.compute_child_statistics(&self.input, args.partition())?,
+ );
Ok(Arc::new(stats.with_fetch(Some(self.fetch), 0, 1)?))
}
@@ -530,6 +535,7 @@ mod tests {
use super::*;
use crate::coalesce_partitions::CoalescePartitionsExec;
use crate::common::collect;
+ use crate::statistics::StatisticsArgs;
use crate::test;
use crate::aggregates::{AggregateExec, AggregateMode, PhysicalGroupBy};
@@ -812,7 +818,9 @@ mod tests {
let offset =
GlobalLimitExec::new(Arc::new(CoalescePartitionsExec::new(csv)), skip, fetch);
- Ok(offset.partition_statistics(None)?.num_rows)
+ Ok(offset
+ .statistics_with_args(&StatisticsArgs::new())?
+ .num_rows)
}
pub fn build_group_by(
@@ -852,7 +860,9 @@ mod tests {
fetch,
);
- Ok(offset.partition_statistics(None)?.num_rows)
+ Ok(offset
+ .statistics_with_args(&StatisticsArgs::new())?
+ .num_rows)
}
async fn row_number_statistics_for_local_limit(
@@ -865,7 +875,9 @@ mod tests {
let offset = LocalLimitExec::new(csv, fetch);
- Ok(offset.partition_statistics(None)?.num_rows)
+ Ok(offset
+ .statistics_with_args(&StatisticsArgs::new())?
+ .num_rows)
}
/// Return a RecordBatch with a single array with row_count sz
diff --git a/datafusion/physical-plan/src/operator_statistics/mod.rs b/datafusion/physical-plan/src/operator_statistics/mod.rs
index 041ef46666..990bb4a682 100644
--- a/datafusion/physical-plan/src/operator_statistics/mod.rs
+++ b/datafusion/physical-plan/src/operator_statistics/mod.rs
@@ -94,6 +94,7 @@ use datafusion_common::stats::Precision;
use datafusion_common::{Result, Statistics};
use crate::ExecutionPlan;
+use crate::statistics::StatisticsArgs;
// ============================================================================
// ExtendedStatistics: Statistics with type-safe extensions
@@ -266,7 +267,7 @@ impl StatisticsProvider for DefaultStatisticsProvider {
plan: &dyn ExecutionPlan,
_child_stats: &[ExtendedStatistics],
) -> Result<StatisticsResult> {
- let base = plan.partition_statistics(None)?;
+ let base = plan.statistics_with_args(&StatisticsArgs::new())?;
Ok(StatisticsResult::Computed(ExtendedStatistics::new_arc(
base,
)))
@@ -358,7 +359,7 @@ impl StatisticsRegistry {
pub fn compute(&self, plan: &dyn ExecutionPlan) -> Result<ExtendedStatistics> {
// Fast path: no providers registered, skip the walk entirely
if self.providers.is_empty() {
- let base = plan.partition_statistics(None)?;
+ let base = plan.statistics_with_args(&StatisticsArgs::new())?;
return Ok(ExtendedStatistics::new_arc(base));
}
@@ -382,7 +383,7 @@ impl StatisticsRegistry {
}
}
// Fallback: use plan's built-in stats
- let base = plan.partition_statistics(None)?;
+ let base = plan.statistics_with_args(&StatisticsArgs::new())?;
Ok(ExtendedStatistics::new_arc(base))
}
@@ -505,7 +506,8 @@ fn computed_with_row_count(
plan: &dyn ExecutionPlan,
num_rows: Precision<usize>,
) -> Result<StatisticsResult> {
- let mut base = Arc::unwrap_or_clone(plan.partition_statistics(None)?);
+ let mut base =
+ Arc::unwrap_or_clone(plan.statistics_with_args(&StatisticsArgs::new())?);
rescale_byte_size(&mut base, num_rows);
Ok(StatisticsResult::Computed(ExtendedStatistics::new(base)))
}
@@ -1023,6 +1025,7 @@ mod tests {
use super::*;
use crate::filter::FilterExec;
use crate::projection::ProjectionExec;
+ use crate::statistics::StatisticsArgs;
use crate::{DisplayAs, DisplayFormatType, PlanProperties};
use arrow::datatypes::{DataType, Field, Schema};
use datafusion_common::stats::Precision;
@@ -1121,9 +1124,9 @@ mod tests {
unimplemented!()
}
- fn partition_statistics(
+ fn statistics_with_args(
&self,
- _partition: Option<usize>,
+ _args: &StatisticsArgs,
) -> Result<Arc<Statistics>> {
Ok(Arc::new(self.stats.clone()))
}
diff --git a/datafusion/physical-plan/src/placeholder_row.rs b/datafusion/physical-plan/src/placeholder_row.rs
index b99f9a9304..64b192d58d 100644
--- a/datafusion/physical-plan/src/placeholder_row.rs
+++ b/datafusion/physical-plan/src/placeholder_row.rs
@@ -33,6 +33,7 @@ use datafusion_common::{Result, assert_or_internal_err};
use datafusion_execution::TaskContext;
use datafusion_physical_expr::EquivalenceProperties;
+use crate::statistics::StatisticsArgs;
use log::trace;
/// Execution plan for empty relation with produce_one_row=true
@@ -164,12 +165,12 @@ impl ExecutionPlan for PlaceholderRowExec {
Ok(Box::pin(cooperative(ms)))
}
- fn partition_statistics(&self, partition: Option<usize>) -> Result<Arc<Statistics>> {
+ fn statistics_with_args(&self, args: &StatisticsArgs) -> Result<Arc<Statistics>> {
let batches = self
.data()
.expect("Create single row placeholder RecordBatch should not fail");
- let batches = match partition {
+ let batches = match args.partition() {
Some(_) => vec![batches],
// entire plan
None => vec![batches; self.partitions],
diff --git a/datafusion/physical-plan/src/projection.rs b/datafusion/physical-plan/src/projection.rs
index ade3a988c7..16b0a5ad7e 100644
--- a/datafusion/physical-plan/src/projection.rs
+++ b/datafusion/physical-plan/src/projection.rs
@@ -33,6 +33,7 @@ use crate::filter_pushdown::{
FilterPushdownPropagation, FilterRemapper, PushedDownPredicate,
};
use crate::joins::utils::{ColumnIndex, JoinFilter, JoinOn, JoinOnRef};
+use crate::statistics::StatisticsArgs;
use crate::{DisplayFormatType, ExecutionPlan, PhysicalExpr, check_if_same_properties};
use std::collections::HashMap;
use std::pin::Pin;
@@ -348,9 +349,10 @@ impl ExecutionPlan for ProjectionExec {
Some(self.metrics.clone_inner())
}
- fn partition_statistics(&self, partition: Option<usize>) -> Result<Arc<Statistics>> {
- let input_stats =
- Arc::unwrap_or_clone(self.input.partition_statistics(partition)?);
+ fn statistics_with_args(&self, args: &StatisticsArgs) -> Result<Arc<Statistics>> {
+ let input_stats = Arc::unwrap_or_clone(
+ args.compute_child_statistics(&self.input, args.partition())?,
+ );
let output_schema = self.schema();
Ok(Arc::new(
self.projector
@@ -1184,6 +1186,7 @@ mod tests {
use crate::common::collect;
use crate::filter_pushdown::PushedDown;
+ use crate::statistics::StatisticsArgs;
use crate::test;
use crate::test::exec::StatisticsExec;
@@ -1374,7 +1377,9 @@ mod tests {
let projection = ProjectionExec::try_new(exprs, input).unwrap();
- let stats = projection.partition_statistics(None).unwrap();
+ let stats = projection
+ .statistics_with_args(&StatisticsArgs::new())
+ .unwrap();
assert_eq!(stats.num_rows, Precision::Exact(10));
assert_eq!(
diff --git a/datafusion/physical-plan/src/repartition/mod.rs b/datafusion/physical-plan/src/repartition/mod.rs
index 3d30dd8276..2298183485 100644
--- a/datafusion/physical-plan/src/repartition/mod.rs
+++ b/datafusion/physical-plan/src/repartition/mod.rs
@@ -39,6 +39,7 @@ use crate::projection::{ProjectionExec, all_columns, make_with_child, update_exp
use crate::sorts::streaming_merge::StreamingMergeBuilder;
use crate::spill::spill_manager::SpillManager;
use crate::spill::spill_pool::{self, SpillPoolWriter};
+use crate::statistics::StatisticsArgs;
use crate::stream::{EmptyRecordBatchStream, RecordBatchStreamAdapter};
use crate::{
DisplayFormatType, ExecutionPlan, Partitioning, PlanProperties, Statistics,
@@ -1361,8 +1362,8 @@ impl ExecutionPlan for RepartitionExec {
Some(self.metrics.clone_inner())
}
- fn partition_statistics(&self, partition: Option<usize>) -> Result<Arc<Statistics>> {
- if let Some(partition) = partition {
+ fn statistics_with_args(&self, args: &StatisticsArgs) -> Result<Arc<Statistics>> {
+ if let Some(partition) = args.partition() {
let partition_count = self.partitioning().partition_count();
if partition_count == 0 {
return Ok(Arc::new(Statistics::new_unknown(&self.schema())));
@@ -1375,7 +1376,8 @@ impl ExecutionPlan for RepartitionExec {
partition_count
);
- let mut stats = Arc::unwrap_or_clone(self.input.partition_statistics(None)?);
+ let mut stats =
+ Arc::unwrap_or_clone(args.compute_child_statistics(&self.input, None)?);
// Distribute statistics across partitions
stats.num_rows = stats
@@ -1398,7 +1400,7 @@ impl ExecutionPlan for RepartitionExec {
Ok(Arc::new(stats))
} else {
- self.input.partition_statistics(None)
+ args.compute_child_statistics(&self.input, None)
}
}
diff --git a/datafusion/physical-plan/src/scalar_subquery.rs b/datafusion/physical-plan/src/scalar_subquery.rs
index 25f7332f95..dd44d09c38 100644
--- a/datafusion/physical-plan/src/scalar_subquery.rs
+++ b/datafusion/physical-plan/src/scalar_subquery.rs
@@ -33,6 +33,7 @@ use datafusion_expr::execution_props::{ScalarSubqueryResults, SubqueryIndex};
use crate::execution_plan::{CardinalityEffect, ExecutionPlan, PlanProperties};
use crate::joins::utils::{OnceAsync, OnceFut};
+use crate::statistics::StatisticsArgs;
use crate::stream::RecordBatchStreamAdapter;
use crate::{DisplayAs, DisplayFormatType, SendableRecordBatchStream};
@@ -235,8 +236,8 @@ impl ExecutionPlan for ScalarSubqueryExec {
vec![false; self.subqueries.len() + 1]
}
- fn partition_statistics(&self, partition: Option<usize>) -> Result<Arc<Statistics>> {
- self.input.partition_statistics(partition)
+ fn statistics_with_args(&self, args: &StatisticsArgs) -> Result<Arc<Statistics>> {
+ args.compute_child_statistics(&self.input, args.partition())
}
fn cardinality_effect(&self) -> CardinalityEffect {
diff --git a/datafusion/physical-plan/src/sorts/partial_sort.rs b/datafusion/physical-plan/src/sorts/partial_sort.rs
index 3bf16af36c..f7b403d943 100644
--- a/datafusion/physical-plan/src/sorts/partial_sort.rs
+++ b/datafusion/physical-plan/src/sorts/partial_sort.rs
@@ -58,6 +58,7 @@ use std::task::{Context, Poll};
use crate::metrics::{BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet};
use crate::sorts::sort::sort_batch;
+use crate::statistics::StatisticsArgs;
use crate::stream::EmptyRecordBatchStream;
use crate::{
DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, ExecutionPlanProperties,
@@ -334,8 +335,8 @@ impl ExecutionPlan for PartialSortExec {
Some(self.metrics_set.clone_inner())
}
- fn partition_statistics(&self, partition: Option<usize>) -> Result<Arc<Statistics>> {
- self.input.partition_statistics(partition)
+ fn statistics_with_args(&self, args: &StatisticsArgs) -> Result<Arc<Statistics>> {
+ args.compute_child_statistics(&self.input, args.partition())
}
}
diff --git a/datafusion/physical-plan/src/sorts/sort.rs b/datafusion/physical-plan/src/sorts/sort.rs
index 929ff4f7df..ccc675c6ef 100644
--- a/datafusion/physical-plan/src/sorts/sort.rs
+++ b/datafusion/physical-plan/src/sorts/sort.rs
@@ -45,6 +45,7 @@ use crate::sorts::streaming_merge::{SortedSpillFile, StreamingMergeBuilder};
use crate::spill::get_record_batch_memory_size;
use crate::spill::in_progress_spill_file::InProgressSpillFile;
use crate::spill::spill_manager::{GetSlicedSize, SpillManager};
+use crate::statistics::StatisticsArgs;
use crate::stream::RecordBatchStreamAdapter;
use crate::stream::ReservationStream;
use crate::topk::TopK;
@@ -1283,13 +1284,14 @@ impl ExecutionPlan for SortExec {
Some(self.metrics_set.clone_inner())
}
- fn partition_statistics(&self, partition: Option<usize>) -> Result<Arc<Statistics>> {
- let p = if !self.preserve_partitioning() {
- None
+ fn statistics_with_args(&self, args: &StatisticsArgs) -> Result<Arc<Statistics>> {
+ let partition = if self.preserve_partitioning() {
+ args.partition()
} else {
- partition
+ None
};
- let stats = Arc::unwrap_or_clone(self.input.partition_statistics(p)?);
+ let child_stats = args.compute_child_statistics(&self.input, partition)?;
+ let stats = Arc::unwrap_or_clone(child_stats);
Ok(Arc::new(stats.with_fetch(self.fetch, 0, 1)?))
}
diff --git a/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs b/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs
index eb9b5f09aa..dcf3a7baad 100644
--- a/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs
+++ b/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs
@@ -24,6 +24,7 @@ use crate::limit::LimitStream;
use crate::metrics::{BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet};
use crate::projection::{ProjectionExec, make_with_child, update_ordering};
use crate::sorts::streaming_merge::StreamingMergeBuilder;
+use crate::statistics::StatisticsArgs;
use crate::{
DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, ExecutionPlanProperties,
Partitioning, PlanProperties, SendableRecordBatchStream, Statistics,
@@ -380,8 +381,8 @@ impl ExecutionPlan for SortPreservingMergeExec {
Some(self.metrics.clone_inner())
}
- fn partition_statistics(&self, _partition: Option<usize>) -> Result<Arc<Statistics>> {
- self.input.partition_statistics(None)
+ fn statistics_with_args(&self, args: &StatisticsArgs) -> Result<Arc<Statistics>> {
+ args.compute_child_statistics(&self.input, None)
}
fn supports_limit_pushdown(&self) -> bool {
diff --git a/datafusion/physical-plan/src/statistics.rs b/datafusion/physical-plan/src/statistics.rs
new file mode 100644
index 0000000000..5ed5558e28
--- /dev/null
+++ b/datafusion/physical-plan/src/statistics.rs
@@ -0,0 +1,229 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Statistics computation for physical plans.
+//!
+//! [`StatisticsArgs`] provides external context to
+//! [`ExecutionPlan::statistics_with_args`].
+
+use crate::ExecutionPlan;
+use datafusion_common::{Result, Statistics, assert_or_internal_err};
+use std::cell::RefCell;
+use std::collections::HashMap;
+use std::rc::Rc;
+use std::sync::Arc;
+
+/// Per-call memoization cache for statistics computation.
+///
+/// Keyed by `(plan node pointer address, partition)`. Shared across
+/// a single statistics walk via [`StatisticsArgs`].
+///
+/// The pointer-based key is safe within a single synchronous walk:
+/// all `Arc<dyn ExecutionPlan>` nodes are held by the plan tree for
+/// the duration of the walk, so addresses cannot be reused.
+#[derive(Debug, Default)]
+struct StatsCache(HashMap<(usize, Option<usize>), Arc<Statistics>>);
+
+impl StatsCache {
+ fn get(
+ &self,
+ plan: &dyn ExecutionPlan,
+ partition: Option<usize>,
+ ) -> Option<&Arc<Statistics>> {
+ let key = (
+ plan as *const dyn ExecutionPlan as *const () as usize,
+ partition,
+ );
+ self.0.get(&key)
+ }
+
+ fn insert(
+ &mut self,
+ plan: &dyn ExecutionPlan,
+ partition: Option<usize>,
+ stats: Arc<Statistics>,
+ ) {
+ let key = (
+ plan as *const dyn ExecutionPlan as *const () as usize,
+ partition,
+ );
+ self.0.insert(key, stats);
+ }
+}
+
+/// Arguments passed to [`ExecutionPlan::statistics_with_args`] carrying
+/// external information that operators can use when computing their
+/// statistics.
+#[derive(Debug, Default)]
+pub struct StatisticsArgs {
+ partition: Option<usize>,
+ /// Shared memoization cache for the current statistics walk.
+ cache: Rc<RefCell<StatsCache>>,
+}
+
+impl StatisticsArgs {
+ /// Creates new statistics arguments with a fresh cache.
+ ///
+ /// By default the partition is set to `None` (statistics should be computed
+ /// for the entire plan).
+ pub fn new() -> Self {
+ Default::default()
+ }
+
+ /// Set the partition to compute statistics
+ ///
+ /// * `None` means statistics should be computed for the entire plan.
+ /// * `Some(idx)` means statistics should be computed for the specified
+ /// partition index.
+ ///
+ /// Changing the partition starts a new statistics walk, so the
+ /// memoization cache is reset to avoid reusing entries computed for a
+ /// different partition.
+ pub fn set_partition(&mut self, partition: Option<usize>) {
+ if self.partition != partition {
+ self.partition = partition;
+ // Drop the previous walk's cache: its entries are keyed by raw
+ // plan pointer and the prior partition, so they must not leak
+ // into the new walk.
+ self.cache = Rc::new(RefCell::new(StatsCache::default()));
+ }
+ }
+
+ /// Builder Style API for [`Self::set_partition`]
+ pub fn with_partition(mut self, partition: Option<usize>) -> Self {
+ self.set_partition(partition);
+ self
+ }
+
+ /// Return the partition to compute statistics
+ pub fn partition(&self) -> Option<usize> {
+ self.partition
+ }
+
+ /// Computes statistics for a child plan, using the shared cache
+ /// to avoid redundant subtree walks.
+ pub fn compute_child_statistics(
+ &self,
+ plan: impl AsRef<dyn ExecutionPlan>,
+ partition: Option<usize>,
+ ) -> Result<Arc<Statistics>> {
+ let plan = plan.as_ref();
+
+ if let Some(idx) = partition {
+ let partition_count = plan.properties().partitioning.partition_count();
+ assert_or_internal_err!(
+ idx < partition_count,
+ "Invalid partition index: {}, the partition count is {}",
+ idx,
+ partition_count
+ );
+ }
+
+ if let Some(cached) = self.cache.borrow().get(plan, partition) {
+ return Ok(Arc::clone(cached));
+ }
+
+ let child_args = StatisticsArgs {
+ partition,
+ cache: Rc::clone(&self.cache),
+ };
+ let result = plan.statistics_with_args(&child_args)?;
+
+ self.cache
+ .borrow_mut()
+ .insert(plan, partition, Arc::clone(&result));
+ Ok(result)
+ }
+}
+
+#[cfg(all(test, feature = "test_utils"))]
+mod tests {
+ use super::*;
+ use crate::coalesce_partitions::CoalescePartitionsExec;
+ use crate::test::exec::StatisticsExec;
+ use arrow::datatypes::{DataType, Field, Schema};
+ use datafusion_common::{ColumnStatistics, stats::Precision};
+
+ fn make_stats_leaf(num_rows: usize) -> Arc<dyn ExecutionPlan> {
+ let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]);
+ let col_stats = vec![ColumnStatistics {
+ null_count: Precision::Exact(0),
+ max_value: Precision::Absent,
+ min_value: Precision::Absent,
+ sum_value: Precision::Absent,
+ distinct_count: Precision::Absent,
+ byte_size: Precision::Absent,
+ }];
+ Arc::new(StatisticsExec::new(
+ Statistics {
+ num_rows: Precision::Exact(num_rows),
+ total_byte_size: Precision::Absent,
+ column_statistics: col_stats,
+ },
+ schema,
+ ))
+ }
+
+ #[test]
+ fn coalesce_returns_overall_stats_for_any_partition() {
+ let leaf = make_stats_leaf(100);
+ let plan: Arc<dyn ExecutionPlan> = Arc::new(CoalescePartitionsExec::new(leaf));
+
+ let args = StatisticsArgs::new().with_partition(Some(0));
+ let stats = plan.statistics_with_args(&args).unwrap();
+ assert_eq!(stats.num_rows, Precision::Exact(100));
+
+ let args_none = StatisticsArgs::new();
+ let stats_none = plan.statistics_with_args(&args_none).unwrap();
+ assert_eq!(stats_none.num_rows, Precision::Exact(100));
+ }
+
+ #[test]
+ fn changing_partition_resets_cache() {
+ let leaf = make_stats_leaf(100);
+
+ // Populate the memoization cache for an initial walk.
+ let mut args = StatisticsArgs::new();
+ let _ = args
+ .compute_child_statistics(Arc::clone(&leaf), Some(0))
+ .unwrap();
+ assert!(
+ !args.cache.borrow().0.is_empty(),
+ "cache should be populated after a statistics walk"
+ );
+
+ // Changing the partition starts a new walk and must reset the cache
+ // so stale, pointer-keyed entries cannot leak across walks.
+ args.set_partition(Some(1));
+ assert!(
+ args.cache.borrow().0.is_empty(),
+ "cache should be cleared when the partition changes"
+ );
+
+ // Setting the partition to its current value is a no-op and retains
+ // the cache (avoids needlessly discarding work mid-walk).
+ let _ = args
+ .compute_child_statistics(Arc::clone(&leaf), Some(0))
+ .unwrap();
+ assert!(!args.cache.borrow().0.is_empty());
+ args.set_partition(Some(1));
+ assert!(
+ !args.cache.borrow().0.is_empty(),
+ "cache should be retained when the partition is unchanged"
+ );
+ }
+}
diff --git a/datafusion/physical-plan/src/test.rs b/datafusion/physical-plan/src/test.rs
index a6e76cebcd..44aacfa87a 100644
--- a/datafusion/physical-plan/src/test.rs
+++ b/datafusion/physical-plan/src/test.rs
@@ -29,6 +29,7 @@ use crate::common;
use crate::execution_plan::{Boundedness, EmissionType};
use crate::memory::MemoryStream;
use crate::metrics::MetricsSet;
+use crate::statistics::StatisticsArgs;
use crate::stream::RecordBatchStreamAdapter;
use crate::streaming::PartitionStream;
use crate::{DisplayAs, DisplayFormatType, PlanProperties};
@@ -164,8 +165,8 @@ impl ExecutionPlan for TestMemoryExec {
unimplemented!()
}
- fn partition_statistics(&self, partition: Option<usize>) -> Result<Arc<Statistics>> {
- if partition.is_some() {
+ fn statistics_with_args(&self, args: &StatisticsArgs) -> Result<Arc<Statistics>> {
+ if args.partition().is_some() {
Ok(Arc::new(Statistics::new_unknown(&self.schema)))
} else {
Ok(Arc::new(self.statistics_inner()?))
diff --git a/datafusion/physical-plan/src/test/exec.rs b/datafusion/physical-plan/src/test/exec.rs
index e162571e32..2bd19ccbeb 100644
--- a/datafusion/physical-plan/src/test/exec.rs
+++ b/datafusion/physical-plan/src/test/exec.rs
@@ -20,7 +20,7 @@
use crate::{
DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, PlanProperties,
RecordBatchStream, SendableRecordBatchStream, Statistics, common,
- execution_plan::Boundedness,
+ execution_plan::Boundedness, statistics::StatisticsArgs,
};
use crate::{
execution_plan::EmissionType,
@@ -249,8 +249,8 @@ impl ExecutionPlan for MockExec {
}
// Panics if one of the batches is an error
- fn partition_statistics(&self, partition: Option<usize>) -> Result<Arc<Statistics>> {
- if partition.is_some() {
+ fn statistics_with_args(&self, args: &StatisticsArgs) -> Result<Arc<Statistics>> {
+ if args.partition().is_some() {
return Ok(Arc::new(Statistics::new_unknown(&self.schema)));
}
let data: Result<Vec<_>> = self
@@ -474,8 +474,8 @@ impl ExecutionPlan for BarrierExec {
Ok(builder.build())
}
- fn partition_statistics(&self, partition: Option<usize>) -> Result<Arc<Statistics>> {
- if partition.is_some() {
+ fn statistics_with_args(&self, args: &StatisticsArgs) -> Result<Arc<Statistics>> {
+ if args.partition().is_some() {
return Ok(Arc::new(Statistics::new_unknown(&self.schema)));
}
Ok(Arc::new(common::compute_record_batch_statistics(
@@ -654,8 +654,8 @@ impl ExecutionPlan for StatisticsExec {
unimplemented!("This plan only serves for testing statistics")
}
- fn partition_statistics(&self, partition: Option<usize>) -> Result<Arc<Statistics>> {
- Ok(Arc::new(if partition.is_some() {
+ fn statistics_with_args(&self, args: &StatisticsArgs) -> Result<Arc<Statistics>> {
+ Ok(Arc::new(if args.partition().is_some() {
Statistics::new_unknown(&self.schema)
} else {
self.stats.clone()
diff --git a/datafusion/physical-plan/src/union.rs b/datafusion/physical-plan/src/union.rs
index 3ea2eb5402..2962428532 100644
--- a/datafusion/physical-plan/src/union.rs
+++ b/datafusion/physical-plan/src/union.rs
@@ -43,6 +43,7 @@ use crate::filter_pushdown::{
};
use crate::metrics::BaselineMetrics;
use crate::projection::{ProjectionExec, make_with_child};
+use crate::statistics::StatisticsArgs;
use crate::stream::ObservedStream;
use arrow::datatypes::{Field, Schema, SchemaRef};
@@ -318,26 +319,35 @@ impl ExecutionPlan for UnionExec {
Some(self.metrics.clone_inner())
}
- fn partition_statistics(&self, partition: Option<usize>) -> Result<Arc<Statistics>> {
- if let Some(partition_idx) = partition {
+ fn statistics_with_args(&self, args: &StatisticsArgs) -> Result<Arc<Statistics>> {
+ if let Some(partition_idx) = args.partition() {
// For a specific partition, find which input it belongs to
let mut remaining_idx = partition_idx;
- for input in &self.inputs {
+ for (i, input) in self.inputs.iter().enumerate() {
let input_partition_count = input.output_partitioning().partition_count();
if remaining_idx < input_partition_count {
- // This partition belongs to this input
- return input.partition_statistics(Some(remaining_idx));
+ // This partition belongs to this input - compute stats
+ // for the specific child at the specific partition
+ let child = &self.inputs[i];
+ return args.compute_child_statistics(child, Some(remaining_idx));
}
remaining_idx -= input_partition_count;
}
// If we get here, the partition index is out of bounds
Ok(Arc::new(Statistics::new_unknown(&self.schema())))
} else {
- let schema = self.schema();
- Ok(Arc::new(merge_input_statistics(
- &self.inputs,
- None,
- schema.as_ref(),
+ // Collect overall stats for each input from the cache
+ let stats = self
+ .inputs
+ .iter()
+ .map(|input| args.compute_child_statistics(input, None))
+ .collect::<Result<Vec<_>>>()?;
+ let stats_refs = stats.iter().map(|s| s.as_ref()).collect::<Vec<_>>();
+
+ Ok(Arc::new(Statistics::try_merge_iter_with_ndv_fallback(
+ stats_refs,
+ self.schema().as_ref(),
+ NdvFallback::Sum,
)?))
}
}
@@ -641,12 +651,20 @@ impl ExecutionPlan for InterleaveExec {
Some(self.metrics.clone_inner())
}
- fn partition_statistics(&self, partition: Option<usize>) -> Result<Arc<Statistics>> {
- let schema = self.schema();
- Ok(Arc::new(merge_input_statistics(
- &self.inputs,
- partition,
- schema.as_ref(),
+ fn statistics_with_args(&self, args: &StatisticsArgs) -> Result<Arc<Statistics>> {
+ let stats = self
+ .inputs
+ .iter()
+ .map(|input| {
+ args.compute_child_statistics(input, args.partition())
+ .map(Arc::unwrap_or_clone)
+ })
+ .collect::<Result<Vec<_>>>()?;
+
+ Ok(Arc::new(Statistics::try_merge_iter_with_ndv_fallback(
+ stats.iter(),
+ self.schema().as_ref(),
+ NdvFallback::Sum,
)?))
}
@@ -806,28 +824,12 @@ impl Stream for CombinedRecordBatchStream {
}
}
-fn merge_input_statistics(
- inputs: &[Arc<dyn ExecutionPlan>],
- partition: Option<usize>,
- schema: &Schema,
-) -> Result<Statistics> {
- let stats = inputs
- .iter()
- .map(|input| {
- input
- .partition_statistics(partition)
- .map(Arc::unwrap_or_clone)
- })
- .collect::<Result<Vec<_>>>()?;
-
- Statistics::try_merge_iter_with_ndv_fallback(stats.iter(), schema, NdvFallback::Sum)
-}
-
#[cfg(test)]
mod tests {
use super::*;
use crate::collect;
use crate::repartition::RepartitionExec;
+ use crate::statistics::StatisticsArgs;
use crate::test::exec::StatisticsExec;
use crate::test::{self, TestMemoryExec};
@@ -1018,7 +1020,7 @@ mod tests {
Arc::new(StatisticsExec::new(right, schema.as_ref().clone()));
let union = UnionExec::try_new(vec![left, right])?;
- let stats = union.partition_statistics(None)?;
+ let stats = union.statistics_with_args(&StatisticsArgs::new())?;
assert_eq!(stats.as_ref(), &expected);
Ok(())
@@ -1035,7 +1037,7 @@ mod tests {
Arc::new(StatisticsExec::new(right, schema.as_ref().clone()));
let union = UnionExec::try_new(vec![left, right])?;
- let stats = union.partition_statistics(None)?;
+ let stats = union.statistics_with_args(&StatisticsArgs::new())?;
assert_eq!(stats.as_ref(), &expected);
Ok(())
@@ -1056,7 +1058,7 @@ mod tests {
)?);
let interleave = InterleaveExec::try_new(vec![left, right])?;
- let stats = interleave.partition_statistics(None)?;
+ let stats = interleave.statistics_with_args(&StatisticsArgs::new())?;
assert_eq!(stats.as_ref(), &expected);
Ok(())
@@ -1078,7 +1080,8 @@ mod tests {
)?);
let interleave = InterleaveExec::try_new(vec![left, right])?;
- let stats = interleave.partition_statistics(Some(0))?;
+ let stats = interleave
+ .statistics_with_args(&StatisticsArgs::new().with_partition(Some(0)))?;
let expected = Statistics::default()
.with_num_rows(Precision::Inexact(5))
diff --git a/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs b/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs
index 6c6b26c9cf..b0a0330441 100644
--- a/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs
+++ b/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs
@@ -28,6 +28,7 @@ use std::task::{Context, Poll};
use super::utils::create_schema;
use crate::metrics::{BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet};
+use crate::statistics::StatisticsArgs;
use crate::stream::EmptyRecordBatchStream;
use crate::windows::{
calc_requirements, get_ordered_partition_by_indices, get_partition_by_sort_exprs,
@@ -377,9 +378,10 @@ impl ExecutionPlan for BoundedWindowAggExec {
Some(self.metrics.clone_inner())
}
- fn partition_statistics(&self, partition: Option<usize>) -> Result<Arc<Statistics>> {
- let input_stat =
- Arc::unwrap_or_clone(self.input.partition_statistics(partition)?);
+ fn statistics_with_args(&self, args: &StatisticsArgs) -> Result<Arc<Statistics>> {
+ let input_stat = Arc::unwrap_or_clone(
+ args.compute_child_statistics(&self.input, args.partition())?,
+ );
Ok(Arc::new(self.statistics_helper(input_stat)?))
}
diff --git a/datafusion/physical-plan/src/windows/window_agg_exec.rs b/datafusion/physical-plan/src/windows/window_agg_exec.rs
index ee3b071fc9..f4bc40cf35 100644
--- a/datafusion/physical-plan/src/windows/window_agg_exec.rs
+++ b/datafusion/physical-plan/src/windows/window_agg_exec.rs
@@ -24,6 +24,7 @@ use std::task::{Context, Poll};
use super::utils::create_schema;
use crate::execution_plan::{CardinalityEffect, EmissionType};
use crate::metrics::{BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet};
+use crate::statistics::StatisticsArgs;
use crate::stream::EmptyRecordBatchStream;
use crate::windows::{
calc_requirements, get_ordered_partition_by_indices, get_partition_by_sort_exprs,
@@ -280,9 +281,10 @@ impl ExecutionPlan for WindowAggExec {
Some(self.metrics.clone_inner())
}
- fn partition_statistics(&self, partition: Option<usize>) -> Result<Arc<Statistics>> {
- let input_stat =
- Arc::unwrap_or_clone(self.input.partition_statistics(partition)?);
+ fn statistics_with_args(&self, args: &StatisticsArgs) -> Result<Arc<Statistics>> {
+ let input_stat = Arc::unwrap_or_clone(
+ args.compute_child_statistics(&self.input, args.partition())?,
+ );
let win_cols = self.window_expr.len();
let input_cols = self.input.schema().fields().len();
// TODO stats: some windowing function will maintain invariants such as min, max...
diff --git a/datafusion/physical-plan/src/work_table.rs b/datafusion/physical-plan/src/work_table.rs
index 28b9c8ddc7..9bf167aa73 100644
--- a/datafusion/physical-plan/src/work_table.rs
+++ b/datafusion/physical-plan/src/work_table.rs
@@ -29,6 +29,7 @@ use crate::{
SendableRecordBatchStream, Statistics,
};
+use crate::statistics::StatisticsArgs;
use arrow::datatypes::SchemaRef;
use arrow::record_batch::RecordBatch;
use datafusion_common::{Result, assert_eq_or_internal_err, internal_datafusion_err};
@@ -227,7 +228,7 @@ impl ExecutionPlan for WorkTableExec {
Some(self.metrics.clone_inner())
}
- fn partition_statistics(&self, _partition: Option<usize>) -> Result<Arc<Statistics>> {
+ fn statistics_with_args(&self, _args: &StatisticsArgs) -> Result<Arc<Statistics>> {
Ok(Arc::new(Statistics::new_unknown(&self.schema())))
}
diff --git a/docs/source/library-user-guide/upgrading/55.0.0.md b/docs/source/library-user-guide/upgrading/55.0.0.md
index 2d4a37474e..6d1f834abf 100644
--- a/docs/source/library-user-guide/upgrading/55.0.0.md
+++ b/docs/source/library-user-guide/upgrading/55.0.0.md
@@ -154,6 +154,77 @@ as a supertrait:
+ pub trait QueryPlanner: Any + Debug
```
+### `ExecutionPlan::partition_statistics` deprecated in favor of `statistics_with_args`
+
+`ExecutionPlan::partition_statistics` is deprecated. A new method
+`statistics_with_args` accepts a `StatisticsArgs` parameter that carries
+the partition index and a shared cache for memoized child statistics lookups.
+
+Existing implementations of `partition_statistics` continue to work unchanged.
+The default `statistics_with_args` delegates to the deprecated method, so no
+migration is required until the deprecated method is removed.
+
+> **Warning:** The delegation is **one-way**: the default `statistics_with_args`
+> calls `partition_statistics`, but the default `partition_statistics` does
+> **not** call `statistics_with_args` — it returns `Statistics::new_unknown`.
+> Nodes that override only `statistics_with_args` will silently return
+> `Statistics::new_unknown` to any caller still using the deprecated
+> `partition_statistics`.
+
+**Who is affected:**
+
+- Users who implement custom `ExecutionPlan` nodes (recommended to migrate)
+- Users who call `partition_statistics` directly (recommended to switch to `statistics_with_args`)
+
+**Migration guide:**
+
+For **implementations**, override `statistics_with_args` instead of
+`partition_statistics`. Leaf nodes that do not have children can ignore
+the args.
+
+Child statistics are looked up via `args.compute_child_statistics(child, partition)`.
+Use `args.partition()` for partition-preserving operators, or `None` for
+partition-merging operators that always need overall stats:
+
+```rust,ignore
+// Before:
+fn partition_statistics(&self, partition: Option<usize>) -> Result<Arc<Statistics>> {
+ let child_stats = self.input.partition_statistics(partition)?;
+ // ... transform child_stats ...
+}
+
+// After (partition-preserving):
+fn statistics_with_args(
+ &self,
+ args: &StatisticsArgs,
+) -> Result<Arc<Statistics>> {
+ let child_stats = args.compute_child_statistics(&self.input, args.partition())?;
+ // ... transform child_stats ...
+}
+
+// After (partition-merging):
+fn statistics_with_args(
+ &self,
+ args: &StatisticsArgs,
+) -> Result<Arc<Statistics>> {
+ let child_stats = args.compute_child_statistics(&self.input, None)?;
+ // ... transform child_stats ...
+}
+```
+
+For **callers**, create a `StatisticsArgs` and call `statistics_with_args`
+directly. The cache is created automatically:
+
+```rust,ignore
+use datafusion_physical_plan::StatisticsArgs;
+
+// Before:
+let stats = plan.partition_statistics(None)?;
+
+// After:
+let stats = plan.statistics_with_args(&StatisticsArgs::new())?;
+```
+
### `DdlStatement::CreateExternalTable` and `CreateFunction` are now boxed
The two largest variants of `datafusion_expr::DdlStatement` are now
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org