1717
1818use arrow:: datatypes:: { DataType , SchemaRef } ;
1919use datafusion_catalog:: Session ;
20- use datafusion_common:: plan_err;
20+ use datafusion_common:: { GetExt , plan_err} ;
2121use datafusion_datasource:: ListingTableUrl ;
22+ use datafusion_datasource:: file_compression_type:: FileCompressionType ;
2223use datafusion_datasource:: file_format:: FileFormat ;
2324use datafusion_expr:: { Partitioning , SortExpr } ;
2425use futures:: StreamExt ;
@@ -33,6 +34,9 @@ pub struct ListingOptions {
3334 /// A suffix on which files should be filtered (leave empty to
3435 /// keep all files on the path)
3536 pub file_extension : String ,
37+ /// Accepted filename suffixes for header-detected compression. An explicit
38+ /// file-extension override disables this filter.
39+ pub auto_file_suffixes : Option < Vec < String > > ,
3640 /// The file format
3741 pub format : Arc < dyn FileFormat > ,
3842 /// The expected partition column names in the folder structure.
@@ -96,14 +100,38 @@ pub struct ListingOptions {
96100 pub output_partitioning : Option < Partitioning > ,
97101}
98102
103+ fn auto_compression_suffixes ( base_extension : & str ) -> Vec < String > {
104+ let base = format ! ( ".{}" , base_extension. trim_start_matches( '.' ) ) ;
105+ [
106+ FileCompressionType :: UNCOMPRESSED ,
107+ FileCompressionType :: GZIP ,
108+ FileCompressionType :: BZIP2 ,
109+ FileCompressionType :: XZ ,
110+ FileCompressionType :: ZSTD ,
111+ FileCompressionType :: DEFLATE ,
112+ ]
113+ . into_iter ( )
114+ . map ( |compression| format ! ( "{base}{}" , compression. get_ext( ) ) )
115+ . collect ( )
116+ }
117+
99118impl ListingOptions {
100119 /// Creates an options instance with the given format
101120 /// Default values:
102121 /// - use default file extension filter
103122 /// - no input partition to discover
104123 pub fn new ( format : Arc < dyn FileFormat > ) -> Self {
124+ let auto = format. compression_type ( ) == Some ( FileCompressionType :: AUTO ) ;
125+ let file_extension = if auto {
126+ String :: new ( )
127+ } else {
128+ format. get_ext ( )
129+ } ;
130+ let auto_file_suffixes =
131+ auto. then ( || auto_compression_suffixes ( & format. get_ext ( ) ) ) ;
105132 Self {
106- file_extension : format. get_ext ( ) ,
133+ file_extension,
134+ auto_file_suffixes,
107135 format,
108136 table_partition_cols : vec ! [ ] ,
109137 file_sort_order : vec ! [ ] ,
@@ -126,6 +154,15 @@ impl ListingOptions {
126154 /// ```
127155 pub fn with_file_extension ( mut self , file_extension : impl Into < String > ) -> Self {
128156 self . file_extension = file_extension. into ( ) ;
157+ self . auto_file_suffixes = None ;
158+ self
159+ }
160+
161+ /// Detect compression while restricting a listing to files of this format.
162+ pub fn with_auto_compression_file_extension ( mut self , base_extension : & str ) -> Self {
163+ self . file_extension . clear ( ) ;
164+ self . auto_file_suffixes = ( !base_extension. is_empty ( ) )
165+ . then ( || auto_compression_suffixes ( base_extension) ) ;
129166 self
130167 }
131168
@@ -150,7 +187,7 @@ impl ListingOptions {
150187 S : Into < String > ,
151188 {
152189 if let Some ( file_extension) = file_extension {
153- self . file_extension = file_extension . into ( ) ;
190+ self = self . with_file_extension ( file_extension ) ;
154191 }
155192 self
156193 }
@@ -280,7 +317,12 @@ impl ListingOptions {
280317 let store = state. runtime_env ( ) . object_store ( table_path) ?;
281318
282319 let all_files: Vec < _ > = table_path
283- . list_all_files ( state, store. as_ref ( ) , & self . file_extension )
320+ . list_all_files_with_suffixes (
321+ state,
322+ store. as_ref ( ) ,
323+ & self . file_extension ,
324+ self . auto_file_suffixes . as_deref ( ) ,
325+ )
284326 . await ?
285327 . try_collect ( )
286328 . await ?;
@@ -375,7 +417,12 @@ impl ListingOptions {
375417 // This can fail to detect inconsistent partition keys
376418 // A DFS traversal approach of the store can help here
377419 let files: Vec < _ > = table_path
378- . list_all_files ( state, store. as_ref ( ) , & self . file_extension )
420+ . list_all_files_with_suffixes (
421+ state,
422+ store. as_ref ( ) ,
423+ & self . file_extension ,
424+ self . auto_file_suffixes . as_deref ( ) ,
425+ )
379426 . await ?
380427 . take ( 10 )
381428 . try_collect ( )
0 commit comments