@@ -36,8 +36,9 @@ use datafusion::datasource::listing::{
3636} ;
3737use datafusion:: datasource:: object_store:: ObjectStoreUrl ;
3838use datafusion:: datasource:: physical_plan:: {
39- ArrowSource , FileGroup , FileOutputMode , FileScanConfigBuilder , FileSinkConfig ,
40- ParquetSource , wrap_partition_type_in_dict, wrap_partition_value_in_dict,
39+ ArrowSource , FileGroup , FileOutputMode , FileScanConfig , FileScanConfigBuilder ,
40+ FileSinkConfig , ParquetSource , wrap_partition_type_in_dict,
41+ wrap_partition_value_in_dict,
4142} ;
4243use datafusion:: datasource:: sink:: DataSinkExec ;
4344use datafusion:: datasource:: source:: DataSourceExec ;
@@ -942,6 +943,59 @@ fn roundtrip_parquet_exec_with_pruning_predicate() -> Result<()> {
942943 roundtrip_test ( DataSourceExec :: from_data_source ( scan_config) )
943944}
944945
946+ #[ test]
947+ fn roundtrip_parquet_exec_attaches_cached_reader_factory_after_roundtrip ( ) -> Result < ( ) > {
948+ let file_schema =
949+ Arc :: new ( Schema :: new ( vec ! [ Field :: new( "col" , DataType :: Utf8 , false ) ] ) ) ;
950+ let file_source = Arc :: new ( ParquetSource :: new ( Arc :: clone ( & file_schema) ) ) ;
951+ let scan_config =
952+ FileScanConfigBuilder :: new ( ObjectStoreUrl :: local_filesystem ( ) , file_source)
953+ . with_file_groups ( vec ! [ FileGroup :: new( vec![ PartitionedFile :: new(
954+ "/path/to/file.parquet" . to_string( ) ,
955+ 1024 ,
956+ ) ] ) ] )
957+ . with_statistics ( Statistics {
958+ num_rows : Precision :: Inexact ( 100 ) ,
959+ total_byte_size : Precision :: Inexact ( 1024 ) ,
960+ column_statistics : Statistics :: unknown_column ( & file_schema) ,
961+ } )
962+ . build ( ) ;
963+ let exec_plan = DataSourceExec :: from_data_source ( scan_config) ;
964+
965+ let ctx = SessionContext :: new ( ) ;
966+ let codec = DefaultPhysicalExtensionCodec { } ;
967+ let proto_converter = DefaultPhysicalProtoConverter { } ;
968+ let roundtripped =
969+ roundtrip_test_and_return ( exec_plan, & ctx, & codec, & proto_converter) ?;
970+
971+ let data_source = roundtripped
972+ . as_any ( )
973+ . downcast_ref :: < DataSourceExec > ( )
974+ . ok_or_else ( || {
975+ internal_datafusion_err ! ( "Expected DataSourceExec after roundtrip" )
976+ } ) ?;
977+ let file_scan = data_source
978+ . data_source ( )
979+ . as_any ( )
980+ . downcast_ref :: < FileScanConfig > ( )
981+ . ok_or_else ( || {
982+ internal_datafusion_err ! ( "Expected FileScanConfig after roundtrip" )
983+ } ) ?;
984+ let parquet_source = file_scan
985+ . file_source ( )
986+ . as_any ( )
987+ . downcast_ref :: < ParquetSource > ( )
988+ . ok_or_else ( || {
989+ internal_datafusion_err ! ( "Expected ParquetSource after roundtrip" )
990+ } ) ?;
991+
992+ assert ! (
993+ parquet_source. parquet_file_reader_factory( ) . is_some( ) ,
994+ "Parquet reader factory should be attached after decoding from protobuf"
995+ ) ;
996+ Ok ( ( ) )
997+ }
998+
945999#[ test]
9461000fn roundtrip_arrow_scan ( ) -> Result < ( ) > {
9471001 let file_schema =
0 commit comments