add partitioning on file reads and parallelize aggregations if partitioning permits

Author: gene-bordegaray

datafusion/core/src/physical_planner.rs

@@ -89,7 +89,7 @@ use datafusion_expr::{
 use datafusion_physical_expr::aggregate::{AggregateExprBuilder, AggregateFunctionExpr};
 use datafusion_physical_expr::expressions::Literal;
 use datafusion_physical_expr::{
-    create_physical_sort_exprs, LexOrdering, PhysicalSortExpr,
+    create_physical_sort_exprs, Distribution, LexOrdering, PhysicalSortExpr,
 };
 use datafusion_physical_optimizer::PhysicalOptimizerRule;
 use datafusion_physical_plan::empty::EmptyExec;
@@ -809,12 +809,19 @@ impl DefaultPhysicalPlanner {
                 // `AggregateFunctionExpr`/`PhysicalSortExpr` objects.
                 let updated_aggregates = initial_aggr.aggr_expr().to_vec();
 
-                let next_partition_mode = if can_repartition {
-                    // construct a second aggregation with 'AggregateMode::FinalPartitioned'
+                // Check if input is already hash-partitioned on a subset of the group-by
+                // columns.
+                let input_already_partitioned = !groups.is_empty()
+                    && initial_aggr.input().output_partitioning().partition_count() > 1
+                    && initial_aggr.input().output_partitioning().satisfy(
+                        &Distribution::HashPartitioned(groups.input_exprs()),
+                        &initial_aggr.input().equivalence_properties(),
+                    );
+
+                let next_partition_mode = if can_repartition || input_already_partitioned
+                {
                     AggregateMode::FinalPartitioned
                 } else {
-                    // construct a second aggregation, keeping the final column name equal to the
-                    // first aggregation and the expressions corresponding to the respective aggregate
                     AggregateMode::Final
                 };
 

datafusion/datasource/src/file_scan_config.rs

@@ -570,6 +570,30 @@ impl DataSource for FileScanConfig {
     }
 
     fn output_partitioning(&self) -> Partitioning {
+        if self.is_properly_partitioned() {
+            // Use the projected schema to find the column index, not the virtual
+            // partition column position.
+            let partition_exprs: Vec<Arc<dyn PhysicalExpr>> = self
+                .table_partition_cols()
+                .iter()
+                .filter_map(|field| {
+                    self.projected_schema()
+                        .fields()
+                        .iter()
+                        .position(|f| f.name() == field.name())
+                        .map(|idx| {
+                            Arc::new(Column::new(field.name(), idx))
+                                as Arc<dyn PhysicalExpr>
+                        })
+                })
+                .collect();
+
+            if partition_exprs.len() == self.table_partition_cols().len() {
+                return Partitioning::Hash(partition_exprs, self.file_groups.len());
+            }
+        }
+
+        // Fall back to UnknownPartitioning
         Partitioning::UnknownPartitioning(self.file_groups.len())
     }
 
@@ -736,6 +760,25 @@ impl FileScanConfig {
         }
     }
 
+    fn is_properly_partitioned(&self) -> bool {
+        if self.file_groups.is_empty() || self.table_partition_cols().is_empty() {
+            return false;
+        }
+
+        self.file_groups.iter().all(|file_group| {
+            let files = file_group.files();
+            if files.is_empty() {
+                return true;
+            }
+
+            let reference_values = &files[0].partition_values;
+
+            files
+                .iter()
+                .all(|file| file.partition_values == *reference_values)
+        })
+    }
+
     pub fn projected_stats(&self) -> Statistics {
         let statistics = self.file_source.statistics().unwrap();
 
@@ -2673,4 +2716,172 @@ mod tests {
         assert_eq!(partition_stats.num_rows, Precision::Exact(100));
         assert_eq!(partition_stats.total_byte_size, Precision::Exact(1024));
     }
+
+    #[test]
+    fn test_is_properly_partitioned() {
+        // All file groups have consistent partition values
+        let file_schema = Arc::new(Schema::new(vec![Field::new(
+            "value",
+            DataType::Int32,
+            false,
+        )]));
+
+        let partition_cols = vec![Field::new(
+            "f_dkey",
+            wrap_partition_type_in_dict(DataType::Utf8),
+            false,
+        )];
+
+        // Case 1: Multiple file groups, each with consistent values
+        let file_groups = vec![
+            FileGroup::new(vec![
+                {
+                    let mut file =
+                        PartitionedFile::new("data/f_dkey=A/file1.parquet", 1024);
+                    file.partition_values =
+                        vec![wrap_partition_value_in_dict(ScalarValue::from("A"))];
+                    file
+                },
+                {
+                    let mut file =
+                        PartitionedFile::new("data/f_dkey=A/file2.parquet", 2048);
+                    file.partition_values =
+                        vec![wrap_partition_value_in_dict(ScalarValue::from("A"))];
+                    file
+                },
+            ]),
+            FileGroup::new(vec![{
+                let mut file = PartitionedFile::new("data/f_dkey=B/file1.parquet", 512);
+                file.partition_values =
+                    vec![wrap_partition_value_in_dict(ScalarValue::from("B"))];
+                file
+            }]),
+            FileGroup::new(vec![{
+                let mut file = PartitionedFile::new("data/f_dkey=C/file1.parquet", 1024);
+                file.partition_values =
+                    vec![wrap_partition_value_in_dict(ScalarValue::from("C"))];
+                file
+            }]),
+        ];
+
+        let config = FileScanConfigBuilder::new(
+            ObjectStoreUrl::parse("test:///").unwrap(),
+            Arc::new(MockSource::new(TableSchema::new(
+                Arc::clone(&file_schema),
+                partition_cols.iter().map(|f| Arc::new(f.clone())).collect(),
+            ))),
+        )
+        .with_file_groups(file_groups)
+        .build();
+
+        assert!(
+            config.is_properly_partitioned(),
+            "Expected true when all file groups have consistent partition values"
+        );
+
+        // Case 2: Mixed partition values within a group (should return false)
+        let mixed_file_groups = vec![FileGroup::new(vec![
+            {
+                let mut file = PartitionedFile::new("data/f_dkey=A/file1.parquet", 1024);
+                file.partition_values =
+                    vec![wrap_partition_value_in_dict(ScalarValue::from("A"))];
+                file
+            },
+            {
+                let mut file = PartitionedFile::new("data/f_dkey=B/file2.parquet", 2048);
+                file.partition_values =
+                    vec![wrap_partition_value_in_dict(ScalarValue::from("B"))];
+                file
+            },
+        ])];
+
+        let config_mixed = FileScanConfigBuilder::new(
+            ObjectStoreUrl::parse("test:///").unwrap(),
+            Arc::new(MockSource::new(TableSchema::new(
+                Arc::clone(&file_schema),
+                partition_cols.iter().map(|f| Arc::new(f.clone())).collect(),
+            ))),
+        )
+        .with_file_groups(mixed_file_groups)
+        .build();
+
+        assert!(
+            !config_mixed.is_properly_partitioned(),
+            "Expected false when file group has mixed partition values"
+        );
+
+        // Case 3: Empty file groups
+        let config_empty = FileScanConfigBuilder::new(
+            ObjectStoreUrl::parse("test:///").unwrap(),
+            Arc::new(MockSource::new(TableSchema::new(
+                Arc::clone(&file_schema),
+                partition_cols.iter().map(|f| Arc::new(f.clone())).collect(),
+            ))),
+        )
+        .with_file_groups(vec![])
+        .build();
+
+        assert!(
+            !config_empty.is_properly_partitioned(),
+            "Expected false for empty file groups"
+        );
+    }
+
+    #[test]
+    fn test_output_partitioning_column_indices() {
+        // Test that output_partitioning returns correct column indices in projected schema
+        let file_schema = Arc::new(Schema::new(vec![
+            Field::new("value", DataType::Int32, false),
+            Field::new("timestamp", DataType::Int64, false),
+        ]));
+
+        let partition_cols = vec![Field::new(
+            "f_dkey",
+            wrap_partition_type_in_dict(DataType::Utf8),
+            false,
+        )];
+
+        let file_groups = vec![FileGroup::new(vec![{
+            let mut file = PartitionedFile::new("data/f_dkey=A/file1.parquet", 1024);
+            file.partition_values =
+                vec![wrap_partition_value_in_dict(ScalarValue::from("A"))];
+            file
+        }])];
+
+        // Full schema is [value@0, timestamp@1, f_dkey@2]
+        // Project indices [0, 2] -> [value@0, f_dkey@1] in projected schema
+        let config = FileScanConfigBuilder::new(
+            ObjectStoreUrl::parse("test:///").unwrap(),
+            Arc::new(MockSource::new(TableSchema::new(
+                Arc::clone(&file_schema),
+                partition_cols.iter().map(|f| Arc::new(f.clone())).collect(),
+            ))),
+        )
+        .with_projection_indices(Some(vec![0, 2]))
+        .with_file_groups(file_groups)
+        .build();
+
+        let partitioning = config.output_partitioning();
+
+        match partitioning {
+            Partitioning::Hash(exprs, count) => {
+                assert_eq!(count, 1, "Expected 1 partition (single file group)");
+                assert_eq!(exprs.len(), 1, "Expected 1 partition expression");
+
+                // The partition expression should reference f_dkey at index 1
+                // in the projected schema [value@0, f_dkey@1]
+                let col = exprs[0]
+                    .as_any()
+                    .downcast_ref::<Column>()
+                    .expect("Expected Column expression");
+                assert_eq!(col.name(), "f_dkey", "Column name should be f_dkey");
+                assert_eq!(
+                    col.index(),
+                    1,
+                    "Column index should be 1 in projected schema [value@0, f_dkey@1]"
+                );
+            }
+            _ => panic!("Expected Hash partitioning, got: {:?}", partitioning),
+        }
+    }
 }