diff --git a/parquet/Cargo.toml b/parquet/Cargo.toml index 1d44c585210b..fa89751d4ab1 100644 --- a/parquet/Cargo.toml +++ b/parquet/Cargo.toml @@ -247,6 +247,11 @@ name = "arrow_reader_row_filter" required-features = ["arrow", "async"] harness = false +[[bench]] +name = "arrow_reader_row_selection_policy" +required-features = ["arrow", "async"] +harness = false + [[bench]] name = "arrow_reader_clickbench" required-features = ["arrow", "async", "object_store"] diff --git a/parquet/benches/arrow_reader_row_selection_policy.rs b/parquet/benches/arrow_reader_row_selection_policy.rs new file mode 100644 index 000000000000..da3f4fc15c10 --- /dev/null +++ b/parquet/benches/arrow_reader_row_selection_policy.rs @@ -0,0 +1,57 @@ +// 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. + +//! Auto-policy benchmark for async Parquet `RowSelection` execution. +//! +//! This benchmark is intended for direct comparisons between `main` and +//! candidate branches. Forced row-selection policies and decode-then-filter +//! diagnostics intentionally belong in separate oracle benchmarks. +//! Page indexes are intentionally disabled so this benchmark isolates +//! row-selection execution from page-level I/O pruning. + +mod row_selection_policy_common; + +use criterion::{Criterion, criterion_group, criterion_main}; +use row_selection_policy_common::cases::{ + ORDER_CASES, SCALE_CASES, SHAPE_CASES, assert_order_cases_are_reverses, +}; +use row_selection_policy_common::register::register_auto_group; +use row_selection_policy_common::shapes::assert_shape_contracts; + +fn benchmark_auto(c: &mut Criterion) { + assert_shape_contracts(); + assert_order_cases_are_reverses(); + + register_auto_group( + c, + "arrow_reader_row_selection_policy/auto/shape", + SHAPE_CASES, + ); + register_auto_group( + c, + "arrow_reader_row_selection_policy/auto/order", + ORDER_CASES, + ); + register_auto_group( + c, + "arrow_reader_row_selection_policy/auto/scale", + SCALE_CASES, + ); +} + +criterion_group!(benches, benchmark_auto); +criterion_main!(benches); diff --git a/parquet/benches/row_selection_policy_common/assertions.rs b/parquet/benches/row_selection_policy_common/assertions.rs new file mode 100644 index 000000000000..057ce355eadf --- /dev/null +++ b/parquet/benches/row_selection_policy_common/assertions.rs @@ -0,0 +1,65 @@ +// 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. + +use parquet::arrow::arrow_reader::RowSelectionPolicy; + +use super::fixture::CaseFixture; +use super::model::{CaseSpec, PAYLOAD_VALUE_MODULUS, ROWS_PER_GROUP}; +use super::runner::run_collect_payload0; +use super::shapes::expand_pattern; + +pub(crate) async fn preflight_auto(case: &CaseSpec, fixture: &CaseFixture) { + let actual = run_collect_payload0(fixture, RowSelectionPolicy::default()).await; + assert_eq!( + actual.row_count, fixture.expected_rows, + "{} returned an unexpected number of rows", + case.name + ); + + let expected = expected_selected_global_rows(case); + assert_eq!(actual.payload0.len(), expected.len(), "{}", case.name); + if let Some((output_row, (actual, expected))) = actual + .payload0 + .iter() + .zip(&expected) + .enumerate() + .find(|(_, (actual, expected))| actual != expected) + { + panic!( + "{} returned the wrong source row at output {output_row}: expected {expected}, got {actual}", + case.name + ); + } +} + +fn expected_selected_global_rows(case: &CaseSpec) -> Vec { + case.row_groups + .iter() + .copied() + .enumerate() + .flat_map(|(row_group_idx, pattern)| { + expand_pattern(pattern, ROWS_PER_GROUP) + .into_iter() + .enumerate() + .filter(|(_, selected)| *selected == 1) + .map(move |(row_idx, _)| { + let global_row = row_group_idx * ROWS_PER_GROUP + row_idx; + global_row.wrapping_rem(PAYLOAD_VALUE_MODULUS) as i32 + }) + }) + .collect() +} diff --git a/parquet/benches/row_selection_policy_common/cases.rs b/parquet/benches/row_selection_policy_common/cases.rs new file mode 100644 index 000000000000..99e9068041f2 --- /dev/null +++ b/parquet/benches/row_selection_policy_common/cases.rs @@ -0,0 +1,121 @@ +// 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. + +use super::model::{CaseSpec, RowGroupPattern}; +use super::shapes::{ + BURSTY_50_SAME_SUMMARY, CLUSTERED_50_RUN128, DENSE_98_44_SKIP1_SELECT63, FRAGMENTED_50_RUN1, + MODERATE_12_5_RUN32, REGULAR_50_RUN32, SPARSE_1_56_RUN32, +}; + +const FOUR_SPARSE: &[RowGroupPattern] = &[RowGroupPattern::Cycle(SPARSE_1_56_RUN32); 4]; + +const FOUR_MODERATE: &[RowGroupPattern] = &[RowGroupPattern::Cycle(MODERATE_12_5_RUN32); 4]; + +const FOUR_FRAGMENTED: &[RowGroupPattern] = &[RowGroupPattern::Cycle(FRAGMENTED_50_RUN1); 4]; + +const FOUR_CLUSTERED: &[RowGroupPattern] = &[RowGroupPattern::Cycle(CLUSTERED_50_RUN128); 4]; + +const FOUR_REGULAR: &[RowGroupPattern] = &[RowGroupPattern::Cycle(REGULAR_50_RUN32); 4]; + +const FOUR_BURSTY: &[RowGroupPattern] = &[RowGroupPattern::Cycle(BURSTY_50_SAME_SUMMARY); 4]; + +const FOUR_DENSE: &[RowGroupPattern] = &[RowGroupPattern::Cycle(DENSE_98_44_SKIP1_SELECT63); 4]; + +const FOUR_ALL_SELECTED: &[RowGroupPattern] = &[RowGroupPattern::AllSelected; 4]; + +const SPARSE_TO_FRAGMENTED: &[RowGroupPattern] = &[ + RowGroupPattern::Cycle(SPARSE_1_56_RUN32), + RowGroupPattern::Cycle(SPARSE_1_56_RUN32), + RowGroupPattern::Cycle(FRAGMENTED_50_RUN1), + RowGroupPattern::Cycle(FRAGMENTED_50_RUN1), +]; + +const FRAGMENTED_TO_SPARSE: &[RowGroupPattern] = &[ + RowGroupPattern::Cycle(FRAGMENTED_50_RUN1), + RowGroupPattern::Cycle(FRAGMENTED_50_RUN1), + RowGroupPattern::Cycle(SPARSE_1_56_RUN32), + RowGroupPattern::Cycle(SPARSE_1_56_RUN32), +]; + +const ONE_FRAGMENTED: &[RowGroupPattern] = &[RowGroupPattern::Cycle(FRAGMENTED_50_RUN1)]; + +const EIGHT_FRAGMENTED: &[RowGroupPattern] = &[RowGroupPattern::Cycle(FRAGMENTED_50_RUN1); 8]; + +pub(crate) fn assert_order_cases_are_reverses() { + assert!( + SPARSE_TO_FRAGMENTED + .iter() + .eq(FRAGMENTED_TO_SPARSE.iter().rev()) + ); +} + +pub(crate) const SHAPE_CASES: &[CaseSpec] = &[ + CaseSpec { + name: "sparse_1_56_run32", + row_groups: FOUR_SPARSE, + }, + CaseSpec { + name: "moderate_12_5_run32", + row_groups: FOUR_MODERATE, + }, + CaseSpec { + name: "fragmented_50_run1", + row_groups: FOUR_FRAGMENTED, + }, + CaseSpec { + name: "clustered_50_run128", + row_groups: FOUR_CLUSTERED, + }, + CaseSpec { + name: "regular_50_run32", + row_groups: FOUR_REGULAR, + }, + CaseSpec { + name: "bursty_50_same_summary", + row_groups: FOUR_BURSTY, + }, + CaseSpec { + name: "dense_98_44_skip1_select63", + row_groups: FOUR_DENSE, + }, + CaseSpec { + name: "all_selected", + row_groups: FOUR_ALL_SELECTED, + }, +]; + +pub(crate) const ORDER_CASES: &[CaseSpec] = &[ + CaseSpec { + name: "sparse2_then_fragmented2", + row_groups: SPARSE_TO_FRAGMENTED, + }, + CaseSpec { + name: "fragmented2_then_sparse2", + row_groups: FRAGMENTED_TO_SPARSE, + }, +]; + +pub(crate) const SCALE_CASES: &[CaseSpec] = &[ + CaseSpec { + name: "fragmented_50_run1/rg01", + row_groups: ONE_FRAGMENTED, + }, + CaseSpec { + name: "fragmented_50_run1/rg08", + row_groups: EIGHT_FRAGMENTED, + }, +]; diff --git a/parquet/benches/row_selection_policy_common/fixture.rs b/parquet/benches/row_selection_policy_common/fixture.rs new file mode 100644 index 000000000000..1470f5b89f3b --- /dev/null +++ b/parquet/benches/row_selection_policy_common/fixture.rs @@ -0,0 +1,163 @@ +// 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. + +use std::ops::Range; +use std::sync::Arc; + +use arrow::array::{ArrayRef, Int32Array, RecordBatch}; +use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; +use bytes::Bytes; +use futures::FutureExt; +use futures::future::BoxFuture; +use parquet::arrow::ArrowWriter; +use parquet::arrow::arrow_reader::ArrowReaderOptions; +use parquet::arrow::async_reader::AsyncFileReader; +use parquet::basic::Compression; +use parquet::errors::Result; +use parquet::file::metadata::{PageIndexPolicy, ParquetMetaData, ParquetMetaDataReader}; +use parquet::file::properties::WriterProperties; + +use super::model::{CaseSpec, PAYLOAD_COLUMNS, PAYLOAD_VALUE_MODULUS, ROWS_PER_GROUP}; +use super::shapes::{expand_pattern, selected_rows}; + +#[derive(Debug)] +pub(crate) struct CaseFixture { + bytes: Bytes, + metadata: Arc, + pub(crate) expected_rows: usize, +} + +impl CaseFixture { + pub(crate) fn reader(&self) -> InMemoryAsyncReader { + InMemoryAsyncReader { + bytes: self.bytes.clone(), + metadata: Arc::clone(&self.metadata), + } + } + + pub(crate) fn schema_descr(&self) -> &parquet::schema::types::SchemaDescriptor { + self.metadata.file_metadata().schema_descr() + } +} + +#[derive(Debug, Clone)] +pub(crate) struct InMemoryAsyncReader { + bytes: Bytes, + metadata: Arc, +} + +impl AsyncFileReader for InMemoryAsyncReader { + fn get_bytes(&mut self, range: Range) -> BoxFuture<'_, Result> { + let bytes = self.bytes.slice(range.start as usize..range.end as usize); + async move { Ok(bytes) }.boxed() + } + + fn get_metadata<'a>( + &'a mut self, + _options: Option<&'a ArrowReaderOptions>, + ) -> BoxFuture<'a, Result>> { + let metadata = Arc::clone(&self.metadata); + async move { Ok(metadata) }.boxed() + } +} + +pub(crate) fn build_fixture(case: &CaseSpec) -> Result { + assert!( + !case.row_groups.is_empty(), + "benchmark case must contain at least one row group" + ); + + let schema = build_schema(); + let properties = WriterProperties::builder() + .set_compression(Compression::UNCOMPRESSED) + .set_dictionary_enabled(false) + .set_max_row_group_row_count(Some(ROWS_PER_GROUP)) + .build(); + + let mut encoded = Vec::new(); + { + let mut writer = ArrowWriter::try_new(&mut encoded, Arc::clone(&schema), Some(properties))?; + for (row_group_idx, pattern) in case.row_groups.iter().copied().enumerate() { + writer.write(&build_row_group_batch( + Arc::clone(&schema), + pattern, + row_group_idx, + )?)?; + } + writer.close()?; + } + + let bytes = Bytes::from(encoded); + let mut metadata_reader = + ParquetMetaDataReader::new().with_page_index_policy(PageIndexPolicy::Skip); + metadata_reader.try_parse(&bytes)?; + let metadata = Arc::new(metadata_reader.finish()?); + + assert_eq!( + metadata.num_row_groups(), + case.row_groups.len(), + "writer did not preserve the requested row-group layout" + ); + for row_group in metadata.row_groups() { + assert_eq!(row_group.num_rows() as usize, ROWS_PER_GROUP); + } + + let expected_rows = case + .row_groups + .iter() + .copied() + .map(|pattern| selected_rows(pattern, ROWS_PER_GROUP)) + .sum(); + + Ok(CaseFixture { + bytes, + metadata, + expected_rows, + }) +} + +fn build_schema() -> SchemaRef { + let mut fields = Vec::with_capacity(PAYLOAD_COLUMNS + 1); + fields.push(Field::new("predicate", DataType::Int32, false)); + fields.extend( + (0..PAYLOAD_COLUMNS) + .map(|column_idx| Field::new(format!("payload_{column_idx}"), DataType::Int32, false)), + ); + Arc::new(Schema::new(fields)) +} + +fn build_row_group_batch( + schema: SchemaRef, + pattern: super::model::RowGroupPattern, + row_group_idx: usize, +) -> Result { + let predicate = expand_pattern(pattern, ROWS_PER_GROUP); + let mut columns = Vec::with_capacity(PAYLOAD_COLUMNS + 1); + columns.push(Arc::new(Int32Array::from(predicate)) as ArrayRef); + + for column_idx in 0..PAYLOAD_COLUMNS { + let values = Int32Array::from_iter_values((0..ROWS_PER_GROUP).map(|row_idx| { + let global_row = row_group_idx * ROWS_PER_GROUP + row_idx; + global_row + .wrapping_add(column_idx * 17) + .wrapping_rem(PAYLOAD_VALUE_MODULUS) as i32 + })); + columns.push(Arc::new(values) as ArrayRef); + } + + Ok(RecordBatch::try_new(schema, columns)?) +} diff --git a/parquet/benches/row_selection_policy_common/mod.rs b/parquet/benches/row_selection_policy_common/mod.rs new file mode 100644 index 000000000000..a4053e8c03b9 --- /dev/null +++ b/parquet/benches/row_selection_policy_common/mod.rs @@ -0,0 +1,24 @@ +// 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. + +pub(crate) mod assertions; +pub(crate) mod cases; +pub(crate) mod fixture; +pub(crate) mod model; +pub(crate) mod register; +pub(crate) mod runner; +pub(crate) mod shapes; diff --git a/parquet/benches/row_selection_policy_common/model.rs b/parquet/benches/row_selection_policy_common/model.rs new file mode 100644 index 000000000000..e44462e713ee --- /dev/null +++ b/parquet/benches/row_selection_policy_common/model.rs @@ -0,0 +1,61 @@ +// 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. + +pub(crate) const ROWS_PER_GROUP: usize = 65_536; +pub(crate) const BATCH_SIZE: usize = 8_192; +pub(crate) const PAYLOAD_COLUMNS: usize = 8; +pub(crate) const PAYLOAD_VALUE_MODULUS: usize = 1_000_003; + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(crate) struct SelectionRun { + pub(crate) selected: bool, + pub(crate) len: usize, +} + +impl SelectionRun { + pub(crate) const fn select(len: usize) -> Self { + Self { + selected: true, + len, + } + } + + pub(crate) const fn skip(len: usize) -> Self { + Self { + selected: false, + len, + } + } +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(crate) enum RowGroupPattern { + Cycle(&'static [SelectionRun]), + AllSelected, +} + +#[derive(Clone, Copy, Debug)] +pub(crate) struct CaseSpec { + pub(crate) name: &'static str, + pub(crate) row_groups: &'static [RowGroupPattern], +} + +impl CaseSpec { + pub(crate) const fn total_rows(self) -> usize { + self.row_groups.len() * ROWS_PER_GROUP + } +} diff --git a/parquet/benches/row_selection_policy_common/register.rs b/parquet/benches/row_selection_policy_common/register.rs new file mode 100644 index 000000000000..7c2e8dc16e42 --- /dev/null +++ b/parquet/benches/row_selection_policy_common/register.rs @@ -0,0 +1,59 @@ +// 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. + +use std::hint; +use std::sync::{Arc, OnceLock}; + +use criterion::{Criterion, Throughput}; + +use super::assertions::preflight_auto; +use super::fixture::{CaseFixture, build_fixture}; +use super::model::CaseSpec; +use super::runner::run_auto; + +pub(crate) fn register_auto_group(c: &mut Criterion, group_name: &str, cases: &'static [CaseSpec]) { + let runtime = Arc::new( + tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(), + ); + let mut group = c.benchmark_group(group_name); + + for case in cases { + group.throughput(Throughput::Elements(case.total_rows() as u64)); + + let fixture: Arc>> = Arc::new(OnceLock::new()); + let preflight = Arc::new(OnceLock::new()); + let runtime = Arc::clone(&runtime); + + group.bench_function(case.name, move |b| { + let fixture = + Arc::clone(fixture.get_or_init(|| Arc::new(build_fixture(case).unwrap()))); + preflight.get_or_init(|| { + runtime.block_on(preflight_auto(case, &fixture)); + }); + + b.iter(|| { + let rows = runtime.block_on(run_auto(&fixture)); + hint::black_box(rows); + }); + }); + } + + group.finish(); +} diff --git a/parquet/benches/row_selection_policy_common/runner.rs b/parquet/benches/row_selection_policy_common/runner.rs new file mode 100644 index 000000000000..bf89fd8a6214 --- /dev/null +++ b/parquet/benches/row_selection_policy_common/runner.rs @@ -0,0 +1,92 @@ +// 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. + +use arrow::array::{Int32Array, RecordBatch}; +use arrow::compute::kernels::cmp::eq; +use arrow::datatypes::Int32Type; +use arrow_array::cast::AsArray; +use futures::StreamExt; +use parquet::arrow::arrow_reader::{ArrowPredicateFn, RowFilter, RowSelectionPolicy}; +use parquet::arrow::{ParquetRecordBatchStreamBuilder, ProjectionMask}; + +use super::fixture::CaseFixture; +use super::model::{BATCH_SIZE, PAYLOAD_COLUMNS}; + +#[derive(Debug, Eq, PartialEq)] +pub(crate) struct RunResult { + pub(crate) row_count: usize, + pub(crate) payload0: Vec, +} + +async fn run_with_consumer( + fixture: &CaseFixture, + policy: RowSelectionPolicy, + mut consume: F, +) -> usize +where + F: FnMut(&RecordBatch), +{ + let predicate_projection = ProjectionMask::roots(fixture.schema_descr(), [0]); + let output_projection = ProjectionMask::roots(fixture.schema_descr(), 1..=PAYLOAD_COLUMNS); + let predicate = ArrowPredicateFn::new(predicate_projection, |batch: RecordBatch| { + eq(batch.column(0), &Int32Array::new_scalar(1)) + }); + let row_filter = RowFilter::new(vec![Box::new(predicate)]); + + let mut stream = ParquetRecordBatchStreamBuilder::new(fixture.reader()) + .await + .unwrap() + .with_batch_size(BATCH_SIZE) + .with_projection(output_projection) + .with_row_filter(row_filter) + .with_row_selection_policy(policy) + .build() + .unwrap(); + + let mut rows = 0; + while let Some(batch) = stream.next().await { + let batch = batch.unwrap(); + rows += batch.num_rows(); + consume(&batch); + } + rows +} + +pub(crate) async fn run(fixture: &CaseFixture, policy: RowSelectionPolicy) -> usize { + run_with_consumer(fixture, policy, |_| {}).await +} + +pub(crate) async fn run_auto(fixture: &CaseFixture) -> usize { + run(fixture, RowSelectionPolicy::default()).await +} + +pub(crate) async fn run_collect_payload0( + fixture: &CaseFixture, + policy: RowSelectionPolicy, +) -> RunResult { + let mut payload0 = Vec::with_capacity(fixture.expected_rows); + let row_count = run_with_consumer(fixture, policy, |batch| { + let values = batch.column(0).as_primitive::(); + payload0.extend(values.values().iter().copied()); + }) + .await; + + RunResult { + row_count, + payload0, + } +} diff --git a/parquet/benches/row_selection_policy_common/shapes.rs b/parquet/benches/row_selection_policy_common/shapes.rs new file mode 100644 index 000000000000..e83667ddb5a8 --- /dev/null +++ b/parquet/benches/row_selection_policy_common/shapes.rs @@ -0,0 +1,150 @@ +// 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. + +use super::model::{RowGroupPattern, SelectionRun}; + +pub(crate) const SPARSE_1_56_RUN32: &[SelectionRun] = + &[SelectionRun::skip(2_016), SelectionRun::select(32)]; + +pub(crate) const MODERATE_12_5_RUN32: &[SelectionRun] = + &[SelectionRun::skip(224), SelectionRun::select(32)]; + +pub(crate) const FRAGMENTED_50_RUN1: &[SelectionRun] = + &[SelectionRun::skip(1), SelectionRun::select(1)]; + +pub(crate) const CLUSTERED_50_RUN128: &[SelectionRun] = + &[SelectionRun::skip(128), SelectionRun::select(128)]; + +pub(crate) const REGULAR_50_RUN32: &[SelectionRun] = + &[SelectionRun::skip(32), SelectionRun::select(32)]; + +pub(crate) const DENSE_98_44_SKIP1_SELECT63: &[SelectionRun] = + &[SelectionRun::skip(1), SelectionRun::select(63)]; + +/// Has the same selectivity, run density, and mean selected/skipped run lengths +/// as [`REGULAR_50_RUN32`], but a different run variance and ordering. +pub(crate) const BURSTY_50_SAME_SUMMARY: &[SelectionRun] = &[ + SelectionRun::skip(1), + SelectionRun::select(1), + SelectionRun::skip(1), + SelectionRun::select(1), + SelectionRun::skip(1), + SelectionRun::select(1), + SelectionRun::skip(125), + SelectionRun::select(125), +]; + +#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] +struct ShapeSummary { + selected_rows: usize, + skipped_rows: usize, + selector_count: usize, + selected_run_count: usize, + skipped_run_count: usize, +} + +impl ShapeSummary { + fn from_cycle(cycle: &[SelectionRun]) -> Self { + cycle.iter().fold(Self::default(), |mut summary, run| { + assert!(run.len > 0, "selection run must be non-empty"); + summary.selector_count += 1; + if run.selected { + summary.selected_rows += run.len; + summary.selected_run_count += 1; + } else { + summary.skipped_rows += run.len; + summary.skipped_run_count += 1; + } + summary + }) + } + + fn total_rows(self) -> usize { + self.selected_rows + self.skipped_rows + } + + fn selected_ratio(self) -> f64 { + self.selected_rows as f64 / self.total_rows() as f64 + } + + fn run_density(self) -> f64 { + self.selector_count as f64 / self.total_rows() as f64 + } + + fn average_selected_run(self) -> f64 { + self.selected_rows as f64 / self.selected_run_count as f64 + } + + fn average_skipped_run(self) -> f64 { + self.skipped_rows as f64 / self.skipped_run_count as f64 + } +} + +pub(crate) fn assert_shape_contracts() { + let regular = ShapeSummary::from_cycle(REGULAR_50_RUN32); + let bursty = ShapeSummary::from_cycle(BURSTY_50_SAME_SUMMARY); + + assert_eq!(regular.selected_ratio(), bursty.selected_ratio()); + assert_eq!(regular.run_density(), bursty.run_density()); + assert_eq!( + regular.average_selected_run(), + bursty.average_selected_run() + ); + assert_eq!(regular.average_skipped_run(), bursty.average_skipped_run()); + + let dense = ShapeSummary::from_cycle(DENSE_98_44_SKIP1_SELECT63); + assert_eq!(dense.selected_rows, 63); + assert_eq!(dense.skipped_rows, 1); + assert_eq!(dense.selector_count, 2); + assert_eq!(dense.average_selected_run(), 63.0); + assert_eq!(dense.average_skipped_run(), 1.0); +} + +pub(crate) fn expand_pattern(pattern: RowGroupPattern, row_count: usize) -> Vec { + match pattern { + RowGroupPattern::AllSelected => vec![1; row_count], + RowGroupPattern::Cycle(cycle) => { + assert!(!cycle.is_empty(), "selection cycle must not be empty"); + let summary = ShapeSummary::from_cycle(cycle); + assert_eq!( + row_count % summary.total_rows(), + 0, + "row group size must be divisible by cycle size" + ); + + let mut values = Vec::with_capacity(row_count); + while values.len() < row_count { + for run in cycle { + values.extend(std::iter::repeat_n(i32::from(run.selected), run.len)); + } + } + assert_eq!(values.len(), row_count); + values + } + } +} + +pub(crate) fn selected_rows(pattern: RowGroupPattern, row_count: usize) -> usize { + match pattern { + RowGroupPattern::AllSelected => row_count, + RowGroupPattern::Cycle(cycle) => { + let summary = ShapeSummary::from_cycle(cycle); + assert_eq!(row_count % summary.total_rows(), 0); + summary.selected_rows * (row_count / summary.total_rows()) + } + } +}