Remove child_stats Vec, use cache for all child statistics lookups

Author: asolimando

datafusion-examples/examples/relation_planner/table_sample.rs

@@ -723,7 +723,7 @@ impl ExecutionPlan for SampleExec {
     }
 
     fn statistics_with_args(&self, args: &StatisticsArgs) -> Result<Arc<Statistics>> {
-        let mut stats = Arc::unwrap_or_clone(Arc::clone(&args.child_stats()[0]));
+        let mut stats = Arc::unwrap_or_clone(args.child_stats_of(self.input.as_ref())?);
         let ratio = self.upper_bound - self.lower_bound;
 
         // Scale statistics by sampling ratio (inexact due to randomness)

datafusion/physical-optimizer/src/output_requirements.rs

@@ -244,7 +244,7 @@ impl ExecutionPlan for OutputRequirementExec {
     }
 
     fn statistics_with_args(&self, args: &StatisticsArgs) -> Result<Arc<Statistics>> {
-        args.child_stats_for(0, self.input.as_ref())
+        args.child_stats_for(self.input.as_ref())
     }
 
     fn try_swapping_with_projection(

datafusion/physical-plan/benches/compute_statistics.rs

@@ -174,12 +174,10 @@ fn compute_statistics_without_shared_cache(
     plan: &dyn ExecutionPlan,
     partition: Option<usize>,
 ) -> Result<Arc<Statistics>> {
-    let child_stats = plan
-        .children()
-        .iter()
-        .map(|child| compute_statistics_without_shared_cache(child.as_ref(), partition))
-        .collect::<Result<Vec<_>>>()?;
-    let args = StatisticsArgs::new(partition, child_stats);
+    for child in plan.children() {
+        compute_statistics_without_shared_cache(child.as_ref(), partition)?;
+    }
+    let args = StatisticsArgs::new(partition);
     plan.statistics_with_args(&args)
 }
 

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

@@ -1645,7 +1645,7 @@ impl ExecutionPlan for AggregateExec {
     }
 
     fn statistics_with_args(&self, args: &StatisticsArgs) -> Result<Arc<Statistics>> {
-        let child_statistics = args.child_stats_for(0, self.input.as_ref())?;
+        let child_statistics = args.child_stats_for(self.input.as_ref())?;
         Ok(Arc::new(self.statistics_inner(&child_statistics)?))
     }
 

datafusion/physical-plan/src/buffer.rs

@@ -246,7 +246,7 @@ impl ExecutionPlan for BufferExec {
     }
 
     fn statistics_with_args(&self, args: &StatisticsArgs) -> Result<Arc<Statistics>> {
-        args.child_stats_for(0, self.input.as_ref())
+        args.child_stats_for(self.input.as_ref())
     }
 
     fn supports_limit_pushdown(&self) -> bool {

datafusion/physical-plan/src/coalesce_batches.rs

@@ -224,7 +224,7 @@ impl ExecutionPlan for CoalesceBatchesExec {
     }
 
     fn statistics_with_args(&self, args: &StatisticsArgs) -> Result<Arc<Statistics>> {
-        let stats = Arc::unwrap_or_clone(args.child_stats_for(0, self.input.as_ref())?);
+        let stats = Arc::unwrap_or_clone(args.child_stats_for(self.input.as_ref())?);
         Ok(Arc::new(stats.with_fetch(self.fetch, 0, 1)?))
     }
 

datafusion/physical-plan/src/coalesce_partitions.rs

@@ -241,7 +241,7 @@ impl ExecutionPlan for CoalescePartitionsExec {
     }
 
     fn statistics_with_args(&self, args: &StatisticsArgs) -> Result<Arc<Statistics>> {
-        let stats = Arc::unwrap_or_clone(Arc::clone(&args.child_stats()[0]));
+        let stats = Arc::unwrap_or_clone(args.child_stats_of(self.input.as_ref())?);
         Ok(Arc::new(stats.with_fetch(self.fetch, 0, 1)?))
     }
 

datafusion/physical-plan/src/coop.rs

@@ -308,7 +308,7 @@ impl ExecutionPlan for CooperativeExec {
     }
 
     fn statistics_with_args(&self, args: &StatisticsArgs) -> Result<Arc<Statistics>> {
-        args.child_stats_for(0, self.input.as_ref())
+        args.child_stats_for(self.input.as_ref())
     }
 
     fn supports_limit_pushdown(&self) -> bool {

datafusion/physical-plan/src/filter.rs

@@ -579,7 +579,7 @@ impl ExecutionPlan for FilterExec {
     /// predicate's selectivity value can be determined for the incoming data.
     fn statistics_with_args(&self, args: &StatisticsArgs) -> Result<Arc<Statistics>> {
         let input_stats =
-            Arc::unwrap_or_clone(args.child_stats_for(0, self.input.as_ref())?);
+            Arc::unwrap_or_clone(args.child_stats_for(self.input.as_ref())?);
         let stats = Self::statistics_helper(
             &self.input.schema(),
             input_stats,

datafusion/physical-plan/src/joins/cross_join.rs

@@ -386,9 +386,9 @@ impl ExecutionPlan for CrossJoinExec {
         // so it always needs overall stats (child_stats provides these).
         // Right side can have multiple partitions, so it needs per-partition
         // stats when a specific partition is requested.
-        let left_stats = Arc::unwrap_or_clone(Arc::clone(&args.child_stats()[0]));
+        let left_stats = Arc::unwrap_or_clone(args.child_stats_of(self.left.as_ref())?);
         let right_stats =
-            Arc::unwrap_or_clone(args.child_stats_for(1, self.right.as_ref())?);
+            Arc::unwrap_or_clone(args.child_stats_for(self.right.as_ref())?);
 
         Ok(Arc::new(stats_cartesian_product(left_stats, right_stats)))
     }