Skip to content

Add bench for data page statistics parquet extraction #10950

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 4 commits into from
Jul 3, 2024
Merged

Add bench for data page statistics parquet extraction #10950

Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
78f4bcc
feat: add data page bench
marvinlanhenke Jun 17, 2024
bcd1407
chore: add comment
marvinlanhenke Jun 17, 2024
1f8f2a4
fix: row_groups + shorten row_group_indices
marvinlanhenke Jun 17, 2024
6f92be1
Merge remote-tracking branch 'apache/main' into add_benchmark_pq_stats
alamb Jul 3, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
168 changes: 134 additions & 34 deletions datafusion/core/benches/parquet_statistic.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,20 +18,26 @@
//! Benchmarks of benchmark for extracting arrow statistics from parquet

use arrow::array::{ArrayRef, DictionaryArray, Float64Array, StringArray, UInt64Array};
use arrow_array::{Int32Array, RecordBatch};
use arrow_array::{Int32Array, Int64Array, RecordBatch};
use arrow_schema::{
DataType::{self, *},
Field, Schema,
};
use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion};
use datafusion::datasource::physical_plan::parquet::StatisticsConverter;
use parquet::arrow::{arrow_reader::ArrowReaderBuilder, ArrowWriter};
use parquet::file::properties::WriterProperties;
use parquet::{
arrow::arrow_reader::ArrowReaderOptions, file::properties::WriterProperties,
};
use parquet::{
arrow::{arrow_reader::ArrowReaderBuilder, ArrowWriter},
file::properties::EnabledStatistics,
};
use std::sync::Arc;
use tempfile::NamedTempFile;
#[derive(Debug, Clone)]
enum TestTypes {
UInt64,
Int64,
F64,
String,
Dictionary,
Expand All @@ -43,18 +49,26 @@ impl fmt::Display for TestTypes {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
match self {
TestTypes::UInt64 => write!(f, "UInt64"),
TestTypes::Int64 => write!(f, "Int64"),
TestTypes::F64 => write!(f, "F64"),
TestTypes::String => write!(f, "String"),
TestTypes::Dictionary => write!(f, "Dictionary(Int32, String)"),
}
}
}

fn create_parquet_file(dtype: TestTypes, row_groups: usize) -> NamedTempFile {
fn create_parquet_file(
dtype: TestTypes,
row_groups: usize,
data_page_row_count_limit: &Option<usize>,
) -> NamedTempFile {
let schema = match dtype {
TestTypes::UInt64 => {
Arc::new(Schema::new(vec![Field::new("col", DataType::UInt64, true)]))
}
TestTypes::Int64 => {
Arc::new(Schema::new(vec![Field::new("col", DataType::Int64, true)]))
}
TestTypes::F64 => Arc::new(Schema::new(vec![Field::new(
"col",
DataType::Float64,
Expand All @@ -70,7 +84,14 @@ fn create_parquet_file(dtype: TestTypes, row_groups: usize) -> NamedTempFile {
)])),
};

let props = WriterProperties::builder().build();
let mut props = WriterProperties::builder().set_max_row_group_size(row_groups);
if let Some(limit) = data_page_row_count_limit {
props = props
.set_data_page_row_count_limit(*limit)
.set_statistics_enabled(EnabledStatistics::Page);
};
let props = props.build();

let file = tempfile::Builder::new()
.suffix(".parquet")
.tempfile()
Expand All @@ -82,11 +103,21 @@ fn create_parquet_file(dtype: TestTypes, row_groups: usize) -> NamedTempFile {
for _ in 0..row_groups {
let batch = match dtype {
TestTypes::UInt64 => make_uint64_batch(),
TestTypes::Int64 => make_int64_batch(),
TestTypes::F64 => make_f64_batch(),
TestTypes::String => make_string_batch(),
TestTypes::Dictionary => make_dict_batch(),
};
writer.write(&batch).unwrap();
if data_page_row_count_limit.is_some() {
// Send batches one at a time. This allows the
// writer to apply the page limit, that is only
// checked on RecordBatch boundaries.
for i in 0..batch.num_rows() {
writer.write(&batch.slice(i, 1)).unwrap();
}
} else {
writer.write(&batch).unwrap();
}
}
writer.close().unwrap();
file
Expand All @@ -109,6 +140,23 @@ fn make_uint64_batch() -> RecordBatch {
.unwrap()
}

fn make_int64_batch() -> RecordBatch {
let array: ArrayRef = Arc::new(Int64Array::from(vec![
Some(1),
Some(2),
Some(3),
Some(4),
Some(5),
]));
RecordBatch::try_new(
Arc::new(arrow::datatypes::Schema::new(vec![
arrow::datatypes::Field::new("col", Int64, false),
])),
vec![array],
)
.unwrap()
}

fn make_f64_batch() -> RecordBatch {
let array: ArrayRef = Arc::new(Float64Array::from(vec![1.0, 2.0, 3.0, 4.0, 5.0]));
RecordBatch::try_new(
Expand Down Expand Up @@ -150,36 +198,88 @@ fn make_dict_batch() -> RecordBatch {
fn criterion_benchmark(c: &mut Criterion) {
let row_groups = 100;
use TestTypes::*;
let types = vec![UInt64, F64, String, Dictionary];
let types = vec![Int64, UInt64, F64, String, Dictionary];
let data_page_row_count_limits = vec![None, Some(1)];

for dtype in types {
let file = create_parquet_file(dtype.clone(), row_groups);
let file = file.reopen().unwrap();
let reader = ArrowReaderBuilder::try_new(file).unwrap();
let metadata = reader.metadata();
let row_groups = metadata.row_groups();

let mut group =
c.benchmark_group(format!("Extract statistics for {}", dtype.clone()));
group.bench_function(
BenchmarkId::new("extract_statistics", dtype.clone()),
|b| {
b.iter(|| {
let converter = StatisticsConverter::try_new(
"col",
reader.schema(),
reader.parquet_schema(),
)
.unwrap();

let _ = converter.row_group_mins(row_groups.iter()).unwrap();
let _ = converter.row_group_maxes(row_groups.iter()).unwrap();
let _ = converter.row_group_null_counts(row_groups.iter()).unwrap();
let _ = converter.row_group_row_counts(row_groups.iter()).unwrap();
})
},
);
group.finish();
for data_page_row_count_limit in &data_page_row_count_limits {
let file =
create_parquet_file(dtype.clone(), row_groups, data_page_row_count_limit);
let file = file.reopen().unwrap();
let options = ArrowReaderOptions::new().with_page_index(true);
let reader = ArrowReaderBuilder::try_new_with_options(file, options).unwrap();
let metadata = reader.metadata();
let row_groups = metadata.row_groups();
let row_group_indices: Vec<_> = (0..row_groups.len()).collect();

let statistic_type = if data_page_row_count_limit.is_some() {
"data page"
} else {
"row group"
};

let mut group = c.benchmark_group(format!(
"Extract {} statistics for {}",
statistic_type,
dtype.clone()
));
group.bench_function(
BenchmarkId::new("extract_statistics", dtype.clone()),
|b| {
b.iter(|| {
let converter = StatisticsConverter::try_new(
"col",
reader.schema(),
reader.parquet_schema(),
)
.unwrap();

if data_page_row_count_limit.is_some() {
let column_page_index = reader
.metadata()
.column_index()
.expect("File should have column page indices");

let column_offset_index = reader
.metadata()
.offset_index()
.expect("File should have column offset indices");

let _ = converter.data_page_mins(
column_page_index,
column_offset_index,
&row_group_indices,
);
let _ = converter.data_page_maxes(
column_page_index,
column_offset_index,
&row_group_indices,
);
let _ = converter.data_page_null_counts(
column_page_index,
column_offset_index,
&row_group_indices,
);
let _ = converter.data_page_row_counts(
column_offset_index,
row_groups,
&row_group_indices,
);
} else {
let _ = converter.row_group_mins(row_groups.iter()).unwrap();
let _ = converter.row_group_maxes(row_groups.iter()).unwrap();
let _ = converter
.row_group_null_counts(row_groups.iter())
.unwrap();
let _ = converter
.row_group_row_counts(row_groups.iter())
.unwrap();
}
})
},
);
group.finish();
}
}
}

Expand Down
Loading