@@ -64,15 +64,25 @@ impl TableProviderFactory for ListingTableFactory {
64
64
. create ( session_state, & cmd. options ) ?;
65
65
66
66
let file_extension = get_extension ( cmd. location . as_str ( ) ) ;
67
+ let mut table_path = ListingTableUrl :: parse ( & cmd. location ) ?;
68
+ let mut options = ListingOptions :: new ( file_format)
69
+ . with_session_config_options ( session_state. config ( ) )
70
+ . with_file_extension ( file_extension) ;
67
71
68
72
let ( provided_schema, table_partition_cols) = if cmd. schema . fields ( ) . is_empty ( ) {
73
+ let part_cols = match cmd. table_partition_cols . is_empty ( ) {
74
+ true => options
75
+ . infer_partitions ( session_state, & table_path)
76
+ . await ?
77
+ . into_iter ( ) ,
78
+ false => cmd. table_partition_cols . clone ( ) . into_iter ( ) ,
79
+ } ;
69
80
(
70
81
None ,
71
- cmd. table_partition_cols
72
- . iter ( )
73
- . map ( |x| {
82
+ part_cols
83
+ . map ( |p| {
74
84
(
75
- x . clone ( ) ,
85
+ p ,
76
86
DataType :: Dictionary (
77
87
Box :: new ( DataType :: UInt16 ) ,
78
88
Box :: new ( DataType :: Utf8 ) ,
@@ -108,13 +118,7 @@ impl TableProviderFactory for ListingTableFactory {
108
118
( Some ( schema) , table_partition_cols)
109
119
} ;
110
120
111
- let mut table_path = ListingTableUrl :: parse ( & cmd. location ) ?;
112
-
113
- let options = ListingOptions :: new ( file_format)
114
- . with_file_extension ( & file_extension)
115
- . with_session_config_options ( session_state. config ( ) )
116
- . with_table_partition_cols ( table_partition_cols) ;
117
-
121
+ options = options. with_table_partition_cols ( table_partition_cols) ;
118
122
options
119
123
. validate_partitions ( session_state, & table_path)
120
124
. await ?;
@@ -189,6 +193,8 @@ fn get_extension(path: &str) -> String {
189
193
mod tests {
190
194
use glob:: Pattern ;
191
195
use std:: collections:: HashMap ;
196
+ use std:: fs;
197
+ use std:: path:: PathBuf ;
192
198
193
199
use super :: * ;
194
200
use crate :: {
@@ -375,4 +381,49 @@ mod tests {
375
381
Pattern :: new( "*.csv" ) . unwrap( )
376
382
) ;
377
383
}
384
+
385
+ #[ tokio:: test]
386
+ async fn test_create_with_hive_partitions ( ) {
387
+ let dir = tempfile:: tempdir ( ) . unwrap ( ) ;
388
+ let mut path = PathBuf :: from ( dir. path ( ) ) ;
389
+ path. extend ( [ "key1=value1" , "key2=value2" ] ) ;
390
+ fs:: create_dir_all ( & path) . unwrap ( ) ;
391
+ path. push ( "data.parquet" ) ;
392
+ fs:: File :: create_new ( & path) . unwrap ( ) ;
393
+
394
+ let factory = ListingTableFactory :: new ( ) ;
395
+ let context = SessionContext :: new ( ) ;
396
+ let state = context. state ( ) ;
397
+ let name = TableReference :: bare ( "foo" ) ;
398
+
399
+ let cmd = CreateExternalTable {
400
+ name,
401
+ location : dir. path ( ) . to_str ( ) . unwrap ( ) . to_string ( ) ,
402
+ file_type : "parquet" . to_string ( ) ,
403
+ schema : Arc :: new ( DFSchema :: empty ( ) ) ,
404
+ table_partition_cols : vec ! [ ] ,
405
+ if_not_exists : false ,
406
+ temporary : false ,
407
+ definition : None ,
408
+ order_exprs : vec ! [ ] ,
409
+ unbounded : false ,
410
+ options : HashMap :: new ( ) ,
411
+ constraints : Constraints :: default ( ) ,
412
+ column_defaults : HashMap :: new ( ) ,
413
+ } ;
414
+ let table_provider = factory. create ( & state, & cmd) . await . unwrap ( ) ;
415
+ let listing_table = table_provider
416
+ . as_any ( )
417
+ . downcast_ref :: < ListingTable > ( )
418
+ . unwrap ( ) ;
419
+
420
+ let listing_options = listing_table. options ( ) ;
421
+ let dtype =
422
+ DataType :: Dictionary ( Box :: new ( DataType :: UInt16 ) , Box :: new ( DataType :: Utf8 ) ) ;
423
+ let expected_cols = vec ! [
424
+ ( String :: from( "key1" ) , dtype. clone( ) ) ,
425
+ ( String :: from( "key2" ) , dtype. clone( ) ) ,
426
+ ] ;
427
+ assert_eq ! ( expected_cols, listing_options. table_partition_cols) ;
428
+ }
378
429
}
0 commit comments