diff --git a/benchmarks/Cargo.toml b/benchmarks/Cargo.toml index 62eea98439ee1..cd222c3a5e5b0 100644 --- a/benchmarks/Cargo.toml +++ b/benchmarks/Cargo.toml @@ -44,7 +44,7 @@ arrow = { workspace = true } async-trait = "0.1" bytes = { workspace = true } clap = { version = "4.6.0", features = ["derive", "env", "string"] } -criterion = { workspace = true, features = ["html_reports"] } +criterion = { workspace = true, features = ["async_tokio", "html_reports"] } datafusion = { workspace = true, default-features = true } datafusion-common = { workspace = true, default-features = true } datafusion-common-runtime = { workspace = true } @@ -71,3 +71,7 @@ tempfile = { workspace = true } [[bench]] harness = false name = "sql" + +[[bench]] +harness = false +name = "parquet_pruning_setup_cache" diff --git a/benchmarks/benches/parquet_pruning_setup_cache.rs b/benchmarks/benches/parquet_pruning_setup_cache.rs new file mode 100644 index 0000000000000..685ad698fbb4f --- /dev/null +++ b/benchmarks/benches/parquet_pruning_setup_cache.rs @@ -0,0 +1,110 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Benchmarks end-to-end Parquet scan cost for a cache-favourable workload. +//! +//! The 128 files share one physical schema and use the same predicate and target +//! partition count, so pruning setup is reusable in a cache-enabled comparison +//! branch. The predicate matches every file, so each scan must adapt the +//! predicate and build (or reuse) pruning setup for all files. +//! +//! This is a focused baseline for comparing cache-disabled and cache-enabled +//! branches. It is not a general DataFusion benchmark and must not be +//! interpreted as a ClickBench performance result. + +use std::{fs::File, sync::Arc}; + +use arrow::{ + array::Int64Array, + datatypes::{DataType, Field, Schema}, + record_batch::RecordBatch, +}; +use criterion::{Criterion, criterion_group, criterion_main}; +use datafusion::prelude::{SessionConfig, SessionContext}; +use parquet::arrow::ArrowWriter; +use tempfile::TempDir; + +const FILES: usize = 128; +const ROWS_PER_FILE: usize = 128; + +fn write_files() -> TempDir { + let directory = tempfile::tempdir().unwrap(); + let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int64, false)])); + + for file_index in 0..FILES { + let start = i64::try_from(file_index * ROWS_PER_FILE).unwrap(); + let values = Int64Array::from_iter_values( + start..start + i64::try_from(ROWS_PER_FILE).unwrap(), + ); + let batch = + RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(values)]).unwrap(); + let file = + File::create(directory.path().join(format!("{file_index}.parquet"))).unwrap(); + let mut writer = ArrowWriter::try_new(file, Arc::clone(&schema), None).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + } + + directory +} + +fn criterion_benchmark(criterion: &mut Criterion) { + let directory = write_files(); + let runtime = tokio::runtime::Runtime::new().unwrap(); + let context = runtime.block_on(async { + let context = SessionContext::new_with_config( + SessionConfig::new().with_target_partitions(1), + ); + context + .register_parquet("t", directory.path().to_str().unwrap(), Default::default()) + .await + .unwrap(); + context + }); + + runtime.block_on(async { + let batches = context + .sql("SELECT id FROM t WHERE id >= 0") + .await + .unwrap() + .collect() + .await + .unwrap(); + assert_eq!( + batches.iter().map(RecordBatch::num_rows).sum::(), + FILES * ROWS_PER_FILE + ); + }); + + criterion.bench_function( + "parquet_pruning_setup_cache/same_schema_files", + |bencher| { + bencher.to_async(&runtime).iter(|| async { + context + .sql("SELECT id FROM t WHERE id >= 0") + .await + .unwrap() + .collect() + .await + .unwrap() + }); + }, + ); +} + +criterion_group!(benches, criterion_benchmark); +criterion_main!(benches);