@@ -81,8 +81,10 @@ enum Scenario {
8181}
8282
8383enum Unit {
84- RowGroup ,
85- Page ,
84+ // pass max row per row_group in parquet writer
85+ RowGroup ( usize ) ,
86+ // pass max row per page in parquet writer
87+ Page ( usize ) ,
8688}
8789
8890/// Test fixture that has an execution context that has an external
@@ -185,13 +187,13 @@ impl ContextWithParquet {
185187 mut config : SessionConfig ,
186188 ) -> Self {
187189 let file = match unit {
188- Unit :: RowGroup => {
190+ Unit :: RowGroup ( row_per_group ) => {
189191 config = config. with_parquet_bloom_filter_pruning ( true ) ;
190- make_test_file_rg ( scenario) . await
192+ make_test_file_rg ( scenario, row_per_group ) . await
191193 }
192- Unit :: Page => {
194+ Unit :: Page ( row_per_page ) => {
193195 config = config. with_parquet_page_index_pruning ( true ) ;
194- make_test_file_page ( scenario) . await
196+ make_test_file_page ( scenario, row_per_page ) . await
195197 }
196198 } ;
197199 let parquet_path = file. path ( ) . to_string_lossy ( ) ;
@@ -880,15 +882,15 @@ fn create_data_batch(scenario: Scenario) -> Vec<RecordBatch> {
880882}
881883
882884/// Create a test parquet file with various data types
883- async fn make_test_file_rg ( scenario : Scenario ) -> NamedTempFile {
885+ async fn make_test_file_rg ( scenario : Scenario , row_per_group : usize ) -> NamedTempFile {
884886 let mut output_file = tempfile:: Builder :: new ( )
885887 . prefix ( "parquet_pruning" )
886888 . suffix ( ".parquet" )
887889 . tempfile ( )
888890 . expect ( "tempfile creation" ) ;
889891
890892 let props = WriterProperties :: builder ( )
891- . set_max_row_group_size ( 5 )
893+ . set_max_row_group_size ( row_per_group )
892894 . set_bloom_filter_enabled ( true )
893895 . build ( ) ;
894896
@@ -906,17 +908,17 @@ async fn make_test_file_rg(scenario: Scenario) -> NamedTempFile {
906908 output_file
907909}
908910
909- async fn make_test_file_page ( scenario : Scenario ) -> NamedTempFile {
911+ async fn make_test_file_page ( scenario : Scenario , row_per_page : usize ) -> NamedTempFile {
910912 let mut output_file = tempfile:: Builder :: new ( )
911913 . prefix ( "parquet_page_pruning" )
912914 . suffix ( ".parquet" )
913915 . tempfile ( )
914916 . expect ( "tempfile creation" ) ;
915917
916- // set row count to 5 , should get same result as rowGroup
918+ // set row count to row_per_page , should get same result as rowGroup
917919 let props = WriterProperties :: builder ( )
918- . set_data_page_row_count_limit ( 5 )
919- . set_write_batch_size ( 5 )
920+ . set_data_page_row_count_limit ( row_per_page )
921+ . set_write_batch_size ( row_per_page )
920922 . build ( ) ;
921923
922924 let batches = create_data_batch ( scenario) ;
0 commit comments