Allow FileSource-specific repartitioning

datafusion/core/src/datasource/data_source.rs

@@ -26,6 +26,8 @@ use crate::datasource::physical_plan::{FileOpener, FileScanConfig};
 
 use arrow::datatypes::SchemaRef;
 use datafusion_common::Statistics;
+use datafusion_datasource::file_groups::FileGroupPartitioner;
+use datafusion_physical_expr::LexOrdering;
 use datafusion_physical_plan::metrics::ExecutionPlanMetricsSet;
 use datafusion_physical_plan::DisplayFormatType;
 
@@ -67,4 +69,33 @@ pub trait FileSource: Send + Sync {
     /// If this returns true, the DataSourceExec may repartition the data
     /// by breaking up the input files into multiple smaller groups.
     fn supports_repartition(&self, config: &FileScanConfig) -> bool;
+
+    /// If supported by the [`FileSource`], redistribute files across partitions according to their size.
+    /// Allows custom file formats to implement their own repartitioning logic.
+    ///
+    /// Provides a default repartitioning behavior, see comments on [`FileGroupPartitioner`] for more detail.
+    fn repartitioned(
+        &self,
+        target_partitions: usize,
+        repartition_file_min_size: usize,
+        output_ordering: Option<LexOrdering>,
+        config: &FileScanConfig,
+    ) -> datafusion_common::Result<Option<FileScanConfig>> {
+        if !self.supports_repartition(config) {
+            return Ok(None);
+        }
+
+        let repartitioned_file_groups_option = FileGroupPartitioner::new()
+            .with_target_partitions(target_partitions)
+            .with_repartition_file_min_size(repartition_file_min_size)
+            .with_preserve_order_within_groups(output_ordering.is_some())
+            .repartition_file_groups(&config.file_groups);
+
+        if let Some(repartitioned_file_groups) = repartitioned_file_groups_option {
+            let mut source = config.clone();
+            source.file_groups = repartitioned_file_groups;
+            return Ok(Some(source));
+        }
+        Ok(None)
+    }
 }

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

@@ -19,8 +19,8 @@
 //! file sources.
 
 use super::{
-    get_projected_output_ordering, statistics::MinMaxStatistics, FileGroupPartitioner,
-    FileGroupsDisplay, FileStream,
+    get_projected_output_ordering, statistics::MinMaxStatistics, FileGroupsDisplay,
+    FileStream,
 };
 use crate::datasource::file_format::file_compression_type::FileCompressionType;
 use crate::datasource::{listing::PartitionedFile, object_store::ObjectStoreUrl};
@@ -203,30 +203,21 @@ impl DataSource for FileScanConfig {
         self.fmt_file_source(t, f)
     }
 
-    /// Redistribute files across partitions according to their size
-    /// See comments on [`FileGroupPartitioner`] for more detail.
+    /// If supported by the underlying [`FileSource`], redistribute files across partitions according to their size.
     fn repartitioned(
         &self,
         target_partitions: usize,
         repartition_file_min_size: usize,
         output_ordering: Option<LexOrdering>,
     ) -> Result<Option<Arc<dyn DataSource>>> {
-        if !self.source.supports_repartition(self) {
-            return Ok(None);
-        }
-
-        let repartitioned_file_groups_option = FileGroupPartitioner::new()
-            .with_target_partitions(target_partitions)
-            .with_repartition_file_min_size(repartition_file_min_size)
-            .with_preserve_order_within_groups(output_ordering.is_some())
-            .repartition_file_groups(&self.file_groups);
-
-        if let Some(repartitioned_file_groups) = repartitioned_file_groups_option {
-            let mut source = self.clone();
-            source.file_groups = repartitioned_file_groups;
-            return Ok(Some(Arc::new(source)));
-        }
-        Ok(None)
+        let source = self.source.repartitioned(
+            target_partitions,
+            repartition_file_min_size,
+            output_ordering,
+            self,
+        )?;
+
+        Ok(source.map(|s| Arc::new(s) as _))
     }
 
     fn output_partitioning(&self) -> Partitioning {