This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-11557-e6425d164080f91eb5246a92d2532e549f2b99fe in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
commit bab25f0a3adf553d9945a823f03357ab7b2baf48 Author: Huaijin <[email protected]> AuthorDate: Thu Oct 15 17:02:54 2026 +0000 bench(parquet): add async row selection coverage with page indexes (#11557) Add async reader benchmarks with required page indexes to `row_selector_boolean_buffer`, preserving existing construction and sync cases. Supports [#11553](https://github.com/apache/arrow-rs/pull/11553). Validated all 72 benchmark cases. --------- Co-authored-by: Ed Seidl <[email protected]> --- parquet/Cargo.toml | 2 +- parquet/benches/row_selector_boolean_buffer.rs | 92 +++++++++++++++----------- 2 files changed, 56 insertions(+), 38 deletions(-) diff --git a/parquet/Cargo.toml b/parquet/Cargo.toml index c36791f571..d71120d9d7 100644 --- a/parquet/Cargo.toml +++ b/parquet/Cargo.toml @@ -332,7 +332,7 @@ required-features = ["arrow"] [[bench]] name = "row_selector_boolean_buffer" harness = false -required-features = ["arrow"] +required-features = ["arrow", "async"] [[bench]] name = "row_group_index_reader" diff --git a/parquet/benches/row_selector_boolean_buffer.rs b/parquet/benches/row_selector_boolean_buffer.rs index d741453735..d0421ae126 100644 --- a/parquet/benches/row_selector_boolean_buffer.rs +++ b/parquet/benches/row_selector_boolean_buffer.rs @@ -20,11 +20,17 @@ use arrow_buffer::{BooleanBuffer, BooleanBufferBuilder}; use arrow_schema::{DataType, Field, Schema}; use bytes::Bytes; use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main}; +use futures::StreamExt; use parquet::arrow::ArrowWriter; -use parquet::arrow::arrow_reader::{ParquetRecordBatchReaderBuilder, RowSelection}; +use parquet::arrow::arrow_reader::{ + ArrowReaderOptions, ParquetRecordBatchReaderBuilder, RowSelection, +}; +use parquet::arrow::async_reader::ParquetRecordBatchStreamBuilder; +use parquet::file::metadata::PageIndexPolicy; use rand::rngs::StdRng; use rand::{RngExt, SeedableRng}; use std::hint; +use std::io::Cursor; use std::sync::Arc; const TOTAL_ROWS: usize = 3_000_000; @@ -78,62 +84,55 @@ fn boolean_array(mask: &BooleanBuffer) -> BooleanArray { fn criterion_benchmark(c: &mut Criterion) { let patterns = build_patterns(); + benchmark_selections(c, "construction", &patterns, |selection| selection); - let mut construction = c.benchmark_group("row_selector_boolean_buffer/construction"); - for case in &patterns { - construction.bench_with_input( - BenchmarkId::new("from_filters", case.label()), - &case.mask, - |b, mask| { - b.iter(|| { - let array = boolean_array(mask); - let selection = RowSelection::from_filters(&[array]); - hint::black_box(selection); - }) - }, - ); + let parquet_data = write_parquet_file(TOTAL_ROWS); + benchmark_selections(c, "reader", &patterns, |selection| { + read_rows(&parquet_data, selection) + }); - construction.bench_with_input( - BenchmarkId::new("from_boolean_buffer", case.label()), - &case.mask, - |b, mask| { - b.iter(|| { - let selection = RowSelection::from_boolean_buffer(mask.clone()); - hint::black_box(selection); - }) - }, - ); - } - construction.finish(); + let runtime = tokio::runtime::Builder::new_current_thread() + .build() + .unwrap(); + benchmark_selections(c, "reader_async", &patterns, |selection| { + runtime.block_on(read_rows_async(&parquet_data, selection)) + }); +} - let parquet_data = write_parquet_file(TOTAL_ROWS); - let mut reader = c.benchmark_group("row_selector_boolean_buffer/reader"); - for case in &patterns { - reader.bench_with_input( +/// Each iteration times constructing a `RowSelection` from a precomputed mask +/// and calling `run` with it. `run` returns the selection for construction-only +/// benchmarks, or reads the selected rows for sync and async reader benchmarks. +/// Generating the masks and writing the Parquet data happen outside the timed loop. +fn benchmark_selections<T>( + c: &mut Criterion, + name: &str, + patterns: &[BenchCase], + mut run: impl FnMut(RowSelection) -> T, +) { + let mut group = c.benchmark_group(format!("row_selector_boolean_buffer/{name}")); + for case in patterns { + group.bench_with_input( BenchmarkId::new("from_filters", case.label()), &case.mask, |b, mask| { b.iter(|| { let selection = RowSelection::from_filters(&[boolean_array(mask)]); - let rows = read_rows(&parquet_data, selection); - hint::black_box(rows); + hint::black_box(run(selection)); }) }, ); - - reader.bench_with_input( + group.bench_with_input( BenchmarkId::new("from_boolean_buffer", case.label()), &case.mask, |b, mask| { b.iter(|| { let selection = RowSelection::from_boolean_buffer(mask.clone()); - let rows = read_rows(&parquet_data, selection); - hint::black_box(rows); + hint::black_box(run(selection)); }) }, ); } - reader.finish(); + group.finish(); } fn build_patterns() -> Vec<BenchCase> { @@ -184,6 +183,25 @@ fn read_rows(parquet_data: &Bytes, selection: RowSelection) -> usize { reader.map(|batch| batch.unwrap().num_rows()).sum::<usize>() } +async fn read_rows_async(parquet_data: &Bytes, selection: RowSelection) -> usize { + let options = ArrowReaderOptions::new().with_page_index_policy(PageIndexPolicy::Required); + let mut reader = ParquetRecordBatchStreamBuilder::new_with_options( + Cursor::new(parquet_data.clone()), + options, + ) + .await + .unwrap() + .with_row_selection(selection) + .build() + .unwrap(); + + let mut rows = 0; + while let Some(batch) = reader.next().await { + rows += batch.unwrap().num_rows(); + } + rows +} + struct BenchCase { pattern: &'static str, selectivity: Selectivity,
