@@ -22,7 +22,7 @@ use std::sync::Arc;
22
22
23
23
use super :: ListingTableUrl ;
24
24
use super :: PartitionedFile ;
25
- use crate :: execution :: context :: SessionState ;
25
+ use datafusion_catalog :: Session ;
26
26
use datafusion_common:: internal_err;
27
27
use datafusion_common:: { HashMap , Result , ScalarValue } ;
28
28
use datafusion_expr:: { BinaryExpr , Operator } ;
@@ -154,7 +154,7 @@ pub fn split_files(
154
154
chunks
155
155
}
156
156
157
- struct Partition {
157
+ pub struct Partition {
158
158
/// The path to the partition, including the table prefix
159
159
path : Path ,
160
160
/// How many path segments below the table prefix `path` contains
@@ -183,7 +183,7 @@ impl Partition {
183
183
}
184
184
185
185
/// Returns a recursive list of the partitions in `table_path` up to `max_depth`
186
- async fn list_partitions (
186
+ pub async fn list_partitions (
187
187
store : & dyn ObjectStore ,
188
188
table_path : & ListingTableUrl ,
189
189
max_depth : usize ,
@@ -364,7 +364,7 @@ fn populate_partition_values<'a>(
364
364
}
365
365
}
366
366
367
- fn evaluate_partition_prefix < ' a > (
367
+ pub fn evaluate_partition_prefix < ' a > (
368
368
partition_cols : & ' a [ ( String , DataType ) ] ,
369
369
filters : & ' a [ Expr ] ,
370
370
) -> Option < Path > {
@@ -405,7 +405,7 @@ fn evaluate_partition_prefix<'a>(
405
405
/// `filters` should only contain expressions that can be evaluated
406
406
/// using only the partition columns.
407
407
pub async fn pruned_partition_list < ' a > (
408
- ctx : & ' a SessionState ,
408
+ ctx : & ' a dyn Session ,
409
409
store : & ' a dyn ObjectStore ,
410
410
table_path : & ' a ListingTableUrl ,
411
411
filters : & ' a [ Expr ] ,
@@ -489,7 +489,7 @@ pub async fn pruned_partition_list<'a>(
489
489
490
490
/// Extract the partition values for the given `file_path` (in the given `table_path`)
491
491
/// associated to the partitions defined by `table_partition_cols`
492
- fn parse_partitions_for_path < ' a , I > (
492
+ pub fn parse_partitions_for_path < ' a , I > (
493
493
table_path : & ListingTableUrl ,
494
494
file_path : & ' a Path ,
495
495
table_partition_cols : I ,
@@ -517,17 +517,36 @@ where
517
517
}
518
518
Some ( part_values)
519
519
}
520
+ /// Describe a partition as a (path, depth, files) tuple for easier assertions
521
+ pub fn describe_partition ( partition : & Partition ) -> ( & str , usize , Vec < & str > ) {
522
+ (
523
+ partition. path . as_ref ( ) ,
524
+ partition. depth ,
525
+ partition
526
+ . files
527
+ . as_ref ( )
528
+ . map ( |f| f. iter ( ) . map ( |f| f. location . filename ( ) . unwrap ( ) ) . collect ( ) )
529
+ . unwrap_or_default ( ) ,
530
+ )
531
+ }
520
532
521
533
#[ cfg( test) ]
522
534
mod tests {
535
+ use async_trait:: async_trait;
536
+ use datafusion_execution:: config:: SessionConfig ;
537
+ use datafusion_execution:: runtime_env:: RuntimeEnv ;
538
+ use futures:: FutureExt ;
539
+ use object_store:: memory:: InMemory ;
540
+ use std:: any:: Any ;
523
541
use std:: ops:: Not ;
524
-
525
- use futures:: StreamExt ;
526
-
527
- use crate :: test:: object_store:: make_test_store_and_state;
528
- use datafusion_expr:: { case, col, lit, Expr } ;
542
+ // use futures::StreamExt;
529
543
530
544
use super :: * ;
545
+ use datafusion_expr:: {
546
+ case, col, lit, AggregateUDF , Expr , LogicalPlan , ScalarUDF , WindowUDF ,
547
+ } ;
548
+ use datafusion_physical_expr_common:: physical_expr:: PhysicalExpr ;
549
+ use datafusion_physical_plan:: ExecutionPlan ;
531
550
532
551
#[ test]
533
552
fn test_split_files ( ) {
@@ -578,7 +597,7 @@ mod tests {
578
597
] ) ;
579
598
let filter = Expr :: eq ( col ( "mypartition" ) , lit ( "val1" ) ) ;
580
599
let pruned = pruned_partition_list (
581
- & state,
600
+ state. as_ref ( ) ,
582
601
store. as_ref ( ) ,
583
602
& ListingTableUrl :: parse ( "file:///tablepath/" ) . unwrap ( ) ,
584
603
& [ filter] ,
@@ -603,7 +622,7 @@ mod tests {
603
622
] ) ;
604
623
let filter = Expr :: eq ( col ( "mypartition" ) , lit ( "val1" ) ) ;
605
624
let pruned = pruned_partition_list (
606
- & state,
625
+ state. as_ref ( ) ,
607
626
store. as_ref ( ) ,
608
627
& ListingTableUrl :: parse ( "file:///tablepath/" ) . unwrap ( ) ,
609
628
& [ filter] ,
@@ -643,7 +662,7 @@ mod tests {
643
662
let filter1 = Expr :: eq ( col ( "part1" ) , lit ( "p1v2" ) ) ;
644
663
let filter2 = Expr :: eq ( col ( "part2" ) , lit ( "p2v1" ) ) ;
645
664
let pruned = pruned_partition_list (
646
- & state,
665
+ state. as_ref ( ) ,
647
666
store. as_ref ( ) ,
648
667
& ListingTableUrl :: parse ( "file:///tablepath/" ) . unwrap ( ) ,
649
668
& [ filter1, filter2] ,
@@ -680,19 +699,6 @@ mod tests {
680
699
) ;
681
700
}
682
701
683
- /// Describe a partition as a (path, depth, files) tuple for easier assertions
684
- fn describe_partition ( partition : & Partition ) -> ( & str , usize , Vec < & str > ) {
685
- (
686
- partition. path . as_ref ( ) ,
687
- partition. depth ,
688
- partition
689
- . files
690
- . as_ref ( )
691
- . map ( |f| f. iter ( ) . map ( |f| f. location . filename ( ) . unwrap ( ) ) . collect ( ) )
692
- . unwrap_or_default ( ) ,
693
- )
694
- }
695
-
696
702
#[ tokio:: test]
697
703
async fn test_list_partition ( ) {
698
704
let ( store, _) = make_test_store_and_state ( & [
@@ -994,4 +1000,74 @@ mod tests {
994
1000
Some ( Path :: from( "a=1970-01-05" ) ) ,
995
1001
) ;
996
1002
}
1003
+
1004
+ pub fn make_test_store_and_state (
1005
+ files : & [ ( & str , u64 ) ] ,
1006
+ ) -> ( Arc < InMemory > , Arc < dyn Session > ) {
1007
+ let memory = InMemory :: new ( ) ;
1008
+
1009
+ for ( name, size) in files {
1010
+ memory
1011
+ . put ( & Path :: from ( * name) , vec ! [ 0 ; * size as usize ] . into ( ) )
1012
+ . now_or_never ( )
1013
+ . unwrap ( )
1014
+ . unwrap ( ) ;
1015
+ }
1016
+
1017
+ ( Arc :: new ( memory) , Arc :: new ( MockSession { } ) )
1018
+ }
1019
+
1020
+ struct MockSession { }
1021
+
1022
+ #[ async_trait]
1023
+ impl Session for MockSession {
1024
+ fn session_id ( & self ) -> & str {
1025
+ unimplemented ! ( )
1026
+ }
1027
+
1028
+ fn config ( & self ) -> & SessionConfig {
1029
+ unimplemented ! ( )
1030
+ }
1031
+
1032
+ async fn create_physical_plan (
1033
+ & self ,
1034
+ _logical_plan : & LogicalPlan ,
1035
+ ) -> Result < Arc < dyn ExecutionPlan > > {
1036
+ unimplemented ! ( )
1037
+ }
1038
+
1039
+ fn create_physical_expr (
1040
+ & self ,
1041
+ _expr : Expr ,
1042
+ _df_schema : & DFSchema ,
1043
+ ) -> Result < Arc < dyn PhysicalExpr > > {
1044
+ unimplemented ! ( )
1045
+ }
1046
+
1047
+ fn scalar_functions ( & self ) -> & std:: collections:: HashMap < String , Arc < ScalarUDF > > {
1048
+ unimplemented ! ( )
1049
+ }
1050
+
1051
+ fn aggregate_functions (
1052
+ & self ,
1053
+ ) -> & std:: collections:: HashMap < String , Arc < AggregateUDF > > {
1054
+ unimplemented ! ( )
1055
+ }
1056
+
1057
+ fn window_functions ( & self ) -> & std:: collections:: HashMap < String , Arc < WindowUDF > > {
1058
+ unimplemented ! ( )
1059
+ }
1060
+
1061
+ fn runtime_env ( & self ) -> & Arc < RuntimeEnv > {
1062
+ unimplemented ! ( )
1063
+ }
1064
+
1065
+ fn execution_props ( & self ) -> & ExecutionProps {
1066
+ unimplemented ! ( )
1067
+ }
1068
+
1069
+ fn as_any ( & self ) -> & dyn Any {
1070
+ unimplemented ! ( )
1071
+ }
1072
+ }
997
1073
}
0 commit comments