Feat: introduce partition statistics API

Author: xudong963

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

@@ -42,7 +42,6 @@ use datafusion_physical_plan::{
 };
 
 use datafusion_datasource::file_groups::FileGroup;
-use datafusion_physical_plan::statistics::PartitionedStatistics;
 use futures::StreamExt;
 use itertools::Itertools;
 use object_store::{GetOptions, GetRange, GetResultPayload, ObjectStore};
@@ -192,8 +191,8 @@ impl ExecutionPlan for ArrowExec {
         self.inner.statistics()
     }
 
-    fn statistics_by_partition(&self) -> Result<PartitionedStatistics> {
-        self.inner.statistics_by_partition()
+    fn partition_statistics(&self, partition: Option<usize>) -> Result<Statistics> {
+        self.inner.partition_statistics(partition)
     }
 
     fn fetch(&self) -> Option<usize> {

datafusion/core/tests/physical_optimizer/partition_statistics.rs

@@ -39,7 +39,9 @@ mod test {
     use datafusion_physical_plan::projection::ProjectionExec;
     use datafusion_physical_plan::sorts::sort::SortExec;
     use datafusion_physical_plan::union::UnionExec;
-    use datafusion_physical_plan::{execute_stream_partitioned, ExecutionPlan};
+    use datafusion_physical_plan::{
+        execute_stream_partitioned, ExecutionPlan, ExecutionPlanProperties,
+    };
     use futures::TryStreamExt;
     use std::sync::Arc;
 
@@ -186,7 +188,9 @@ mod test {
     #[tokio::test]
     async fn test_statistics_by_partition_of_data_source() -> Result<()> {
         let scan = create_scan_exec_with_statistics(None, Some(2)).await;
-        let statistics = scan.statistics_by_partition()?;
+        let statistics = (0..scan.output_partitioning().partition_count())
+            .map(|idx| scan.partition_statistics(Some(idx)))
+            .collect::<Result<Vec<_>>>()?;
         let expected_statistic_partition_1 =
             create_partition_statistics(2, 110, 3, 4, true);
         let expected_statistic_partition_2 =
@@ -212,8 +216,11 @@ mod test {
         // Add projection execution plan
         let exprs: Vec<(Arc<dyn PhysicalExpr>, String)> =
             vec![(Arc::new(Column::new("id", 0)), "id".to_string())];
-        let projection = ProjectionExec::try_new(exprs, scan)?;
-        let statistics = projection.statistics_by_partition()?;
+        let projection: Arc<dyn ExecutionPlan> =
+            Arc::new(ProjectionExec::try_new(exprs, scan)?);
+        let statistics = (0..projection.output_partitioning().partition_count())
+            .map(|idx| projection.partition_statistics(Some(idx)))
+            .collect::<Result<Vec<_>>>()?;
         let expected_statistic_partition_1 =
             create_partition_statistics(2, 8, 3, 4, false);
         let expected_statistic_partition_2 =
@@ -225,7 +232,7 @@ mod test {
 
         // Check the statistics_by_partition with real results
         let expected_stats = vec![(3, 4, 2), (1, 2, 2)];
-        validate_statistics_with_data(Arc::new(projection), expected_stats, 0).await?;
+        validate_statistics_with_data(projection, expected_stats, 0).await?;
         Ok(())
     }
 
@@ -243,8 +250,10 @@ mod test {
             }]),
             scan_1,
         );
-        let sort_exec = Arc::new(sort.clone());
-        let statistics = sort_exec.statistics_by_partition()?;
+        let sort_exec: Arc<dyn ExecutionPlan> = Arc::new(sort.clone());
+        let statistics = (0..sort_exec.output_partitioning().partition_count())
+            .map(|idx| sort_exec.partition_statistics(Some(idx)))
+            .collect::<Result<Vec<_>>>()?;
         let expected_statistic_partition =
             create_partition_statistics(4, 220, 1, 4, true);
         assert_eq!(statistics.len(), 1);
@@ -256,7 +265,7 @@ mod test {
         // Sort with preserve_partitioning
         let scan_2 = create_scan_exec_with_statistics(None, Some(2)).await;
         // Add sort execution plan
-        let sort_exec = Arc::new(
+        let sort_exec: Arc<dyn ExecutionPlan> = Arc::new(
             SortExec::new(
                 LexOrdering::new(vec![PhysicalSortExpr {
                     expr: Arc::new(Column::new("id", 0)),
@@ -273,7 +282,9 @@ mod test {
             create_partition_statistics(2, 110, 3, 4, true);
         let expected_statistic_partition_2 =
             create_partition_statistics(2, 110, 1, 2, true);
-        let statistics = sort_exec.statistics_by_partition()?;
+        let statistics = (0..sort_exec.output_partitioning().partition_count())
+            .map(|idx| sort_exec.partition_statistics(Some(idx)))
+            .collect::<Result<Vec<_>>>()?;
         assert_eq!(statistics.len(), 2);
         assert_eq!(statistics[0], expected_statistic_partition_1);
         assert_eq!(statistics[1], expected_statistic_partition_2);
@@ -296,7 +307,7 @@ mod test {
         )?;
         let filter: Arc<dyn ExecutionPlan> =
             Arc::new(FilterExec::try_new(predicate, scan)?);
-        let full_statistics = filter.statistics()?;
+        let full_statistics = filter.partition_statistics(None)?;
         let expected_full_statistic = Statistics {
             num_rows: Precision::Inexact(0),
             total_byte_size: Precision::Inexact(0),
@@ -319,7 +330,9 @@ mod test {
         };
         assert_eq!(full_statistics, expected_full_statistic);
 
-        let statistics = filter.statistics_by_partition()?;
+        let statistics = (0..filter.output_partitioning().partition_count())
+            .map(|idx| filter.partition_statistics(Some(idx)))
+            .collect::<Result<Vec<_>>>()?;
         assert_eq!(statistics.len(), 2);
         assert_eq!(statistics[0], expected_full_statistic);
         assert_eq!(statistics[1], expected_full_statistic);
@@ -329,8 +342,11 @@ mod test {
     #[tokio::test]
     async fn test_statistic_by_partition_of_union() -> Result<()> {
         let scan = create_scan_exec_with_statistics(None, Some(2)).await;
-        let union_exec = Arc::new(UnionExec::new(vec![scan.clone(), scan]));
-        let statistics = union_exec.statistics_by_partition()?;
+        let union_exec: Arc<dyn ExecutionPlan> =
+            Arc::new(UnionExec::new(vec![scan.clone(), scan]));
+        let statistics = (0..union_exec.output_partitioning().partition_count())
+            .map(|idx| union_exec.partition_statistics(Some(idx)))
+            .collect::<Result<Vec<_>>>()?;
         // Check that we have 4 partitions (2 from each scan)
         assert_eq!(statistics.len(), 4);
         let expected_statistic_partition_1 =
@@ -360,8 +376,11 @@ mod test {
                                                 WITH ORDER (id ASC);";
         let right_scan =
             create_scan_exec_with_statistics(Some(right_create_table_sql), Some(2)).await;
-        let cross_join = CrossJoinExec::new(left_scan, right_scan);
-        let statistics = cross_join.statistics_by_partition()?;
+        let cross_join: Arc<dyn ExecutionPlan> =
+            Arc::new(CrossJoinExec::new(left_scan, right_scan));
+        let statistics = (0..cross_join.output_partitioning().partition_count())
+            .map(|idx| cross_join.partition_statistics(Some(idx)))
+            .collect::<Result<Vec<_>>>()?;
         // Check that we have 2 partitions
         assert_eq!(statistics.len(), 2);
         let mut expected_statistic_partition_1 =
@@ -391,52 +410,60 @@ mod test {
 
         // Check the statistics_by_partition with real results
         let expected_stats = vec![(1, 4, 8), (1, 4, 8)];
-        validate_statistics_with_data(Arc::new(cross_join), expected_stats, 0).await?;
+        validate_statistics_with_data(cross_join, expected_stats, 0).await?;
         Ok(())
     }
 
     #[tokio::test]
     async fn test_statistic_by_partition_of_coalesce_batches() -> Result<()> {
         let scan = create_scan_exec_with_statistics(None, Some(2)).await;
-        let coalesce_batches = CoalesceBatchesExec::new(scan, 2);
+        dbg!(scan.partition_statistics(Some(0))?);
+        let coalesce_batches: Arc<dyn ExecutionPlan> =
+            Arc::new(CoalesceBatchesExec::new(scan, 2));
         let expected_statistic_partition_1 =
             create_partition_statistics(2, 110, 3, 4, true);
         let expected_statistic_partition_2 =
             create_partition_statistics(2, 110, 1, 2, true);
-        let statistics = coalesce_batches.statistics_by_partition()?;
+        let statistics = (0..coalesce_batches.output_partitioning().partition_count())
+            .map(|idx| coalesce_batches.partition_statistics(Some(idx)))
+            .collect::<Result<Vec<_>>>()?;
         assert_eq!(statistics.len(), 2);
         assert_eq!(statistics[0], expected_statistic_partition_1);
         assert_eq!(statistics[1], expected_statistic_partition_2);
 
         // Check the statistics_by_partition with real results
         let expected_stats = vec![(3, 4, 2), (1, 2, 2)];
-        validate_statistics_with_data(Arc::new(coalesce_batches), expected_stats, 0)
-            .await?;
+        validate_statistics_with_data(coalesce_batches, expected_stats, 0).await?;
         Ok(())
     }
 
     #[tokio::test]
     async fn test_statistic_by_partition_of_coalesce_partitions() -> Result<()> {
         let scan = create_scan_exec_with_statistics(None, Some(2)).await;
-        let coalesce_partitions = CoalescePartitionsExec::new(scan);
+        let coalesce_partitions: Arc<dyn ExecutionPlan> =
+            Arc::new(CoalescePartitionsExec::new(scan));
         let expected_statistic_partition =
             create_partition_statistics(4, 220, 1, 4, true);
-        let statistics = coalesce_partitions.statistics_by_partition()?;
+        let statistics = (0..coalesce_partitions.output_partitioning().partition_count())
+            .map(|idx| coalesce_partitions.partition_statistics(Some(idx)))
+            .collect::<Result<Vec<_>>>()?;
         assert_eq!(statistics.len(), 1);
         assert_eq!(statistics[0], expected_statistic_partition);
 
         // Check the statistics_by_partition with real results
         let expected_stats = vec![(1, 4, 4)];
-        validate_statistics_with_data(Arc::new(coalesce_partitions), expected_stats, 0)
-            .await?;
+        validate_statistics_with_data(coalesce_partitions, expected_stats, 0).await?;
         Ok(())
     }
 
     #[tokio::test]
     async fn test_statistic_by_partition_of_local_limit() -> Result<()> {
         let scan = create_scan_exec_with_statistics(None, Some(2)).await;
-        let local_limit = LocalLimitExec::new(scan.clone(), 1);
-        let statistics = local_limit.statistics_by_partition()?;
+        let local_limit: Arc<dyn ExecutionPlan> =
+            Arc::new(LocalLimitExec::new(scan.clone(), 1));
+        let statistics = (0..local_limit.output_partitioning().partition_count())
+            .map(|idx| local_limit.partition_statistics(Some(idx)))
+            .collect::<Result<Vec<_>>>()?;
         assert_eq!(statistics.len(), 2);
         let schema = scan.schema();
         let mut expected_statistic_partition = Statistics::new_unknown(&schema);
@@ -449,8 +476,11 @@ mod test {
     #[tokio::test]
     async fn test_statistic_by_partition_of_global_limit_partitions() -> Result<()> {
         let scan = create_scan_exec_with_statistics(None, Some(2)).await;
-        let global_limit = GlobalLimitExec::new(scan.clone(), 0, Some(2));
-        let statistics = global_limit.statistics_by_partition()?;
+        let global_limit: Arc<dyn ExecutionPlan> =
+            Arc::new(GlobalLimitExec::new(scan.clone(), 0, Some(2)));
+        let statistics = (0..global_limit.output_partitioning().partition_count())
+            .map(|idx| global_limit.partition_statistics(Some(idx)))
+            .collect::<Result<Vec<_>>>()?;
         assert_eq!(statistics.len(), 1);
         let mut expected_statistic_partition = Statistics::new_unknown(&scan.schema());
         expected_statistic_partition.num_rows = Precision::Exact(2);

datafusion/datasource-avro/src/source.rs

@@ -40,7 +40,6 @@ use datafusion_physical_plan::{
     DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties,
 };
 
-use datafusion_physical_plan::statistics::PartitionedStatistics;
 use object_store::ObjectStore;
 
 /// Execution plan for scanning Avro data source
@@ -142,8 +141,8 @@ impl ExecutionPlan for AvroExec {
         self.inner.statistics()
     }
 
-    fn statistics_by_partition(&self) -> Result<PartitionedStatistics> {
-        self.inner.statistics_by_partition()
+    fn partition_statistics(&self, partition: Option<usize>) -> Result<Statistics> {
+        self.inner.partition_statistics(partition)
     }
 
     fn metrics(&self) -> Option<MetricsSet> {

datafusion/datasource-csv/src/source.rs

@@ -51,7 +51,6 @@ use datafusion_physical_plan::{
 
 use crate::file_format::CsvDecoder;
 use datafusion_datasource::file_groups::FileGroup;
-use datafusion_physical_plan::statistics::PartitionedStatistics;
 use futures::{StreamExt, TryStreamExt};
 use object_store::buffered::BufWriter;
 use object_store::{GetOptions, GetResultPayload, ObjectStore};
@@ -382,8 +381,8 @@ impl ExecutionPlan for CsvExec {
         self.inner.statistics()
     }
 
-    fn statistics_by_partition(&self) -> Result<PartitionedStatistics> {
-        self.inner.statistics_by_partition()
+    fn partition_statistics(&self, partition: Option<usize>) -> Result<Statistics> {
+        self.inner.partition_statistics(partition)
     }
 
     fn metrics(&self) -> Option<MetricsSet> {

datafusion/datasource/src/file_groups.rs

@@ -420,11 +420,6 @@ impl FileGroup {
         self.files.push(file);
     }
 
-    /// Get the statistics for this group
-    pub fn statistics(&self) -> &Option<Arc<Statistics>> {
-        &self.statistics
-    }
-
     /// Get the statistics for this group
     pub fn statistics_ref(&self) -> Option<&Statistics> {
         self.statistics.as_deref()

datafusion/datasource/src/source.rs

@@ -25,7 +25,6 @@ use std::sync::Arc;
 use datafusion_physical_plan::execution_plan::{Boundedness, EmissionType};
 use datafusion_physical_plan::metrics::{ExecutionPlanMetricsSet, MetricsSet};
 use datafusion_physical_plan::projection::ProjectionExec;
-use datafusion_physical_plan::statistics::PartitionedStatistics;
 use datafusion_physical_plan::{
     DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties,
 };
@@ -189,24 +188,22 @@ impl ExecutionPlan for DataSourceExec {
         self.data_source.statistics()
     }
 
-    fn statistics_by_partition(&self) -> Result<PartitionedStatistics> {
-        let mut statistics = {
-            let mut v =
-                Vec::with_capacity(self.properties().partitioning.partition_count());
-            (0..self.properties().partitioning.partition_count())
-                .for_each(|_| v.push(Arc::new(Statistics::new_unknown(&self.schema()))));
-            v
-        };
-        if let Some(file_config) =
-            self.data_source.as_any().downcast_ref::<FileScanConfig>()
-        {
-            for (idx, file_group) in file_config.file_groups.iter().enumerate() {
-                if let Some(stat) = file_group.statistics() {
-                    statistics[idx] = Arc::clone(stat);
+    fn partition_statistics(&self, partition: Option<usize>) -> Result<Statistics> {
+        if let Some(partition) = partition {
+            let mut statistics = Statistics::new_unknown(&self.schema());
+            if let Some(file_config) =
+                self.data_source.as_any().downcast_ref::<FileScanConfig>()
+            {
+                if let Some(file_group) = file_config.file_groups.get(partition) {
+                    if let Some(stat) = file_group.statistics_ref() {
+                        statistics = stat.clone();
+                    }
                 }
             }
+            Ok(statistics)
+        } else {
+            Ok(self.data_source.statistics()?)
         }
-        Ok(PartitionedStatistics::new(statistics))
     }
 
     fn with_fetch(&self, limit: Option<usize>) -> Option<Arc<dyn ExecutionPlan>> {

datafusion/datasource/src/statistics.rs

@@ -476,7 +476,7 @@ pub fn compute_all_files_statistics(
     // Then summary statistics across all file groups
     let file_groups_statistics = file_groups_with_stats
         .iter()
-        .filter_map(|file_group| file_group.statistics().as_deref());
+        .filter_map(|file_group| file_group.statistics_ref());
 
     let mut statistics =
         Statistics::try_merge_iter(file_groups_statistics, &table_schema)?;

datafusion/physical-plan/src/coalesce_batches.rs

@@ -24,7 +24,6 @@ use std::task::{Context, Poll};
 
 use super::metrics::{BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet};
 use super::{DisplayAs, ExecutionPlanProperties, PlanProperties, Statistics};
-use crate::statistics::PartitionedStatistics;
 use crate::{
     DisplayFormatType, ExecutionPlan, RecordBatchStream, SendableRecordBatchStream,
 };
@@ -200,24 +199,11 @@ impl ExecutionPlan for CoalesceBatchesExec {
         Statistics::with_fetch(self.input.statistics()?, self.schema(), self.fetch, 0, 1)
     }
 
-    fn statistics_by_partition(&self) -> Result<PartitionedStatistics> {
-        let input_stats = self.input.statistics_by_partition()?;
-
-        let stats: Result<Vec<Arc<Statistics>>> = input_stats
-            .iter()
-            .map(|stat| {
-                let fetched_stat = Statistics::with_fetch(
-                    stat.clone(),
-                    self.schema(),
-                    self.fetch,
-                    0,
-                    1,
-                )?;
-                Ok(Arc::new(fetched_stat))
-            })
-            .collect();
-
-        Ok(PartitionedStatistics::new(stats?))
+    fn partition_statistics(&self, partition: Option<usize>) -> Result<Statistics> {
+        let input_stats = self.input.partition_statistics(partition)?;
+        let fetched_stat =
+            Statistics::with_fetch(input_stats.clone(), self.schema(), self.fetch, 0, 1)?;
+        Ok(fetched_stat)
     }
 
     fn with_fetch(&self, limit: Option<usize>) -> Option<Arc<dyn ExecutionPlan>> {

datafusion/physical-plan/src/coalesce_partitions.rs

@@ -31,7 +31,6 @@ use crate::execution_plan::CardinalityEffect;
 use crate::projection::{make_with_child, ProjectionExec};
 use crate::{DisplayFormatType, ExecutionPlan, Partitioning};
 
-use crate::statistics::PartitionedStatistics;
 use datafusion_common::{internal_err, Result};
 use datafusion_execution::TaskContext;
 
@@ -200,10 +199,8 @@ impl ExecutionPlan for CoalescePartitionsExec {
         Statistics::with_fetch(self.input.statistics()?, self.schema(), self.fetch, 0, 1)
     }
 
-    fn statistics_by_partition(&self) -> Result<PartitionedStatistics> {
-        Ok(PartitionedStatistics::new(vec![Arc::new(
-            self.statistics()?,
-        )]))
+    fn partition_statistics(&self, _partition: Option<usize>) -> Result<Statistics> {
+        Statistics::with_fetch(self.input.statistics()?, self.schema(), self.fetch, 0, 1)
     }
 
     fn supports_limit_pushdown(&self) -> bool {