Skip to content

Commit dea4a1b

Browse files
authored
[minor] make parquet prune tests more readable (#10112)
* [minor] make parquet prune tests more readable * typo
1 parent 4ad4f90 commit dea4a1b

File tree

3 files changed

+73
-22
lines changed

3 files changed

+73
-22
lines changed

datafusion/core/tests/parquet/mod.rs

Lines changed: 14 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -81,8 +81,10 @@ enum Scenario {
8181
}
8282

8383
enum 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

Comments
 (0)