save

Author: xudong963

datafusion/core/src/datasource/physical_plan/arrow_file.rs

@@ -191,7 +191,7 @@ impl ExecutionPlan for ArrowExec {
         self.inner.statistics()
     }
 
-    fn statistics_by_partition(&self) -> Result<Vec<Statistics>> {
+    fn statistics_by_partition(&self) -> Result<PartitionedStatistics> {
         self.inner.statistics_by_partition()
     }
 

datafusion/core/tests/physical_optimizer/partition_statistics.rs

@@ -39,15 +39,15 @@ mod test {
     ///
     /// This function:
     /// - Creates an external table from './tests/data/test_statistics_per_partition'
-    /// - If we set the `target_partition` to `2, the data contains 2 partitions, each with 2 rows
+    /// - If we set the `target_partition` to 2, the data contains 2 partitions, each with 2 rows
     /// - Each partition has an "id" column (INT) with the following values:
     ///   - First partition: [3, 4]
     ///   - Second partition: [1, 2]
     /// - Each row is 110 bytes in size
     ///
     /// @param target_partition Optional parameter to set the target partitions
     /// @return ExecutionPlan representing the scan of the table with statistics
-    async fn generate_listing_table_with_statistics(
+    async fn create_scan_exec_with_statistics(
         target_partition: Option<usize>,
     ) -> Arc<dyn ExecutionPlan> {
         let mut session_config = SessionConfig::new().with_collect_statistics(true);
@@ -111,7 +111,7 @@ mod test {
     #[tokio::test]
     async fn test_statistics_by_partition_of_data_source() -> datafusion_common::Result<()>
     {
-        let scan = generate_listing_table_with_statistics(Some(2)).await;
+        let scan = create_scan_exec_with_statistics(Some(2)).await;
         let statistics = scan.statistics_by_partition()?;
         let expected_statistic_partition_1 =
             create_partition_statistics(2, 110, 3, 4, true);
@@ -127,7 +127,7 @@ mod test {
     #[tokio::test]
     async fn test_statistics_by_partition_of_projection() -> datafusion_common::Result<()>
     {
-        let scan = generate_listing_table_with_statistics(Some(2)).await;
+        let scan = create_scan_exec_with_statistics(Some(2)).await;
         // Add projection execution plan
         let exprs: Vec<(Arc<dyn PhysicalExpr>, String)> =
             vec![(Arc::new(Column::new("id", 0)), "id".to_string())];
@@ -146,7 +146,7 @@ mod test {
 
     #[tokio::test]
     async fn test_statistics_by_partition_of_sort() -> datafusion_common::Result<()> {
-        let scan = generate_listing_table_with_statistics(Some(2)).await;
+        let scan = create_scan_exec_with_statistics(Some(2)).await;
         // Add sort execution plan
         let sort = SortExec::new(
             LexOrdering::new(vec![PhysicalSortExpr {
@@ -179,7 +179,7 @@ mod test {
 
     #[tokio::test]
     async fn test_statistics_by_partition_of_filter() -> datafusion_common::Result<()> {
-        let scan = generate_listing_table_with_statistics(Some(2)).await;
+        let scan = create_scan_exec_with_statistics(Some(2)).await;
         let schema = Schema::new(vec![Field::new("id", DataType::Int32, false)]);
         let predicate = binary(
             Arc::new(Column::new("id", 0)),
@@ -221,7 +221,7 @@ mod test {
 
     #[tokio::test]
     async fn test_statistic_by_partition_of_union() -> datafusion_common::Result<()> {
-        let scan = generate_listing_table_with_statistics(Some(2)).await;
+        let scan = create_scan_exec_with_statistics(Some(2)).await;
         let union_exec = Arc::new(UnionExec::new(vec![scan.clone(), scan]));
         let statistics = union_exec.statistics_by_partition()?;
         // Check that we have 4 partitions (2 from each scan)

datafusion/datasource-avro/src/source.rs

@@ -141,7 +141,7 @@ impl ExecutionPlan for AvroExec {
         self.inner.statistics()
     }
 
-    fn statistics_by_partition(&self) -> Result<Vec<Statistics>> {
+    fn statistics_by_partition(&self) -> Result<PartitionedStatistics> {
         self.inner.statistics_by_partition()
     }
 

datafusion/datasource-csv/src/source.rs

@@ -381,7 +381,7 @@ impl ExecutionPlan for CsvExec {
         self.inner.statistics()
     }
 
-    fn statistics_by_partition(&self) -> Result<Vec<Statistics>> {
+    fn statistics_by_partition(&self) -> Result<PartitionedStatistics> {
         self.inner.statistics_by_partition()
     }
 

datafusion/physical-plan/src/coalesce_batches.rs

@@ -195,8 +195,12 @@ impl ExecutionPlan for CoalesceBatchesExec {
         Statistics::with_fetch(self.input.statistics()?, self.schema(), self.fetch, 0, 1)
     }
 
-    fn statistics_by_partition(&self) -> Result<Vec<Statistics>> {
-        Ok(vec![self.statistics()?])
+    fn statistics_by_partition(&self) -> Result<PartitionedStatistics> {
+        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/coalesce_partitions.rs

@@ -193,7 +193,7 @@ impl ExecutionPlan for CoalescePartitionsExec {
         Statistics::with_fetch(self.input.statistics()?, self.schema(), self.fetch, 0, 1)
     }
 
-    fn statistics_by_partition(&self) -> Result<Vec<Statistics>> {
+    fn statistics_by_partition(&self) -> Result<PartitionedStatistics> {
         Ok(vec![self.statistics()?])
     }
 

datafusion/physical-plan/src/execution_plan.rs

@@ -52,6 +52,7 @@ use datafusion_physical_expr::{EquivalenceProperties, LexOrdering};
 use datafusion_physical_expr_common::sort_expr::LexRequirement;
 
 use futures::stream::{StreamExt, TryStreamExt};
+use crate::statistics::PartitionedStatistics;
 
 /// Represent nodes in the DataFusion Physical Plan.
 ///
@@ -430,11 +431,11 @@ pub trait ExecutionPlan: Debug + DisplayAs + Send + Sync {
     /// Returns statistics for each partition of this `ExecutionPlan` node.
     /// If statistics are not available, returns an array of
     /// [`Statistics::new_unknown`] for each partition.
-    fn statistics_by_partition(&self) -> Result<Vec<Statistics>> {
-        Ok(vec![
-            Statistics::new_unknown(&self.schema());
+    fn statistics_by_partition(&self) -> Result<PartitionedStatistics> {
+        Ok(PartitionedStatistics::new(vec![
+            Arc::new(Statistics::new_unknown(&self.schema()));
             self.properties().partitioning.partition_count()
-        ])
+        ]))
     }
 
     /// Returns `true` if a limit can be safely pushed down through this

datafusion/physical-plan/src/filter.rs

@@ -409,7 +409,7 @@ impl ExecutionPlan for FilterExec {
         Ok(stats.project(self.projection.as_ref()))
     }
 
-    fn statistics_by_partition(&self) -> Result<Vec<Statistics>> {
+    fn statistics_by_partition(&self) -> Result<PartitionedStatistics> {
         self.input
             .statistics_by_partition()?
             .into_iter()

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

@@ -345,7 +345,7 @@ impl ExecutionPlan for CrossJoinExec {
         ))
     }
 
-    fn statistics_by_partition(&self) -> Result<Vec<Statistics>> {
+    fn statistics_by_partition(&self) -> Result<PartitionedStatistics> {
         let left_stats = self.left.statistics_by_partition()?;
         let right_stats = self.right.statistics_by_partition()?;
 

datafusion/physical-plan/src/limit.rs

@@ -202,7 +202,7 @@ impl ExecutionPlan for GlobalLimitExec {
         )
     }
 
-    fn statistics_by_partition(&self) -> Result<Vec<Statistics>> {
+    fn statistics_by_partition(&self) -> Result<PartitionedStatistics> {
         Ok(vec![self.statistics()?])
     }
 
@@ -347,7 +347,7 @@ impl ExecutionPlan for LocalLimitExec {
         )
     }
 
-    fn statistics_by_partition(&self) -> Result<Vec<Statistics>> {
+    fn statistics_by_partition(&self) -> Result<PartitionedStatistics> {
         self.input
             .statistics_by_partition()?
             .into_iter()

datafusion/physical-plan/src/projection.rs

@@ -251,7 +251,7 @@ impl ExecutionPlan for ProjectionExec {
         ))
     }
 
-    fn statistics_by_partition(&self) -> Result<Vec<Statistics>> {
+    fn statistics_by_partition(&self) -> Result<PartitionedStatistics> {
         Ok(self
             .input
             .statistics_by_partition()?

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

@@ -1273,7 +1273,7 @@ impl ExecutionPlan for SortExec {
         Statistics::with_fetch(self.input.statistics()?, self.schema(), self.fetch, 0, 1)
     }
 
-    fn statistics_by_partition(&self) -> Result<Vec<Statistics>> {
+    fn statistics_by_partition(&self) -> Result<PartitionedStatistics> {
         if !self.preserve_partitioning() {
             return Ok(vec![self.statistics()?]);
         }

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

@@ -346,7 +346,7 @@ impl ExecutionPlan for SortPreservingMergeExec {
         self.input.statistics()
     }
 
-    fn statistics_by_partition(&self) -> Result<Vec<Statistics>> {
+    fn statistics_by_partition(&self) -> Result<PartitionedStatistics> {
         Ok(vec![self.statistics()?])
     }
 

datafusion/physical-plan/src/union.rs

@@ -270,7 +270,7 @@ impl ExecutionPlan for UnionExec {
             .unwrap_or_else(|| Statistics::new_unknown(&self.schema())))
     }
 
-    fn statistics_by_partition(&self) -> Result<Vec<Statistics>> {
+    fn statistics_by_partition(&self) -> Result<PartitionedStatistics> {
         Ok(self
             .inputs
             .iter()

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

@@ -364,7 +364,7 @@ impl ExecutionPlan for BoundedWindowAggExec {
         self.statistics_helper(input_stat)
     }
 
-    fn statistics_by_partition(&self) -> Result<Vec<Statistics>> {
+    fn statistics_by_partition(&self) -> Result<PartitionedStatistics> {
         self.input
             .statistics_by_partition()?
             .into_iter()