functional way

Author: xudong963

datafusion/physical-plan/src/filter.rs

@@ -410,18 +410,19 @@ impl ExecutionPlan for FilterExec {
     }
 
     fn statistics_by_partition(&self) -> Result<Vec<Statistics>> {
-        let input_stats = self.input.statistics_by_partition()?;
-        let mut stats = Vec::with_capacity(input_stats.len());
-        for input_stat in input_stats {
-            let stat = Self::statistics_helper(
-                self.schema(),
-                input_stat,
-                self.predicate(),
-                self.default_selectivity,
-            )?;
-            stats.push(stat.project(self.projection.as_ref()));
-        }
-        Ok(stats)
+        self.input
+            .statistics_by_partition()?
+            .into_iter()
+            .map(|input_stat| {
+                Self::statistics_helper(
+                    self.schema(),
+                    input_stat,
+                    self.predicate(),
+                    self.default_selectivity,
+                )
+                .map(|stat| stat.project(self.projection.as_ref()))
+            })
+            .collect()
     }
 
     fn cardinality_effect(&self) -> CardinalityEffect {

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

@@ -355,16 +355,15 @@ impl ExecutionPlan for CrossJoinExec {
 
         // Summarize the `left_stats`
         let statistics = compute_summary_statistics(
-            left_stats.iter(),
+            left_stats.into_iter(),
             self.schema.fields().len(),
             |stats| Some(stats),
         );
 
-        let mut stats = Vec::new();
-        for right in right_stats.iter() {
-            stats.push(stats_cartesian_product(statistics.clone(), right.clone()));
-        }
-        Ok(stats)
+        Ok(right_stats
+            .into_iter()
+            .map(|right| stats_cartesian_product(statistics.clone(), right))
+            .collect())
     }
 
     /// Tries to swap the projection with its input [`CrossJoinExec`]. If it can be done,

datafusion/physical-plan/src/limit.rs

@@ -348,19 +348,13 @@ impl ExecutionPlan for LocalLimitExec {
     }
 
     fn statistics_by_partition(&self) -> Result<Vec<Statistics>> {
-        let input_stats = self.input.statistics_by_partition()?;
-        let mut stats = Vec::with_capacity(input_stats.len());
-        for input_stat in input_stats {
-            let stat = Statistics::with_fetch(
-                input_stat,
-                self.schema(),
-                Some(self.fetch),
-                0,
-                1,
-            )?;
-            stats.push(stat);
-        }
-        Ok(stats)
+        self.input
+            .statistics_by_partition()?
+            .into_iter()
+            .map(|input_stat| {
+                Statistics::with_fetch(input_stat, self.schema(), Some(self.fetch), 0, 1)
+            })
+            .collect()
     }
 
     fn fetch(&self) -> Option<usize> {

datafusion/physical-plan/src/projection.rs

@@ -252,17 +252,18 @@ impl ExecutionPlan for ProjectionExec {
     }
 
     fn statistics_by_partition(&self) -> Result<Vec<Statistics>> {
-        let input_stats = self.input.statistics_by_partition()?;
-        let mut stats = Vec::with_capacity(input_stats.len());
-        for input_stat in input_stats {
-            let stat = stats_projection(
-                input_stat.clone(),
-                self.expr.iter().map(|(e, _)| Arc::clone(e)),
-                Arc::clone(&self.schema),
-            );
-            stats.push(stat);
-        }
-        Ok(stats)
+        Ok(self
+            .input
+            .statistics_by_partition()?
+            .into_iter()
+            .map(|input_stats| {
+                stats_projection(
+                    input_stats,
+                    self.expr.iter().map(|(e, _)| Arc::clone(e)),
+                    Arc::clone(&self.schema),
+                )
+            })
+            .collect())
     }
 
     fn supports_limit_pushdown(&self) -> bool {

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

@@ -680,10 +680,6 @@ impl ExecutionPlan for RepartitionExec {
         self.input.statistics()
     }
 
-    fn statistics_by_partition(&self) -> Result<Vec<Statistics>> {
-        todo!()
-    }
-
     fn cardinality_effect(&self) -> CardinalityEffect {
         CardinalityEffect::Equal
     }

datafusion/physical-plan/src/sorts/sort.rs

@@ -1274,18 +1274,11 @@ impl ExecutionPlan for SortExec {
     }
 
     fn statistics_by_partition(&self) -> Result<Vec<Statistics>> {
-        let input_stats = self.input.statistics_by_partition()?;
-        let mut stats = Vec::with_capacity(input_stats.len());
-        for stat in input_stats {
-            stats.push(Statistics::with_fetch(
-                stat,
-                self.schema(),
-                self.fetch,
-                0,
-                1,
-            )?);
-        }
-        Ok(stats)
+        self.input
+            .statistics_by_partition()?
+            .into_iter()
+            .map(|stat| Statistics::with_fetch(stat, self.schema(), self.fetch, 0, 1))
+            .collect()
     }
 
     fn with_fetch(&self, limit: Option<usize>) -> Option<Arc<dyn ExecutionPlan>> {

datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs

@@ -365,13 +365,11 @@ impl ExecutionPlan for BoundedWindowAggExec {
     }
 
     fn statistics_by_partition(&self) -> Result<Vec<Statistics>> {
-        let input_stats = self.input.statistics_by_partition()?;
-        let mut output_stats = Vec::with_capacity(input_stats.len());
-        for stat in input_stats {
-            let output_stat = self.statistics_helper(stat.clone())?;
-            output_stats.push(output_stat);
-        }
-        Ok(output_stats)
+        self.input
+            .statistics_by_partition()?
+            .into_iter()
+            .map(|stat| self.statistics_helper(stat))
+            .collect()
     }
 }