@@ -22,7 +22,7 @@ use std::sync::Arc;
2222
2323use super :: ListingTableUrl ;
2424use super :: PartitionedFile ;
25- use crate :: execution :: context :: SessionState ;
25+ use datafusion_catalog :: Session ;
2626use datafusion_common:: internal_err;
2727use datafusion_common:: { HashMap , Result , ScalarValue } ;
2828use datafusion_expr:: { BinaryExpr , Operator } ;
@@ -154,7 +154,7 @@ pub fn split_files(
154154 chunks
155155}
156156
157- struct Partition {
157+ pub struct Partition {
158158 /// The path to the partition, including the table prefix
159159 path : Path ,
160160 /// How many path segments below the table prefix `path` contains
@@ -183,7 +183,7 @@ impl Partition {
183183}
184184
185185/// 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 (
187187 store : & dyn ObjectStore ,
188188 table_path : & ListingTableUrl ,
189189 max_depth : usize ,
@@ -364,7 +364,7 @@ fn populate_partition_values<'a>(
364364 }
365365}
366366
367- fn evaluate_partition_prefix < ' a > (
367+ pub fn evaluate_partition_prefix < ' a > (
368368 partition_cols : & ' a [ ( String , DataType ) ] ,
369369 filters : & ' a [ Expr ] ,
370370) -> Option < Path > {
@@ -405,7 +405,7 @@ fn evaluate_partition_prefix<'a>(
405405/// `filters` should only contain expressions that can be evaluated
406406/// using only the partition columns.
407407pub async fn pruned_partition_list < ' a > (
408- ctx : & ' a SessionState ,
408+ ctx : & ' a dyn Session ,
409409 store : & ' a dyn ObjectStore ,
410410 table_path : & ' a ListingTableUrl ,
411411 filters : & ' a [ Expr ] ,
@@ -489,7 +489,7 @@ pub async fn pruned_partition_list<'a>(
489489
490490/// Extract the partition values for the given `file_path` (in the given `table_path`)
491491/// 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 > (
493493 table_path : & ListingTableUrl ,
494494 file_path : & ' a Path ,
495495 table_partition_cols : I ,
@@ -517,17 +517,36 @@ where
517517 }
518518 Some ( part_values)
519519}
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+ }
520532
521533#[ cfg( test) ]
522534mod 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 ;
523541 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;
529543
530544 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 ;
531550
532551 #[ test]
533552 fn test_split_files ( ) {
@@ -578,7 +597,7 @@ mod tests {
578597 ] ) ;
579598 let filter = Expr :: eq ( col ( "mypartition" ) , lit ( "val1" ) ) ;
580599 let pruned = pruned_partition_list (
581- & state,
600+ state. as_ref ( ) ,
582601 store. as_ref ( ) ,
583602 & ListingTableUrl :: parse ( "file:///tablepath/" ) . unwrap ( ) ,
584603 & [ filter] ,
@@ -603,7 +622,7 @@ mod tests {
603622 ] ) ;
604623 let filter = Expr :: eq ( col ( "mypartition" ) , lit ( "val1" ) ) ;
605624 let pruned = pruned_partition_list (
606- & state,
625+ state. as_ref ( ) ,
607626 store. as_ref ( ) ,
608627 & ListingTableUrl :: parse ( "file:///tablepath/" ) . unwrap ( ) ,
609628 & [ filter] ,
@@ -643,7 +662,7 @@ mod tests {
643662 let filter1 = Expr :: eq ( col ( "part1" ) , lit ( "p1v2" ) ) ;
644663 let filter2 = Expr :: eq ( col ( "part2" ) , lit ( "p2v1" ) ) ;
645664 let pruned = pruned_partition_list (
646- & state,
665+ state. as_ref ( ) ,
647666 store. as_ref ( ) ,
648667 & ListingTableUrl :: parse ( "file:///tablepath/" ) . unwrap ( ) ,
649668 & [ filter1, filter2] ,
@@ -680,19 +699,6 @@ mod tests {
680699 ) ;
681700 }
682701
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-
696702 #[ tokio:: test]
697703 async fn test_list_partition ( ) {
698704 let ( store, _) = make_test_store_and_state ( & [
@@ -994,4 +1000,74 @@ mod tests {
9941000 Some ( Path :: from( "a=1970-01-05" ) ) ,
9951001 ) ;
9961002 }
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+ }
9971073}
0 commit comments