rebase main and fix tests

Author: xudong963

datafusion/core/src/datasource/listing/table.rs

@@ -1154,9 +1154,18 @@ impl ListingTable {
             .map_schema(self.file_schema.as_ref())?;
         // Use schema_mapper to map each file-level column statistics to table-level column statistics
         file_groups.iter_mut().try_for_each(|file_group| {
-            if let Some(stat) = file_group.statistics_mut() {
-                stat.column_statistics =
-                    schema_mapper.map_column_statistics(&stat.column_statistics)?;
+            // Update each file's statistics's column statistics in file_group
+            for idx in 0..file_group.len() {
+                if let Some(stat) = file_group[idx].statistics.as_ref() {
+                    let column_statistics =
+                        schema_mapper.map_column_statistics(&stat.column_statistics)?;
+                    // Update the file's statistics with the mapped column statistics
+                    file_group[idx].statistics = Some(Arc::new(Statistics {
+                        num_rows: stat.num_rows.clone(),
+                        total_byte_size: stat.total_byte_size.clone(),
+                        column_statistics,
+                    }));
+                }
             }
             Ok::<_, DataFusionError>(())
         })?;

datafusion/core/tests/physical_optimizer/partition_statistics.rs

@@ -476,14 +476,15 @@ 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;
+        // Skip 2 rows
         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);
+        let expected_statistic_partition =
+            create_partition_statistics(2, 110, 3, 4, true);
         assert_eq!(statistics[0], expected_statistic_partition);
         Ok(())
     }

datafusion/physical-plan/src/union.rs

@@ -262,16 +262,32 @@ impl ExecutionPlan for UnionExec {
     }
 
     fn partition_statistics(&self, partition: Option<usize>) -> Result<Statistics> {
-        let stats = self
-            .inputs
-            .iter()
-            .map(|input_exec| input_exec.partition_statistics(partition))
-            .collect::<Result<Vec<_>>>()?;
+        if let Some(partition_idx) = partition {
+            // For a specific partition, find which input it belongs to
+            let mut remaining_idx = partition_idx;
+            for input in &self.inputs {
+                let input_partition_count = input.output_partitioning().partition_count();
+                if remaining_idx < input_partition_count {
+                    // This partition belongs to this input
+                    return input.partition_statistics(Some(remaining_idx));
+                }
+                remaining_idx -= input_partition_count;
+            }
+            // If we get here, the partition index is out of bounds
+            return Ok(Statistics::new_unknown(&self.schema()));
+        } else {
+            // Collect statistics from all inputs
+            let stats = self
+                .inputs
+                .iter()
+                .map(|input_exec| input_exec.partition_statistics(None))
+                .collect::<Result<Vec<_>>>()?;
 
-        Ok(stats
-            .into_iter()
-            .reduce(stats_union)
-            .unwrap_or_else(|| Statistics::new_unknown(&self.schema())))
+            Ok(stats
+                .into_iter()
+                .reduce(stats_union)
+                .unwrap_or_else(|| Statistics::new_unknown(&self.schema())))
+        }
     }
 
     fn benefits_from_input_partitioning(&self) -> Vec<bool> {