update

Author: xudong963

datafusion/catalog-listing/src/helpers.rs

@@ -17,7 +17,6 @@
 
 //! Helper functions for the table implementation
 
-use std::mem;
 use std::sync::Arc;
 
 use datafusion_catalog::Session;
@@ -40,7 +39,6 @@ use log::{debug, trace};
 
 use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion};
 use datafusion_common::{Column, DFSchema, DataFusionError};
-use datafusion_datasource::file_groups::FileGroup;
 use datafusion_expr::{Expr, Volatility};
 use datafusion_physical_expr::create_physical_expr;
 use object_store::path::Path;
@@ -122,39 +120,6 @@ pub fn expr_applicable_for_cols(col_names: &[&str], expr: &Expr) -> bool {
 /// The maximum number of concurrent listing requests
 const CONCURRENCY_LIMIT: usize = 100;
 
-/// Partition the list of files into `n` groups
-pub fn split_files(mut file_group: FileGroup, n: usize) -> Vec<FileGroup> {
-    if file_group.is_empty() {
-        return vec![];
-    }
-
-    // ObjectStore::list does not guarantee any consistent order and for some
-    // implementations such as LocalFileSystem, it may be inconsistent. Thus
-    // Sort files by path to ensure consistent plans when run more than once.
-    file_group.files.sort_by(|a, b| a.path().cmp(b.path()));
-
-    // effectively this is div with rounding up instead of truncating
-    let chunk_size = file_group.len().div_ceil(n);
-    let mut chunks = Vec::with_capacity(n);
-    let mut current_chunk = Vec::with_capacity(chunk_size);
-    for file in file_group.files.drain(..) {
-        current_chunk.push(file);
-        if current_chunk.len() == chunk_size {
-            let full_chunk = FileGroup::new(mem::replace(
-                &mut current_chunk,
-                Vec::with_capacity(chunk_size),
-            ));
-            chunks.push(full_chunk);
-        }
-    }
-
-    if !current_chunk.is_empty() {
-        chunks.push(FileGroup::new(current_chunk))
-    }
-
-    chunks
-}
-
 pub struct Partition {
     /// The path to the partition, including the table prefix
     path: Path,
@@ -541,6 +506,7 @@ mod tests {
     use object_store::memory::InMemory;
     use std::any::Any;
     use std::ops::Not;
+    use datafusion_datasource::file_groups::FileGroup;
     // use futures::StreamExt;
 
     use super::*;
@@ -561,32 +527,33 @@ mod tests {
             new_partitioned_file("e"),
         ]);
 
-        let chunks = split_files(files.clone(), 1);
+        let chunks = files.clone().split_files(1);
         assert_eq!(1, chunks.len());
         assert_eq!(5, chunks[0].len());
 
-        let chunks = split_files(files.clone(), 2);
+        let chunks = files.clone().split_files(2);
         assert_eq!(2, chunks.len());
         assert_eq!(3, chunks[0].len());
         assert_eq!(2, chunks[1].len());
 
-        let chunks = split_files(files.clone(), 5);
+        let chunks = files.clone().split_files(5);
         assert_eq!(5, chunks.len());
         assert_eq!(1, chunks[0].len());
         assert_eq!(1, chunks[1].len());
         assert_eq!(1, chunks[2].len());
         assert_eq!(1, chunks[3].len());
         assert_eq!(1, chunks[4].len());
 
-        let chunks = split_files(files, 123);
+        let chunks = files.clone().split_files(123);
         assert_eq!(5, chunks.len());
         assert_eq!(1, chunks[0].len());
         assert_eq!(1, chunks[1].len());
         assert_eq!(1, chunks[2].len());
         assert_eq!(1, chunks[3].len());
         assert_eq!(1, chunks[4].len());
 
-        let chunks = split_files(FileGroup::default(), 2);
+        let mut empty_group = FileGroup::default();
+        let chunks = empty_group.split_files(2);
         assert_eq!(0, chunks.len());
     }
 

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

@@ -20,7 +20,7 @@
 use std::collections::HashMap;
 use std::{any::Any, str::FromStr, sync::Arc};
 
-use super::helpers::{expr_applicable_for_cols, pruned_partition_list, split_files};
+use super::helpers::{expr_applicable_for_cols, pruned_partition_list};
 use super::{ListingTableUrl, PartitionedFile};
 
 use crate::datasource::{
@@ -1128,7 +1128,7 @@ impl ListingTable {
             .boxed()
             .buffer_unordered(ctx.config_options().execution.meta_fetch_concurrency);
 
-        let (files, statistics) = get_statistics_with_limit(
+        let (mut files, statistics) = get_statistics_with_limit(
             files,
             self.schema(),
             limit,
@@ -1137,7 +1137,7 @@ impl ListingTable {
         .await?;
 
         Ok((
-            split_files(files, self.options.target_partitions),
+            files.split_files(self.options.target_partitions),
             statistics,
         ))
     }

datafusion/datasource/src/file_groups.rs

@@ -23,6 +23,7 @@ use itertools::Itertools;
 use std::cmp::min;
 use std::collections::BinaryHeap;
 use std::iter::repeat_with;
+use std::mem;
 use std::ops::{Index, IndexMut};
 
 /// Repartition input files into `target_partitions` partitions, if total file size exceed
@@ -407,6 +408,39 @@ impl FileGroup {
     pub fn push(&mut self, file: PartitionedFile) {
         self.files.push(file);
     }
+
+    /// Partition the list of files into `n` groups
+    pub fn split_files(&mut self, n: usize) -> Vec<FileGroup> {
+        if self.is_empty() {
+            return vec![];
+        }
+
+        // ObjectStore::list does not guarantee any consistent order and for some
+        // implementations such as LocalFileSystem, it may be inconsistent. Thus
+        // Sort files by path to ensure consistent plans when run more than once.
+        self.files.sort_by(|a, b| a.path().cmp(b.path()));
+
+        // effectively this is div with rounding up instead of truncating
+        let chunk_size = self.len().div_ceil(n);
+        let mut chunks = Vec::with_capacity(n);
+        let mut current_chunk = Vec::with_capacity(chunk_size);
+        for file in self.files.drain(..) {
+            current_chunk.push(file);
+            if current_chunk.len() == chunk_size {
+                let full_chunk = FileGroup::new(mem::replace(
+                    &mut current_chunk,
+                    Vec::with_capacity(chunk_size),
+                ));
+                chunks.push(full_chunk);
+            }
+        }
+
+        if !current_chunk.is_empty() {
+            chunks.push(FileGroup::new(current_chunk))
+        }
+
+        chunks
+    }
 }
 
 impl Index<usize> for FileGroup {

datafusion/datasource/src/file_scan_config.rs

@@ -72,7 +72,7 @@ use crate::{
 /// # use object_store::ObjectStore;
 /// # use datafusion_common::Statistics;
 /// # use datafusion_datasource::file::FileSource;
-/// use datafusion_datasource::file_groups::FileGroup;
+/// # use datafusion_datasource::file_groups::FileGroup;
 /// # use datafusion_datasource::PartitionedFile;
 /// # use datafusion_datasource::file_scan_config::FileScanConfig;
 /// # use datafusion_datasource::file_stream::FileOpener;