save

Author: xudong963

datafusion/datasource/src/source.rs

@@ -175,6 +175,23 @@ impl ExecutionPlan for DataSourceExec {
         self.data_source.statistics()
     }
 
+    fn statistics_by_partition(&self) -> datafusion_common::Result<Vec<Statistics>> {
+        let mut statistics = vec![
+            Statistics::new_unknown(&self.schema());
+            self.properties().partitioning.partition_count()
+        ];
+        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] = stat.clone();
+                }
+            }
+        }
+        Ok(statistics)
+    }
+
     fn with_fetch(&self, limit: Option<usize>) -> Option<Arc<dyn ExecutionPlan>> {
         let data_source = self.data_source.with_fetch(limit)?;
         let cache = self.cache.clone();

datafusion/physical-plan/src/execution_plan.rs

@@ -427,6 +427,16 @@ pub trait ExecutionPlan: Debug + DisplayAs + Send + Sync {
         Ok(Statistics::new_unknown(&self.schema()))
     }
 
+    /// 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());
+            self.properties().partitioning.partition_count()
+        ])
+    }
+
     /// Returns `true` if a limit can be safely pushed down through this
     /// `ExecutionPlan` node.
     ///

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

@@ -680,6 +680,10 @@ 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

@@ -1273,6 +1273,10 @@ impl ExecutionPlan for SortExec {
         Statistics::with_fetch(self.input.statistics()?, self.schema(), self.fetch, 0, 1)
     }
 
+    fn statistics_by_partition(&self) -> Result<Vec<Statistics>> {
+        todo!()
+    }
+
     fn with_fetch(&self, limit: Option<usize>) -> Option<Arc<dyn ExecutionPlan>> {
         Some(Arc::new(SortExec::with_fetch(self, limit)))
     }

datafusion/physical-plan/src/union.rs

@@ -270,6 +270,17 @@ impl ExecutionPlan for UnionExec {
             .unwrap_or_else(|| Statistics::new_unknown(&self.schema())))
     }
 
+    fn statistics_by_partition(&self) -> Result<Vec<Statistics>> {
+        Ok(self
+            .inputs
+            .iter()
+            .map(|child| child.statistics_by_partition())
+            .collect::<Result<Vec<_>>>()?
+            .into_iter()
+            .flatten()
+            .collect())
+    }
+
     fn benefits_from_input_partitioning(&self) -> Vec<bool> {
         vec![false; self.children().len()]
     }