Replace partition_statistics_with_context with statistics_with_args

Introduce StatisticsArgs struct combining partition, child_stats, and
cache. Adds child_stats_for helper to eliminate per-operator boilerplate.
Removes StatisticsContext (flattened into StatisticsArgs).

Author: asolimando

datafusion-examples/examples/relation_planner/table_sample.rs

@@ -108,7 +108,7 @@ use datafusion::{
     },
     physical_expr::EquivalenceProperties,
     physical_plan::{
-        DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, StatisticsContext,
+        DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, StatisticsArgs,
         metrics::{BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet, RecordOutput},
     },
     physical_planner::{DefaultPhysicalPlanner, ExtensionPlanner, PhysicalPlanner},
@@ -722,12 +722,8 @@ impl ExecutionPlan for SampleExec {
         Some(self.metrics.clone_inner())
     }
 
-    fn partition_statistics_with_context(
-        &self,
-        _partition: Option<usize>,
-        ctx: &StatisticsContext,
-    ) -> Result<Arc<Statistics>> {
-        let mut stats = Arc::unwrap_or_clone(Arc::clone(&ctx.child_stats()[0]));
+    fn statistics_with_args(&self, args: &StatisticsArgs) -> Result<Arc<Statistics>> {
+        let mut stats = Arc::unwrap_or_clone(Arc::clone(&args.child_stats()[0]));
         let ratio = self.upper_bound - self.lower_bound;
 
         // Scale statistics by sampling ratio (inexact due to randomness)

datafusion/core/tests/custom_sources_cases/mod.rs

@@ -41,7 +41,7 @@ use datafusion_common::stats::Precision;
 use datafusion_common::tree_node::TreeNodeRecursion;
 use datafusion_physical_expr::EquivalenceProperties;
 use datafusion_physical_plan::PlanProperties;
-use datafusion_physical_plan::StatisticsContext;
+use datafusion_physical_plan::StatisticsArgs;
 use datafusion_physical_plan::execution_plan::{Boundedness, EmissionType};
 use datafusion_physical_plan::placeholder_row::PlaceholderRowExec;
 
@@ -180,12 +180,8 @@ impl ExecutionPlan for CustomExecutionPlan {
         Ok(Box::pin(TestCustomRecordBatchStream { nb_batch: 1 }))
     }
 
-    fn partition_statistics_with_context(
-        &self,
-        partition: Option<usize>,
-        _ctx: &StatisticsContext,
-    ) -> 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();

datafusion/core/tests/custom_sources_cases/statistics.rs

@@ -36,7 +36,7 @@ use datafusion_catalog::Session;
 use datafusion_common::tree_node::TreeNodeRecursion;
 use datafusion_common::{project_schema, stats::Precision};
 use datafusion_physical_expr::EquivalenceProperties;
-use datafusion_physical_plan::StatisticsContext;
+use datafusion_physical_plan::StatisticsArgs;
 use datafusion_physical_plan::compute_statistics;
 use datafusion_physical_plan::execution_plan::{Boundedness, EmissionType};
 
@@ -176,12 +176,8 @@ impl ExecutionPlan for StatisticsValidation {
         unimplemented!("This plan only serves for testing statistics")
     }
 
-    fn partition_statistics_with_context(
-        &self,
-        partition: Option<usize>,
-        _ctx: &StatisticsContext,
-    ) -> 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()))

datafusion/core/tests/physical_optimizer/join_selection.rs

@@ -45,7 +45,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, StatisticsContext,
+    DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, StatisticsArgs,
     compute_statistics,
     execution_plan::{Boundedness, EmissionType},
 };
@@ -1140,12 +1140,8 @@ impl ExecutionPlan for StatisticsExec {
         unimplemented!("This plan only serves for testing statistics")
     }
 
-    fn partition_statistics_with_context(
-        &self,
-        partition: Option<usize>,
-        _ctx: &StatisticsContext,
-    ) -> 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()

datafusion/core/tests/physical_optimizer/test_utils.rs

@@ -70,7 +70,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, StatisticsContext, displayable,
+    PlanProperties, SortOrderPushdownResult, StatisticsArgs, displayable,
 };
 
 /// Create a non sorted parquet exec
@@ -983,11 +983,7 @@ impl ExecutionPlan for TestScan {
         internal_err!("TestScan is for testing optimizer only, not for execution")
     }
 
-    fn partition_statistics_with_context(
-        &self,
-        _partition: Option<usize>,
-        _ctx: &StatisticsContext,
-    ) -> Result<Arc<Statistics>> {
+    fn statistics_with_args(&self, _args: &StatisticsArgs) -> Result<Arc<Statistics>> {
         Ok(Arc::new(Statistics::new_unknown(&self.schema)))
     }
 

datafusion/datasource/src/source.rs

@@ -46,7 +46,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::StatisticsContext;
+use datafusion_physical_plan::StatisticsArgs;
 use datafusion_physical_plan::filter_pushdown::{
     ChildPushdownResult, FilterPushdownPhase, FilterPushdownPropagation, PushedDown,
 };
@@ -452,12 +452,8 @@ impl ExecutionPlan for DataSourceExec {
         Some(metrics)
     }
 
-    fn partition_statistics_with_context(
-        &self,
-        partition: Option<usize>,
-        _ctx: &StatisticsContext,
-    ) -> 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>> {

datafusion/physical-optimizer/src/limit_pushdown.rs

@@ -71,6 +71,7 @@ use datafusion_common::stats::Precision;
 use datafusion_common::tree_node::{Transformed, TreeNodeRecursion};
 use datafusion_common::utils::combine_limit;
 use datafusion_physical_plan::coalesce_partitions::CoalescePartitionsExec;
+use datafusion_physical_plan::compute_statistics;
 use datafusion_physical_plan::empty::EmptyExec;
 use datafusion_physical_plan::limit::{GlobalLimitExec, LocalLimitExec};
 use datafusion_physical_plan::placeholder_row::PlaceholderRowExec;
@@ -351,7 +352,7 @@ fn limit_eliminable_exact_num_rows(
     }
 
     if matches!(
-        current.partition_statistics(None)?.num_rows,
+        compute_statistics(current.as_ref(), None)?.num_rows,
         Precision::Exact(0)
     ) {
         return Ok(Some(0));

datafusion/physical-optimizer/src/output_requirements.rs

@@ -34,13 +34,13 @@ 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::StatisticsContext;
 use datafusion_physical_plan::execution_plan::Boundedness;
 use datafusion_physical_plan::projection::{
     ProjectionExec, make_with_child, update_expr, update_ordering_requirement,
 };
 use datafusion_physical_plan::sorts::sort::SortExec;
 use datafusion_physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec;
+use datafusion_physical_plan::statistics_context::StatisticsArgs;
 use datafusion_physical_plan::{
     DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, PlanProperties,
     SendableRecordBatchStream,
@@ -243,15 +243,8 @@ impl ExecutionPlan for OutputRequirementExec {
         unreachable!();
     }
 
-    fn partition_statistics_with_context(
-        &self,
-        partition: Option<usize>,
-        ctx: &StatisticsContext,
-    ) -> Result<Arc<Statistics>> {
-        match partition {
-            Some(_) => ctx.compute_child_statistics(self.input.as_ref(), partition),
-            None => Ok(Arc::clone(&ctx.child_stats()[0])),
-        }
+    fn statistics_with_args(&self, args: &StatisticsArgs) -> Result<Arc<Statistics>> {
+        args.child_stats_for(0, self.input.as_ref())
     }
 
     fn try_swapping_with_projection(

datafusion/physical-plan/benches/compute_statistics.rs

@@ -41,9 +41,7 @@ use datafusion_physical_plan::execution_plan::{
     Boundedness, EmissionType, ExecutionPlan, PlanProperties,
 };
 use datafusion_physical_plan::joins::CrossJoinExec;
-use datafusion_physical_plan::statistics_context::{
-    StatisticsContext, compute_statistics,
-};
+use datafusion_physical_plan::statistics_context::{StatisticsArgs, compute_statistics};
 use datafusion_physical_plan::{
     DisplayAs, DisplayFormatType, Partitioning, SendableRecordBatchStream,
 };
@@ -119,11 +117,7 @@ impl ExecutionPlan for BenchLeaf {
         unimplemented!()
     }
 
-    fn partition_statistics_with_context(
-        &self,
-        _partition: Option<usize>,
-        _ctx: &StatisticsContext,
-    ) -> Result<Arc<Statistics>> {
+    fn statistics_with_args(&self, _args: &StatisticsArgs) -> Result<Arc<Statistics>> {
         Ok(Arc::new(Statistics::new_unknown(&self.schema)))
     }
 }
@@ -166,8 +160,8 @@ fn compute_statistics_without_shared_cache(
         .iter()
         .map(|child| compute_statistics_without_shared_cache(child.as_ref(), partition))
         .collect::<Result<Vec<_>>>()?;
-    let ctx = StatisticsContext::new(child_stats);
-    plan.partition_statistics_with_context(partition, &ctx)
+    let args = StatisticsArgs::new(partition, child_stats);
+    plan.statistics_with_args(&args)
 }
 
 fn bench_compute_statistics(c: &mut Criterion) {

datafusion/physical-plan/src/aggregates/mod.rs

@@ -31,7 +31,7 @@ use crate::filter_pushdown::{
     FilterPushdownPropagation, PushedDownPredicate,
 };
 use crate::metrics::{ExecutionPlanMetricsSet, MetricsSet};
-use crate::statistics_context::StatisticsContext;
+use crate::statistics_context::StatisticsArgs;
 use crate::{
     DisplayFormatType, Distribution, ExecutionPlan, InputOrderMode,
     SendableRecordBatchStream, Statistics, check_if_same_properties,
@@ -1644,15 +1644,8 @@ impl ExecutionPlan for AggregateExec {
         Some(self.metrics.clone_inner())
     }
 
-    fn partition_statistics_with_context(
-        &self,
-        partition: Option<usize>,
-        ctx: &StatisticsContext,
-    ) -> Result<Arc<Statistics>> {
-        let child_statistics = match partition {
-            Some(_) => ctx.compute_child_statistics(self.input.as_ref(), partition)?,
-            None => Arc::clone(&ctx.child_stats()[0]),
-        };
+    fn statistics_with_args(&self, args: &StatisticsArgs) -> Result<Arc<Statistics>> {
+        let child_statistics = args.child_stats_for(0, self.input.as_ref())?;
         Ok(Arc::new(self.statistics_inner(&child_statistics)?))
     }
 
@@ -2801,12 +2794,8 @@ mod tests {
             Ok(Box::pin(stream))
         }
 
-        fn partition_statistics_with_context(
-            &self,
-            partition: Option<usize>,
-            _ctx: &StatisticsContext,
-        ) -> 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();