refactor(physical-plan): externalize statistics traversal into StatisticsContext (#23051)

## Which issue does this PR close?

- Closes #22958.

## Rationale for this change

Follow-up to #21815. Per @2010YOUY01's suggestion there, this decouples
statistics
traversal/caching from each operator's computation: operators now
express only local
propagation, while the walk and cache live in an external
`StatisticsContext`.

## What changes are included in this PR?

- `StatisticsContext` owns the bottom-up walk + per-walk cache;
`compute(plan, args)`
  resolves children, then calls the operator.
- `ExecutionPlan::statistics_from_inputs(input_stats, args)` — stateless
local
propagation from pre-resolved child stats; default delegates to
deprecated
  `partition_statistics`.
- `ExecutionPlan::child_stats_requests(partition) -> Vec<ChildStats>` —
per-child
directive (`At(Option<usize>)` / `Skip`) for which partition each child
is computed
at (e.g. broadcast joins request the build side overall; `UnionExec`
skips
  non-owning inputs).
- Removes the unreleased `statistics_with_args`; `partition_statistics`
stays
  deprecated.
- Migrates all operators/callers; updates the 55.0.0 upgrade guide.

## Are these changes tested?

Yes — existing statistics tests now run through
`StatisticsContext::compute`, plus new
unit tests for the cache and the `compute_statistics` benchmark.

## Are there any user-facing changes?

Yes (public `ExecutionPlan` API): `partition_statistics` deprecated in
favor of
`statistics_from_inputs`; new `StatisticsContext` / `ChildStats`;
unreleased
`statistics_with_args` removed. Upgrade guide updated.

---

Reviewer note: `statistics_from_inputs` takes `args: &StatisticsArgs`
(currently just
`partition`) as a non-breaking extension seam rather than a bare
`partition` — happy
to change if preferred (see #22958).

----

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: xudong.w <wxd963996380@gmail.com>

Author: asolimando

datafusion-examples/examples/relation_planner/table_sample.rs

@@ -108,7 +108,8 @@ use datafusion::{
     },
     physical_expr::EquivalenceProperties,
     physical_plan::{
-        DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, StatisticsArgs,
+        ChildStats, DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties,
+        StatisticsArgs,
         metrics::{BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet, RecordOutput},
     },
     physical_planner::{DefaultPhysicalPlanner, ExtensionPlanner, PhysicalPlanner},
@@ -722,10 +723,16 @@ impl ExecutionPlan for SampleExec {
         Some(self.metrics.clone_inner())
     }
 
-    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())?,
-        );
+    fn child_stats_requests(&self, partition: Option<usize>) -> Vec<ChildStats> {
+        vec![ChildStats::At(partition)]
+    }
+
+    fn statistics_from_inputs(
+        &self,
+        input_stats: &[Arc<Statistics>],
+        _args: &StatisticsArgs,
+    ) -> Result<Arc<Statistics>> {
+        let mut stats = input_stats[0].as_ref().clone();
         let ratio = self.upper_bound - self.lower_bound;
 
         // Scale statistics by sampling ratio (inexact due to randomness)

datafusion/core/src/datasource/file_format/csv.rs

@@ -45,7 +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::statistics::{StatisticsArgs, StatisticsContext};
     use datafusion_physical_plan::{ExecutionPlan, collect};
 
     use arrow::array::{
@@ -217,11 +217,14 @@ mod tests {
 
         // test metadata
         assert_eq!(
-            exec.statistics_with_args(&StatisticsArgs::new())?.num_rows,
+            StatisticsContext::new()
+                .compute(exec.as_ref(), &StatisticsArgs::new())?
+                .num_rows,
             Precision::Absent
         );
         assert_eq!(
-            exec.statistics_with_args(&StatisticsArgs::new())?
+            StatisticsContext::new()
+                .compute(exec.as_ref(), &StatisticsArgs::new())?
                 .total_byte_size,
             Precision::Absent
         );

datafusion/core/src/datasource/file_format/json.rs

@@ -36,7 +36,7 @@ mod tests {
         BatchDeserializer, DecoderDeserializer, DeserializerOutput,
     };
     use datafusion_datasource::file_format::FileFormat;
-    use datafusion_physical_plan::statistics::StatisticsArgs;
+    use datafusion_physical_plan::statistics::{StatisticsArgs, StatisticsContext};
     use datafusion_physical_plan::{ExecutionPlan, collect};
 
     use arrow::compute::concat_batches;
@@ -119,11 +119,14 @@ mod tests {
 
         // test metadata
         assert_eq!(
-            exec.statistics_with_args(&StatisticsArgs::new())?.num_rows,
+            StatisticsContext::new()
+                .compute(exec.as_ref(), &StatisticsArgs::new())?
+                .num_rows,
             Precision::Absent
         );
         assert_eq!(
-            exec.statistics_with_args(&StatisticsArgs::new())?
+            StatisticsContext::new()
+                .compute(exec.as_ref(), &StatisticsArgs::new())?
                 .total_byte_size,
             Precision::Absent
         );

datafusion/core/src/datasource/file_format/parquet.rs

@@ -141,7 +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::statistics::{StatisticsArgs, StatisticsContext};
     use datafusion_physical_plan::stream::RecordBatchStreamAdapter;
     use datafusion_physical_plan::{ExecutionPlan, collect};
 
@@ -716,12 +716,15 @@ mod tests {
 
         // test metadata
         assert_eq!(
-            exec.statistics_with_args(&StatisticsArgs::new())?.num_rows,
+            StatisticsContext::new()
+                .compute(exec.as_ref(), &StatisticsArgs::new())?
+                .num_rows,
             Precision::Exact(8)
         );
         // TODO correct byte size: https://github.com/apache/datafusion/issues/14936
         assert_eq!(
-            exec.statistics_with_args(&StatisticsArgs::new())?
+            StatisticsContext::new()
+                .compute(exec.as_ref(), &StatisticsArgs::new())?
                 .total_byte_size,
             Precision::Absent,
         );
@@ -766,11 +769,14 @@ mod tests {
 
         // note: even if the limit is set, the executor rounds up to the batch size
         assert_eq!(
-            exec.statistics_with_args(&StatisticsArgs::new())?.num_rows,
+            StatisticsContext::new()
+                .compute(exec.as_ref(), &StatisticsArgs::new())?
+                .num_rows,
             Precision::Exact(8)
         );
         assert_eq!(
-            exec.statistics_with_args(&StatisticsArgs::new())?
+            StatisticsContext::new()
+                .compute(exec.as_ref(), &StatisticsArgs::new())?
                 .total_byte_size,
             Precision::Absent,
         );

datafusion/core/src/datasource/listing/table.rs

@@ -145,7 +145,7 @@ mod tests {
     use datafusion_physical_expr::expressions::{Column, 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::statistics::{StatisticsArgs, StatisticsContext};
     use datafusion_physical_plan::{
         ExecutionPlanProperties, Partitioning, RangePartitioning, SplitPoint, collect,
     };
@@ -266,11 +266,14 @@ mod tests {
 
         // test metadata
         assert_eq!(
-            exec.statistics_with_args(&StatisticsArgs::new())?.num_rows,
+            StatisticsContext::new()
+                .compute(exec.as_ref(), &StatisticsArgs::new())?
+                .num_rows,
             Precision::Exact(8)
         );
         assert_eq!(
-            exec.statistics_with_args(&StatisticsArgs::new())?
+            StatisticsContext::new()
+                .compute(exec.as_ref(), &StatisticsArgs::new())?
                 .total_byte_size,
             Precision::Absent,
         );
@@ -1612,16 +1615,16 @@ mod tests {
 
         let exec_default = table_default.scan(&state, None, &[], None).await?;
         assert_eq!(
-            exec_default
-                .statistics_with_args(&StatisticsArgs::new())?
+            StatisticsContext::new()
+                .compute(exec_default.as_ref(), &StatisticsArgs::new())?
                 .num_rows,
             Precision::Exact(8)
         );
 
         // TODO correct byte size: https://github.com/apache/datafusion/issues/14936
         assert_eq!(
-            exec_default
-                .statistics_with_args(&StatisticsArgs::new())?
+            StatisticsContext::new()
+                .compute(exec_default.as_ref(), &StatisticsArgs::new())?
                 .total_byte_size,
             Precision::Absent
         );
@@ -1638,14 +1641,14 @@ mod tests {
 
         let exec_disabled = table_disabled.scan(&state, None, &[], None).await?;
         assert_eq!(
-            exec_disabled
-                .statistics_with_args(&StatisticsArgs::new())?
+            StatisticsContext::new()
+                .compute(exec_disabled.as_ref(), &StatisticsArgs::new())?
                 .num_rows,
             Precision::Absent
         );
         assert_eq!(
-            exec_disabled
-                .statistics_with_args(&StatisticsArgs::new())?
+            StatisticsContext::new()
+                .compute(exec_disabled.as_ref(), &StatisticsArgs::new())?
                 .total_byte_size,
             Precision::Absent
         );
@@ -1662,15 +1665,15 @@ mod tests {
 
         let exec_enabled = table_enabled.scan(&state, None, &[], None).await?;
         assert_eq!(
-            exec_enabled
-                .statistics_with_args(&StatisticsArgs::new())?
+            StatisticsContext::new()
+                .compute(exec_enabled.as_ref(), &StatisticsArgs::new())?
                 .num_rows,
             Precision::Exact(8)
         );
         // TODO correct byte size: https://github.com/apache/datafusion/issues/14936
         assert_eq!(
-            exec_enabled
-                .statistics_with_args(&StatisticsArgs::new())?
+            StatisticsContext::new()
+                .compute(exec_enabled.as_ref(), &StatisticsArgs::new())?
                 .total_byte_size,
             Precision::Absent,
         );

datafusion/core/tests/custom_sources_cases/mod.rs

@@ -179,7 +179,11 @@ impl ExecutionPlan for CustomExecutionPlan {
         Ok(Box::pin(TestCustomRecordBatchStream { nb_batch: 1 }))
     }
 
-    fn statistics_with_args(&self, args: &StatisticsArgs) -> Result<Arc<Statistics>> {
+    fn statistics_from_inputs(
+        &self,
+        _input_stats: &[Arc<Statistics>],
+        args: &StatisticsArgs,
+    ) -> Result<Arc<Statistics>> {
         if args.partition().is_some() {
             return Ok(Arc::new(Statistics::new_unknown(&self.schema())));
         }

datafusion/core/tests/custom_sources_cases/statistics.rs

@@ -35,8 +35,8 @@ 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 datafusion_physical_plan::{StatisticsArgs, StatisticsContext};
 
 use async_trait::async_trait;
 
@@ -174,7 +174,11 @@ impl ExecutionPlan for StatisticsValidation {
         unimplemented!("This plan only serves for testing statistics")
     }
 
-    fn statistics_with_args(&self, args: &StatisticsArgs) -> Result<Arc<Statistics>> {
+    fn statistics_from_inputs(
+        &self,
+        _input_stats: &[Arc<Statistics>],
+        args: &StatisticsArgs,
+    ) -> Result<Arc<Statistics>> {
         if args.partition().is_some() {
             Ok(Arc::new(Statistics::new_unknown(&self.schema)))
         } else {
@@ -233,7 +237,8 @@ async fn sql_basic() -> Result<()> {
     // the statistics should be those of the source
     assert_eq!(
         stats,
-        *physical_plan.statistics_with_args(&StatisticsArgs::new())?
+        *StatisticsContext::new()
+            .compute(physical_plan.as_ref(), &StatisticsArgs::new())?
     );
 
     Ok(())
@@ -250,7 +255,8 @@ async fn sql_filter() -> Result<()> {
         .unwrap();
 
     let physical_plan = df.create_physical_plan().await.unwrap();
-    let stats = physical_plan.statistics_with_args(&StatisticsArgs::new())?;
+    let stats = StatisticsContext::new()
+        .compute(physical_plan.as_ref(), &StatisticsArgs::new())?;
     assert_eq!(stats.num_rows, Precision::Inexact(7));
 
     Ok(())
@@ -265,7 +271,8 @@ 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.statistics_with_args(&StatisticsArgs::new())?;
+    let limit_stats = StatisticsContext::new()
+        .compute(physical_plan.as_ref(), &StatisticsArgs::new())?;
     assert_eq!(limit_stats.num_rows, Precision::Exact(5));
     // c1: NDV=2 stays at 2 (already below limit of 5)
     assert_eq!(
@@ -286,7 +293,8 @@ async fn sql_limit() -> Result<()> {
     // when the limit is larger than the original number of lines, statistics remain unchanged
     assert_eq!(
         stats,
-        *physical_plan.statistics_with_args(&StatisticsArgs::new())?
+        *StatisticsContext::new()
+            .compute(physical_plan.as_ref(), &StatisticsArgs::new())?
     );
 
     Ok(())
@@ -304,7 +312,8 @@ async fn sql_window() -> Result<()> {
 
     let physical_plan = df.create_physical_plan().await.unwrap();
 
-    let result = physical_plan.statistics_with_args(&StatisticsArgs::new())?;
+    let result = StatisticsContext::new()
+        .compute(physical_plan.as_ref(), &StatisticsArgs::new())?;
 
     assert_eq!(stats.num_rows, result.num_rows);
     let col_stats = &result.column_statistics;

datafusion/core/tests/parquet/file_statistics.rs

@@ -45,7 +45,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 datafusion_physical_plan::statistics::{StatisticsArgs, StatisticsContext};
 use tempfile::tempdir;
 
 #[tokio::test]
@@ -65,7 +65,8 @@ 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.statistics_with_args(&StatisticsArgs::new())
+        StatisticsContext::new()
+            .compute(exec.as_ref(), &StatisticsArgs::new())
             .unwrap()
             .num_rows,
         Precision::Exact(8),
@@ -99,8 +100,8 @@ async fn check_stats_precision_with_filter_pushdown() {
     );
     // Scan with filter pushdown, stats are inexact
     assert_eq!(
-        optimized_exec
-            .statistics_with_args(&StatisticsArgs::new())
+        StatisticsContext::new()
+            .compute(optimized_exec.as_ref(), &StatisticsArgs::new())
             .unwrap()
             .num_rows,
         Precision::Inexact(8),
@@ -135,15 +136,15 @@ async fn load_table_stats_with_session_level_cache() {
     let exec1 = table1.scan(&state1, None, &[], None).await.unwrap();
 
     assert_eq!(
-        exec1
-            .statistics_with_args(&StatisticsArgs::new())
+        StatisticsContext::new()
+            .compute(exec1.as_ref(), &StatisticsArgs::new())
             .unwrap()
             .num_rows,
         Precision::Exact(8)
     );
     assert_eq!(
-        exec1
-            .statistics_with_args(&StatisticsArgs::new())
+        StatisticsContext::new()
+            .compute(exec1.as_ref(), &StatisticsArgs::new())
             .unwrap()
             .total_byte_size,
         // Byte size is absent because we cannot estimate the output size
@@ -157,15 +158,15 @@ 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
-            .statistics_with_args(&StatisticsArgs::new())
+        StatisticsContext::new()
+            .compute(exec2.as_ref(), &StatisticsArgs::new())
             .unwrap()
             .num_rows,
         Precision::Exact(8)
     );
     assert_eq!(
-        exec2
-            .statistics_with_args(&StatisticsArgs::new())
+        StatisticsContext::new()
+            .compute(exec2.as_ref(), &StatisticsArgs::new())
             .unwrap()
             .total_byte_size,
         // Absent because the data contains variable length columns
@@ -178,15 +179,15 @@ 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
-            .statistics_with_args(&StatisticsArgs::new())
+        StatisticsContext::new()
+            .compute(exec3.as_ref(), &StatisticsArgs::new())
             .unwrap()
             .num_rows,
         Precision::Exact(8)
     );
     assert_eq!(
-        exec3
-            .statistics_with_args(&StatisticsArgs::new())
+        StatisticsContext::new()
+            .compute(exec3.as_ref(), &StatisticsArgs::new())
             .unwrap()
             .total_byte_size,
         // Absent because the data contains variable length columns
@@ -252,7 +253,9 @@ async fn anonymous_parquet_stats_cache_with_explicit_wider_schema() {
         .await
         .unwrap();
 
-    let stats = plan.statistics_with_args(&StatisticsArgs::new()).unwrap();
+    let stats = StatisticsContext::new()
+        .compute(plan.as_ref(), &StatisticsArgs::new())
+        .unwrap();
     assert_eq!(stats.column_statistics.len(), 2);
     assert_eq!(stats.column_statistics[1].null_count, Precision::Exact(1));