[HSTACK] - consolidate partitions if we have many small files

Author: aditanase

datafusion/datasource/src/file_groups.rs

@@ -213,7 +213,13 @@ impl FileGroupPartitioner {
             .iter()
             .map(|f| f.object_meta.size as i64)
             .sum::<i64>();
-        if total_size < (repartition_file_min_size as i64) || total_size == 0 {
+
+        // bail if we are asked to *split* a set of files that are already too small
+        // if we are being asked to consolidate, we proceed
+        if (total_size < (repartition_file_min_size as i64)
+            && target_partitions >= file_groups.len())
+            || total_size == 0
+        {
             return None;
         }
 
@@ -231,6 +237,18 @@ impl FileGroupPartitioner {
                     let mut produced_files = vec![];
                     let mut range_start = 0;
                     while range_start < source_file.object_meta.size {
+                        // Skip splitting files smaller than repartition_file_min_size
+                        // This may result in a number of partitions slightly smaller than requested
+                        if source_file.object_meta.size < repartition_file_min_size {
+                            state.1 += source_file.object_meta.size;
+                            if state.1 > target_partition_size {
+                                state.0 += 1;
+                                state.1 = 0;
+                            }
+                            produced_files.push((state.0, source_file.clone()));
+                            break;
+                        }
+
                         let range_end = min(
                             range_start + (target_partition_size - state.1),
                             source_file.object_meta.size,