diff --git a/parquet/Cargo.toml b/parquet/Cargo.toml index ce4212e8e3ea..58c1f4980ca4 100644 --- a/parquet/Cargo.toml +++ b/parquet/Cargo.toml @@ -259,6 +259,16 @@ name = "arrow_reader_row_selection_policy" required-features = ["arrow", "async"] harness = false +[[bench]] +name = "arrow_reader_row_selection_policy_heterogeneous" +required-features = ["arrow", "async"] +harness = false + +[[bench]] +name = "arrow_reader_row_selection_policy_sampler" +required-features = ["arrow", "async"] +harness = false + [[bench]] name = "arrow_reader_clickbench" required-features = ["arrow", "async"] diff --git a/parquet/benches/arrow_reader_row_selection_policy_heterogeneous.rs b/parquet/benches/arrow_reader_row_selection_policy_heterogeneous.rs new file mode 100644 index 000000000000..3ff1be808c13 --- /dev/null +++ b/parquet/benches/arrow_reader_row_selection_policy_heterogeneous.rs @@ -0,0 +1,43 @@ +// 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. + +//! Baseline benchmark for row-selection execution over a heterogeneous +//! projection. The fixture combines fixed-width, variable-width, dictionary, +//! and fixed-length byte-array decoding while keeping the logical selection +//! identical for every output column. +//! +//! 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::HETEROGENEOUS_CASES; +use row_selection_policy_common::register::register_heterogeneous_group; +use row_selection_policy_common::shapes::assert_shape_contracts; + +fn benchmark_heterogeneous(c: &mut Criterion) { + assert_shape_contracts(); + register_heterogeneous_group( + c, + "arrow_reader_row_selection_policy/heterogeneous", + HETEROGENEOUS_CASES, + ); +} + +criterion_group!(benches, benchmark_heterogeneous); +criterion_main!(benches); diff --git a/parquet/benches/arrow_reader_row_selection_policy_sampler.rs b/parquet/benches/arrow_reader_row_selection_policy_sampler.rs new file mode 100644 index 000000000000..0e727cef9496 --- /dev/null +++ b/parquet/benches/arrow_reader_row_selection_policy_sampler.rs @@ -0,0 +1,32 @@ +// 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. + +//! Offline paired sampler for the per-column row-selection cost model. +//! +//! This is intentionally separate from Criterion benchmarks: it emits raw, +//! resumable observations for offline analysis instead of producing a single +//! benchmark summary. + +mod row_selection_policy_common; +mod row_selection_policy_sampler; + +fn main() { + if let Err(error) = row_selection_policy_sampler::run() { + eprintln!("row-selection sampler failed: {error}"); + std::process::exit(1); + } +} diff --git a/parquet/benches/row_selection_policy_common/assertions.rs b/parquet/benches/row_selection_policy_common/assertions.rs index 057ce355eadf..8b005a67e5c7 100644 --- a/parquet/benches/row_selection_policy_common/assertions.rs +++ b/parquet/benches/row_selection_policy_common/assertions.rs @@ -15,11 +15,19 @@ // specific language governing permissions and limitations // under the License. +use arrow::array::StringArray; +use arrow::datatypes::{DataType, Int32Type}; +use arrow_array::ArrayAccessor; +use arrow_array::cast::AsArray; use parquet::arrow::arrow_reader::RowSelectionPolicy; -use super::fixture::CaseFixture; +use super::fixture::{ + CaseFixture, HETEROGENEOUS_FIXED_BINARY_WIDTH, heterogeneous_dictionary_key, + heterogeneous_dictionary_value, heterogeneous_fixed_binary_value, heterogeneous_int32_value, + heterogeneous_string_value, +}; use super::model::{CaseSpec, PAYLOAD_VALUE_MODULUS, ROWS_PER_GROUP}; -use super::runner::run_collect_payload0; +use super::runner::{run_collect_payload0, run_with_consumer}; use super::shapes::expand_pattern; pub(crate) async fn preflight_auto(case: &CaseSpec, fixture: &CaseFixture) { @@ -35,9 +43,13 @@ pub(crate) async fn preflight_auto(case: &CaseSpec, fixture: &CaseFixture) { if let Some((output_row, (actual, expected))) = actual .payload0 .iter() - .zip(&expected) + .zip( + expected + .iter() + .map(|row| row.wrapping_rem(PAYLOAD_VALUE_MODULUS) as i32), + ) .enumerate() - .find(|(_, (actual, expected))| actual != expected) + .find(|(_, (actual, expected))| **actual != *expected) { panic!( "{} returned the wrong source row at output {output_row}: expected {expected}, got {actual}", @@ -46,7 +58,139 @@ pub(crate) async fn preflight_auto(case: &CaseSpec, fixture: &CaseFixture) { } } -fn expected_selected_global_rows(case: &CaseSpec) -> Vec { +pub(crate) async fn preflight_heterogeneous(case: &CaseSpec, fixture: &CaseFixture) { + let expected = expected_selected_global_rows(case); + for (policy_name, policy) in [ + ("auto", RowSelectionPolicy::default()), + ("auto_per_column", RowSelectionPolicy::AutoPerColumn), + ("selectors", RowSelectionPolicy::Selectors), + ("mask", RowSelectionPolicy::Mask), + ] { + let mut output_offset = 0; + let row_count = run_with_consumer(fixture, policy, |batch| { + let batch_end = output_offset + batch.num_rows(); + assert!( + batch_end <= expected.len(), + "{} ({policy_name}) returned too many rows", + case.name + ); + assert_heterogeneous_batch( + case, + policy_name, + batch, + &expected[output_offset..batch_end], + ); + output_offset = batch_end; + }) + .await; + + assert_eq!( + row_count, fixture.expected_rows, + "{} ({policy_name}) returned an unexpected number of rows", + case.name + ); + assert_eq!( + output_offset, + expected.len(), + "{} ({policy_name}) did not return every expected row", + case.name + ); + } +} + +fn assert_heterogeneous_batch( + case: &CaseSpec, + policy_name: &str, + batch: &arrow::record_batch::RecordBatch, + expected_rows: &[usize], +) { + let expected_types = [ + DataType::Int32, + DataType::Int32, + DataType::Utf8View, + DataType::Utf8View, + DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)), + DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)), + DataType::FixedSizeBinary(HETEROGENEOUS_FIXED_BINARY_WIDTH as i32), + DataType::FixedSizeBinary(HETEROGENEOUS_FIXED_BINARY_WIDTH as i32), + ]; + assert_eq!(batch.num_columns(), expected_types.len()); + for (column_idx, expected_type) in expected_types.iter().enumerate() { + assert_eq!( + batch.column(column_idx).data_type(), + expected_type, + "{} ({policy_name}) returned the wrong type for payload_{column_idx}", + case.name + ); + } + + let int32 = [ + batch.column(0).as_primitive::(), + batch.column(1).as_primitive::(), + ]; + let strings = [ + batch.column(2).as_string_view(), + batch.column(3).as_string_view(), + ]; + let dictionaries = [ + batch + .column(4) + .as_dictionary::() + .downcast_dict::() + .unwrap(), + batch + .column(5) + .as_dictionary::() + .downcast_dict::() + .unwrap(), + ]; + let fixed_binary = [ + batch.column(6).as_fixed_size_binary(), + batch.column(7).as_fixed_size_binary(), + ]; + + for (batch_row, global_row) in expected_rows.iter().copied().enumerate() { + for (column_idx, values) in int32.iter().enumerate() { + assert_eq!( + values.value(batch_row), + heterogeneous_int32_value(column_idx, global_row), + "{} ({policy_name}) returned the wrong payload_{column_idx} value at source row {global_row}", + case.name + ); + } + for (offset, values) in strings.iter().enumerate() { + let column_idx = offset + 2; + assert_eq!( + values.value(batch_row), + heterogeneous_string_value(column_idx, global_row), + "{} ({policy_name}) returned the wrong payload_{column_idx} value at source row {global_row}", + case.name + ); + } + for (offset, values) in dictionaries.iter().enumerate() { + let column_idx = offset + 4; + let key = heterogeneous_dictionary_key(column_idx, global_row); + assert_eq!( + values.value(batch_row), + heterogeneous_dictionary_value(column_idx, key), + "{} ({policy_name}) returned the wrong payload_{column_idx} value at source row {global_row}", + case.name + ); + } + for (offset, values) in fixed_binary.iter().enumerate() { + let column_idx = offset + 6; + let expected = heterogeneous_fixed_binary_value(column_idx, global_row); + assert_eq!( + values.value(batch_row), + expected.as_slice(), + "{} ({policy_name}) returned the wrong payload_{column_idx} value at source row {global_row}", + case.name + ); + } + } +} + +fn expected_selected_global_rows(case: &CaseSpec) -> Vec { case.row_groups .iter() .copied() @@ -56,10 +200,7 @@ fn expected_selected_global_rows(case: &CaseSpec) -> Vec { .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 - }) + .map(move |(row_idx, _)| row_group_idx * ROWS_PER_GROUP + row_idx) }) .collect() } diff --git a/parquet/benches/row_selection_policy_common/cases.rs b/parquet/benches/row_selection_policy_common/cases.rs index 99e9068041f2..40cb869e81e4 100644 --- a/parquet/benches/row_selection_policy_common/cases.rs +++ b/parquet/benches/row_selection_policy_common/cases.rs @@ -18,7 +18,7 @@ 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, + MODERATE_12_5_RUN32, REGULAR_50_RUN8, REGULAR_50_RUN32, SPARSE_1_56_RUN32, }; const FOUR_SPARSE: &[RowGroupPattern] = &[RowGroupPattern::Cycle(SPARSE_1_56_RUN32); 4]; @@ -31,6 +31,8 @@ const FOUR_CLUSTERED: &[RowGroupPattern] = &[RowGroupPattern::Cycle(CLUSTERED_50 const FOUR_REGULAR: &[RowGroupPattern] = &[RowGroupPattern::Cycle(REGULAR_50_RUN32); 4]; +const FOUR_BOUNDARY: &[RowGroupPattern] = &[RowGroupPattern::Cycle(REGULAR_50_RUN8); 4]; + const FOUR_BURSTY: &[RowGroupPattern] = &[RowGroupPattern::Cycle(BURSTY_50_SAME_SUMMARY); 4]; const FOUR_DENSE: &[RowGroupPattern] = &[RowGroupPattern::Cycle(DENSE_98_44_SKIP1_SELECT63); 4]; @@ -119,3 +121,26 @@ pub(crate) const SCALE_CASES: &[CaseSpec] = &[ row_groups: EIGHT_FRAGMENTED, }, ]; + +pub(crate) const HETEROGENEOUS_CASES: &[CaseSpec] = &[ + CaseSpec { + name: "boundary_50_run8", + row_groups: FOUR_BOUNDARY, + }, + CaseSpec { + name: "sparse_1_56_run32", + row_groups: FOUR_SPARSE, + }, + CaseSpec { + name: "fragmented_50_run1", + row_groups: FOUR_FRAGMENTED, + }, + CaseSpec { + name: "clustered_50_run128", + row_groups: FOUR_CLUSTERED, + }, + CaseSpec { + name: "sparse2_then_fragmented2", + row_groups: SPARSE_TO_FRAGMENTED, + }, +]; diff --git a/parquet/benches/row_selection_policy_common/fixture.rs b/parquet/benches/row_selection_policy_common/fixture.rs index 1470f5b89f3b..0c4415a86253 100644 --- a/parquet/benches/row_selection_policy_common/fixture.rs +++ b/parquet/benches/row_selection_policy_common/fixture.rs @@ -18,22 +18,30 @@ use std::ops::Range; use std::sync::Arc; -use arrow::array::{ArrayRef, Int32Array, RecordBatch}; -use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; +use arrow::array::{ + ArrayRef, DictionaryArray, FixedSizeBinaryArray, Int32Array, RecordBatch, StringArray, + StringViewArray, +}; +use arrow::datatypes::{DataType, Field, Int32Type, 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::basic::{Compression, Encoding, Type}; use parquet::errors::Result; use parquet::file::metadata::{PageIndexPolicy, ParquetMetaData, ParquetMetaDataReader}; use parquet::file::properties::WriterProperties; +use parquet::schema::types::ColumnPath; use super::model::{CaseSpec, PAYLOAD_COLUMNS, PAYLOAD_VALUE_MODULUS, ROWS_PER_GROUP}; use super::shapes::{expand_pattern, selected_rows}; +pub(crate) const HETEROGENEOUS_STRING_WIDTH: usize = 64; +pub(crate) const HETEROGENEOUS_DICTIONARY_CARDINALITY: usize = 1_024; +pub(crate) const HETEROGENEOUS_FIXED_BINARY_WIDTH: usize = 32; + #[derive(Debug)] pub(crate) struct CaseFixture { bytes: Bytes, @@ -101,6 +109,43 @@ pub(crate) fn build_fixture(case: &CaseSpec) -> Result { writer.close()?; } + finish_fixture(case, encoded) +} + +pub(crate) fn build_heterogeneous_fixture(case: &CaseSpec) -> Result { + assert!( + !case.row_groups.is_empty(), + "benchmark case must contain at least one row group" + ); + + let schema = build_heterogeneous_schema(); + let properties = WriterProperties::builder() + .set_compression(Compression::UNCOMPRESSED) + .set_dictionary_enabled(false) + .set_column_dictionary_enabled(ColumnPath::from("payload_4"), true) + .set_column_dictionary_enabled(ColumnPath::from("payload_5"), true) + .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_heterogeneous_row_group_batch( + Arc::clone(&schema), + pattern, + row_group_idx, + )?)?; + } + writer.close()?; + } + + let fixture = finish_fixture(case, encoded)?; + validate_heterogeneous_metadata(&fixture.metadata); + Ok(fixture) +} + +fn finish_fixture(case: &CaseSpec, encoded: Vec) -> Result { let bytes = Bytes::from(encoded); let mut metadata_reader = ParquetMetaDataReader::new().with_page_index_policy(PageIndexPolicy::Skip); @@ -140,6 +185,29 @@ fn build_schema() -> SchemaRef { Arc::new(Schema::new(fields)) } +fn build_heterogeneous_schema() -> SchemaRef { + let mut fields = Vec::with_capacity(PAYLOAD_COLUMNS + 1); + fields.push(Field::new("predicate", DataType::Int32, false)); + fields.extend( + [ + DataType::Int32, + DataType::Int32, + DataType::Utf8View, + DataType::Utf8View, + DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)), + DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)), + DataType::FixedSizeBinary(HETEROGENEOUS_FIXED_BINARY_WIDTH as i32), + DataType::FixedSizeBinary(HETEROGENEOUS_FIXED_BINARY_WIDTH as i32), + ] + .into_iter() + .enumerate() + .map(|(column_idx, data_type)| { + Field::new(format!("payload_{column_idx}"), data_type, false) + }), + ); + Arc::new(Schema::new(fields)) +} + fn build_row_group_batch( schema: SchemaRef, pattern: super::model::RowGroupPattern, @@ -161,3 +229,140 @@ fn build_row_group_batch( Ok(RecordBatch::try_new(schema, columns)?) } + +fn build_heterogeneous_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..2 { + let values = Int32Array::from_iter_values((0..ROWS_PER_GROUP).map(|row_idx| { + heterogeneous_int32_value(column_idx, global_row(row_group_idx, row_idx)) + })); + columns.push(Arc::new(values) as ArrayRef); + } + + for column_idx in 2..4 { + let values = StringViewArray::from_iter_values((0..ROWS_PER_GROUP).map(|row_idx| { + heterogeneous_string_value(column_idx, global_row(row_group_idx, row_idx)) + })); + columns.push(Arc::new(values) as ArrayRef); + } + + for column_idx in 4..6 { + let keys = Int32Array::from_iter_values((0..ROWS_PER_GROUP).map(|row_idx| { + heterogeneous_dictionary_key(column_idx, global_row(row_group_idx, row_idx)) as i32 + })); + let values = StringArray::from_iter_values( + (0..HETEROGENEOUS_DICTIONARY_CARDINALITY) + .map(|key| heterogeneous_dictionary_value(column_idx, key)), + ); + let dictionary = DictionaryArray::::try_new(keys, Arc::new(values))?; + columns.push(Arc::new(dictionary) as ArrayRef); + } + + for column_idx in 6..8 { + let values = (0..ROWS_PER_GROUP) + .flat_map(|row_idx| { + heterogeneous_fixed_binary_value(column_idx, global_row(row_group_idx, row_idx)) + }) + .collect::>(); + let values = FixedSizeBinaryArray::try_new( + HETEROGENEOUS_FIXED_BINARY_WIDTH as i32, + values.into(), + None, + )?; + columns.push(Arc::new(values) as ArrayRef); + } + + Ok(RecordBatch::try_new(schema, columns)?) +} + +fn validate_heterogeneous_metadata(metadata: &ParquetMetaData) { + let expected_types = [ + Type::INT32, + Type::INT32, + Type::INT32, + Type::BYTE_ARRAY, + Type::BYTE_ARRAY, + Type::BYTE_ARRAY, + Type::BYTE_ARRAY, + Type::FIXED_LEN_BYTE_ARRAY, + Type::FIXED_LEN_BYTE_ARRAY, + ]; + + for row_group in metadata.row_groups() { + assert_eq!(row_group.num_columns(), expected_types.len()); + for (column_idx, (column, expected_type)) in + row_group.columns().iter().zip(expected_types).enumerate() + { + assert_eq!( + column.column_type(), + expected_type, + "unexpected physical type for column {column_idx}" + ); + + let dictionary_encoded = column.encodings_mask().is_set(Encoding::RLE_DICTIONARY); + assert_eq!( + dictionary_encoded, + matches!(column_idx, 5 | 6), + "unexpected dictionary encoding for column {column_idx}" + ); + } + } +} + +fn global_row(row_group_idx: usize, row_idx: usize) -> usize { + row_group_idx * ROWS_PER_GROUP + row_idx +} + +pub(crate) fn heterogeneous_int32_value(column_idx: usize, global_row: usize) -> i32 { + global_row + .wrapping_add(column_idx * 17) + .wrapping_rem(PAYLOAD_VALUE_MODULUS) as i32 +} + +pub(crate) fn heterogeneous_string_value(column_idx: usize, global_row: usize) -> String { + let mixed = mix64((global_row as u64) ^ ((column_idx as u64) << 48)); + let remixed = mix64(mixed ^ 0x9e37_79b9_7f4a_7c15); + let value = format!("payload_{column_idx}:{global_row:016x}:{mixed:016x}:{remixed:016x}:end"); + debug_assert_eq!(value.len(), HETEROGENEOUS_STRING_WIDTH); + value +} + +pub(crate) fn heterogeneous_dictionary_key(column_idx: usize, global_row: usize) -> usize { + global_row + .wrapping_mul(31) + .wrapping_add(column_idx * 101) + .wrapping_rem(HETEROGENEOUS_DICTIONARY_CARDINALITY) +} + +pub(crate) fn heterogeneous_dictionary_value(column_idx: usize, key: usize) -> String { + let mixed = mix64((key as u64) ^ ((column_idx as u64) << 48)); + format!("dict_{column_idx}:{key:04x}:{mixed:016x}") +} + +pub(crate) fn heterogeneous_fixed_binary_value( + column_idx: usize, + global_row: usize, +) -> [u8; HETEROGENEOUS_FIXED_BINARY_WIDTH] { + let seed = (global_row as u64) ^ ((column_idx as u64) << 48); + let mut value = [0; HETEROGENEOUS_FIXED_BINARY_WIDTH]; + for (lane, chunk) in value.chunks_exact_mut(8).enumerate() { + let mixed = mix64(seed ^ (lane as u64).wrapping_mul(0x9e37_79b9_7f4a_7c15)); + chunk.copy_from_slice(&mixed.to_le_bytes()); + } + value +} + +fn mix64(mut value: u64) -> u64 { + value ^= value >> 30; + value = value.wrapping_mul(0xbf58_476d_1ce4_e5b9); + value ^= value >> 27; + value = value.wrapping_mul(0x94d0_49bb_1331_11eb); + value ^ (value >> 31) +} diff --git a/parquet/benches/row_selection_policy_common/mod.rs b/parquet/benches/row_selection_policy_common/mod.rs index a4053e8c03b9..c636cf32fb4f 100644 --- a/parquet/benches/row_selection_policy_common/mod.rs +++ b/parquet/benches/row_selection_policy_common/mod.rs @@ -15,6 +15,10 @@ // specific language governing permissions and limitations // under the License. +// This support module is shared by benchmark binaries with intentionally +// disjoint case and fixture registrations. +#![allow(dead_code)] + pub(crate) mod assertions; pub(crate) mod cases; pub(crate) mod fixture; diff --git a/parquet/benches/row_selection_policy_common/register.rs b/parquet/benches/row_selection_policy_common/register.rs index 7c2e8dc16e42..576ac16e73be 100644 --- a/parquet/benches/row_selection_policy_common/register.rs +++ b/parquet/benches/row_selection_policy_common/register.rs @@ -20,10 +20,11 @@ use std::sync::{Arc, OnceLock}; use criterion::{Criterion, Throughput}; -use super::assertions::preflight_auto; -use super::fixture::{CaseFixture, build_fixture}; +use super::assertions::{preflight_auto, preflight_heterogeneous}; +use super::fixture::{CaseFixture, build_fixture, build_heterogeneous_fixture}; use super::model::CaseSpec; -use super::runner::run_auto; +use super::runner::{run, run_auto}; +use parquet::arrow::arrow_reader::RowSelectionPolicy; pub(crate) fn register_auto_group(c: &mut Criterion, group_name: &str, cases: &'static [CaseSpec]) { let runtime = Arc::new( @@ -57,3 +58,51 @@ pub(crate) fn register_auto_group(c: &mut Criterion, group_name: &str, cases: &' group.finish(); } + +pub(crate) fn register_heterogeneous_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()); + + for (policy_name, policy) in [ + ("auto", RowSelectionPolicy::default()), + ("auto_per_column", RowSelectionPolicy::AutoPerColumn), + ("selectors", RowSelectionPolicy::Selectors), + ("mask", RowSelectionPolicy::Mask), + ] { + let fixture = Arc::clone(&fixture); + let preflight = Arc::clone(&preflight); + let runtime = Arc::clone(&runtime); + + group.bench_function(format!("{}/{policy_name}", case.name), move |b| { + let fixture = Arc::clone( + fixture.get_or_init(|| Arc::new(build_heterogeneous_fixture(case).unwrap())), + ); + preflight.get_or_init(|| { + runtime.block_on(preflight_heterogeneous(case, &fixture)); + }); + + b.iter(|| { + let rows = runtime.block_on(run(&fixture, policy)); + 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 index bf89fd8a6214..095a0c6812b0 100644 --- a/parquet/benches/row_selection_policy_common/runner.rs +++ b/parquet/benches/row_selection_policy_common/runner.rs @@ -32,7 +32,7 @@ pub(crate) struct RunResult { pub(crate) payload0: Vec, } -async fn run_with_consumer( +pub(crate) async fn run_with_consumer( fixture: &CaseFixture, policy: RowSelectionPolicy, mut consume: F, diff --git a/parquet/benches/row_selection_policy_common/shapes.rs b/parquet/benches/row_selection_policy_common/shapes.rs index e83667ddb5a8..a1975870215d 100644 --- a/parquet/benches/row_selection_policy_common/shapes.rs +++ b/parquet/benches/row_selection_policy_common/shapes.rs @@ -32,6 +32,11 @@ pub(crate) const CLUSTERED_50_RUN128: &[SelectionRun] = pub(crate) const REGULAR_50_RUN32: &[SelectionRun] = &[SelectionRun::skip(32), SelectionRun::select(32)]; +/// Refinement boundary where cheap and wide column decoders choose different +/// row-selection strategies. +pub(crate) const REGULAR_50_RUN8: &[SelectionRun] = + &[SelectionRun::skip(8), SelectionRun::select(8)]; + pub(crate) const DENSE_98_44_SKIP1_SELECT63: &[SelectionRun] = &[SelectionRun::skip(1), SelectionRun::select(63)]; diff --git a/parquet/benches/row_selection_policy_sampler/README.md b/parquet/benches/row_selection_policy_sampler/README.md new file mode 100644 index 000000000000..69713c4b49cb --- /dev/null +++ b/parquet/benches/row_selection_policy_sampler/README.md @@ -0,0 +1,155 @@ + + +# Row-selection policy sampler + +This bench-only binary collects paired observations for two related questions: + +- `Mask` versus `Selectors` for one projected column at a time +- end-to-end `Auto` versus `AutoPerColumn` over a heterogeneous projection + +It does not change the production `AutoPerColumn` planner. + +Fixtures and Parquet metadata are built before measurement. Each observation +times reader construction through complete batch consumption against preloaded +bytes; fixture generation, correctness checks, checksums, and storage I/O are +outside the timer. The two forced policies are run serially in randomized A/B +order, and their decoded `RecordBatch` values must match before a point is +sampled. + +Run a short smoke pass: + +```shell +cargo bench -p parquet --features arrow,async \ + --bench arrow_reader_row_selection_policy_sampler -- \ + --stage smoke --output /tmp/row-selection-smoke.jsonl +``` + +Run a time-bounded Pilot and resume it later: + +```shell +cargo bench -p parquet --features arrow,async \ + --bench arrow_reader_row_selection_policy_sampler -- \ + --stage pilot --budget-seconds 600 \ + --output /tmp/row-selection-pilot.jsonl + +cargo bench -p parquet --features arrow,async \ + --bench arrow_reader_row_selection_policy_sampler -- \ + --stage pilot --budget-seconds 600 \ + --output /tmp/row-selection-pilot.jsonl --resume +``` + +After the Pilot identifies a coarse boundary, scan intermediate run lengths +and selectivities before fitting a model: + +```shell +cargo bench -p parquet --features arrow,async \ + --bench arrow_reader_row_selection_policy_sampler -- \ + --stage refinement --budget-seconds 600 \ + --output /tmp/row-selection-refinement.jsonl +``` + +The refinement defaults use a 3% practical decision band. This prevents tiny, +unstable differences near the boundary from being treated as decisive wins. + +Validate the complete policy decision over the same heterogeneous cases used +by the Criterion benchmark: + +```shell +cargo bench -p parquet --features arrow,async \ + --bench arrow_reader_row_selection_policy_sampler -- \ + --stage policy-validation \ + --output /tmp/row-selection-policy-validation.jsonl +``` + +Use repeatable `--case` filters to isolate a workload, for example +`--case boundary_50_run8`. This stage measures `Auto` and `AutoPerColumn` in +randomized two-pair blocks, so each block contains both execution orders. Its +defaults use 20 warmup pairs, 20--60 measured pairs, 16 complete scans per +timed observation, a 3% practical decision band, and an `Auto` control after +every four measured pairs. + +Each policy-validation record includes the raw pairs, periodic controls, +bootstrap confidence interval, execution-order effect, and the actual +Mask/Selectors/fallback decision counts made by `AutoPerColumn`. The final +`run_end.validation_passed` is false if any point exceeds 10% control drift or +a 5% execution-order effect. `run_end.promotion_eligible` additionally requires +that no point remains statistically inconclusive; a practical tie is a valid +conclusion and is not considered inconclusive. + +Inspect the decisions and the final gate with: + +```shell +jq -c 'select(.record_type == "experiment") | + {case: .experiment.case, summary, decisions: .policy_decisions}' \ + /tmp/row-selection-policy-validation.jsonl + +jq 'select(.record_type == "run_end") | + {validation_passed, promotion_eligible, validation_warning_points, + inconclusive_points, remaining_experiments}' \ + /tmp/row-selection-policy-validation.jsonl +``` + +Validate sparse page loading through the public push decoder: + +```shell +cargo bench -p parquet --features arrow,async \ + --bench arrow_reader_row_selection_policy_sampler -- \ + --stage page-validation \ + --output /tmp/row-selection-pages.jsonl +``` + +This stage asks the empty-buffer decoder which ranges it needs under both +forced policies and fails unless each requests a strict subset of the encoded +column. Timed observations still use preloaded buffers, so they measure decode +latency rather than filesystem latency. + +## Output and restart rules + +The versioned JSONL stream contains a manifest, start/end control points, raw +paired observations, summaries, and a final run record. Each observation also +carries the sampling settings used for it, so a resumed run remains +self-describing if its time budget or sampling granularity is adjusted. A Pilot +always runs its small mandatory type coverage before the wall-clock budget can +stop the randomized expansion set. + +With `--resume`, complete and unsupported experiment IDs are skipped while +incomplete timeout records are retried. Resume rejects a different schema, +manifest, stage, seed, or machine signature. A mandatory fixture error aborts +the run; an optional fixture error is recorded as `unsupported` so the rest of +the manifest can continue. The sampler never silently substitutes one policy +for the other. Policy-validation stability warnings are retained across resume +and continue to make the final validation gate fail. + +## Promoting measurements into the planner + +Treat effects inside the refinement stage's 3% decision band as ties. Before a +threshold changes production planning, repeat the same manifest on a warm +machine, require control drift below 10%, and check that the practical winner +is stable. A stable single-column refinement is necessary but not sufficient: +the corresponding `policy-validation` run must also finish with +`validation_passed: true`. Sequential Criterion policy groups are useful for +throughput reporting, but are not promotion evidence because they do not +cancel cross-window drift. Rules apply only to sampled Arrow/encoding families; +unmodeled, nested, or ambiguous metadata must retain the compatibility +fallback. + +Use `--help` for sampling, filtering, confidence, and checkpoint options. Raw +JSONL files are machine-specific experiment artifacts and are not intended to +be committed to the repository. diff --git a/parquet/benches/row_selection_policy_sampler/cli.rs b/parquet/benches/row_selection_policy_sampler/cli.rs new file mode 100644 index 000000000000..5d7c8e993831 --- /dev/null +++ b/parquet/benches/row_selection_policy_sampler/cli.rs @@ -0,0 +1,435 @@ +// 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::collections::HashSet; +use std::ffi::OsString; +use std::path::PathBuf; +use std::time::{Duration, SystemTime, UNIX_EPOCH}; + +use serde_json::{Value, json}; + +use super::model::{FixtureKind, Stage}; + +pub(crate) enum ParseOutcome { + Help, + Run(Cli), +} + +pub(crate) struct Cli { + pub(crate) stage: Stage, + pub(crate) budget: Duration, + pub(crate) seed: u64, + pub(crate) output: PathBuf, + pub(crate) resume: bool, + pub(crate) kinds: Vec, + pub(crate) cases: Vec, + pub(crate) min_pairs: usize, + pub(crate) max_pairs: usize, + pub(crate) warmup_pairs: usize, + pub(crate) bootstrap_samples: usize, + pub(crate) decision_band: f64, + pub(crate) target_ci_width: f64, + pub(crate) point_timeout: Duration, + pub(crate) inner_iterations: usize, + pub(crate) control_interval_pairs: usize, + pub(crate) ephemeral_output: bool, +} + +impl Cli { + pub(crate) fn parse() -> Result { + let mut args = std::env::args_os().skip(1); + let mut stage = None; + let mut budget_seconds = None; + let mut seed = 0x5eed_c057_2026_0817u64; + let mut output = None; + let mut resume = false; + let mut kinds = Vec::new(); + let mut cases = Vec::new(); + let mut min_pairs = None; + let mut max_pairs = None; + let mut warmup_pairs = None; + let mut bootstrap_samples = None; + let mut decision_band = None; + let mut target_ci_width = None; + let mut point_timeout_seconds = 60.0; + let mut inner_iterations = None; + let mut control_interval_pairs = None; + let mut ephemeral_output = false; + + while let Some(arg) = args.next() { + let arg = arg + .into_string() + .map_err(|_| "sampler arguments must be valid UTF-8".to_string())?; + match arg.as_str() { + "-h" | "--help" => { + print_help(); + return Ok(ParseOutcome::Help); + } + "--test" => ephemeral_output = true, + "--smoke" => stage = Some(Stage::Smoke), + // Cargo appends this libtest-compatible flag to bench + // binaries even when `harness = false`. + "--bench" => {} + "--stage" => { + stage = Some(Stage::parse(&next_string(&mut args, "--stage")?)?); + } + "--budget-seconds" => { + budget_seconds = Some(parse_f64( + &next_string(&mut args, "--budget-seconds")?, + "--budget-seconds", + )?); + } + "--seed" => { + seed = parse_u64(&next_string(&mut args, "--seed")?, "--seed")?; + } + "--output" => { + output = Some(PathBuf::from(next_os(&mut args, "--output")?)); + } + "--resume" => resume = true, + "--kind" => { + kinds.push(FixtureKind::parse(&next_string(&mut args, "--kind")?)?); + } + "--case" => cases.push(next_string(&mut args, "--case")?), + "--min-pairs" => { + min_pairs = Some(parse_usize( + &next_string(&mut args, "--min-pairs")?, + "--min-pairs", + )?); + } + "--max-pairs" => { + max_pairs = Some(parse_usize( + &next_string(&mut args, "--max-pairs")?, + "--max-pairs", + )?); + } + "--warmup-pairs" => { + warmup_pairs = Some(parse_usize( + &next_string(&mut args, "--warmup-pairs")?, + "--warmup-pairs", + )?); + } + "--bootstrap-samples" => { + bootstrap_samples = Some(parse_usize( + &next_string(&mut args, "--bootstrap-samples")?, + "--bootstrap-samples", + )?); + } + "--decision-band" => { + decision_band = Some(parse_f64( + &next_string(&mut args, "--decision-band")?, + "--decision-band", + )?); + } + "--target-ci-width" => { + target_ci_width = Some(parse_f64( + &next_string(&mut args, "--target-ci-width")?, + "--target-ci-width", + )?); + } + "--point-timeout-seconds" => { + point_timeout_seconds = parse_f64( + &next_string(&mut args, "--point-timeout-seconds")?, + "--point-timeout-seconds", + )?; + } + "--inner-iterations" => { + inner_iterations = Some(parse_usize( + &next_string(&mut args, "--inner-iterations")?, + "--inner-iterations", + )?); + } + "--control-interval-pairs" => { + control_interval_pairs = Some(parse_usize( + &next_string(&mut args, "--control-interval-pairs")?, + "--control-interval-pairs", + )?); + } + _ => return Err(format!("unknown argument '{arg}', use --help")), + } + } + + if ephemeral_output && stage.is_some_and(|stage| stage != Stage::Smoke) { + return Err("--test cannot be combined with a non-smoke --stage".into()); + } + let stage = stage.unwrap_or(Stage::Smoke); + let defaults = SamplingDefaults::for_stage(stage); + let min_pairs = min_pairs.unwrap_or(defaults.min_pairs); + let max_pairs = max_pairs.unwrap_or(defaults.max_pairs); + let warmup_pairs = warmup_pairs.unwrap_or(defaults.warmup_pairs); + let bootstrap_samples = bootstrap_samples.unwrap_or(defaults.bootstrap_samples); + let budget_seconds = budget_seconds.unwrap_or(defaults.budget_seconds); + let decision_band = decision_band.unwrap_or(defaults.decision_band); + let target_ci_width = target_ci_width.unwrap_or(defaults.target_ci_width); + let inner_iterations = inner_iterations.unwrap_or(defaults.inner_iterations); + let control_interval_pairs = + control_interval_pairs.unwrap_or(defaults.control_interval_pairs); + + if !budget_seconds.is_finite() || budget_seconds <= 0.0 { + return Err("--budget-seconds must be finite and greater than zero".into()); + } + if min_pairs == 0 || max_pairs < min_pairs { + return Err("pair counts must satisfy 0 < min-pairs <= max-pairs".into()); + } + if bootstrap_samples == 0 { + return Err("--bootstrap-samples must be greater than zero".into()); + } + if !decision_band.is_finite() || decision_band < 0.0 { + return Err("--decision-band must be finite and non-negative".into()); + } + if !target_ci_width.is_finite() || target_ci_width <= 0.0 { + return Err("--target-ci-width must be finite and greater than zero".into()); + } + if !point_timeout_seconds.is_finite() || point_timeout_seconds <= 0.0 { + return Err("--point-timeout-seconds must be finite and greater than zero".into()); + } + if inner_iterations == 0 { + return Err("--inner-iterations must be greater than zero".into()); + } + if control_interval_pairs == 0 { + return Err("--control-interval-pairs must be greater than zero".into()); + } + let mut seen = HashSet::new(); + kinds.retain(|kind| seen.insert(*kind)); + let mut seen = HashSet::new(); + cases.retain(|case| seen.insert(case.clone())); + + if stage == Stage::PolicyValidation && !kinds.is_empty() { + return Err("--kind is not supported by policy-validation; use --case".into()); + } + if stage != Stage::PolicyValidation && !cases.is_empty() { + return Err("--case is only supported by policy-validation".into()); + } + if stage == Stage::PolicyValidation + && (!min_pairs.is_multiple_of(2) + || !max_pairs.is_multiple_of(2) + || !control_interval_pairs.is_multiple_of(2)) + { + return Err( + "policy-validation pair counts and control interval must be even to preserve balanced execution-order blocks" + .into(), + ); + } + + if ephemeral_output && output.is_some() { + return Err("--test cannot be combined with --output".into()); + } + if ephemeral_output && resume { + return Err("--test cannot be combined with --resume".into()); + } + let output = output.unwrap_or_else(|| default_output(stage, ephemeral_output)); + + Ok(ParseOutcome::Run(Self { + stage, + budget: Duration::from_secs_f64(budget_seconds), + seed, + output, + resume, + kinds, + cases, + min_pairs, + max_pairs, + warmup_pairs, + bootstrap_samples, + decision_band, + target_ci_width, + point_timeout: Duration::from_secs_f64(point_timeout_seconds), + inner_iterations, + control_interval_pairs, + ephemeral_output, + })) + } + + pub(crate) fn sampling_json(&self) -> Value { + json!({ + "budget_seconds": self.budget.as_secs_f64(), + "min_pairs": self.min_pairs, + "max_pairs": self.max_pairs, + "warmup_pairs": self.warmup_pairs, + "bootstrap_samples": self.bootstrap_samples, + "decision_band": self.decision_band, + "target_ci_width": self.target_ci_width, + "point_timeout_seconds": self.point_timeout.as_secs_f64(), + "inner_iterations": self.inner_iterations, + "control_interval_pairs": self.control_interval_pairs, + "kind_filters": self.kinds.iter().map(|kind| kind.as_str()).collect::>(), + "case_filters": &self.cases, + }) + } +} + +struct SamplingDefaults { + budget_seconds: f64, + min_pairs: usize, + max_pairs: usize, + warmup_pairs: usize, + bootstrap_samples: usize, + decision_band: f64, + target_ci_width: f64, + inner_iterations: usize, + control_interval_pairs: usize, +} + +impl SamplingDefaults { + fn for_stage(stage: Stage) -> Self { + match stage { + Stage::Smoke => Self { + budget_seconds: 30.0, + min_pairs: 2, + max_pairs: 2, + warmup_pairs: 1, + bootstrap_samples: 100, + decision_band: 0.0, + target_ci_width: 0.05, + inner_iterations: 1, + control_interval_pairs: 4, + }, + Stage::Pilot => Self { + budget_seconds: 300.0, + min_pairs: 6, + max_pairs: 30, + warmup_pairs: 2, + bootstrap_samples: 400, + decision_band: 0.0, + target_ci_width: 0.05, + inner_iterations: 1, + control_interval_pairs: 4, + }, + Stage::Refinement => Self { + budget_seconds: 300.0, + min_pairs: 8, + max_pairs: 40, + warmup_pairs: 2, + bootstrap_samples: 600, + decision_band: 0.03, + target_ci_width: 0.04, + inner_iterations: 1, + control_interval_pairs: 4, + }, + Stage::PolicyValidation => Self { + budget_seconds: 300.0, + min_pairs: 20, + max_pairs: 60, + warmup_pairs: 20, + bootstrap_samples: 1_000, + decision_band: 0.03, + target_ci_width: 0.03, + inner_iterations: 16, + control_interval_pairs: 4, + }, + Stage::PageValidation => Self { + budget_seconds: 60.0, + min_pairs: 4, + max_pairs: 12, + warmup_pairs: 1, + bootstrap_samples: 200, + decision_band: 0.0, + target_ci_width: 0.05, + inner_iterations: 1, + control_interval_pairs: 4, + }, + } + } +} + +fn default_output(stage: Stage, ephemeral: bool) -> PathBuf { + let now = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_secs(); + let name = format!( + "row-selection-{}-{now}-{}.jsonl", + stage.as_str(), + std::process::id() + ); + if ephemeral { + return std::env::temp_dir().join(name); + } + let target = std::env::var_os("CARGO_TARGET_DIR") + .map(PathBuf::from) + .unwrap_or_else(|| { + PathBuf::from(env!("CARGO_MANIFEST_DIR")) + .parent() + .expect("parquet crate must be in a workspace") + .join("target") + }); + target.join("row-selection-cost-samples").join(name) +} + +fn next_os(args: &mut impl Iterator, flag: &str) -> Result { + args.next() + .ok_or_else(|| format!("{flag} requires a value")) +} + +fn next_string(args: &mut impl Iterator, flag: &str) -> Result { + next_os(args, flag)? + .into_string() + .map_err(|_| format!("{flag} requires a UTF-8 value")) +} + +fn parse_usize(value: &str, flag: &str) -> Result { + value + .parse() + .map_err(|_| format!("{flag} requires a positive integer, got '{value}'")) +} + +fn parse_u64(value: &str, flag: &str) -> Result { + value + .parse() + .map_err(|_| format!("{flag} requires an unsigned integer, got '{value}'")) +} + +fn parse_f64(value: &str, flag: &str) -> Result { + value + .parse() + .map_err(|_| format!("{flag} requires a number, got '{value}'")) +} + +fn print_help() { + println!( + r"Offline paired row-selection cost sampler + +Usage: + cargo bench -p parquet --features arrow,async \ + --bench arrow_reader_row_selection_policy_sampler -- [OPTIONS] + +Options: + --stage + Sampling stage (default: smoke) + --smoke Alias for --stage smoke + --test Ephemeral smoke run used by validation + --budget-seconds Expansion sampling wall-clock budget + --seed Deterministic manifest/order seed + --output JSONL output path + --resume Resume an existing matching JSONL file + --kind Repeatable type filter: int32, + string-view, dictionary, fixed-binary + --case Repeatable policy-validation case filter + --min-pairs Minimum measured pairs per point + (even for policy-validation) + --max-pairs Maximum measured pairs per point + (even for policy-validation) + --warmup-pairs Untimed warmup pairs per point + --bootstrap-samples Bootstrap resamples for the median CI + --decision-band Early-stop band around zero + --target-ci-width Precision early-stop threshold + --point-timeout-seconds Mark a slow point incomplete + --inner-iterations Full scans per timed observation + --control-interval-pairs Even policy-validation control cadence + -h, --help Print this help +" + ); +} diff --git a/parquet/benches/row_selection_policy_sampler/fixture.rs b/parquet/benches/row_selection_policy_sampler/fixture.rs new file mode 100644 index 000000000000..4bc4b3147731 --- /dev/null +++ b/parquet/benches/row_selection_policy_sampler/fixture.rs @@ -0,0 +1,388 @@ +// 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::collections::HashMap; +use std::sync::Arc; + +use arrow_array::builder::{FixedSizeBinaryBuilder, StringBuilder, StringViewBuilder}; +use arrow_array::types::Int32Type; +use arrow_array::{ + Array, ArrayRef, Date32Array, Decimal128Array, DictionaryArray, Float64Array, Int32Array, + Int64Array, RecordBatch, StringArray, +}; +use arrow_cast::display::array_value_to_string; +use arrow_schema::{DataType, Field, Schema}; +use bytes::Bytes; +use parquet::arrow::ArrowWriter; +use parquet::arrow::arrow_reader::{ArrowReaderMetadata, ArrowReaderOptions}; +use parquet::basic::Compression; +use parquet::file::metadata::PageIndexPolicy; +use parquet::file::properties::{EnabledStatistics, WriterProperties}; +use serde_json::{Value, json}; + +use super::model::{FixtureKind, FixtureSpec, stable_hash}; + +const FIXTURE_CACHE_CAPACITY: usize = 8; + +pub(crate) struct Fixture { + pub(crate) bytes: Bytes, + pub(crate) metadata: ArrowReaderMetadata, +} + +impl Fixture { + fn try_new(spec: &FixtureSpec) -> Result { + validate_spec(spec)?; + let schema = Arc::new(Schema::new(vec![Field::new( + "payload", + data_type(spec), + spec.nullable, + )])); + let properties = WriterProperties::builder() + .set_compression(Compression::UNCOMPRESSED) + .set_dictionary_enabled(matches!( + spec.kind, + FixtureKind::Dictionary | FixtureKind::DictStringView + )) + .set_max_row_group_row_count(Some(spec.rows)) + .set_data_page_row_count_limit(spec.page_rows) + .set_statistics_enabled(EnabledStatistics::Page) + .build(); + let batch = build_batch(spec, Arc::clone(&schema))?; + + let mut encoded = Vec::new(); + { + let mut writer = ArrowWriter::try_new(&mut encoded, schema, Some(properties)) + .map_err(|error| error.to_string())?; + for offset in (0..spec.rows).step_by(spec.page_rows) { + let len = spec.page_rows.min(spec.rows - offset); + writer + .write(&batch.slice(offset, len)) + .map_err(|error| error.to_string())?; + } + writer.close().map_err(|error| error.to_string())?; + } + + let bytes = Bytes::from(encoded); + let options = ArrowReaderOptions::new().with_page_index_policy(PageIndexPolicy::Required); + let metadata = ArrowReaderMetadata::load(&bytes, options) + .map_err(|error| format!("failed to load generated metadata: {error}"))?; + if metadata.metadata().num_row_groups() != 1 + || metadata.metadata().row_group(0).num_rows() as usize != spec.rows + { + return Err("writer did not preserve the requested single row group".into()); + } + Ok(Self { bytes, metadata }) + } + + pub(crate) fn metadata_json(&self) -> Value { + let row_group = self.metadata.metadata().row_group(0); + let column = row_group.column(0); + let offset_index = self + .metadata + .metadata() + .offset_index() + .and_then(|row_groups| row_groups.first()) + .and_then(|columns| columns.first()); + let page_rows = offset_index + .map(|index| { + let pages = index.page_locations(); + pages + .iter() + .enumerate() + .map(|(idx, page)| { + let end = pages + .get(idx + 1) + .map(|next| next.first_row_index) + .unwrap_or_else(|| row_group.num_rows()); + end - page.first_row_index + }) + .collect::>() + }) + .unwrap_or_default(); + json!({ + "arrow_type": format!("{:?}", self.metadata.schema().field(0).data_type()), + "physical_type": format!("{:?}", column.column_type()), + "encodings": column.encodings().map(|encoding| format!("{encoding:?}")).collect::>(), + "compression": format!("{:?}", column.compression()), + "num_values": column.num_values(), + "compressed_bytes": column.compressed_size(), + "uncompressed_bytes": column.uncompressed_size(), + "null_count": column.statistics().and_then(|stats| stats.null_count_opt()), + "distinct_count": column.statistics().and_then(|stats| stats.distinct_count_opt()), + "data_page_count": offset_index.map_or(0, |index| index.page_locations().len()), + "data_page_rows": page_rows, + "file_bytes": self.bytes.len(), + }) + } +} + +pub(crate) struct FixtureCache { + fixtures: HashMap>, +} + +impl FixtureCache { + pub(crate) fn new() -> Self { + Self { + fixtures: HashMap::new(), + } + } + + pub(crate) fn get(&mut self, spec: &FixtureSpec) -> Result, String> { + if let Some(fixture) = self.fixtures.get(spec) { + return Ok(Arc::clone(fixture)); + } + if self.fixtures.len() >= FIXTURE_CACHE_CAPACITY { + self.fixtures.clear(); + } + let fixture = Arc::new(Fixture::try_new(spec)?); + self.fixtures.insert(spec.clone(), Arc::clone(&fixture)); + Ok(fixture) + } +} + +pub(crate) fn logical_checksum(batches: &[RecordBatch]) -> Result { + let mut encoded = Vec::new(); + for batch in batches { + encoded.extend_from_slice(&(batch.num_rows() as u64).to_le_bytes()); + encoded.extend_from_slice(&(batch.num_columns() as u64).to_le_bytes()); + for column in batch.columns() { + encoded.extend_from_slice(format!("{:?}", column.data_type()).as_bytes()); + for row in 0..column.len() { + if column.is_null(row) { + encoded.push(0xff); + } else { + encoded.push(0x00); + let value = array_value_to_string(column.as_ref(), row) + .map_err(|error| error.to_string())?; + encoded.extend_from_slice(value.as_bytes()); + encoded.push(0x00); + } + } + } + } + Ok(stable_hash(&encoded)) +} + +fn validate_spec(spec: &FixtureSpec) -> Result<(), String> { + if spec.rows == 0 || spec.page_rows == 0 { + return Err("row and page sizes must be non-zero".into()); + } + if spec.nullable != spec.null_every.is_some() { + return Err("nullable fixtures must specify null_every, required fixtures must not".into()); + } + if spec.null_every == Some(0) { + return Err("null_every must be non-zero".into()); + } + match spec.kind { + FixtureKind::Int32 if spec.value_width != 4 => { + Err("Int32 fixtures require value_width=4".into()) + } + FixtureKind::Int64 | FixtureKind::Float64 if spec.value_width != 8 => { + Err("64-bit primitive fixtures require value_width=8".into()) + } + FixtureKind::Decimal128 if spec.value_width != 16 => { + Err("Decimal128 fixtures require value_width=16".into()) + } + FixtureKind::Date32 if spec.value_width != 4 => { + Err("Date32 fixtures require value_width=4".into()) + } + FixtureKind::Dictionary | FixtureKind::DictStringView + if spec.dictionary_cardinality == 0 => + { + Err("dictionary fixtures require a non-zero cardinality".into()) + } + FixtureKind::String + | FixtureKind::StringView + | FixtureKind::DictStringView + | FixtureKind::Dictionary + | FixtureKind::FixedBinary + if spec.value_width == 0 => + { + Err("byte-oriented fixtures require a non-zero value width".into()) + } + _ => Ok(()), + } +} + +fn data_type(spec: &FixtureSpec) -> DataType { + match spec.kind { + FixtureKind::Int32 => DataType::Int32, + FixtureKind::Int64 => DataType::Int64, + FixtureKind::Float64 => DataType::Float64, + FixtureKind::Decimal128 => DataType::Decimal128(DECIMAL128_PRECISION, DECIMAL128_SCALE), + FixtureKind::Date32 => DataType::Date32, + FixtureKind::String => DataType::Utf8, + // Same Arrow type as `StringView`; the writer property is what makes + // this one dictionary-encoded in Parquet. + FixtureKind::StringView | FixtureKind::DictStringView => DataType::Utf8View, + FixtureKind::Dictionary => { + DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)) + } + FixtureKind::FixedBinary => DataType::FixedSizeBinary(spec.value_width as i32), + } +} + +fn build_batch(spec: &FixtureSpec, schema: Arc) -> Result { + let array: ArrayRef = match spec.kind { + FixtureKind::Int32 => Arc::new(Int32Array::from_iter( + (0..spec.rows).map(|row| (!is_null(spec, row)).then_some(int32_value(row))), + )), + FixtureKind::Int64 => Arc::new(Int64Array::from_iter( + (0..spec.rows).map(|row| (!is_null(spec, row)).then_some(int64_value(row))), + )), + FixtureKind::Float64 => Arc::new(Float64Array::from_iter( + (0..spec.rows).map(|row| (!is_null(spec, row)).then_some(float64_value(row))), + )), + FixtureKind::Decimal128 => Arc::new( + Decimal128Array::from_iter( + (0..spec.rows).map(|row| (!is_null(spec, row)).then_some(decimal128_value(row))), + ) + .with_precision_and_scale(DECIMAL128_PRECISION, DECIMAL128_SCALE) + .map_err(|error| error.to_string())?, + ), + FixtureKind::String => { + let mut builder = StringBuilder::with_capacity(spec.rows, spec.rows * spec.value_width); + for row in 0..spec.rows { + if is_null(spec, row) { + builder.append_null(); + } else { + builder.append_value(string_value(row, spec.value_width, 0x51)); + } + } + Arc::new(builder.finish()) + } + FixtureKind::Date32 => Arc::new(Date32Array::from_iter( + (0..spec.rows).map(|row| (!is_null(spec, row)).then_some(date32_value(row))), + )), + FixtureKind::StringView => { + let mut builder = StringViewBuilder::with_capacity(spec.rows); + for row in 0..spec.rows { + if is_null(spec, row) { + builder.append_null(); + } else { + builder.append_value(string_value(row, spec.value_width, 0x51)); + } + } + Arc::new(builder.finish()) + } + FixtureKind::DictStringView => { + // Values are drawn from a small pool so the writer actually picks + // dictionary encoding, which is the case the planner rejects today. + let mut builder = StringViewBuilder::with_capacity(spec.rows); + for row in 0..spec.rows { + if is_null(spec, row) { + builder.append_null(); + } else { + let key = row.wrapping_mul(31) % spec.dictionary_cardinality; + builder.append_value(string_value(key, spec.value_width, 0xd1)); + } + } + Arc::new(builder.finish()) + } + FixtureKind::Dictionary => { + let keys = Int32Array::from_iter((0..spec.rows).map(|row| { + (!is_null(spec, row)) + .then_some((row.wrapping_mul(31) % spec.dictionary_cardinality) as i32) + })); + let values = StringArray::from_iter_values( + (0..spec.dictionary_cardinality) + .map(|key| string_value(key, spec.value_width, 0xd1)), + ); + Arc::new( + DictionaryArray::::try_new(keys, Arc::new(values)) + .map_err(|error| error.to_string())?, + ) + } + FixtureKind::FixedBinary => { + let mut builder = + FixedSizeBinaryBuilder::with_capacity(spec.rows, spec.value_width as i32); + for row in 0..spec.rows { + if is_null(spec, row) { + builder.append_null(); + } else { + builder + .append_value(binary_value(row, spec.value_width)) + .map_err(|error| error.to_string())?; + } + } + Arc::new(builder.finish()) + } + }; + RecordBatch::try_new(schema, vec![array]).map_err(|error| error.to_string()) +} + +fn is_null(spec: &FixtureSpec, row: usize) -> bool { + spec.null_every + .is_some_and(|every| row.is_multiple_of(every)) +} + +fn int32_value(row: usize) -> i32 { + row.wrapping_mul(31).wrapping_add(17) as i32 +} + +/// Matches the TPC-DS money columns (`Decimal128(7, 2)`), which are the widest +/// unmodeled slice of that schema. +const DECIMAL128_PRECISION: u8 = 7; +const DECIMAL128_SCALE: i8 = 2; + +/// Stays inside `DECIMAL128_PRECISION`: the modulus is the largest prime below +/// `10^7`, so every value has at most seven decimal digits. +fn decimal128_value(row: usize) -> i128 { + (int64_value(row) % 9_999_991) as i128 +} + +/// Days since the epoch, kept inside a plausible calendar range so the values +/// look like the date dimensions the model is meant to cover. +fn date32_value(row: usize) -> i32 { + (int32_value(row).rem_euclid(36_524)) + 10_957 +} + +fn int64_value(row: usize) -> i64 { + (row as u64) + .wrapping_mul(0x9e37_79b9_7f4a_7c15) + .wrapping_add(17) as i64 +} + +/// Finite by construction. A NaN would decode correctly but compare unequal to +/// itself, which would fail the sampler's paired value check. +fn float64_value(row: usize) -> f64 { + (int64_value(row) % 1_000_003) as f64 / 3.0 +} + +fn string_value(row: usize, width: usize, salt: u8) -> String { + let bytes = binary_value(row ^ usize::from(salt), width); + bytes + .into_iter() + .map(|byte| b'a' + byte % 26) + .map(char::from) + .collect() +} + +fn binary_value(row: usize, width: usize) -> Vec { + let mut value = Vec::with_capacity(width); + let mut state = (row as u64) ^ 0x9e37_79b9_7f4a_7c15; + while value.len() < width { + state ^= state >> 30; + state = state.wrapping_mul(0xbf58_476d_1ce4_e5b9); + state ^= state >> 27; + state = state.wrapping_mul(0x94d0_49bb_1331_11eb); + state ^= state >> 31; + value.extend_from_slice(&state.to_le_bytes()); + } + value.truncate(width); + value +} diff --git a/parquet/benches/row_selection_policy_sampler/mod.rs b/parquet/benches/row_selection_policy_sampler/mod.rs new file mode 100644 index 000000000000..dc801575d3db --- /dev/null +++ b/parquet/benches/row_selection_policy_sampler/mod.rs @@ -0,0 +1,189 @@ +// 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. + +mod cli; +mod fixture; +mod model; +mod output; +mod policy_validation; +mod sampling; + +use std::time::Instant; + +use serde_json::{Value, json}; + +use self::cli::{Cli, ParseOutcome}; +use self::fixture::FixtureCache; +use self::model::{ExperimentManifest, Stage}; +use self::output::JsonlOutput; +use self::sampling::{SampledPoint, sample_point}; + +pub(crate) fn run() -> Result<(), String> { + let ParseOutcome::Run(cli) = Cli::parse()? else { + return Ok(()); + }; + if cli.stage == Stage::PolicyValidation { + return policy_validation::run(cli); + } + let manifest = ExperimentManifest::generate(cli.stage, &cli.kinds, cli.seed)?; + let mut output = JsonlOutput::open( + &cli, + &manifest.id, + manifest.experiments.len(), + manifest.mandatory_count, + )?; + println!( + "row-selection sampler: stage={}, manifest={}, experiments={}, output={}", + cli.stage, + manifest.id, + manifest.experiments.len(), + cli.output.display() + ); + if output.resumed_records != 0 { + println!( + "resuming after {} completed/unsupported experiment records", + output.resumed_records + ); + } + + let started = Instant::now(); + let mut fixtures = FixtureCache::new(); + let control_experiment = manifest + .experiments + .first() + .expect("a non-empty manifest was validated"); + let control_fixture = fixtures.get(&control_experiment.fixture).map_err(|error| { + format!( + "failed to build mandatory control fixture {}: {error}", + control_experiment.id + ) + })?; + let start_control = sample_point(&cli, control_experiment, &control_fixture, true)?; + output.write(&control_record(&start_control, control_experiment, "start"))?; + + let mut completed_this_run = 0usize; + let mut unsupported_this_run = 0usize; + let mut incomplete_this_run = 0usize; + let mut skipped_resume = 0usize; + let mut budget_exhausted = false; + + for (idx, experiment) in manifest.experiments.iter().enumerate() { + if output.is_completed(&experiment.id) { + skipped_resume += 1; + continue; + } + if started.elapsed() >= cli.budget && !experiment.mandatory { + budget_exhausted = true; + break; + } + println!( + "[{}/{}] {} {} / {} / {}", + idx + 1, + manifest.experiments.len(), + experiment.id, + experiment.fixture.kind, + experiment.selection.name, + experiment.batch_size + ); + + let fixture = match fixtures.get(&experiment.fixture) { + Ok(fixture) => fixture, + Err(error) if !experiment.mandatory => { + output.write(&json!({ + "record_type": "experiment", + "status": "unsupported", + "experiment": experiment.to_json(), + "sampling": cli.sampling_json(), + "reason": error, + }))?; + output.mark_completed(&experiment.id); + unsupported_this_run += 1; + continue; + } + Err(error) => { + return Err(format!( + "mandatory experiment {} is unsupported: {error}", + experiment.id + )); + } + }; + let point = sample_point(&cli, experiment, &fixture, false)?; + output.write(&point.to_json(experiment, "experiment"))?; + if point.is_complete() { + output.mark_completed(&experiment.id); + completed_this_run += 1; + } else { + incomplete_this_run += 1; + } + } + + let end_control = sample_point(&cli, control_experiment, &control_fixture, true)?; + output.write(&control_record(&end_control, control_experiment, "end"))?; + let control_drift = end_control.baseline_ns() / start_control.baseline_ns() - 1.0; + let remaining = manifest + .experiments + .iter() + .filter(|experiment| !output.is_completed(&experiment.id)) + .count(); + output.write(&json!({ + "record_type": "run_end", + "stage": cli.stage.as_str(), + "manifest_id": manifest.id, + "elapsed_seconds": started.elapsed().as_secs_f64(), + "completed_this_run": completed_this_run, + "unsupported_this_run": unsupported_this_run, + "incomplete_this_run": incomplete_this_run, + "skipped_from_resume": skipped_resume, + "remaining_experiments": remaining, + "budget_exhausted": budget_exhausted, + "control_baseline_drift": control_drift, + "control_drift_warning": control_drift.abs() > 0.10, + "sampling": cli.sampling_json(), + }))?; + + println!( + "sampler finished in {:.3}s: completed={}, unsupported={}, incomplete={}, remaining={}, control drift={:+.2}%", + started.elapsed().as_secs_f64(), + completed_this_run, + unsupported_this_run, + incomplete_this_run, + remaining, + control_drift * 100.0 + ); + + let ephemeral_output = cli.ephemeral_output; + let output_path = cli.output.clone(); + drop(output); + if ephemeral_output { + std::fs::remove_file(&output_path).map_err(|error| { + format!( + "failed to remove ephemeral output {}: {error}", + output_path.display() + ) + })?; + } + Ok(()) +} + +fn control_record(point: &SampledPoint, experiment: &model::Experiment, position: &str) -> Value { + let mut value = point.to_json(experiment, "control"); + value + .as_object_mut() + .expect("sample records are JSON objects") + .insert("position".into(), json!(position)); + value +} diff --git a/parquet/benches/row_selection_policy_sampler/model.rs b/parquet/benches/row_selection_policy_sampler/model.rs new file mode 100644 index 000000000000..3b714d1ba48b --- /dev/null +++ b/parquet/benches/row_selection_policy_sampler/model.rs @@ -0,0 +1,799 @@ +// 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::collections::HashSet; +use std::fmt::{Display, Formatter}; + +use arrow_buffer::BooleanBuffer; +use parquet::arrow::arrow_reader::{RowSelection, RowSelector}; +use rand::{RngExt, SeedableRng, rngs::StdRng}; +use serde_json::{Value, json}; + +pub(crate) const OUTPUT_SCHEMA_VERSION: u64 = 1; +const MAX_PILOT_EXPANSION_POINTS: usize = 192; +const MAX_REFINEMENT_EXPANSION_POINTS: usize = 256; + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(crate) enum Stage { + Smoke, + Pilot, + Refinement, + PolicyValidation, + PageValidation, +} + +impl Stage { + pub(crate) fn parse(value: &str) -> Result { + match value { + "smoke" => Ok(Self::Smoke), + "pilot" => Ok(Self::Pilot), + "refinement" => Ok(Self::Refinement), + "policy-validation" => Ok(Self::PolicyValidation), + "page-validation" => Ok(Self::PageValidation), + _ => Err(format!( + "unknown stage '{value}', expected smoke, pilot, refinement, policy-validation, or page-validation" + )), + } + } + + pub(crate) fn as_str(self) -> &'static str { + match self { + Self::Smoke => "smoke", + Self::Pilot => "pilot", + Self::Refinement => "refinement", + Self::PolicyValidation => "policy-validation", + Self::PageValidation => "page-validation", + } + } +} + +impl Display for Stage { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + f.write_str(self.as_str()) + } +} + +#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)] +pub(crate) enum FixtureKind { + Int32, + Int64, + Float64, + Decimal128, + Date32, + String, + StringView, + DictStringView, + Dictionary, + FixedBinary, +} + +impl FixtureKind { + pub(crate) const ALL: [Self; 10] = [ + Self::Int32, + Self::Int64, + Self::Float64, + Self::Decimal128, + Self::Date32, + Self::String, + Self::StringView, + Self::DictStringView, + Self::Dictionary, + Self::FixedBinary, + ]; + + pub(crate) fn parse(value: &str) -> Result { + match value { + "int32" => Ok(Self::Int32), + "int64" => Ok(Self::Int64), + "float64" => Ok(Self::Float64), + "decimal128" => Ok(Self::Decimal128), + "date32" => Ok(Self::Date32), + "string" => Ok(Self::String), + "string-view" => Ok(Self::StringView), + "dict-string-view" => Ok(Self::DictStringView), + "dictionary" => Ok(Self::Dictionary), + "fixed-binary" => Ok(Self::FixedBinary), + _ => Err(format!( + "unknown kind '{value}', expected int32, int64, float64, decimal128, \ + date32, string, string-view, dict-string-view, dictionary, or \ + fixed-binary" + )), + } + } + + pub(crate) fn as_str(self) -> &'static str { + match self { + Self::Int32 => "int32", + Self::Int64 => "int64", + Self::Float64 => "float64", + Self::Decimal128 => "decimal128", + Self::Date32 => "date32", + Self::String => "string", + Self::StringView => "string-view", + Self::DictStringView => "dict-string-view", + Self::Dictionary => "dictionary", + Self::FixedBinary => "fixed-binary", + } + } +} + +impl Display for FixtureKind { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + f.write_str(self.as_str()) + } +} + +#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)] +pub(crate) enum SelectionBacking { + Selectors, + Mask, +} + +impl SelectionBacking { + fn as_str(self) -> &'static str { + match self { + Self::Selectors => "selectors", + Self::Mask => "mask", + } + } +} + +#[derive(Clone, Debug, Eq, Hash, PartialEq)] +pub(crate) struct FixtureSpec { + pub(crate) kind: FixtureKind, + pub(crate) rows: usize, + pub(crate) nullable: bool, + pub(crate) null_every: Option, + pub(crate) value_width: usize, + pub(crate) dictionary_cardinality: usize, + pub(crate) page_rows: usize, +} + +impl FixtureSpec { + fn default_for(kind: FixtureKind, nullable: bool) -> Self { + let (value_width, dictionary_cardinality) = match kind { + FixtureKind::Int32 => (4, 0), + FixtureKind::Int64 | FixtureKind::Float64 => (8, 0), + FixtureKind::Decimal128 => (16, 0), + FixtureKind::Date32 => (4, 0), + FixtureKind::String | FixtureKind::StringView => (64, 0), + FixtureKind::DictStringView => (32, 256), + FixtureKind::Dictionary => (32, 256), + FixtureKind::FixedBinary => (32, 0), + }; + Self { + kind, + rows: 16_384, + nullable, + null_every: nullable.then_some(4), + value_width, + dictionary_cardinality, + page_rows: 512, + } + } + + pub(crate) fn canonical(&self) -> String { + format!( + "kind={};rows={};nullable={};null_every={};width={};cardinality={};page_rows={}", + self.kind, + self.rows, + self.nullable, + self.null_every.unwrap_or(0), + self.value_width, + self.dictionary_cardinality, + self.page_rows + ) + } + + pub(crate) fn to_json(&self) -> Value { + json!({ + "kind": self.kind.as_str(), + "rows": self.rows, + "nullable": self.nullable, + "null_every": self.null_every, + "value_width": self.value_width, + "dictionary_cardinality": self.dictionary_cardinality, + "page_rows": self.page_rows, + "compression": "UNCOMPRESSED", + }) + } +} + +#[derive(Clone, Debug, Eq, Hash, PartialEq)] +pub(crate) struct SelectionSpec { + pub(crate) name: &'static str, + pub(crate) skip_run: usize, + pub(crate) select_run: usize, + pub(crate) offset: usize, + pub(crate) backing: SelectionBacking, +} + +impl SelectionSpec { + fn canonical(&self) -> String { + format!( + "selection={};skip={};select={};offset={};backing={}", + self.name, + self.skip_run, + self.select_run, + self.offset, + self.backing.as_str() + ) + } + + pub(crate) fn to_json(&self) -> Value { + json!({ + "name": self.name, + "skip_run": self.skip_run, + "select_run": self.select_run, + "offset": self.offset, + "source_backing": self.backing.as_str(), + }) + } + + pub(crate) fn materialize(&self, rows: usize, batch_size: usize) -> SelectionMaterialization { + let mut bits = vec![false; rows]; + let mut cursor = self.offset.min(rows); + while cursor < rows { + cursor = cursor.saturating_add(self.skip_run).min(rows); + let end = cursor.saturating_add(self.select_run).min(rows); + bits[cursor..end].fill(true); + cursor = end; + } + + let selectors = selectors_from_bits(&bits); + let selection = match self.backing { + SelectionBacking::Selectors => RowSelection::from(selectors), + SelectionBacking::Mask => { + RowSelection::from_boolean_buffer(BooleanBuffer::from(bits.clone())) + } + }; + let stats = SelectionStats::new(&bits, batch_size); + SelectionMaterialization { selection, stats } + } +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(crate) enum ExecutionMode { + SyncOracle, + PageValidation, +} + +impl ExecutionMode { + fn as_str(self) -> &'static str { + match self { + Self::SyncOracle => "sync-oracle", + Self::PageValidation => "page-validation", + } + } +} + +#[derive(Clone, Debug)] +pub(crate) struct Experiment { + pub(crate) id: String, + pub(crate) fixture: FixtureSpec, + pub(crate) selection: SelectionSpec, + pub(crate) batch_size: usize, + pub(crate) mode: ExecutionMode, + pub(crate) mandatory: bool, +} + +impl Experiment { + fn new( + fixture: FixtureSpec, + selection: SelectionSpec, + batch_size: usize, + mode: ExecutionMode, + mandatory: bool, + ) -> Self { + let canonical = format!( + "{};{};batch_size={};mode={}", + fixture.canonical(), + selection.canonical(), + batch_size, + mode.as_str() + ); + Self { + id: format!("{:016x}", stable_hash(canonical.as_bytes())), + fixture, + selection, + batch_size, + mode, + mandatory, + } + } + + pub(crate) fn to_json(&self) -> Value { + json!({ + "experiment_id": self.id, + "mandatory": self.mandatory, + "execution_mode": self.mode.as_str(), + "batch_size": self.batch_size, + "fixture": self.fixture.to_json(), + "selection": self.selection.to_json(), + }) + } +} + +pub(crate) struct ExperimentManifest { + pub(crate) id: String, + pub(crate) experiments: Vec, + pub(crate) mandatory_count: usize, +} + +impl ExperimentManifest { + pub(crate) fn generate(stage: Stage, kinds: &[FixtureKind], seed: u64) -> Result { + let kinds = if kinds.is_empty() { + FixtureKind::ALL.to_vec() + } else { + kinds.to_vec() + }; + let mut experiments = match stage { + Stage::Smoke => smoke_experiments(&kinds), + Stage::Pilot => pilot_experiments(&kinds, seed), + Stage::Refinement => refinement_experiments(&kinds, seed), + Stage::PolicyValidation => { + return Err("policy-validation uses its dedicated heterogeneous manifest".into()); + } + Stage::PageValidation => page_validation_experiments(&kinds), + }; + if experiments.is_empty() { + return Err("the selected stage and kind filters produced no experiments".into()); + } + + let mut seen = HashSet::new(); + experiments.retain(|experiment| seen.insert(experiment.id.clone())); + let mandatory_count = experiments + .iter() + .filter(|experiment| experiment.mandatory) + .count(); + let mut manifest_key = format!( + "schema={OUTPUT_SCHEMA_VERSION};stage={};seed={seed}", + stage.as_str() + ); + for experiment in &experiments { + manifest_key.push(';'); + manifest_key.push_str(&experiment.id); + } + Ok(Self { + id: format!("{:016x}", stable_hash(manifest_key.as_bytes())), + experiments, + mandatory_count, + }) + } +} + +pub(crate) struct SelectionMaterialization { + pub(crate) selection: RowSelection, + pub(crate) stats: SelectionStats, +} + +#[derive(Debug)] +pub(crate) struct SelectionStats { + total_rows: usize, + selected_rows: usize, + selected_runs: usize, + skipped_runs: usize, + transitions: usize, + first_selected: Option, + last_selected: Option, + selected_span_rows: usize, + mask_decode_rows_without_page_pruning: usize, + output_batches: usize, + mean_selected_run: f64, + mean_skipped_run: f64, + max_selected_run: usize, + max_skipped_run: usize, +} + +impl SelectionStats { + fn new(bits: &[bool], batch_size: usize) -> Self { + let selected_positions = bits + .iter() + .enumerate() + .filter_map(|(idx, selected)| selected.then_some(idx)) + .collect::>(); + let selected_rows = selected_positions.len(); + let first_selected = selected_positions.first().copied(); + let last_selected = selected_positions.last().copied(); + let selected_span_rows = first_selected + .zip(last_selected) + .map(|(first, last)| last - first + 1) + .unwrap_or(0); + let mask_decode_rows_without_page_pruning = selected_positions + .chunks(batch_size) + .map(|chunk| chunk[chunk.len() - 1] - chunk[0] + 1) + .sum(); + + let execution_bits = last_selected.map_or(&bits[..0], |last| &bits[..=last]); + let mut selected_lengths = Vec::new(); + let mut skipped_lengths = Vec::new(); + for selector in selectors_from_bits(execution_bits) { + if selector.skip { + skipped_lengths.push(selector.row_count); + } else { + selected_lengths.push(selector.row_count); + } + } + let selected_runs = selected_lengths.len(); + let skipped_runs = skipped_lengths.len(); + let transitions = selected_runs.saturating_add(skipped_runs).saturating_sub(1); + + Self { + total_rows: bits.len(), + selected_rows, + selected_runs, + skipped_runs, + transitions, + first_selected, + last_selected, + selected_span_rows, + mask_decode_rows_without_page_pruning, + output_batches: selected_rows.div_ceil(batch_size), + mean_selected_run: mean(&selected_lengths), + mean_skipped_run: mean(&skipped_lengths), + max_selected_run: selected_lengths.into_iter().max().unwrap_or(0), + max_skipped_run: skipped_lengths.into_iter().max().unwrap_or(0), + } + } + + pub(crate) fn selected_rows(&self) -> usize { + self.selected_rows + } + + pub(crate) fn to_json(&self) -> Value { + json!({ + "total_rows": self.total_rows, + "selected_rows": self.selected_rows, + "selectivity": self.selected_rows as f64 / self.total_rows as f64, + "selected_runs": self.selected_runs, + "skipped_runs": self.skipped_runs, + "transitions": self.transitions, + "first_selected": self.first_selected, + "last_selected": self.last_selected, + "selected_span_rows": self.selected_span_rows, + "mask_decode_rows_without_page_pruning": self.mask_decode_rows_without_page_pruning, + "output_batches": self.output_batches, + "mean_selected_run": self.mean_selected_run, + "mean_skipped_run": self.mean_skipped_run, + "max_selected_run": self.max_selected_run, + "max_skipped_run": self.max_skipped_run, + }) + } +} + +fn selectors_from_bits(bits: &[bool]) -> Vec { + let mut selectors = Vec::new(); + let Some((&first, tail)) = bits.split_first() else { + return selectors; + }; + let mut selected = first; + let mut rows = 1usize; + for &next in tail { + if next == selected { + rows += 1; + } else { + selectors.push(if selected { + RowSelector::select(rows) + } else { + RowSelector::skip(rows) + }); + selected = next; + rows = 1; + } + } + selectors.push(if selected { + RowSelector::select(rows) + } else { + RowSelector::skip(rows) + }); + selectors +} + +fn mean(values: &[usize]) -> f64 { + if values.is_empty() { + 0.0 + } else { + values.iter().sum::() as f64 / values.len() as f64 + } +} + +fn selection_anchors() -> Vec { + [ + ("sparse", 63, 1, 0), + ("fragmented", 1, 1, 0), + ("clustered", 128, 128, 17), + ("dense", 1, 63, 0), + ] + .into_iter() + .flat_map(|(name, skip_run, select_run, offset)| { + [SelectionBacking::Selectors, SelectionBacking::Mask].map(|backing| SelectionSpec { + name, + skip_run, + select_run, + offset, + backing, + }) + }) + .collect() +} + +fn smoke_experiments(kinds: &[FixtureKind]) -> Vec { + let selections = selection_anchors(); + kinds + .iter() + .copied() + .enumerate() + .map(|(idx, kind)| { + let nullable = idx % 2 == 1; + let mut fixture = FixtureSpec::default_for(kind, nullable); + fixture.rows = 8_192; + fixture.page_rows = 256; + Experiment::new( + fixture, + selections[idx % selections.len()].clone(), + 1_024, + ExecutionMode::SyncOracle, + true, + ) + }) + .collect() +} + +fn pilot_experiments(kinds: &[FixtureKind], seed: u64) -> Vec { + let selections = selection_anchors(); + let mut mandatory = Vec::new(); + for &kind in kinds { + for (idx, selection) in [ + SelectionSpec { + name: "sparse", + skip_run: 63, + select_run: 1, + offset: 0, + backing: SelectionBacking::Selectors, + }, + SelectionSpec { + name: "fragmented", + skip_run: 1, + select_run: 1, + offset: 0, + backing: SelectionBacking::Mask, + }, + ] + .into_iter() + .enumerate() + { + let fixture = FixtureSpec::default_for(kind, idx == 1); + mandatory.push(Experiment::new( + fixture, + selection, + if idx == 0 { 1_024 } else { 8_192 }, + ExecutionMode::SyncOracle, + true, + )); + } + } + + let mut expansion = Vec::new(); + for fixture in fixture_variants(kinds) { + for selection in &selections { + for batch_size in [256, 1_024, 8_192] { + if batch_size <= fixture.rows { + expansion.push(Experiment::new( + fixture.clone(), + selection.clone(), + batch_size, + ExecutionMode::SyncOracle, + false, + )); + } + } + } + } + seeded_shuffle(&mut expansion, seed ^ 0x51ec_7100_c057_0001); + expansion.truncate(MAX_PILOT_EXPANSION_POINTS); + mandatory.extend(expansion); + mandatory +} + +fn fixture_variants(kinds: &[FixtureKind]) -> Vec { + let mut fixtures = Vec::new(); + for &kind in kinds { + let widths: &[usize] = match kind { + FixtureKind::Int32 => &[4], + FixtureKind::Int64 | FixtureKind::Float64 => &[8], + FixtureKind::Decimal128 => &[16], + FixtureKind::Date32 => &[4], + FixtureKind::String | FixtureKind::StringView => &[16, 64, 256], + FixtureKind::DictStringView => &[16, 64], + FixtureKind::Dictionary => &[16, 64], + FixtureKind::FixedBinary => &[8, 32], + }; + let cardinalities: &[usize] = match kind { + FixtureKind::Dictionary | FixtureKind::DictStringView => &[16, 1_024], + _ => &[0], + }; + for &rows in &[16_384, 65_536] { + for &page_rows in &[256, 2_048] { + for &nullable in &[false, true] { + for &value_width in widths { + for &dictionary_cardinality in cardinalities { + fixtures.push(FixtureSpec { + kind, + rows, + nullable, + null_every: nullable.then_some(4), + value_width, + dictionary_cardinality, + page_rows, + }); + } + } + } + } + } + } + fixtures +} + +fn refinement_experiments(kinds: &[FixtureKind], seed: u64) -> Vec { + let shapes = refinement_shapes(); + let mut mandatory = Vec::with_capacity(kinds.len() * shapes.len()); + for (kind_idx, &kind) in kinds.iter().enumerate() { + for (shape_idx, &(name, skip_run, select_run)) in shapes.iter().enumerate() { + let backing = if (kind_idx + shape_idx).is_multiple_of(2) { + SelectionBacking::Selectors + } else { + SelectionBacking::Mask + }; + mandatory.push(Experiment::new( + FixtureSpec::default_for(kind, false), + SelectionSpec { + name, + skip_run, + select_run, + offset: 0, + backing, + }, + 1_024, + ExecutionMode::SyncOracle, + true, + )); + } + } + + let selections = shapes + .into_iter() + .flat_map(|(name, skip_run, select_run)| { + [SelectionBacking::Selectors, SelectionBacking::Mask].map(|backing| SelectionSpec { + name, + skip_run, + select_run, + offset: 0, + backing, + }) + }) + .collect::>(); + let mut expansion = Vec::new(); + for fixture in fixture_variants(kinds) { + for selection in &selections { + for batch_size in [256, 1_024, 8_192] { + if batch_size <= fixture.rows { + expansion.push(Experiment::new( + fixture.clone(), + selection.clone(), + batch_size, + ExecutionMode::SyncOracle, + false, + )); + } + } + } + } + seeded_shuffle(&mut expansion, seed ^ 0x7ef1_6e00_c057_0001); + expansion.truncate(MAX_REFINEMENT_EXPANSION_POINTS); + mandatory.extend(expansion); + mandatory +} + +/// Selection shapes around and between the coarse Pilot anchors. These vary +/// transition density independently from selectivity so the measured data, +/// rather than the legacy threshold, locates the decision boundary. +fn refinement_shapes() -> Vec<(&'static str, usize, usize)> { + let mut shapes = Vec::new(); + shapes.extend( + // Dense through 16: every sampled Arrow family crosses between 2 and + // 16, so unit resolution there is what separates one type's threshold + // from another's. Beyond 16 the choice is uniformly Selectors and a + // coarse grid is enough to confirm it. + [ + 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 18, 20, 24, 28, 32, 40, 48, 64, + ] + .into_iter() + .map(|run| ("balanced-grid", run, run)), + ); + shapes.extend( + [15, 31, 47, 63, 95] + .into_iter() + .map(|skip_run| ("sparse-grid", skip_run, 1)), + ); + shapes.extend( + [15, 31, 47, 63, 95] + .into_iter() + .map(|select_run| ("dense-grid", 1, select_run)), + ); + shapes.extend( + [4, 8, 12, 16, 24, 32] + .into_iter() + .map(|run| ("quarter-grid", run * 3, run)), + ); + shapes.extend( + [4, 8, 12, 16, 24, 32] + .into_iter() + .map(|run| ("three-quarter-grid", run, run * 3)), + ); + shapes +} + +fn page_validation_experiments(kinds: &[FixtureKind]) -> Vec { + if !kinds.contains(&FixtureKind::Int32) { + return Vec::new(); + } + let fixture = FixtureSpec { + kind: FixtureKind::Int32, + rows: 8_192, + nullable: false, + null_every: None, + value_width: 4, + dictionary_cardinality: 0, + page_rows: 128, + }; + [SelectionBacking::Selectors, SelectionBacking::Mask] + .into_iter() + .map(|backing| { + Experiment::new( + fixture.clone(), + SelectionSpec { + name: "page-sparse", + skip_run: 511, + select_run: 1, + offset: 0, + backing, + }, + 512, + ExecutionMode::PageValidation, + true, + ) + }) + .collect() +} + +fn seeded_shuffle(values: &mut [T], seed: u64) { + let mut rng = StdRng::seed_from_u64(seed); + for idx in (1..values.len()).rev() { + let swap_idx = rng.random_range(0..=idx); + values.swap(idx, swap_idx); + } +} + +pub(crate) fn stable_hash(bytes: &[u8]) -> u64 { + let mut hash = 0xcbf2_9ce4_8422_2325u64; + for byte in bytes { + hash ^= u64::from(*byte); + hash = hash.wrapping_mul(0x0000_0100_0000_01b3); + } + hash +} diff --git a/parquet/benches/row_selection_policy_sampler/output.rs b/parquet/benches/row_selection_policy_sampler/output.rs new file mode 100644 index 000000000000..1159472c1272 --- /dev/null +++ b/parquet/benches/row_selection_policy_sampler/output.rs @@ -0,0 +1,306 @@ +// 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::collections::HashSet; +use std::fs::{File, OpenOptions}; +use std::io::{BufRead, BufReader, BufWriter, Write}; +use std::process::Command; +use std::time::{SystemTime, UNIX_EPOCH}; + +use serde_json::{Value, json}; +use sysinfo::System; + +use super::cli::Cli; +use super::model::{OUTPUT_SCHEMA_VERSION, stable_hash}; + +pub(crate) struct JsonlOutput { + writer: BufWriter, + completed: HashSet, + validation_warnings: HashSet, + validation_inconclusive: HashSet, + pub(crate) resumed_records: usize, +} + +impl JsonlOutput { + pub(crate) fn open( + cli: &Cli, + manifest_id: &str, + experiment_count: usize, + mandatory_count: usize, + ) -> Result { + let machine = machine_info(); + if let Some(parent) = cli.output.parent() + && !parent.as_os_str().is_empty() + { + std::fs::create_dir_all(parent).map_err(|error| { + format!( + "failed to create output directory {}: {error}", + parent.display() + ) + })?; + } + + if cli.resume { + let state = read_resume_state(cli, manifest_id, &machine.signature)?; + let resumed_records = state.completed.len(); + let file = OpenOptions::new() + .append(true) + .open(&cli.output) + .map_err(|error| format!("failed to append {}: {error}", cli.output.display()))?; + return Ok(Self { + writer: BufWriter::new(file), + completed: state.completed, + validation_warnings: state.validation_warnings, + validation_inconclusive: state.validation_inconclusive, + resumed_records, + }); + } + + let file = OpenOptions::new() + .write(true) + .create_new(true) + .open(&cli.output) + .map_err(|error| { + format!( + "failed to create {}: {error}; use --resume for an existing file", + cli.output.display() + ) + })?; + let mut output = Self { + writer: BufWriter::new(file), + completed: HashSet::new(), + validation_warnings: HashSet::new(), + validation_inconclusive: HashSet::new(), + resumed_records: 0, + }; + output.write(&json!({ + "record_type": "manifest", + "schema_version": OUTPUT_SCHEMA_VERSION, + "manifest_id": manifest_id, + "stage": cli.stage.as_str(), + "seed": cli.seed, + "experiment_count": experiment_count, + "mandatory_count": mandatory_count, + "sampling": cli.sampling_json(), + "machine_signature": machine.signature, + "machine": machine.details, + "crate_version": env!("CARGO_PKG_VERSION"), + "created_unix_seconds": unix_seconds(), + "git": git_info(), + }))?; + Ok(output) + } + + pub(crate) fn is_completed(&self, experiment_id: &str) -> bool { + self.completed.contains(experiment_id) + } + + pub(crate) fn mark_completed(&mut self, experiment_id: &str) { + self.completed.insert(experiment_id.to_string()); + } + + pub(crate) fn mark_validation_warning(&mut self, experiment_id: &str) { + self.validation_warnings.insert(experiment_id.to_string()); + } + + pub(crate) fn validation_warning_count(&self) -> usize { + self.validation_warnings.len() + } + + pub(crate) fn mark_validation_inconclusive(&mut self, experiment_id: &str) { + self.validation_inconclusive + .insert(experiment_id.to_string()); + } + + pub(crate) fn validation_inconclusive_count(&self) -> usize { + self.validation_inconclusive.len() + } + + pub(crate) fn write(&mut self, value: &Value) -> Result<(), String> { + serde_json::to_writer(&mut self.writer, value).map_err(|error| error.to_string())?; + self.writer.write_all(b"\n").map_err(to_string)?; + self.writer.flush().map_err(to_string) + } +} + +#[derive(Default)] +struct ResumeState { + completed: HashSet, + validation_warnings: HashSet, + validation_inconclusive: HashSet, +} + +struct MachineInfo { + signature: String, + details: Value, +} + +fn machine_info() -> MachineInfo { + let mut system = System::new(); + system.refresh_cpu_all(); + let cpu = system.cpus().first(); + let os = System::long_os_version().or_else(System::name); + let kernel = System::kernel_version(); + let cpu_brand = cpu.map(|cpu| cpu.brand().to_string()); + let logical_cpus = system.cpus().len(); + let signature_source = format!( + "arch={};os={os:?};kernel={kernel:?};cpu={cpu_brand:?};logical_cpus={logical_cpus}", + std::env::consts::ARCH + ); + MachineInfo { + signature: format!("{:016x}", stable_hash(signature_source.as_bytes())), + details: json!({ + "arch": std::env::consts::ARCH, + "os_family": std::env::consts::OS, + "os_version": os, + "kernel_version": kernel, + "cpu_brand": cpu_brand, + "cpu_frequency_mhz_at_start": cpu.map(|cpu| cpu.frequency()), + "logical_cpus": logical_cpus, + "available_parallelism": std::thread::available_parallelism().ok().map(|value| value.get()), + "rustc": command_output("rustc", &["--version"]), + "power_state": null, + }), + } +} + +fn git_info() -> Value { + let manifest_dir = env!("CARGO_MANIFEST_DIR"); + let commit = command_output("git", &["-C", manifest_dir, "rev-parse", "HEAD"]); + let dirty = Command::new("git") + .args(["-C", manifest_dir, "status", "--porcelain"]) + .output() + .ok() + .filter(|output| output.status.success()) + .is_some_and(|output| !output.stdout.is_empty()); + json!({"commit": commit, "dirty": dirty}) +} + +fn command_output(command: &str, args: &[&str]) -> Option { + let output = Command::new(command).args(args).output().ok()?; + output + .status + .success() + .then(|| String::from_utf8_lossy(&output.stdout).trim().to_string()) +} + +fn read_resume_state( + cli: &Cli, + manifest_id: &str, + machine_signature: &str, +) -> Result { + let file = File::open(&cli.output) + .map_err(|error| format!("failed to read {}: {error}", cli.output.display()))?; + let mut state = ResumeState::default(); + let mut header_seen = false; + for (line_idx, line) in BufReader::new(file).lines().enumerate() { + let line = line.map_err(to_string)?; + if line.trim().is_empty() { + continue; + } + let value: Value = serde_json::from_str(&line) + .map_err(|error| format!("invalid JSONL at line {}: {error}", line_idx + 1))?; + match value.get("record_type").and_then(Value::as_str) { + Some("manifest") if !header_seen => { + validate_header(cli, manifest_id, machine_signature, &value)?; + header_seen = true; + } + Some("experiment") => { + let status = value.get("status").and_then(Value::as_str); + if matches!(status, Some("complete" | "unsupported")) + && let Some(id) = value + .pointer("/experiment/experiment_id") + .and_then(Value::as_str) + { + state.completed.insert(id.to_string()); + if value + .pointer("/summary/stability_warning") + .and_then(Value::as_bool) + .unwrap_or(false) + { + state.validation_warnings.insert(id.to_string()); + } + if value + .pointer("/summary/practical_decision") + .and_then(Value::as_str) + == Some("inconclusive") + { + state.validation_inconclusive.insert(id.to_string()); + } + } + } + _ => {} + } + } + if !header_seen { + return Err(format!( + "{} does not contain a sampler manifest header", + cli.output.display() + )); + } + Ok(state) +} + +fn validate_header( + cli: &Cli, + manifest_id: &str, + machine_signature: &str, + header: &Value, +) -> Result<(), String> { + let checks = [ + ( + "schema_version", + header.get("schema_version").cloned(), + json!(OUTPUT_SCHEMA_VERSION), + ), + ( + "manifest_id", + header.get("manifest_id").cloned(), + json!(manifest_id), + ), + ( + "stage", + header.get("stage").cloned(), + json!(cli.stage.as_str()), + ), + ("seed", header.get("seed").cloned(), json!(cli.seed)), + ( + "machine_signature", + header.get("machine_signature").cloned(), + json!(machine_signature), + ), + ]; + for (name, actual, expected) in checks { + if actual.as_ref() != Some(&expected) { + return Err(format!( + "cannot resume: manifest field {name} is {actual:?}, expected {expected}" + )); + } + } + Ok(()) +} + +fn unix_seconds() -> u64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_secs() +} + +fn to_string(error: impl std::fmt::Display) -> String { + error.to_string() +} diff --git a/parquet/benches/row_selection_policy_sampler/policy_validation.rs b/parquet/benches/row_selection_policy_sampler/policy_validation.rs new file mode 100644 index 000000000000..8d9b0ef5cdd3 --- /dev/null +++ b/parquet/benches/row_selection_policy_sampler/policy_validation.rs @@ -0,0 +1,718 @@ +// 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. + +//! Paired end-to-end validation of `Auto` against `AutoPerColumn`. + +use std::hint::black_box; +use std::time::Instant; + +use arrow::array::{Int32Array, RecordBatch}; +use arrow::compute::kernels::cmp::eq; +use futures::StreamExt; +use parquet::arrow::arrow_reader::metrics::ArrowReaderMetrics; +use parquet::arrow::arrow_reader::{ArrowPredicateFn, RowFilter, RowSelectionPolicy}; +use parquet::arrow::{ParquetRecordBatchStreamBuilder, ProjectionMask}; +use rand::{RngExt, SeedableRng, rngs::StdRng}; +use serde_json::{Value, json}; +use tokio::runtime::Runtime; + +use super::cli::Cli; +use super::fixture::logical_checksum; +use super::model::{OUTPUT_SCHEMA_VERSION, Stage, stable_hash}; +use super::output::JsonlOutput; +use super::sampling::{bootstrap_median_ci, median}; +use crate::row_selection_policy_common::cases::HETEROGENEOUS_CASES; +use crate::row_selection_policy_common::fixture::{ + CaseFixture, InMemoryAsyncReader, build_heterogeneous_fixture, +}; +use crate::row_selection_policy_common::model::{ + BATCH_SIZE, CaseSpec, PAYLOAD_COLUMNS, ROWS_PER_GROUP, RowGroupPattern, +}; + +const CONTROL_DRIFT_LIMIT: f64 = 0.10; +const ORDER_BIAS_LIMIT: f64 = 0.05; + +#[derive(Clone)] +struct PolicyExperiment { + id: String, + case: &'static CaseSpec, +} + +impl PolicyExperiment { + fn to_json(&self) -> Value { + json!({ + "experiment_id": self.id, + "mandatory": true, + "execution_mode": "async-policy-pair", + "case": self.case.name, + "batch_size": BATCH_SIZE, + "row_group_count": self.case.row_groups.len(), + "rows_per_group": ROWS_PER_GROUP, + "total_rows": self.case.total_rows(), + "payload_columns": PAYLOAD_COLUMNS, + }) + } +} + +struct PolicyManifest { + id: String, + experiments: Vec, +} + +impl PolicyManifest { + fn generate(case_filters: &[String], seed: u64) -> Result { + for case_filter in case_filters { + if !HETEROGENEOUS_CASES + .iter() + .any(|case| case.name == case_filter) + { + let known = HETEROGENEOUS_CASES + .iter() + .map(|case| case.name) + .collect::>() + .join(", "); + return Err(format!( + "unknown policy-validation case '{case_filter}', expected one of: {known}" + )); + } + } + + let experiments = HETEROGENEOUS_CASES + .iter() + .filter(|case| case_filters.is_empty() || case_filters.iter().any(|v| v == case.name)) + .map(|case| { + let canonical = format!( + "policy-validation-v1;case={};patterns={};rows_per_group={ROWS_PER_GROUP};batch_size={BATCH_SIZE};payload_columns={PAYLOAD_COLUMNS}", + case.name, + case_patterns_canonical(case) + ); + let id = format!("{:016x}", stable_hash(canonical.as_bytes())); + PolicyExperiment { id, case } + }) + .collect::>(); + + let mut manifest_key = format!( + "schema={OUTPUT_SCHEMA_VERSION};stage={};seed={seed}", + Stage::PolicyValidation.as_str() + ); + for experiment in &experiments { + manifest_key.push(';'); + manifest_key.push_str(&experiment.id); + } + Ok(Self { + id: format!("{:016x}", stable_hash(manifest_key.as_bytes())), + experiments, + }) + } +} + +fn case_patterns_canonical(case: &CaseSpec) -> String { + case.row_groups + .iter() + .map(|pattern| match pattern { + RowGroupPattern::AllSelected => "all".to_string(), + RowGroupPattern::Cycle(runs) => runs + .iter() + .map(|run| format!("{}{}", if run.selected { 's' } else { 'k' }, run.len)) + .collect::>() + .join("-"), + }) + .collect::>() + .join("|") +} + +#[derive(Debug)] +struct PolicyPairSample { + pair_index: usize, + auto_first: bool, + auto_ns: u64, + auto_per_column_ns: u64, + log_ratio: f64, +} + +impl PolicyPairSample { + fn to_json(&self) -> Value { + json!({ + "pair_index": self.pair_index, + "order": if self.auto_first { "auto-first" } else { "auto-per-column-first" }, + "auto_ns": self.auto_ns, + "auto_per_column_ns": self.auto_per_column_ns, + "log_auto_per_column_over_auto": self.log_ratio, + }) + } +} + +#[derive(Debug)] +struct ControlSample { + position: &'static str, + after_pairs: usize, + auto_ns: u64, + drift_from_start: f64, +} + +impl ControlSample { + fn to_json(&self) -> Value { + json!({ + "position": self.position, + "after_pairs": self.after_pairs, + "auto_ns": self.auto_ns, + "drift_from_start": self.drift_from_start, + }) + } +} + +#[derive(Debug)] +struct DecisionCounts { + mask: usize, + selectors: usize, + fallback: usize, +} + +impl DecisionCounts { + fn from_metrics(metrics: &ArrowReaderMetrics) -> Self { + Self { + mask: metrics.row_selection_mask_decisions().unwrap_or_default(), + selectors: metrics + .row_selection_selector_decisions() + .unwrap_or_default(), + fallback: metrics + .row_selection_fallback_decisions() + .unwrap_or_default(), + } + } + + fn to_json(&self) -> Value { + json!({ + "mask": self.mask, + "selectors": self.selectors, + "fallback": self.fallback, + }) + } +} + +struct PolicyPoint { + samples: Vec, + controls: Vec, + rows: usize, + checksum: u64, + decisions: DecisionCounts, + median_auto_ns: f64, + median_auto_per_column_ns: f64, + median_log_ratio: f64, + ci_low: f64, + ci_high: f64, + order_effect: Option, + max_abs_control_drift: f64, + end_control_drift: f64, + stability_warning: bool, + decision: &'static str, + stop_reason: &'static str, + elapsed_ms: u64, + complete: bool, +} + +impl PolicyPoint { + fn to_json(&self, experiment: &PolicyExperiment, cli: &Cli) -> Value { + json!({ + "record_type": "experiment", + "status": if self.complete { "complete" } else { "incomplete" }, + "experiment": experiment.to_json(), + "preflight": { + "rows": self.rows, + "checksum": format!("{:016x}", self.checksum), + }, + "policy_decisions": { + "auto_per_column": self.decisions.to_json(), + }, + "sampling": cli.sampling_json(), + "samples": self.samples.iter().map(PolicyPairSample::to_json).collect::>(), + "controls": self.controls.iter().map(ControlSample::to_json).collect::>(), + "summary": { + "pairs": self.samples.len(), + "median_auto_ns": self.median_auto_ns, + "median_auto_per_column_ns": self.median_auto_per_column_ns, + "median_log_auto_per_column_over_auto": self.median_log_ratio, + "bootstrap_ci95": [self.ci_low, self.ci_high], + "practical_decision": self.decision, + "decision_band": cli.decision_band, + "order_effect_log_ratio": self.order_effect, + "order_bias_limit": ORDER_BIAS_LIMIT, + "max_abs_control_drift": self.max_abs_control_drift, + "end_control_drift": self.end_control_drift, + "control_drift_limit": CONTROL_DRIFT_LIMIT, + "stability_warning": self.stability_warning, + "stop_reason": self.stop_reason, + }, + "elapsed_ms": self.elapsed_ms, + }) + } +} + +pub(crate) fn run(cli: Cli) -> Result<(), String> { + let manifest = PolicyManifest::generate(&cli.cases, cli.seed)?; + let mut output = JsonlOutput::open( + &cli, + &manifest.id, + manifest.experiments.len(), + manifest.experiments.len(), + )?; + println!( + "row-selection sampler: stage={}, manifest={}, experiments={}, output={}", + cli.stage, + manifest.id, + manifest.experiments.len(), + cli.output.display() + ); + if output.resumed_records != 0 { + println!( + "resuming after {} completed policy-validation records", + output.resumed_records + ); + } + + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .map_err(|error| error.to_string())?; + let started = Instant::now(); + let mut completed_this_run = 0usize; + let mut incomplete_this_run = 0usize; + let mut skipped_resume = 0usize; + + for (idx, experiment) in manifest.experiments.iter().enumerate() { + if output.is_completed(&experiment.id) { + skipped_resume += 1; + continue; + } + println!( + "[{}/{}] {} {}", + idx + 1, + manifest.experiments.len(), + experiment.id, + experiment.case.name + ); + let fixture = build_heterogeneous_fixture(experiment.case).map_err(|error| { + format!( + "failed to build mandatory policy-validation fixture {}: {error}", + experiment.case.name + ) + })?; + let point = sample_policy_point(&cli, experiment, &fixture, &runtime)?; + output.write(&point.to_json(experiment, &cli))?; + if point.complete { + output.mark_completed(&experiment.id); + if point.stability_warning { + output.mark_validation_warning(&experiment.id); + } + if point.decision == "inconclusive" { + output.mark_validation_inconclusive(&experiment.id); + } + completed_this_run += 1; + } else { + incomplete_this_run += 1; + } + } + + let remaining = manifest + .experiments + .iter() + .filter(|experiment| !output.is_completed(&experiment.id)) + .count(); + let validation_warnings = output.validation_warning_count(); + let inconclusive_points = output.validation_inconclusive_count(); + let validation_passed = remaining == 0 && validation_warnings == 0; + let promotion_eligible = validation_passed && inconclusive_points == 0; + output.write(&json!({ + "record_type": "run_end", + "stage": cli.stage.as_str(), + "manifest_id": manifest.id, + "elapsed_seconds": started.elapsed().as_secs_f64(), + "completed_this_run": completed_this_run, + "unsupported_this_run": 0, + "incomplete_this_run": incomplete_this_run, + "skipped_from_resume": skipped_resume, + "remaining_experiments": remaining, + "budget_exhausted": false, + "validation_warning_points": validation_warnings, + "inconclusive_points": inconclusive_points, + "validation_passed": validation_passed, + "promotion_eligible": promotion_eligible, + "sampling": cli.sampling_json(), + }))?; + println!( + "policy validation finished in {:.3}s: completed={}, incomplete={}, remaining={}, stability warnings={}, inconclusive={}, validation passed={}, promotion eligible={}", + started.elapsed().as_secs_f64(), + completed_this_run, + incomplete_this_run, + remaining, + validation_warnings, + inconclusive_points, + validation_passed, + promotion_eligible + ); + + let ephemeral_output = cli.ephemeral_output; + let output_path = cli.output.clone(); + drop(output); + if ephemeral_output { + std::fs::remove_file(&output_path).map_err(|error| { + format!( + "failed to remove ephemeral output {}: {error}", + output_path.display() + ) + })?; + } + Ok(()) +} + +fn sample_policy_point( + cli: &Cli, + experiment: &PolicyExperiment, + fixture: &CaseFixture, + runtime: &Runtime, +) -> Result { + let started = Instant::now(); + let auto_batches = runtime.block_on(capture_batches( + fixture, + RowSelectionPolicy::default(), + None, + ))?; + let metrics = ArrowReaderMetrics::enabled(); + let auto_per_column_batches = runtime.block_on(capture_batches( + fixture, + RowSelectionPolicy::AutoPerColumn, + Some(metrics.clone()), + ))?; + if auto_batches != auto_per_column_batches { + return Err(format!( + "correctness failure for {}: Auto and AutoPerColumn returned different RecordBatches", + experiment.case.name + )); + } + let rows = auto_batches + .iter() + .map(RecordBatch::num_rows) + .sum::(); + if rows != fixture.expected_rows { + return Err(format!( + "correctness failure for {}: expected {} rows, got {rows}", + experiment.case.name, fixture.expected_rows + )); + } + let checksum = logical_checksum(&auto_batches)?; + let decisions = DecisionCounts::from_metrics(&metrics); + if decisions.mask + decisions.selectors == 0 { + return Err(format!( + "policy-validation fixture {} recorded no AutoPerColumn decisions", + experiment.case.name + )); + } + drop(auto_batches); + drop(auto_per_column_batches); + + let seed = cli.seed ^ stable_hash(experiment.id.as_bytes()) ^ 0xa070_c011_2026_0817; + let mut rng = StdRng::seed_from_u64(seed); + for _ in 0..cli.warmup_pairs { + run_pair(fixture, runtime, rng.random_bool(0.5), cli.inner_iterations)?; + } + + let start_control_ns = measure_policy( + fixture, + runtime, + RowSelectionPolicy::default(), + cli.inner_iterations, + )?; + let mut controls = vec![ControlSample { + position: "start", + after_pairs: 0, + auto_ns: start_control_ns, + drift_from_start: 0.0, + }]; + let mut samples = Vec::with_capacity(cli.max_pairs); + let mut stop_reason = "max_pairs"; + let mut block_auto_first = false; + + while samples.len() < cli.max_pairs { + if samples.len().is_multiple_of(2) { + block_auto_first = rng.random_bool(0.5); + } + let auto_first = if samples.len().is_multiple_of(2) { + block_auto_first + } else { + !block_auto_first + }; + let (auto_ns, auto_per_column_ns) = + run_pair(fixture, runtime, auto_first, cli.inner_iterations)?; + samples.push(PolicyPairSample { + pair_index: samples.len(), + auto_first, + auto_ns, + auto_per_column_ns, + log_ratio: (auto_per_column_ns as f64 / auto_ns as f64).ln(), + }); + + if samples.len().is_multiple_of(cli.control_interval_pairs) { + let auto_ns = measure_policy( + fixture, + runtime, + RowSelectionPolicy::default(), + cli.inner_iterations, + )?; + controls.push(ControlSample { + position: "periodic", + after_pairs: samples.len(), + auto_ns, + drift_from_start: auto_ns as f64 / start_control_ns as f64 - 1.0, + }); + } + + // Only stop after a complete two-pair block so both execution orders + // always contribute equally to a completed point. + if samples.len() >= cli.min_pairs && samples.len().is_multiple_of(2) { + let log_ratios = samples + .iter() + .map(|sample| sample.log_ratio) + .collect::>(); + let (_, ci_low, ci_high) = bootstrap_median_ci( + &log_ratios, + cli.bootstrap_samples, + seed ^ samples.len() as u64, + ); + if ci_low > cli.decision_band || ci_high < -cli.decision_band { + stop_reason = "direction_confident"; + break; + } + if ci_high - ci_low <= cli.target_ci_width { + stop_reason = "target_precision"; + break; + } + } + if started.elapsed() >= cli.point_timeout { + stop_reason = "point_timeout"; + break; + } + } + + let end_control_ns = measure_policy( + fixture, + runtime, + RowSelectionPolicy::default(), + cli.inner_iterations, + )?; + let end_control_drift = end_control_ns as f64 / start_control_ns as f64 - 1.0; + controls.push(ControlSample { + position: "end", + after_pairs: samples.len(), + auto_ns: end_control_ns, + drift_from_start: end_control_drift, + }); + + let log_ratios = samples + .iter() + .map(|sample| sample.log_ratio) + .collect::>(); + let (median_log_ratio, ci_low, ci_high) = + bootstrap_median_ci(&log_ratios, cli.bootstrap_samples, seed ^ 0xb007_57a9); + let order_effect = order_effect(&samples); + let max_abs_control_drift = controls + .iter() + .map(|control| control.drift_from_start.abs()) + .fold(0.0, f64::max); + let stability_warning = max_abs_control_drift > CONTROL_DRIFT_LIMIT + || order_effect.is_some_and(|effect| effect.abs() > ORDER_BIAS_LIMIT); + let decision = practical_decision(stability_warning, ci_low, ci_high, cli.decision_band); + let complete = stop_reason != "point_timeout" || samples.len() == cli.max_pairs; + + Ok(PolicyPoint { + median_auto_ns: median(samples.iter().map(|sample| sample.auto_ns as f64).collect()), + median_auto_per_column_ns: median( + samples + .iter() + .map(|sample| sample.auto_per_column_ns as f64) + .collect(), + ), + samples, + controls, + rows, + checksum, + decisions, + median_log_ratio, + ci_low, + ci_high, + order_effect, + max_abs_control_drift, + end_control_drift, + stability_warning, + decision, + stop_reason, + elapsed_ms: millis(started.elapsed()), + complete, + }) +} + +fn run_pair( + fixture: &CaseFixture, + runtime: &Runtime, + auto_first: bool, + inner_iterations: usize, +) -> Result<(u64, u64), String> { + if auto_first { + Ok(( + measure_policy( + fixture, + runtime, + RowSelectionPolicy::default(), + inner_iterations, + )?, + measure_policy( + fixture, + runtime, + RowSelectionPolicy::AutoPerColumn, + inner_iterations, + )?, + )) + } else { + let auto_per_column = measure_policy( + fixture, + runtime, + RowSelectionPolicy::AutoPerColumn, + inner_iterations, + )?; + let auto = measure_policy( + fixture, + runtime, + RowSelectionPolicy::default(), + inner_iterations, + )?; + Ok((auto, auto_per_column)) + } +} + +fn measure_policy( + fixture: &CaseFixture, + runtime: &Runtime, + policy: RowSelectionPolicy, + inner_iterations: usize, +) -> Result { + let started = Instant::now(); + let mut rows = 0usize; + for _ in 0..inner_iterations { + let observed = runtime.block_on(run_rows(fixture, policy))?; + if observed != fixture.expected_rows { + return Err(format!( + "policy-validation timed run expected {} rows, got {observed}", + fixture.expected_rows + )); + } + rows = rows.saturating_add(observed); + } + black_box(rows); + let nanos = started.elapsed().as_nanos() / inner_iterations as u128; + Ok(u64::try_from(nanos.max(1)).unwrap_or(u64::MAX)) +} + +async fn capture_batches( + fixture: &CaseFixture, + policy: RowSelectionPolicy, + metrics: Option, +) -> Result, String> { + let mut stream = stream_builder(fixture, policy, metrics) + .await? + .build() + .map_err(|error| error.to_string())?; + let mut batches = Vec::new(); + while let Some(batch) = stream.next().await { + batches.push(batch.map_err(|error| error.to_string())?); + } + Ok(batches) +} + +async fn run_rows(fixture: &CaseFixture, policy: RowSelectionPolicy) -> Result { + let mut stream = stream_builder(fixture, policy, None) + .await? + .build() + .map_err(|error| error.to_string())?; + let mut rows = 0usize; + while let Some(batch) = stream.next().await { + let batch = batch.map_err(|error| error.to_string())?; + rows += black_box(batch.num_rows()); + black_box(batch); + } + Ok(rows) +} + +async fn stream_builder( + fixture: &CaseFixture, + policy: RowSelectionPolicy, + metrics: Option, +) -> Result, String> { + 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 builder = ParquetRecordBatchStreamBuilder::new(fixture.reader()) + .await + .map_err(|error| error.to_string())? + .with_batch_size(BATCH_SIZE) + .with_projection(output_projection) + .with_row_filter(row_filter) + .with_row_selection_policy(policy); + if let Some(metrics) = metrics { + builder = builder.with_metrics(metrics); + } + Ok(builder) +} + +fn order_effect(samples: &[PolicyPairSample]) -> Option { + let auto_first = samples + .iter() + .filter(|sample| sample.auto_first) + .map(|sample| sample.log_ratio) + .collect::>(); + let auto_per_column_first = samples + .iter() + .filter(|sample| !sample.auto_first) + .map(|sample| sample.log_ratio) + .collect::>(); + (!auto_first.is_empty() && !auto_per_column_first.is_empty()) + .then(|| median(auto_first) - median(auto_per_column_first)) +} + +fn practical_decision( + stability_warning: bool, + ci_low: f64, + ci_high: f64, + decision_band: f64, +) -> &'static str { + if stability_warning { + "unstable" + } else if ci_high < -decision_band { + "auto-per-column-faster" + } else if ci_low > decision_band { + "auto-faster" + } else if ci_low >= -decision_band && ci_high <= decision_band { + "practical-tie" + } else { + "inconclusive" + } +} + +fn millis(duration: std::time::Duration) -> u64 { + u64::try_from(duration.as_millis()).unwrap_or(u64::MAX) +} diff --git a/parquet/benches/row_selection_policy_sampler/sampling.rs b/parquet/benches/row_selection_policy_sampler/sampling.rs new file mode 100644 index 000000000000..1a0bbb6a754d --- /dev/null +++ b/parquet/benches/row_selection_policy_sampler/sampling.rs @@ -0,0 +1,547 @@ +// 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::black_box; +use std::ops::Range; +use std::time::Instant; + +use arrow_array::RecordBatch; +use parquet::DecodeResult; +use parquet::arrow::ProjectionMask; +use parquet::arrow::arrow_reader::{ + ParquetRecordBatchReaderBuilder, RowSelection, RowSelectionPolicy, +}; +use parquet::arrow::push_decoder::{ParquetPushDecoderBuilder, PushBuffers}; +use rand::{RngExt, SeedableRng, rngs::StdRng}; +use serde_json::{Value, json}; + +use super::cli::Cli; +use super::fixture::{Fixture, logical_checksum}; +use super::model::{ExecutionMode, Experiment, stable_hash}; + +#[derive(Debug)] +struct PairSample { + pair_index: usize, + mask_first: bool, + mask_ns: u64, + selectors_ns: u64, + log_ratio: f64, +} + +impl PairSample { + fn to_json(&self) -> Value { + json!({ + "pair_index": self.pair_index, + "order": if self.mask_first { "mask-first" } else { "selectors-first" }, + "mask_ns": self.mask_ns, + "selectors_ns": self.selectors_ns, + "log_mask_over_selectors": self.log_ratio, + }) + } +} + +#[derive(Debug)] +struct SampleSummary { + median_mask_ns: f64, + median_selectors_ns: f64, + median_log_ratio: f64, + ci_low: f64, + ci_high: f64, + stop_reason: &'static str, +} + +impl SampleSummary { + fn to_json(&self, pairs: usize) -> Value { + json!({ + "pairs": pairs, + "median_mask_ns": self.median_mask_ns, + "median_selectors_ns": self.median_selectors_ns, + "median_log_mask_over_selectors": self.median_log_ratio, + "bootstrap_ci95": [self.ci_low, self.ci_high], + "stop_reason": self.stop_reason, + }) + } + + pub(crate) fn baseline_ns(&self) -> f64 { + f64::midpoint(self.median_mask_ns, self.median_selectors_ns) + } +} + +pub(crate) struct SampledPoint { + samples: Vec, + summary: SampleSummary, + rows: usize, + checksum: u64, + selection_stats: Value, + metadata: Value, + diagnostic: Value, + sampling: Value, + elapsed_ms: u64, + complete: bool, +} + +impl SampledPoint { + pub(crate) fn baseline_ns(&self) -> f64 { + self.summary.baseline_ns() + } + + pub(crate) fn to_json(&self, experiment: &Experiment, record_type: &str) -> Value { + json!({ + "record_type": record_type, + "status": if self.complete { "complete" } else { "incomplete" }, + "experiment": experiment.to_json(), + "selection_stats": self.selection_stats, + "column_metadata": self.metadata, + "preflight": { + "rows": self.rows, + "checksum": format!("{:016x}", self.checksum), + }, + "diagnostic": self.diagnostic, + "sampling": self.sampling, + "samples": self.samples.iter().map(PairSample::to_json).collect::>(), + "summary": self.summary.to_json(self.samples.len()), + "elapsed_ms": self.elapsed_ms, + }) + } + + pub(crate) fn is_complete(&self) -> bool { + self.complete + } +} + +pub(crate) fn sample_point( + cli: &Cli, + experiment: &Experiment, + fixture: &Fixture, + control: bool, +) -> Result { + let started = Instant::now(); + let materialized = experiment + .selection + .materialize(experiment.fixture.rows, experiment.batch_size); + if materialized.stats.selected_rows() == 0 { + return Err(format!( + "experiment {} generated an empty selection", + experiment.id + )); + } + let projection = ProjectionMask::roots(fixture.metadata.parquet_schema(), [0]); + let prepared = PreparedRun::new( + experiment.mode, + fixture, + projection, + materialized.selection, + experiment.batch_size, + )?; + let diagnostic = prepared.diagnostic()?; + + let mask_batches = prepared.capture(RowSelectionPolicy::Mask)?; + let selector_batches = prepared.capture(RowSelectionPolicy::Selectors)?; + if mask_batches != selector_batches { + return Err(format!( + "correctness failure for {}: Mask and Selectors returned different RecordBatches", + experiment.id + )); + } + let rows = mask_batches + .iter() + .map(RecordBatch::num_rows) + .sum::(); + if rows != materialized.stats.selected_rows() { + return Err(format!( + "correctness failure for {}: expected {} selected rows, got {rows}", + experiment.id, + materialized.stats.selected_rows() + )); + } + let checksum = logical_checksum(&mask_batches)?; + + let seed = + cli.seed ^ stable_hash(experiment.id.as_bytes()) ^ if control { 0xc017_7010 } else { 0 }; + let mut rng = StdRng::seed_from_u64(seed); + for _ in 0..cli.warmup_pairs { + if rng.random_bool(0.5) { + prepared.run_rows(RowSelectionPolicy::Mask)?; + prepared.run_rows(RowSelectionPolicy::Selectors)?; + } else { + prepared.run_rows(RowSelectionPolicy::Selectors)?; + prepared.run_rows(RowSelectionPolicy::Mask)?; + } + } + + let mut samples = Vec::with_capacity(cli.max_pairs); + let mut stop_reason = "max_pairs"; + while samples.len() < cli.max_pairs { + let mask_first = rng.random_bool(0.5); + let (mask_ns, selectors_ns) = if mask_first { + ( + measure(&prepared, RowSelectionPolicy::Mask, cli.inner_iterations)?, + measure( + &prepared, + RowSelectionPolicy::Selectors, + cli.inner_iterations, + )?, + ) + } else { + let selectors = measure( + &prepared, + RowSelectionPolicy::Selectors, + cli.inner_iterations, + )?; + let mask = measure(&prepared, RowSelectionPolicy::Mask, cli.inner_iterations)?; + (mask, selectors) + }; + samples.push(PairSample { + pair_index: samples.len(), + mask_first, + mask_ns, + selectors_ns, + log_ratio: (mask_ns as f64 / selectors_ns as f64).ln(), + }); + + if samples.len() >= cli.min_pairs { + let log_ratios = samples + .iter() + .map(|sample| sample.log_ratio) + .collect::>(); + let (_, ci_low, ci_high) = bootstrap_median_ci( + &log_ratios, + cli.bootstrap_samples, + seed ^ samples.len() as u64, + ); + if ci_low > cli.decision_band || ci_high < -cli.decision_band { + stop_reason = "direction_confident"; + break; + } + if ci_high - ci_low <= cli.target_ci_width { + stop_reason = "target_precision"; + break; + } + } + if started.elapsed() >= cli.point_timeout { + stop_reason = "point_timeout"; + break; + } + } + + let log_ratios = samples + .iter() + .map(|sample| sample.log_ratio) + .collect::>(); + let (median_log_ratio, ci_low, ci_high) = + bootstrap_median_ci(&log_ratios, cli.bootstrap_samples, seed ^ 0xb007_57a9); + let summary = SampleSummary { + median_mask_ns: median(samples.iter().map(|sample| sample.mask_ns as f64).collect()), + median_selectors_ns: median( + samples + .iter() + .map(|sample| sample.selectors_ns as f64) + .collect(), + ), + median_log_ratio, + ci_low, + ci_high, + stop_reason, + }; + let complete = stop_reason != "point_timeout" || samples.len() == cli.max_pairs; + Ok(SampledPoint { + samples, + summary, + rows, + checksum, + selection_stats: materialized.stats.to_json(), + metadata: fixture.metadata_json(), + diagnostic, + sampling: cli.sampling_json(), + elapsed_ms: millis(started.elapsed()), + complete, + }) +} + +struct PreparedRun<'a> { + mode: ExecutionMode, + fixture: &'a Fixture, + projection: ProjectionMask, + selection: RowSelection, + batch_size: usize, + full_buffers: Option, +} + +impl<'a> PreparedRun<'a> { + fn new( + mode: ExecutionMode, + fixture: &'a Fixture, + projection: ProjectionMask, + selection: RowSelection, + batch_size: usize, + ) -> Result { + let full_buffers = if mode == ExecutionMode::PageValidation { + let len = fixture.bytes.len() as u64; + let mut buffers = PushBuffers::new(len); + buffers + .push_range(0..len, fixture.bytes.clone()) + .map_err(|error| error.to_string())?; + Some(buffers) + } else { + None + }; + Ok(Self { + mode, + fixture, + projection, + selection, + batch_size, + full_buffers, + }) + } + + fn capture(&self, policy: RowSelectionPolicy) -> Result, String> { + match self.mode { + ExecutionMode::SyncOracle => self.capture_sync(policy), + ExecutionMode::PageValidation => self.capture_push(policy), + } + } + + fn run_rows(&self, policy: RowSelectionPolicy) -> Result { + match self.mode { + ExecutionMode::SyncOracle => self.run_sync_rows(policy), + ExecutionMode::PageValidation => self.run_push_rows(policy), + } + } + + fn capture_sync(&self, policy: RowSelectionPolicy) -> Result, String> { + ParquetRecordBatchReaderBuilder::new_with_metadata( + self.fixture.bytes.clone(), + self.fixture.metadata.clone(), + ) + .with_projection(self.projection.clone()) + .with_batch_size(self.batch_size) + .with_row_selection(self.selection.clone()) + .with_row_selection_policy(policy) + .build() + .map_err(|error| error.to_string())? + .collect::, _>>() + .map_err(|error| error.to_string()) + } + + fn run_sync_rows(&self, policy: RowSelectionPolicy) -> Result { + let reader = ParquetRecordBatchReaderBuilder::new_with_metadata( + self.fixture.bytes.clone(), + self.fixture.metadata.clone(), + ) + .with_projection(self.projection.clone()) + .with_batch_size(self.batch_size) + .with_row_selection(self.selection.clone()) + .with_row_selection_policy(policy) + .build() + .map_err(|error| error.to_string())?; + let mut rows = 0usize; + for batch in reader { + let batch = batch.map_err(|error| error.to_string())?; + rows += black_box(batch.num_rows()); + black_box(batch); + } + Ok(rows) + } + + fn capture_push(&self, policy: RowSelectionPolicy) -> Result, String> { + let mut decoder = self + .push_builder(policy, true)? + .build() + .map_err(to_string)?; + let mut batches = Vec::new(); + loop { + match decoder.try_decode().map_err(to_string)? { + DecodeResult::Data(batch) => batches.push(batch), + DecodeResult::NeedsData(ranges) => { + return Err(format!( + "prefetched push decoder unexpectedly requested {ranges:?}" + )); + } + DecodeResult::Finished => return Ok(batches), + } + } + } + + fn run_push_rows(&self, policy: RowSelectionPolicy) -> Result { + let mut decoder = self + .push_builder(policy, true)? + .build() + .map_err(to_string)?; + let mut rows = 0usize; + loop { + match decoder.try_decode().map_err(to_string)? { + DecodeResult::Data(batch) => { + rows += black_box(batch.num_rows()); + black_box(batch); + } + DecodeResult::NeedsData(ranges) => { + return Err(format!( + "prefetched push decoder unexpectedly requested {ranges:?}" + )); + } + DecodeResult::Finished => return Ok(rows), + } + } + } + + fn diagnostic(&self) -> Result { + if self.mode != ExecutionMode::PageValidation { + return Ok(json!({"requested_ranges": []})); + } + + let selectors = self.requested_ranges(RowSelectionPolicy::Selectors)?; + let mask = self.requested_ranges(RowSelectionPolicy::Mask)?; + let column_compressed_bytes = u64::try_from( + self.fixture + .metadata + .metadata() + .row_group(0) + .column(0) + .compressed_size(), + ) + .map_err(|_| "generated column has a negative compressed size".to_string())?; + for (policy, ranges) in [("Selectors", &selectors), ("Mask", &mask)] { + let bytes = requested_bytes(ranges); + if ranges.is_empty() || bytes >= column_compressed_bytes { + return Err(format!( + "page-pruning validation failed for {policy}: requested {} ranges / {bytes} bytes from a {column_compressed_bytes}-byte column", + ranges.len() + )); + } + } + + Ok(json!({ + "page_pruning_validated": true, + "column_compressed_bytes": column_compressed_bytes, + "policies_requested_identical_ranges": selectors == mask, + "by_policy": { + "selectors": ranges_json(&selectors), + "mask": ranges_json(&mask), + }, + })) + } + + fn requested_ranges(&self, policy: RowSelectionPolicy) -> Result>, String> { + let mut decoder = self + .push_builder(policy, false)? + .build() + .map_err(to_string)?; + match decoder.try_decode().map_err(to_string)? { + DecodeResult::NeedsData(ranges) => Ok(ranges), + DecodeResult::Data(batch) => Err(format!( + "empty-buffer diagnostic unexpectedly decoded {} rows", + batch.num_rows() + )), + DecodeResult::Finished => Err("empty-buffer diagnostic unexpectedly finished".into()), + } + } + + fn push_builder( + &self, + policy: RowSelectionPolicy, + prefetched: bool, + ) -> Result { + let builder = ParquetPushDecoderBuilder::new_with_metadata(self.fixture.metadata.clone()) + .with_projection(self.projection.clone()) + .with_batch_size(self.batch_size) + .with_row_selection(self.selection.clone()) + .with_row_selection_policy(policy); + if prefetched { + Ok(builder.with_buffers( + self.full_buffers + .as_ref() + .ok_or_else(|| "page validation is missing prefetched buffers".to_string())? + .clone(), + )) + } else { + Ok(builder) + } + } +} + +fn measure( + prepared: &PreparedRun<'_>, + policy: RowSelectionPolicy, + inner_iterations: usize, +) -> Result { + let started = Instant::now(); + let mut rows = 0usize; + for _ in 0..inner_iterations { + rows = rows.saturating_add(prepared.run_rows(policy)?); + } + black_box(rows); + let nanos = started.elapsed().as_nanos() / inner_iterations as u128; + Ok(u64::try_from(nanos.max(1)).unwrap_or(u64::MAX)) +} + +pub(crate) fn bootstrap_median_ci( + observed: &[f64], + bootstrap_samples: usize, + seed: u64, +) -> (f64, f64, f64) { + let estimate = median(observed.to_vec()); + let mut rng = StdRng::seed_from_u64(seed); + let mut bootstrapped = Vec::with_capacity(bootstrap_samples); + for _ in 0..bootstrap_samples { + let values = (0..observed.len()) + .map(|_| observed[rng.random_range(0..observed.len())]) + .collect(); + bootstrapped.push(median(values)); + } + bootstrapped.sort_by(f64::total_cmp); + let low_idx = ((bootstrap_samples as f64 * 0.025).floor() as usize) + .min(bootstrap_samples.saturating_sub(1)); + let high_idx = ((bootstrap_samples as f64 * 0.975).ceil() as usize) + .saturating_sub(1) + .min(bootstrap_samples.saturating_sub(1)); + (estimate, bootstrapped[low_idx], bootstrapped[high_idx]) +} + +pub(crate) fn median(mut values: Vec) -> f64 { + values.sort_by(f64::total_cmp); + let middle = values.len() / 2; + if values.len().is_multiple_of(2) { + f64::midpoint(values[middle - 1], values[middle]) + } else { + values[middle] + } +} + +fn ranges_json(ranges: &[Range]) -> Value { + json!({ + "requested_range_count": ranges.len(), + "requested_bytes": requested_bytes(ranges), + "requested_ranges": ranges.iter().map(|range| json!({ + "start": range.start, + "end": range.end, + "bytes": range.end - range.start, + })).collect::>(), + }) +} + +fn requested_bytes(ranges: &[Range]) -> u64 { + ranges.iter().map(|range| range.end - range.start).sum() +} + +fn millis(duration: std::time::Duration) -> u64 { + u64::try_from(duration.as_millis()).unwrap_or(u64::MAX) +} + +fn to_string(error: impl std::fmt::Display) -> String { + error.to_string() +} diff --git a/parquet/src/arrow/array_reader/builder.rs b/parquet/src/arrow/array_reader/builder.rs index f89a789647af..03c322b39738 100644 --- a/parquet/src/arrow/array_reader/builder.rs +++ b/parquet/src/arrow/array_reader/builder.rs @@ -17,7 +17,7 @@ use std::sync::{Arc, RwLock}; -use arrow_schema::{DataType, Fields, SchemaBuilder}; +use arrow_schema::{DataType, FieldRef, Fields, SchemaBuilder}; use crate::arrow::ProjectionMask; use crate::arrow::array_reader::byte_view_array::make_byte_view_array_reader; @@ -197,6 +197,52 @@ impl<'a> ArrayReaderBuilder<'a> { Ok(reader) } + /// Build one reader for each projected top-level Arrow field. + /// + /// Nested fields remain a single reader subtree so definition and + /// repetition levels are never split across row-selection strategies. + pub(crate) fn build_top_level_array_readers( + &self, + field: Option<&ParquetField>, + mask: &ProjectionMask, + ) -> Result)>> { + let Some(field) = field else { + return Ok(Vec::new()); + }; + let DataType::Struct(arrow_fields) = &field.arrow_type else { + return Err(general_err!( + "Internal Error: top-level Parquet field must be a struct" + )); + }; + let children = field.children().ok_or_else(|| { + general_err!("Internal Error: top-level Parquet field has no children") + })?; + if arrow_fields.len() != children.len() { + return Err(general_err!( + "Internal Error: Arrow/Parquet top-level field count mismatch" + )); + } + + let mut readers = Vec::with_capacity(children.len()); + for (arrow_field, parquet_field) in arrow_fields.iter().zip(children) { + let Some(reader) = self.build_reader(ReaderArgs { + field: parquet_field, + mask, + padding_threshold: None, + })? else { + continue; + }; + let field = Arc::new( + arrow_field + .as_ref() + .clone() + .with_data_type(reader.get_data_type().clone()), + ); + readers.push((field, reader)); + } + Ok(readers) + } + /// Return the total number of rows fn num_rows(&self) -> usize { self.row_groups.num_rows() diff --git a/parquet/src/arrow/arrow_reader/metrics.rs b/parquet/src/arrow/arrow_reader/metrics.rs index b36d79586bb3..b21725401d29 100644 --- a/parquet/src/arrow/arrow_reader/metrics.rs +++ b/parquet/src/arrow/arrow_reader/metrics.rs @@ -17,8 +17,9 @@ //! [ArrowReaderMetrics] for collecting metrics about the Arrow reader +use super::selection::RowSelectionStrategy; use std::sync::Arc; -use std::sync::atomic::AtomicUsize; +use std::sync::atomic::{AtomicUsize, Ordering}; /// This enum represents the state of Arrow reader metrics collection. /// @@ -90,6 +91,35 @@ impl ArrowReaderMetrics { } } + /// Number of row-selection decisions using mask execution. + /// + /// One decision is recorded per row group and projected top-level Arrow + /// field, for both predicate projections and the final output projection. + /// + /// Returns `None` if metrics are disabled. + pub fn row_selection_mask_decisions(&self) -> Option { + self.load(|inner| &inner.row_selection_mask_decisions) + } + + /// Number of row-selection decisions using selector execution. + /// + /// One decision is recorded per row group and projected top-level Arrow + /// field, for both predicate projections and the final output projection. + /// + /// Returns `None` if metrics are disabled. + pub fn row_selection_selector_decisions(&self) -> Option { + self.load(|inner| &inner.row_selection_selector_decisions) + } + + /// Number of decisions made by the compatibility fallback. + /// + /// This is a subset of the mask and selector decision counters. + /// + /// Returns `None` if metrics are disabled. + pub fn row_selection_fallback_decisions(&self) -> Option { + self.load(|inner| &inner.row_selection_fallback_decisions) + } + /// Increments the count of records read from the inner reader pub(crate) fn increment_inner_reads(&self, count: usize) { let Self::Enabled(inner) = self else { @@ -110,6 +140,49 @@ impl ArrowReaderMetrics { .records_read_from_cache .fetch_add(count, std::sync::atomic::Ordering::Relaxed); } + + /// Records `count` identical decisions at once, for the case where every + /// projected column shares a threshold and the per-column loop is skipped. + pub(crate) fn record_shared_row_selection_decision( + &self, + strategy: RowSelectionStrategy, + fallback: bool, + count: usize, + ) { + for _ in 0..count { + self.record_row_selection_decision(strategy, fallback); + } + } + + pub(crate) fn record_row_selection_decision( + &self, + strategy: RowSelectionStrategy, + fallback: bool, + ) { + let Self::Enabled(inner) = self else { + return; + }; + let counter = match strategy { + RowSelectionStrategy::Mask => &inner.row_selection_mask_decisions, + RowSelectionStrategy::Selectors => &inner.row_selection_selector_decisions, + }; + counter.fetch_add(1, Ordering::Relaxed); + if fallback { + inner + .row_selection_fallback_decisions + .fetch_add(1, Ordering::Relaxed); + } + } + + fn load( + &self, + counter: impl FnOnce(&ArrowReaderMetricsInner) -> &AtomicUsize, + ) -> Option { + match self { + Self::Disabled => None, + Self::Enabled(inner) => Some(counter(inner).load(Ordering::Relaxed)), + } + } } /// Holds the actual metrics for the Arrow reader. @@ -122,6 +195,12 @@ pub struct ArrowReaderMetricsInner { records_read_from_inner: AtomicUsize, /// Total number of records read from previously cached pages records_read_from_cache: AtomicUsize, + /// Per-column row-selection decisions using masks. + row_selection_mask_decisions: AtomicUsize, + /// Per-column row-selection decisions using selectors. + row_selection_selector_decisions: AtomicUsize, + /// Decisions made by the compatibility fallback. + row_selection_fallback_decisions: AtomicUsize, } impl ArrowReaderMetricsInner { @@ -130,6 +209,9 @@ impl ArrowReaderMetricsInner { Self { records_read_from_inner: AtomicUsize::new(0), records_read_from_cache: AtomicUsize::new(0), + row_selection_mask_decisions: AtomicUsize::new(0), + row_selection_selector_decisions: AtomicUsize::new(0), + row_selection_fallback_decisions: AtomicUsize::new(0), } } } diff --git a/parquet/src/arrow/arrow_reader/mod.rs b/parquet/src/arrow/arrow_reader/mod.rs index 113faeb3f546..9da5f25e8511 100644 --- a/parquet/src/arrow/arrow_reader/mod.rs +++ b/parquet/src/arrow/arrow_reader/mod.rs @@ -57,6 +57,7 @@ pub use read_plan::{PredicateOptions, ReadPlan, ReadPlanBuilder}; mod filter; pub mod metrics; +pub(crate) mod per_column; mod read_plan; pub(crate) mod selection; pub mod statistics; @@ -1237,26 +1238,77 @@ impl ParquetRecordBatchReaderBuilder { let mut cache_projection = predicate.projection().clone(); cache_projection.intersect(&projection); - let array_reader = ArrayReaderBuilder::new(&reader, &metrics) + let array_reader_builder = ArrayReaderBuilder::new(&reader, &metrics) .with_batch_size(batch_size) - .with_parquet_metadata(&reader.metadata) - .build_array_reader(fields.as_deref(), predicate.projection())?; - - plan_builder = plan_builder.with_predicate(array_reader, predicate.as_mut())?; + .with_parquet_metadata(&reader.metadata); + + if matches!(row_selection_policy, RowSelectionPolicy::AutoPerColumn) { + let predicate_reader = match per_column::PerColumnReader::try_new( + &reader, + &array_reader_builder, + fields.as_deref(), + predicate.projection(), + &plan_builder, + batch_size, + &metrics, + )? { + per_column::PerColumnDecision::Engaged(reader) => { + ParquetRecordBatchReader::new_per_column(reader, batch_size) + } + per_column::PerColumnDecision::Fallback(strategy) => { + let array_reader = array_reader_builder + .build_array_reader(fields.as_deref(), predicate.projection())?; + ParquetRecordBatchReader::new( + array_reader, + plan_builder + .clone() + .with_row_selection_policy(strategy.into_policy()) + .build(), + ) + } + }; + plan_builder = + plan_builder.with_predicate_reader(predicate_reader, predicate.as_mut())?; + } else { + let array_reader = array_reader_builder + .build_array_reader(fields.as_deref(), predicate.projection())?; + plan_builder = plan_builder.with_predicate(array_reader, predicate.as_mut())?; + } } } - let array_reader = ArrayReaderBuilder::new(&reader, &metrics) - .with_batch_size(batch_size) - .with_parquet_metadata(&reader.metadata) - .build_array_reader(fields.as_deref(), &projection)?; - - let read_plan = plan_builder + let mut plan_builder = plan_builder .limited(reader.num_rows()) .with_offset(offset) .with_limit(limit) - .build_limited() - .build(); + .build_limited(); + + let array_reader_builder = ArrayReaderBuilder::new(&reader, &metrics) + .with_batch_size(batch_size) + .with_parquet_metadata(&reader.metadata); + + if matches!(row_selection_policy, RowSelectionPolicy::AutoPerColumn) { + match per_column::PerColumnReader::try_new( + &reader, + &array_reader_builder, + fields.as_deref(), + &projection, + &plan_builder, + batch_size, + &metrics, + )? { + per_column::PerColumnDecision::Engaged(reader) => { + return Ok(ParquetRecordBatchReader::new_per_column(reader, batch_size)); + } + per_column::PerColumnDecision::Fallback(strategy) => { + plan_builder = plan_builder.with_row_selection_policy(strategy.into_policy()); + } + } + } + + let array_reader = + array_reader_builder.build_array_reader(fields.as_deref(), &projection)?; + let read_plan = plan_builder.build(); Ok(ParquetRecordBatchReader::new(array_reader, read_plan)) } @@ -1649,6 +1701,10 @@ impl ParquetRecordBatchReader { } } + pub(crate) fn new_per_column(reader: per_column::PerColumnReader, batch_size: usize) -> Self { + Self::new(Box::new(reader), ReadPlanBuilder::new(batch_size).build()) + } + #[inline(always)] pub(crate) fn batch_size(&self) -> usize { self.read_plan.batch_size() @@ -1669,9 +1725,10 @@ pub(crate) mod tests { use rand::{Rng, RngExt, SeedableRng, random, rng}; use tempfile::tempfile; + use crate::arrow::arrow_reader::metrics::ArrowReaderMetrics; use crate::arrow::arrow_reader::{ ArrowPredicateFn, ArrowReaderMetadata, ArrowReaderOptions, ParquetRecordBatchReader, - ParquetRecordBatchReaderBuilder, RowFilter, RowSelection, RowSelector, + ParquetRecordBatchReaderBuilder, RowFilter, RowSelection, RowSelectionPolicy, RowSelector, }; use crate::arrow::schema::{ add_encoded_arrow_schema_to_metadata, @@ -4673,6 +4730,364 @@ pub(crate) mod tests { assert_eq!(reader.read_plan.batch_size(), num_rows as usize); } + #[test] + fn test_auto_per_column_defers_when_columns_share_a_threshold() { + let batch = RecordBatch::try_from_iter([ + ( + "left", + Arc::new(Int32Array::from_iter_values(0..256)) as ArrayRef, + ), + ( + "right", + Arc::new(Int32Array::from_iter_values(1000..1256)) as ArrayRef, + ), + ]) + .unwrap(); + let props = WriterProperties::builder() + .set_max_row_group_row_count(Some(128)) + .build(); + let mut data = Vec::new(); + let mut writer = ArrowWriter::try_new(&mut data, batch.schema(), Some(props)).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + let data = Bytes::from(data); + + // Row group 0 has two long runs (Selectors); row group 1 alternates + // every row (Mask). With batch size 31 the first + // output batch crosses the row-group/strategy boundary. + let mut selectors = vec![RowSelector::select(1), RowSelector::skip(127)]; + for _ in 0..64 { + selectors.push(RowSelector::select(1)); + selectors.push(RowSelector::skip(1)); + } + let selection = RowSelection::from(selectors); + + let expected = ParquetRecordBatchReaderBuilder::try_new(data.clone()) + .unwrap() + .with_batch_size(31) + .with_row_selection(selection.clone()) + .with_row_selection_policy(RowSelectionPolicy::Selectors) + .build() + .unwrap() + .collect::, _>>() + .unwrap(); + + let metrics = ArrowReaderMetrics::enabled(); + let actual = ParquetRecordBatchReaderBuilder::try_new(data.clone()) + .unwrap() + .with_batch_size(31) + .with_metrics(metrics.clone()) + .with_row_selection(selection.clone()) + .with_row_selection_policy(RowSelectionPolicy::AutoPerColumn) + .build() + .unwrap() + .collect::, _>>() + .unwrap(); + + assert_eq!(actual, expected); + // Both columns are `Int32`, so they resolve to the same threshold and + // can never disagree. Splitting the plan would then only separate one + // row group from the other, which the split executor cannot exploit, + // so the policy defers to the global decision instead of paying for a + // per-row-group pass over the selection. All four decisions (two + // columns times two row groups) are therefore the same. + assert_eq!(metrics.row_selection_selector_decisions(), Some(0)); + assert_eq!(metrics.row_selection_mask_decisions(), Some(4)); + assert_eq!(metrics.row_selection_fallback_decisions(), Some(0)); + + let read_filtered = |policy| { + let builder = ParquetRecordBatchReaderBuilder::try_new(data.clone()).unwrap(); + let predicate_projection = ProjectionMask::leaves(builder.parquet_schema(), [0]); + let predicate = ArrowPredicateFn::new(predicate_projection, |batch: RecordBatch| { + let values = batch.column(0).as_primitive::(); + Ok(BooleanArray::from( + values + .values() + .iter() + .map(|value| value % 3 == 0) + .collect::>(), + )) + }); + builder + .with_batch_size(31) + .with_row_selection(selection.clone()) + .with_row_filter(RowFilter::new(vec![Box::new(predicate)])) + .with_row_selection_policy(policy) + .build() + .unwrap() + .collect::, _>>() + .unwrap() + }; + assert_eq!( + read_filtered(RowSelectionPolicy::AutoPerColumn), + read_filtered(RowSelectionPolicy::Selectors) + ); + } + + #[test] + fn test_auto_per_column_metadata_model_can_mix_columns() { + let rows = 256; + let mut wide = FixedSizeBinaryBuilder::with_capacity(rows, 32); + for row in 0..rows { + wide.append_value([row as u8; 32]).unwrap(); + } + let mut narrow_view = StringViewBuilder::with_capacity(rows); + let mut wide_view = StringViewBuilder::with_capacity(rows); + let narrow_value = "n".repeat(16); + let wide_value = "w".repeat(64); + for _ in 0..rows { + narrow_view.append_value(&narrow_value); + wide_view.append_value(&wide_value); + } + let batch = RecordBatch::try_from_iter([ + ( + "cheap", + Arc::new(Int32Array::from_iter_values(0..rows as i32)) as ArrayRef, + ), + ("wide", Arc::new(wide.finish()) as ArrayRef), + ("narrow_view", Arc::new(narrow_view.finish()) as ArrayRef), + ("wide_view", Arc::new(wide_view.finish()) as ArrayRef), + ]) + .unwrap(); + let props = WriterProperties::builder() + .set_dictionary_enabled(false) + .build(); + let mut data = Vec::new(); + let mut writer = ArrowWriter::try_new(&mut data, batch.schema(), Some(props)).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + let data = Bytes::from(data); + + // Average run length 12 is below the sampled Int32 and narrow Utf8View + // thresholds, but above the FixedSizeBinary(32) and wide Utf8View + // thresholds. + let selection = RowSelection::from( + (0..10) + .flat_map(|_| [RowSelector::skip(12), RowSelector::select(12)]) + .collect::>(), + ); + let expected = ParquetRecordBatchReaderBuilder::try_new(data.clone()) + .unwrap() + .with_batch_size(64) + .with_row_selection(selection.clone()) + .with_row_selection_policy(RowSelectionPolicy::Selectors) + .build() + .unwrap() + .collect::, _>>() + .unwrap(); + + let metrics = ArrowReaderMetrics::enabled(); + let actual = ParquetRecordBatchReaderBuilder::try_new(data) + .unwrap() + .with_batch_size(64) + .with_metrics(metrics.clone()) + .with_row_selection(selection) + .with_row_selection_policy(RowSelectionPolicy::AutoPerColumn) + .build() + .unwrap() + .collect::, _>>() + .unwrap(); + + assert_eq!(actual, expected); + assert_eq!(metrics.row_selection_mask_decisions(), Some(2)); + assert_eq!(metrics.row_selection_selector_decisions(), Some(2)); + assert_eq!(metrics.row_selection_fallback_decisions(), Some(0)); + } + + /// Dictionary-encoded `Utf8View` is modelled rather than deferred to the + /// compatibility fallback: it is the largest single slice of real string + /// columns, so leaving it unmodelled meant the planner never ran on them. + #[test] + fn test_auto_per_column_dictionary_encoded_view_is_modelled() { + let rows = 256; + let mut values = StringViewBuilder::with_capacity(rows); + let value = "dictionary-value".repeat(4); + for _ in 0..rows { + values.append_value(&value); + } + let batch = + RecordBatch::try_from_iter([("view", Arc::new(values.finish()) as ArrayRef)]).unwrap(); + let mut data = Vec::new(); + let mut writer = ArrowWriter::try_new(&mut data, batch.schema(), None).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + let data = Bytes::from(data); + let selection = RowSelection::from( + (0..16) + .flat_map(|_| [RowSelector::skip(8), RowSelector::select(8)]) + .collect::>(), + ); + + let metrics = ArrowReaderMetrics::enabled(); + let actual = ParquetRecordBatchReaderBuilder::try_new(data.clone()) + .unwrap() + .with_metrics(metrics.clone()) + .with_row_selection(selection.clone()) + .with_row_selection_policy(RowSelectionPolicy::AutoPerColumn) + .build() + .unwrap() + .collect::, _>>() + .unwrap(); + let expected = ParquetRecordBatchReaderBuilder::try_new(data) + .unwrap() + .with_row_selection(selection) + .with_row_selection_policy(RowSelectionPolicy::Selectors) + .build() + .unwrap() + .collect::, _>>() + .unwrap(); + + assert_eq!(actual, expected); + // An average run of 8 is below the dictionary-view threshold, so the + // model picks Mask on its own rather than inheriting it from the + // legacy global threshold. + assert_eq!(metrics.row_selection_mask_decisions(), Some(1)); + assert_eq!(metrics.row_selection_fallback_decisions(), Some(0)); + } + + #[test] + fn test_auto_per_column_keeps_nested_field_as_one_subtree() { + let mut list_builder = ListBuilder::new(Int32Builder::new()); + for value in 0..256 { + if value % 7 == 0 { + list_builder.append_null(); + } else { + list_builder.append_value([Some(value), Some(value + 1)]); + } + } + let batch = RecordBatch::try_from_iter([ + ("nested", Arc::new(list_builder.finish()) as ArrayRef), + ( + "value", + Arc::new(Int32Array::from_iter_values(0..256)) as ArrayRef, + ), + ]) + .unwrap(); + let props = WriterProperties::builder() + .set_max_row_group_row_count(Some(128)) + .build(); + let mut data = Vec::new(); + let mut writer = ArrowWriter::try_new(&mut data, batch.schema(), Some(props)).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + let data = Bytes::from(data); + + let mut selectors = vec![RowSelector::select(1), RowSelector::skip(127)]; + for _ in 0..64 { + selectors.push(RowSelector::select(1)); + selectors.push(RowSelector::skip(1)); + } + let selection = RowSelection::from(selectors); + let read = |policy| { + ParquetRecordBatchReaderBuilder::try_new(data.clone()) + .unwrap() + .with_batch_size(31) + .with_row_selection(selection.clone()) + .with_row_selection_policy(policy) + .build() + .unwrap() + .collect::, _>>() + .unwrap() + }; + + assert_eq!( + read(RowSelectionPolicy::AutoPerColumn), + read(RowSelectionPolicy::Selectors) + ); + } + + #[test] + fn test_auto_per_column_random_differential() { + let strings = + StringArray::from_iter_values((0..320).map(|value| format!("value-{value:03}"))); + let batch = RecordBatch::try_from_iter([ + ( + "left", + Arc::new(Int32Array::from_iter_values(0..320)) as ArrayRef, + ), + ("text", Arc::new(strings) as ArrayRef), + ( + "right", + Arc::new(Int64Array::from_iter_values(10_000..10_320)) as ArrayRef, + ), + ]) + .unwrap(); + let props = WriterProperties::builder() + .set_max_row_group_row_count(Some(64)) + .build(); + let mut data = Vec::new(); + let mut writer = ArrowWriter::try_new(&mut data, batch.schema(), Some(props)).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + let data = Bytes::from(data); + let schema_builder = ParquetRecordBatchReaderBuilder::try_new(data.clone()).unwrap(); + let mut rng = StdRng::seed_from_u64(0xA170_C011); + + for case in 0..32 { + let mut bits = Vec::with_capacity(320); + let mut selected = rng.random_range(0..2) == 0; + while bits.len() < 320 { + let run = rng.random_range(1..80).min(320 - bits.len()); + bits.extend(std::iter::repeat_n(selected, run)); + selected = !selected; + } + // Keep both the empty-tail and non-empty-tail cases in the corpus, + // but ensure every selection has at least one output row. + let forced_selected = case % bits.len(); + bits[forced_selected] = true; + + let selection = if case % 2 == 0 { + RowSelection::from(BooleanBuffer::from(bits.clone())) + } else { + let mut selectors = Vec::new(); + let mut start = 0; + while start < bits.len() { + let value = bits[start]; + let mut end = start + 1; + while end < bits.len() && bits[end] == value { + end += 1; + } + selectors.push(RowSelector { + row_count: end - start, + skip: !value, + }); + start = end; + } + RowSelection::from(selectors) + }; + let selected_rows = bits.iter().filter(|&&value| value).count(); + let offset = rng.random_range(0..selected_rows + 6); + let limit = rng.random_range(0..selected_rows + 6); + let batch_size = rng.random_range(1..48); + let projected = match case % 3 { + 0 => vec![0, 1, 2], + 1 => vec![0, 2], + _ => vec![1], + }; + let projection = ProjectionMask::leaves(schema_builder.parquet_schema(), projected); + + let read = |policy| { + ParquetRecordBatchReaderBuilder::try_new(data.clone()) + .unwrap() + .with_projection(projection.clone()) + .with_batch_size(batch_size) + .with_row_selection(selection.clone()) + .with_row_selection_policy(policy) + .with_offset(offset) + .with_limit(limit) + .build() + .unwrap() + .collect::, _>>() + .unwrap() + }; + assert_eq!( + read(RowSelectionPolicy::AutoPerColumn), + read(RowSelectionPolicy::Selectors), + "case {case}, batch_size {batch_size}, offset {offset}, limit {limit}" + ); + } + } + #[test] fn test_read_with_page_index_enabled() { let testdata = arrow::util::test_util::parquet_test_data(); diff --git a/parquet/src/arrow/arrow_reader/per_column.rs b/parquet/src/arrow/arrow_reader/per_column.rs new file mode 100644 index 000000000000..be27bd2ab390 --- /dev/null +++ b/parquet/src/arrow/arrow_reader/per_column.rs @@ -0,0 +1,1158 @@ +// 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. + +//! Internal per-column row-selection planning and execution. +//! +//! Planning resolves one [`RowSelectionStrategy`] per row group and projected +//! top-level Arrow field, from the selection's run statistics and a +//! metadata-only cost model. If every decision agrees, planning declines and +//! the caller keeps the existing global execution path. +//! +//! Execution gives each top-level field its own reader and replays the same +//! compiled chunks through all of them, one lane per field. Lanes issue +//! different reads and skips but must advance in lockstep; see the +//! [`window`](super::selection) module for that invariant. + +use super::metrics::ArrowReaderMetrics; +use super::selection::{ + BatchWindow, ColumnInstruction, DEFAULT_ROW_SELECTION_THRESHOLD, LoadedRowRanges, + RowSelectionExecutionPlan, RowSelectionStrategy, mask_run_count, +}; +use super::{ReadPlanBuilder, RowSelection, RowSelectionPolicy}; +use crate::arrow::ProjectionMask; +use crate::arrow::array_reader::{ArrayReader, ArrayReaderBuilder, RowGroups}; +use crate::arrow::schema::{ParquetField, ParquetFieldType}; +use crate::basic::Encoding; +use crate::errors::{ParquetError, Result}; +use crate::file::metadata::RowGroupMetaData; +use crate::file::page_index::offset_index::OffsetIndexMetaData; +use arrow_array::{Array, ArrayRef, BooleanArray, StructArray, new_empty_array}; +use arrow_buffer::BooleanBufferBuilder; +use arrow_schema::{DataType, Fields}; +use arrow_select::filter::filter; +use std::any::Any; +use std::sync::Arc; + +pub(crate) struct ColumnSelectionContext<'a> { + /// Projected top-level field, including its nested Parquet leaves. + pub(crate) field: &'a ParquetField, + pub(crate) row_group: &'a RowGroupMetaData, + /// Selection statistics for this row group, computed once and shared by + /// every projected column. + pub(crate) selection: RowSelectionStatistics, +} + +#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] +pub(crate) struct RowSelectionStatistics { + physical_rows: usize, + run_count: usize, +} + +impl RowSelectionStatistics { + #[inline] + fn auto_selection_strategy(self, threshold: usize) -> RowSelectionStrategy { + debug_assert!(threshold >= MIN_RUN_THRESHOLD); + if self.run_count == 0 || self.physical_rows < self.run_count.saturating_mul(threshold) { + RowSelectionStrategy::Mask + } else { + RowSelectionStrategy::Selectors + } + } +} + +/// Internal seam for the metadata cost model. Returning `None` applies the +/// column-local compatibility fallback. +pub(crate) trait ColumnSelectionPlanner { + fn strategy(&self, context: ColumnSelectionContext<'_>) -> Option; +} + +/// Exclusive run-length thresholds calibrated by the refinement sampler. +/// +/// `RowSelection::auto_selection_strategy` selects Mask when the average run +/// is below the threshold. The values below deliberately stop at the last +/// repeatable Mask win outside the sampler's 3% practical-equivalence band. +/// +/// The values below come from a refinement pass on x86_64 (AMD EPYC): five +/// seeds, unit resolution over runs 1..=16, restricted to the sampler's +/// `balanced-grid` shape family (`skip == select`), which is the only family +/// that varies run length without also varying selectivity. +/// +/// **These are hardware-specific.** The same fixtures sampled on aarch64 put +/// every crossover far lower — `Int64` at 15 against 28 here, `Int32` at 15 +/// against 24, `Date32` at 16 against 28 — because Mask keeps its edge over +/// Selectors across a much wider run-length range on EPYC. The gap is not +/// academic: on TPC-DS SF10 the x86 values are worth -14.5% on q36 and -13.0% +/// on q75, while the aarch64 values on the same machine leave most of that on +/// the table. Anyone calibrating for a different target should rerun +/// `arrow_reader_row_selection_policy_sampler --stage refinement` there rather +/// than carrying these numbers over. +const WIDE_BYTE_RUN_THRESHOLD: usize = 8; +const FIXED_BINARY_RUN_THRESHOLD: usize = 9; +/// Dictionary-encoded `Utf8View` used to fall through to the compatibility +/// fallback. It is the single largest unmodeled slice of both TPC-DS (31 +/// columns) and ClickBench (28), and it samples well clear of the plain +/// `Utf8View` threshold, so it gets its own value rather than sharing one. +const DICTIONARY_UTF8_VIEW_RUN_THRESHOLD: usize = 16; +/// Only the INT32-backed precision range was sampled; wider decimals keep the +/// fallback. +const DECIMAL128_INT32_RUN_THRESHOLD: usize = 18; +const DECIMAL128_INT32_MAX_PRECISION: u8 = 9; +const INT32_RUN_THRESHOLD: usize = 24; +const INT64_RUN_THRESHOLD: usize = 28; +const DATE32_RUN_THRESHOLD: usize = 28; +/// Inherited from the earlier x86_64 sampling: neither `Dictionary(Int32, +/// Utf8)` nor narrow `Utf8View` occurs in TPC-DS or ClickBench, so this pass +/// did not resample them. +const DICTIONARY_UTF8_RUN_THRESHOLD: usize = 17; +const NARROW_BYTE_RUN_THRESHOLD: usize = 13; +/// Smallest value any arm can return, used to guard the strategy helper. +const MIN_RUN_THRESHOLD: usize = WIDE_BYTE_RUN_THRESHOLD; +const NARROW_UTF8_VIEW_BYTES: usize = 32; + +struct MetadataColumnSelectionPlanner; + +impl ColumnSelectionPlanner for MetadataColumnSelectionPlanner { + fn strategy(&self, context: ColumnSelectionContext<'_>) -> Option { + let threshold = metadata_run_threshold(context.field, context.row_group)?; + Some(context.selection.auto_selection_strategy(threshold)) + } +} + +/// Returns `None` for columns outside the sampled model. The caller preserves +/// compatibility by applying the legacy global threshold to those columns. +fn metadata_run_threshold(field: &ParquetField, row_group: &RowGroupMetaData) -> Option { + let ParquetFieldType::Primitive { col_idx, .. } = &field.field_type else { + return None; + }; + match &field.arrow_type { + DataType::Int32 => Some(INT32_RUN_THRESHOLD), + DataType::Int64 => Some(INT64_RUN_THRESHOLD), + DataType::Date32 => Some(DATE32_RUN_THRESHOLD), + DataType::Decimal128(precision, _) if *precision <= DECIMAL128_INT32_MAX_PRECISION => { + Some(DECIMAL128_INT32_RUN_THRESHOLD) + } + DataType::Dictionary(key, value) + if key.as_ref() == &DataType::Int32 && value.as_ref() == &DataType::Utf8 => + { + Some(DICTIONARY_UTF8_RUN_THRESHOLD) + } + DataType::FixedSizeBinary(width) if *width > 0 => Some(FIXED_BINARY_RUN_THRESHOLD), + DataType::Utf8View => { + let column = row_group.columns().get(*col_idx)?; + if column.encodings().any(|encoding| { + matches!( + encoding, + Encoding::PLAIN_DICTIONARY | Encoding::RLE_DICTIONARY + ) + }) { + return Some(DICTIONARY_UTF8_VIEW_RUN_THRESHOLD); + } + let values = usize::try_from(column.num_values()).ok()?; + let bytes = usize::try_from(column.uncompressed_size()).ok()?; + let narrow_bytes = values.checked_mul(NARROW_UTF8_VIEW_BYTES)?; + if values == 0 { + None + } else if bytes <= narrow_bytes { + Some(NARROW_BYTE_RUN_THRESHOLD) + } else { + Some(WIDE_BYTE_RUN_THRESHOLD) + } + } + _ => None, + } +} + +pub(crate) enum PerColumnDecision { + Fallback(RowSelectionStrategy), + Engaged(PerColumnReader), +} + +/// Computes the sparse page ranges available to each projected top-level +/// Arrow field. Nested leaves are intersected so a Mask fragment never asks a +/// shared reader subtree to decode through a missing page. +pub(crate) fn loaded_row_ranges_for_top_level_fields( + fields: Option<&ParquetField>, + projection: &ProjectionMask, + selection: Option<&RowSelection>, + offset_index: Option<&[OffsetIndexMetaData]>, + total_rows: usize, +) -> Vec> { + projected_top_level_fields(fields, projection) + .into_iter() + .map(|field| { + let (Some(selection), Some(offset_index)) = (selection, offset_index) else { + return None; + }; + let mut leaves = Vec::new(); + collect_projected_leaves(field, projection, &mut leaves); + leaves + .into_iter() + .filter_map(|leaf_idx| { + let pages = &offset_index.get(leaf_idx)?.page_locations; + (!pages.is_empty()).then(|| { + RowSelection::from_consecutive_ranges( + selection + .row_ranges_for_selected_pages(pages, total_rows) + .into_iter(), + total_rows, + ) + }) + }) + .reduce(|loaded, leaf| loaded.intersection(&leaf)) + .filter(|loaded| loaded.skipped_row_count() != 0) + .map(LoadedRowRanges::from_selection) + }) + .collect() +} + +fn projected_top_level_fields<'a>( + fields: Option<&'a ParquetField>, + projection: &ProjectionMask, +) -> Vec<&'a ParquetField> { + fields + .and_then(ParquetField::children) + .into_iter() + .flatten() + .filter(|field| field_is_projected(field, projection)) + .collect() +} + +fn field_is_projected(field: &ParquetField, projection: &ProjectionMask) -> bool { + match &field.field_type { + ParquetFieldType::Primitive { col_idx, .. } => projection.leaf_included(*col_idx), + ParquetFieldType::Group { children } => children + .iter() + .any(|child| field_is_projected(child, projection)), + ParquetFieldType::Virtual(_) => true, + } +} + +fn collect_projected_leaves( + field: &ParquetField, + projection: &ProjectionMask, + leaves: &mut Vec, +) { + match &field.field_type { + ParquetFieldType::Primitive { col_idx, .. } => { + if projection.leaf_included(*col_idx) { + leaves.push(*col_idx); + } + } + ParquetFieldType::Group { children } => { + for child in children { + collect_projected_leaves(child, projection, leaves); + } + } + ParquetFieldType::Virtual(_) => {} + } +} + +struct TopLevelColumnReader { + reader: Box, +} + +pub(crate) struct PerColumnReader { + columns: Vec, + fields: Fields, + data_type: DataType, + plan: Arc, + next_batch: usize, + buffered: Option, +} + +impl PerColumnReader { + pub(super) fn try_new( + row_groups: &dyn RowGroups, + array_reader_builder: &ArrayReaderBuilder<'_>, + fields: Option<&ParquetField>, + projection: &ProjectionMask, + plan_builder: &ReadPlanBuilder, + batch_size: usize, + metrics: &ArrowReaderMetrics, + ) -> Result { + Self::try_new_with_planner( + row_groups, + array_reader_builder, + fields, + projection, + plan_builder, + batch_size, + metrics, + None, + &MetadataColumnSelectionPlanner, + ) + } + + #[allow(clippy::too_many_arguments)] + pub(crate) fn try_new_with_loaded_ranges( + row_groups: &dyn RowGroups, + array_reader_builder: &ArrayReaderBuilder<'_>, + fields: Option<&ParquetField>, + projection: &ProjectionMask, + plan_builder: &ReadPlanBuilder, + batch_size: usize, + metrics: &ArrowReaderMetrics, + loaded_row_ranges: Vec>, + ) -> Result { + Self::try_new_with_planner( + row_groups, + array_reader_builder, + fields, + projection, + plan_builder, + batch_size, + metrics, + Some(loaded_row_ranges), + &MetadataColumnSelectionPlanner, + ) + } + + #[allow(clippy::too_many_arguments)] + fn try_new_with_planner( + row_groups: &dyn RowGroups, + array_reader_builder: &ArrayReaderBuilder<'_>, + fields: Option<&ParquetField>, + projection: &ProjectionMask, + plan_builder: &ReadPlanBuilder, + batch_size: usize, + metrics: &ArrowReaderMetrics, + loaded_row_ranges: Option>>, + planner: &dyn ColumnSelectionPlanner, + ) -> Result { + if !matches!( + plan_builder.row_selection_policy(), + RowSelectionPolicy::AutoPerColumn + ) { + return Ok(PerColumnDecision::Fallback( + plan_builder.resolve_selection_strategy(), + )); + } + let Some(selection) = plan_builder.selection() else { + return Ok(PerColumnDecision::Fallback(RowSelectionStrategy::Selectors)); + }; + if batch_size == 0 || !selection.selects_any() { + return Ok(PerColumnDecision::Fallback( + plan_builder.resolve_selection_strategy(), + )); + } + + let projected_fields = projected_top_level_fields(fields, projection); + let column_count = projected_fields.len(); + if column_count == 0 { + return Ok(PerColumnDecision::Fallback( + plan_builder.resolve_selection_strategy(), + )); + } + + // Columns that resolve to the same threshold can never disagree, and + // disagreement is the only thing the split executor is here to + // exploit. Thresholds come from column metadata, so settling this up + // front avoids the per-row-group pass over the selection entirely — + // that pass is what made `AutoPerColumn` cost more than the global + // policy on wide, uniformly typed scans. + if let Some((threshold, unmodelled, row_group_count)) = + shared_threshold(row_groups, &projected_fields) + { + let strategy = selection.auto_selection_strategy(threshold); + metrics.record_shared_row_selection_decision( + strategy, + unmodelled, + row_group_count * column_count, + ); + return Ok(PerColumnDecision::Fallback(strategy)); + } + + let (row_group_rows, strategies, uniform) = + plan_strategies(row_groups, selection, &projected_fields, metrics, planner)?; + if let Some(strategy) = uniform { + return Ok(PerColumnDecision::Fallback(strategy)); + } + + let readers = array_reader_builder.build_top_level_array_readers(fields, projection)?; + if readers.len() != column_count { + return Err(general_err!( + "Internal Error: planned {column_count} top-level columns but built {}", + readers.len() + )); + } + let (output_fields, columns): (Vec<_>, Vec<_>) = readers + .into_iter() + .map(|(field, reader)| (field, TopLevelColumnReader { reader })) + .unzip(); + let fields = Fields::from(output_fields); + let data_type = DataType::Struct(fields.clone()); + let plan = RowSelectionExecutionPlan::try_new( + selection.clone(), + &row_group_rows, + column_count, + strategies, + loaded_row_ranges, + batch_size, + )?; + Ok(PerColumnDecision::Engaged(Self { + columns, + fields, + data_type, + plan: Arc::new(plan), + next_batch: 0, + buffered: None, + })) + } + + fn next_array(&mut self) -> Result> { + let Self { + columns, + fields, + plan, + next_batch, + .. + } = self; + let Some(batch) = plan.batch(*next_batch) else { + return Ok(None); + }; + *next_batch += 1; + + let arrays = columns + .iter_mut() + .enumerate() + .map(|(column_idx, column)| { + read_column_batch(column.reader.as_mut(), plan, batch, column_idx) + }) + .collect::>>()?; + let array = StructArray::try_new(fields.clone(), arrays, None)?; + if array.len() != batch.selected_rows { + return Err(general_err!( + "Internal Error: per-column batch produced {} rows, expected {}", + array.len(), + batch.selected_rows + )); + } + Ok(Some(Arc::new(array))) + } + + fn empty_array(&self) -> ArrayRef { + let arrays = self + .fields + .iter() + .map(|field| new_empty_array(field.data_type())) + .collect(); + Arc::new(StructArray::new(self.fields.clone(), arrays, None)) + } +} + +impl ArrayReader for PerColumnReader { + fn as_any(&self) -> &dyn Any { + self + } + + fn get_data_type(&self) -> &DataType { + &self.data_type + } + + fn read_records(&mut self, batch_size: usize) -> Result { + if self.buffered.is_some() { + return Err(general_err!( + "Internal Error: per-column batch must be consumed before reading again" + )); + } + if batch_size == 0 { + self.buffered = Some(self.empty_array()); + return Ok(0); + } + let Some(array) = self.next_array()? else { + self.buffered = Some(self.empty_array()); + return Ok(0); + }; + if array.len() > batch_size { + return Err(general_err!( + "Internal Error: per-column reader produced {} rows for batch size {batch_size}", + array.len() + )); + } + let rows = array.len(); + self.buffered = Some(array); + Ok(rows) + } + + fn consume_batch(&mut self) -> Result { + Ok(self.buffered.take().unwrap_or_else(|| self.empty_array())) + } + + fn skip_records(&mut self, _num_records: usize) -> Result { + Err(general_err!( + "Internal Error: per-column root reader does not support outer row skipping" + )) + } + + fn get_def_levels(&self) -> Option<&[i16]> { + None + } + + fn get_rep_levels(&self) -> Option<&[i16]> { + None + } +} + +/// Returns the threshold every projected column resolves to in every row +/// group, plus whether any of them got there through the compatibility +/// fallback and how many row groups were inspected. `None` means the columns +/// disagree and real per-column planning is needed. +/// +/// Reads column metadata only; it never walks the selection. +fn shared_threshold( + row_groups: &dyn RowGroups, + fields: &[&ParquetField], +) -> Option<(usize, bool, usize)> { + let mut shared: Option = None; + let mut unmodelled = false; + let mut row_group_count = 0usize; + for row_group in row_groups.row_groups() { + row_group_count += 1; + for field in fields { + let threshold = match metadata_run_threshold(field, row_group) { + Some(threshold) => threshold, + None => { + unmodelled = true; + DEFAULT_ROW_SELECTION_THRESHOLD + } + }; + match shared { + Some(previous) if previous != threshold => return None, + Some(_) => {} + None => shared = Some(threshold), + } + } + } + shared.map(|threshold| (threshold, unmodelled, row_group_count)) +} + +fn plan_strategies( + row_groups: &dyn RowGroups, + selection: &RowSelection, + fields: &[&ParquetField], + metrics: &ArrowReaderMetrics, + planner: &dyn ColumnSelectionPlanner, +) -> Result<( + Vec, + Vec, + Option, +)> { + let column_count = fields.len(); + let row_groups = row_groups.row_groups().collect::>(); + let row_group_rows = row_groups + .iter() + .map(|row_group| row_group.num_rows() as usize) + .collect::>(); + let selection_statistics = selection_statistics(selection, &row_group_rows); + + let mut strategies = vec![RowSelectionStrategy::Selectors; row_groups.len() * column_count]; + let mut first_active = None; + let mut uniform = true; + // Tracked apart from `uniform`: a plan whose columns all agree inside every + // row group buys nothing from the split executor, even when two row groups + // reach different strategies. Only genuine column-level disagreement is + // worth the split path. + let mut columns_disagree = false; + + for (row_group_index, (row_group, &statistics)) in + row_groups.iter().zip(&selection_statistics).enumerate() + { + if statistics.physical_rows == 0 { + continue; + } + + let fallback = statistics.auto_selection_strategy(DEFAULT_ROW_SELECTION_THRESHOLD); + let mut row_group_strategy: Option = None; + for (column_index, field) in fields.iter().enumerate() { + let context = ColumnSelectionContext { + field, + row_group, + selection: statistics, + }; + let planned = planner.strategy(context); + let strategy = planned.unwrap_or(fallback); + metrics.record_row_selection_decision(strategy, planned.is_none()); + strategies[row_group_index * column_count + column_index] = strategy; + match row_group_strategy { + Some(previous) if previous != strategy => columns_disagree = true, + Some(_) => {} + None => row_group_strategy = Some(strategy), + } + match first_active { + Some(first) if first != strategy => uniform = false, + None => first_active = Some(strategy), + _ => {} + } + } + } + + let first_active = first_active.unwrap_or(RowSelectionStrategy::Selectors); + for (row_group_index, statistics) in selection_statistics.iter().enumerate() { + if statistics.physical_rows == 0 { + strategies[row_group_index * column_count..(row_group_index + 1) * column_count] + .fill(first_active); + } + } + + Ok(( + row_group_rows, + strategies, + (uniform || !columns_disagree).then_some(first_active), + )) +} + +/// Computes row-group-local run statistics in one pass over either selection +/// backing. Selector run counts saturate once every supported threshold must +/// choose Mask, avoiding unnecessary arithmetic on highly fragmented input. +fn selection_statistics( + selection: &RowSelection, + row_group_rows: &[usize], +) -> Vec { + match selection.as_mask() { + Some(mask) => { + let mut offset = 0usize; + row_group_rows + .iter() + .map(|&row_group_rows| { + let remaining = mask.len().saturating_sub(offset); + let physical_rows = row_group_rows.min(remaining); + let run_count = mask_run_count(&mask.slice(offset, physical_rows)); + offset = offset.saturating_add(physical_rows); + RowSelectionStatistics { + physical_rows, + run_count, + } + }) + .collect() + } + None => { + let mut statistics = vec![RowSelectionStatistics::default(); row_group_rows.len()]; + let mut row_group_index = 0; + let mut rows_left_in_group = row_group_rows.first().copied().unwrap_or_default(); + + for selector in selection.iter() { + let mut rows_left_in_run = selector.row_count; + while rows_left_in_run != 0 { + while rows_left_in_group == 0 { + row_group_index += 1; + let Some(&rows) = row_group_rows.get(row_group_index) else { + return statistics; + }; + rows_left_in_group = rows; + } + + let rows = rows_left_in_run.min(rows_left_in_group); + let current = &mut statistics[row_group_index]; + current.physical_rows += rows; + + // Once this boundary is reached, even the smallest model + // threshold chooses Mask. Keep a lower bound instead of + // counting the remainder of this fragmented row group. + let saturation = row_group_rows[row_group_index] + .checked_div(WIDE_BYTE_RUN_THRESHOLD) + .and_then(|runs| runs.checked_add(1)) + .unwrap_or(usize::MAX); + if current.run_count < saturation { + current.run_count += 1; + } + + rows_left_in_run -= rows; + rows_left_in_group -= rows; + } + } + + statistics + } + } +} + +fn read_column_batch( + reader: &mut dyn ArrayReader, + plan: &RowSelectionExecutionPlan, + batch: &BatchWindow, + column_idx: usize, +) -> Result { + let chunks = plan.chunks(batch); + let needs_filter = chunks + .iter() + .any(|chunk| plan.strategy(chunk, column_idx) == RowSelectionStrategy::Mask); + let mut filter_mask = needs_filter.then(|| BooleanBufferBuilder::new(batch.selected_rows)); + + for chunk in chunks { + let mut consumed_rows = 0usize; + for instruction in plan.lower_chunk(chunk, column_idx)? { + match instruction { + ColumnInstruction::Skip(rows) => { + exact_skip(reader, rows)?; + consumed_rows += rows; + } + ColumnInstruction::Read { rows, mask } => { + exact_read(reader, rows)?; + consumed_rows += rows; + if let Some(filter_mask) = filter_mask.as_mut() { + match mask { + Some(mask) => filter_mask.append_buffer(&mask), + None => filter_mask.append_n(rows, true), + } + } + } + } + } + // Lanes that advance by different amounts misalign the output columns + // without necessarily changing any column's length, so check here + // rather than relying on the batch-level length checks alone. + if consumed_rows != chunk.physical_rows.len() { + return Err(general_err!( + "Internal Error: column {column_idx} consumed {consumed_rows} rows of a {}-row chunk", + chunk.physical_rows.len() + )); + } + } + + let array = reader.consume_batch()?; + let array = match filter_mask { + Some(mut filter_mask) => { + let filter_mask = BooleanArray::from(filter_mask.finish()); + if filter_mask.len() != array.len() { + return Err(general_err!( + "Internal Error: per-column filter has {} rows for an array of {} rows", + filter_mask.len(), + array.len() + )); + } + filter(array.as_ref(), &filter_mask)? + } + None => array, + }; + if array.len() != batch.selected_rows { + return Err(general_err!( + "Internal Error: per-column reader produced {} rows, expected {}", + array.len(), + batch.selected_rows + )); + } + Ok(array) +} + +fn exact_skip(reader: &mut dyn ArrayReader, rows: usize) -> Result<()> { + if rows == 0 { + return Ok(()); + } + let skipped = reader.skip_records(rows)?; + if skipped != rows { + return Err(general_err!( + "failed to skip rows, expected {rows}, got {skipped}" + )); + } + Ok(()) +} + +fn exact_read(reader: &mut dyn ArrayReader, rows: usize) -> Result<()> { + if rows == 0 { + return Ok(()); + } + let read = reader.read_records(rows)?; + if read != rows { + return Err(general_err!( + "failed to read rows, expected {rows}, got {read}" + )); + } + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::arrow::arrow_reader::RowSelector; + use crate::arrow::schema::VirtualColumnType; + use crate::basic::Type as PhysicalType; + use crate::column::page::PageIterator; + use crate::file::metadata::ParquetMetaData; + use crate::schema::types::{SchemaDescriptor, Type}; + use arrow_array::{Array, Int32Array, RecordBatch}; + use arrow_buffer::BooleanBuffer; + use arrow_schema::{DataType, Field}; + use std::any::Any; + + struct TestArrayReader { + values: Vec, + position: usize, + buffered: Vec, + data_type: DataType, + } + + impl TestArrayReader { + fn new(values: impl IntoIterator) -> Self { + Self { + values: values.into_iter().collect(), + position: 0, + buffered: Vec::new(), + data_type: DataType::Int32, + } + } + } + + impl ArrayReader for TestArrayReader { + fn as_any(&self) -> &dyn Any { + self + } + + fn get_data_type(&self) -> &DataType { + &self.data_type + } + + fn read_records(&mut self, batch_size: usize) -> Result { + let end = (self.position + batch_size).min(self.values.len()); + self.buffered + .extend_from_slice(&self.values[self.position..end]); + let read = end - self.position; + self.position = end; + Ok(read) + } + + fn consume_batch(&mut self) -> Result { + Ok(Arc::new(Int32Array::from(std::mem::take( + &mut self.buffered, + )))) + } + + fn skip_records(&mut self, num_records: usize) -> Result { + let end = (self.position + num_records).min(self.values.len()); + let skipped = end - self.position; + self.position = end; + Ok(skipped) + } + + fn get_def_levels(&self) -> Option<&[i16]> { + None + } + + fn get_rep_levels(&self) -> Option<&[i16]> { + None + } + } + + fn empty_row_group() -> RowGroupMetaData { + let schema = Arc::new(SchemaDescriptor::new(Arc::new( + Type::group_type_builder("schema") + .with_fields(Vec::new()) + .build() + .unwrap(), + ))); + RowGroupMetaData::builder(schema) + .set_num_rows(0) + .set_total_byte_size(0) + .set_column_metadata(Vec::new()) + .build() + .unwrap() + } + + fn primitive_field(arrow_type: DataType) -> ParquetField { + let primitive_type = Arc::new( + Type::primitive_type_builder("payload", PhysicalType::INT32) + .build() + .unwrap(), + ); + ParquetField { + rep_level: 0, + def_level: 0, + nullable: false, + arrow_type, + field_type: ParquetFieldType::Primitive { + col_idx: 0, + primitive_type, + }, + } + } + + #[test] + fn metadata_model_thresholds_and_fallback_are_explicit() { + let row_group = empty_row_group(); + assert_eq!( + metadata_run_threshold(&primitive_field(DataType::Int32), &row_group), + Some(INT32_RUN_THRESHOLD) + ); + assert_eq!( + metadata_run_threshold( + &primitive_field(DataType::Dictionary( + Box::new(DataType::Int32), + Box::new(DataType::Utf8), + )), + &row_group, + ), + Some(DICTIONARY_UTF8_RUN_THRESHOLD) + ); + assert_eq!( + metadata_run_threshold(&primitive_field(DataType::FixedSizeBinary(8)), &row_group,), + Some(FIXED_BINARY_RUN_THRESHOLD) + ); + assert_eq!( + metadata_run_threshold(&primitive_field(DataType::FixedSizeBinary(32)), &row_group,), + Some(FIXED_BINARY_RUN_THRESHOLD) + ); + assert_eq!( + metadata_run_threshold(&primitive_field(DataType::Int64), &row_group), + Some(INT64_RUN_THRESHOLD) + ); + assert_eq!( + metadata_run_threshold(&primitive_field(DataType::Date32), &row_group), + Some(DATE32_RUN_THRESHOLD) + ); + assert_eq!( + metadata_run_threshold(&primitive_field(DataType::Decimal128(7, 2)), &row_group), + Some(DECIMAL128_INT32_RUN_THRESHOLD) + ); + // Beyond the INT32-backed precision range Parquet switches physical + // storage, which this pass did not sample. + assert_eq!( + metadata_run_threshold(&primitive_field(DataType::Decimal128(20, 2)), &row_group), + None + ); + + let unmodeled = ParquetField { + rep_level: 0, + def_level: 0, + nullable: false, + arrow_type: DataType::Int64, + field_type: ParquetFieldType::Virtual(VirtualColumnType::RowNumber), + }; + assert_eq!(metadata_run_threshold(&unmodeled, &row_group), None); + } + + #[test] + fn row_group_statistics_match_split_selection_decisions() { + let selectors = vec![ + RowSelector::skip(2), + RowSelector::select(5), + RowSelector::skip(1), + RowSelector::select(7), + ]; + let row_group_rows = [3, 0, 5, 7]; + let selector_selection = RowSelection::from(selectors.clone()); + let mask = selectors + .iter() + .flat_map(|selector| std::iter::repeat_n(!selector.skip, selector.row_count)) + .collect::(); + let mask_selection = RowSelection::from_boolean_buffer(mask); + + for selection in [&selector_selection, &mask_selection] { + let statistics = selection_statistics(selection, &row_group_rows); + let mut remaining = selection.clone(); + for (&rows, statistics) in row_group_rows.iter().zip(statistics) { + let row_group_selection = remaining.split_off(rows); + for threshold in [ + WIDE_BYTE_RUN_THRESHOLD, + INT32_RUN_THRESHOLD, + DICTIONARY_UTF8_RUN_THRESHOLD, + DEFAULT_ROW_SELECTION_THRESHOLD, + ] { + assert_eq!( + statistics.auto_selection_strategy(threshold), + row_group_selection.auto_selection_strategy(threshold) + ); + } + } + } + } + + #[test] + fn mixed_columns_stay_aligned_when_a_batch_crosses_row_groups() { + let selection = RowSelection::from(vec![ + RowSelector::skip(1), + RowSelector::select(1), + RowSelector::skip(1), + RowSelector::select(2), + RowSelector::skip(1), + RowSelector::select(1), + RowSelector::skip(1), + RowSelector::select(2), + RowSelector::skip(1), + RowSelector::select(1), + ]); + let plan = RowSelectionExecutionPlan::try_new( + selection, + &[6, 6], + 2, + vec![ + RowSelectionStrategy::Selectors, + RowSelectionStrategy::Mask, + RowSelectionStrategy::Mask, + RowSelectionStrategy::Selectors, + ], + None, + 4, + ) + .unwrap(); + let fields = vec![ + Arc::new(Field::new("left", DataType::Int32, false)), + Arc::new(Field::new("right", DataType::Int32, false)), + ]; + let fields = Fields::from(fields); + let mut reader = PerColumnReader { + columns: vec![ + TopLevelColumnReader { + reader: Box::new(TestArrayReader::new(0..12)), + }, + TopLevelColumnReader { + reader: Box::new(TestArrayReader::new(100..112)), + }, + ], + data_type: DataType::Struct(fields.clone()), + fields, + plan: Arc::new(plan), + next_batch: 0, + buffered: None, + }; + + let mut next_batch = |batch_size| { + let rows = reader.read_records(batch_size).unwrap(); + let array = reader.consume_batch().unwrap(); + let array = array.as_any().downcast_ref::().unwrap(); + let batch = RecordBatch::from(array); + assert_eq!(batch.num_rows(), rows); + batch + }; + let first = next_batch(4); + let second = next_batch(4); + assert_eq!(next_batch(4).num_rows(), 0); + + let values = |batch: &RecordBatch, column| { + batch + .column(column) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec() + }; + assert_eq!(values(&first, 0), vec![1, 3, 4, 6]); + assert_eq!(values(&first, 1), vec![101, 103, 104, 106]); + assert_eq!(values(&second, 0), vec![8, 9, 11]); + assert_eq!(values(&second, 1), vec![108, 109, 111]); + } + + struct StrategyRowGroups { + row_groups: Vec, + } + + impl RowGroups for StrategyRowGroups { + fn num_rows(&self) -> usize { + self.row_groups + .iter() + .map(|row_group| row_group.num_rows() as usize) + .sum() + } + + fn column_chunks(&self, _i: usize) -> Result> { + unreachable!("strategy planning does not open column chunks") + } + + fn row_groups(&self) -> Box + '_> { + Box::new(self.row_groups.iter()) + } + + fn metadata(&self) -> &ParquetMetaData { + unreachable!("strategy planning does not access file metadata") + } + } + + struct ForcedMixedPlanner; + + impl ColumnSelectionPlanner for ForcedMixedPlanner { + fn strategy(&self, context: ColumnSelectionContext<'_>) -> Option { + assert_eq!(context.field.arrow_type, DataType::Int32); + Some(match &context.field.field_type { + ParquetFieldType::Virtual(VirtualColumnType::RowNumber) => { + RowSelectionStrategy::Selectors + } + ParquetFieldType::Virtual(VirtualColumnType::RowGroupIndex) => { + RowSelectionStrategy::Mask + } + _ => unreachable!(), + }) + } + } + + #[test] + fn planner_receives_top_level_field_and_records_forced_decisions() { + let schema = Arc::new(SchemaDescriptor::new(Arc::new( + Type::group_type_builder("schema") + .with_fields(Vec::new()) + .build() + .unwrap(), + ))); + let row_group = || { + RowGroupMetaData::builder(Arc::clone(&schema)) + .set_num_rows(6) + .set_total_byte_size(0) + .set_column_metadata(Vec::new()) + .build() + .unwrap() + }; + let row_groups = StrategyRowGroups { + row_groups: vec![row_group(), row_group()], + }; + let fields = [ + ParquetField { + rep_level: 0, + def_level: 0, + nullable: false, + arrow_type: DataType::Int32, + field_type: ParquetFieldType::Virtual(VirtualColumnType::RowNumber), + }, + ParquetField { + rep_level: 0, + def_level: 0, + nullable: false, + arrow_type: DataType::Int32, + field_type: ParquetFieldType::Virtual(VirtualColumnType::RowGroupIndex), + }, + ]; + let field_refs = fields.iter().collect::>(); + let selection = RowSelection::from(vec![ + RowSelector::select(2), + RowSelector::skip(2), + RowSelector::select(4), + RowSelector::skip(1), + RowSelector::select(3), + ]); + let metrics = ArrowReaderMetrics::enabled(); + + let (rows, strategies, uniform) = plan_strategies( + &row_groups, + &selection, + &field_refs, + &metrics, + &ForcedMixedPlanner, + ) + .unwrap(); + assert_eq!(rows, vec![6, 6]); + assert_eq!( + strategies, + vec![ + RowSelectionStrategy::Selectors, + RowSelectionStrategy::Mask, + RowSelectionStrategy::Selectors, + RowSelectionStrategy::Mask, + ] + ); + assert_eq!(uniform, None); + assert_eq!(metrics.row_selection_selector_decisions(), Some(2)); + assert_eq!(metrics.row_selection_mask_decisions(), Some(2)); + assert_eq!(metrics.row_selection_fallback_decisions(), Some(0)); + } +} diff --git a/parquet/src/arrow/arrow_reader/read_plan.rs b/parquet/src/arrow/arrow_reader/read_plan.rs index 63e1b1ce351c..ef14d149543c 100644 --- a/parquet/src/arrow/arrow_reader/read_plan.rs +++ b/parquet/src/arrow/arrow_reader/read_plan.rs @@ -20,7 +20,8 @@ use crate::arrow::array_reader::ArrayReader; use crate::arrow::arrow_reader::selection::{ - LoadedRowRanges, RowSelectionInner, RowSelectionPolicy, RowSelectionStrategy, + DEFAULT_ROW_SELECTION_THRESHOLD, LoadedRowRanges, RowSelectionInner, RowSelectionPolicy, + RowSelectionStrategy, }; use crate::arrow::arrow_reader::{ ArrowPredicate, ParquetRecordBatchReader, RowSelection, RowSelectionCursor, RowSelector, @@ -168,6 +169,14 @@ impl ReadPlanBuilder { selection.auto_selection_strategy(threshold) } + // Per-column planning intercepts this policy before a global + // ReadPlan is built. Unsupported projections retain the legacy + // Auto-32 behavior as a compatibility fallback. + RowSelectionPolicy::AutoPerColumn => self + .selection + .as_ref() + .map(|selection| selection.auto_selection_strategy(DEFAULT_ROW_SELECTION_THRESHOLD)) + .unwrap_or(RowSelectionStrategy::Selectors), } } @@ -194,7 +203,7 @@ impl ReadPlanBuilder { /// Like [`Self::with_predicate`], but allows additional options such as a /// match-count limit for early termination (see /// [`PredicateOptions::with_limit`]). - pub fn with_predicate_options(mut self, options: PredicateOptions<'_>) -> Result { + pub fn with_predicate_options(self, options: PredicateOptions<'_>) -> Result { let PredicateOptions { array_reader, predicate, @@ -202,6 +211,40 @@ impl ReadPlanBuilder { total_rows, } = options; + let reader = ParquetRecordBatchReader::new(array_reader, self.clone().build()); + self.apply_predicate(reader, predicate, limit, total_rows) + } + + /// Evaluates a predicate using a pre-planned record batch reader. + /// + /// This is the internal entry point used by per-column row-selection + /// planning. Predicate evaluation and selection composition remain shared + /// with the existing single-reader path. + pub(crate) fn with_predicate_reader( + self, + reader: ParquetRecordBatchReader, + predicate: &mut dyn ArrowPredicate, + ) -> Result { + self.apply_predicate(reader, predicate, None, 0) + } + + pub(crate) fn with_predicate_reader_options( + self, + reader: ParquetRecordBatchReader, + predicate: &mut dyn ArrowPredicate, + limit: Option, + total_rows: usize, + ) -> Result { + self.apply_predicate(reader, predicate, limit, total_rows) + } + + fn apply_predicate( + mut self, + reader: ParquetRecordBatchReader, + predicate: &mut dyn ArrowPredicate, + limit: Option, + total_rows: usize, + ) -> Result { // Target length for the concatenated filter output: // - Prior selection ⇒ the reader yields that many rows; `and_then` // below requires the filter output to match. @@ -213,7 +256,6 @@ impl ReadPlanBuilder { None => limit.map(|_| total_rows), }; - let reader = ParquetRecordBatchReader::new(array_reader, self.clone().build()); let mut filters = vec![]; let mut processed_rows: usize = 0; let mut matched_rows: usize = 0; @@ -745,6 +787,53 @@ mod tests { assert_eq!(selection, &expected); } + #[test] + fn auto_per_column_predicate_output_matches_the_default_policy() { + use crate::arrow::ProjectionMask; + use crate::arrow::array_reader::StructArrayReader; + use crate::arrow::array_reader::test_util::make_int32_page_reader; + use crate::arrow::arrow_reader::ArrowPredicateFn; + use arrow_schema::{DataType as ArrowType, Field, Fields}; + + // Per-column planning may decline after the predicate has run, so its + // predicate output must cost no more than the default policy's. It also + // does not benefit from mask backing: compiling the execution windows + // always consumes selectors, while the mask is built lazily and only + // when some column actually picks that strategy. + let selection_for = |policy| { + let data: Vec = (0..6).collect(); + let levels = vec![0; data.len()]; + let leaf = make_int32_page_reader(&data, &levels, &levels, 0, 0, None); + let struct_type = ArrowType::Struct(Fields::from(vec![Field::new( + "c0", + ArrowType::Int32, + false, + )])); + let struct_reader = StructArrayReader::new(struct_type, vec![leaf], 0, 0, false, None); + let mut predicate = ArrowPredicateFn::new(ProjectionMask::all(), |_| { + Ok(BooleanArray::from(vec![ + true, false, true, false, true, false, + ])) + }); + ReadPlanBuilder::new(16) + .with_row_selection_policy(policy) + .with_predicate_options(PredicateOptions::new( + Box::new(struct_reader), + &mut predicate, + )) + .unwrap() + .selection() + .cloned() + .unwrap() + }; + + let per_column = selection_for(RowSelectionPolicy::AutoPerColumn); + let default = selection_for(RowSelectionPolicy::default()); + assert_eq!(per_column.row_count(), 3); + assert_eq!(per_column, default); + assert_eq!(per_column.as_mask().is_some(), default.as_mask().is_some()); + } + #[test] fn with_predicate_options_limit_handles_null_filters() { use crate::arrow::ProjectionMask; diff --git a/parquet/src/arrow/arrow_reader/selection/boolean.rs b/parquet/src/arrow/arrow_reader/selection/boolean.rs index 1196fac2196b..9076a2dda497 100644 --- a/parquet/src/arrow/arrow_reader/selection/boolean.rs +++ b/parquet/src/arrow/arrow_reader/selection/boolean.rs @@ -243,6 +243,42 @@ pub(super) fn mask_has_at_least_runs(mask: &BooleanBuffer, min_runs: usize) -> b run_count + usize::from(last_end < total_rows) >= min_runs } +/// Counts alternating set and unset runs using word-sized transition masks. +/// +/// Unlike [`MaskRunIter`], this does not enumerate every set slice, which is +/// important when a highly fragmented mask is used for cost-model statistics. +pub(crate) fn mask_run_count(mask: &BooleanBuffer) -> usize { + if mask.is_empty() { + return 0; + } + + let chunks = mask.bit_chunks(); + let mut previous_bit = None; + let mut transitions = 0usize; + let mut count_word = |word: u64, bits: usize| { + if bits == 0 { + return; + } + // Bits are packed least-significant first. Supplying the first bit as + // its own predecessor excludes the start of the mask from the + // transition count; subsequent words use the preceding word's tail. + let preceding = previous_bit.unwrap_or(word & 1); + let valid = if bits == 64 { + u64::MAX + } else { + (1_u64 << bits) - 1 + }; + transitions += ((word ^ ((word << 1) | preceding)) & valid).count_ones() as usize; + previous_bit = Some((word >> (bits - 1)) & 1); + }; + + for word in chunks.iter() { + count_word(word, 64); + } + count_word(chunks.remainder_bits(), chunks.remainder_len()); + transitions + 1 +} + /// Split a mask into `(head, tail)` at `row_count`, preserving an empty mask tail /// when the split point is past the end. pub(super) fn split_off_mask( @@ -753,6 +789,7 @@ mod tests { fn test_mask_has_at_least_runs() { fn assert_run_count(bits: Vec, expected_runs: usize) { let mask = BooleanBuffer::from(bits); + assert_eq!(mask_run_count(&mask), expected_runs); for min_runs in 0..=expected_runs + 2 { assert_eq!( mask_has_at_least_runs(&mask, min_runs), @@ -771,11 +808,34 @@ mod tests { // Exercise the unaligned iterator path as mask-backed selections can be slices. let mask = BooleanBuffer::from(vec![true, false, false, true, true, false, true, true]) .slice(1, 6); + assert_eq!(mask_run_count(&mask), 4); for min_runs in 0..=6 { assert_eq!(mask_has_at_least_runs(&mask, min_runs), 4 >= min_runs); } } + #[test] + fn test_mask_run_count_fuzzes_word_boundaries_and_offsets() { + let mut rand = rng(); + for _ in 0..500 { + let prefix = rand.random_range(0..16); + let len = rand.random_range(0..512); + let bits = (0..prefix + len) + .map(|_| rand.random_bool(0.5)) + .collect::>(); + let expected = if len == 0 { + 0 + } else { + 1 + bits[prefix..prefix + len] + .windows(2) + .filter(|pair| pair[0] != pair[1]) + .count() + }; + let mask = BooleanBuffer::from(bits).slice(prefix, len); + assert_eq!(mask_run_count(&mask), expected); + } + } + #[test] fn test_trim_mask_fuzz_equivalence() { let mut rand = rng(); diff --git a/parquet/src/arrow/arrow_reader/selection/cursor.rs b/parquet/src/arrow/arrow_reader/selection/cursor.rs index 9a6caad24b52..acb4aaa904d8 100644 --- a/parquet/src/arrow/arrow_reader/selection/cursor.rs +++ b/parquet/src/arrow/arrow_reader/selection/cursor.rs @@ -23,7 +23,7 @@ //! the selection itself stays immutable. use super::boolean::boolean_mask_from_selectors; -use super::{RowSelection, RowSelector}; +use super::{DEFAULT_ROW_SELECTION_THRESHOLD, RowSelection, RowSelector}; use crate::errors::ParquetError; use arrow_array::BooleanArray; use arrow_buffer::BooleanBuffer; @@ -33,6 +33,7 @@ use std::sync::Arc; /// Policy for picking a strategy to materialize [`RowSelection`] during execution. #[derive(Clone, Copy, Debug, Eq, PartialEq)] +#[non_exhaustive] pub enum RowSelectionPolicy { /// Use a queue of [`RowSelector`] values Selectors, @@ -43,11 +44,23 @@ pub enum RowSelectionPolicy { /// Average selector length below which masks are preferred threshold: usize, }, + /// Choose the execution strategy independently for each top-level + /// projected Arrow field in each row group. Nested children use their + /// top-level field's strategy. + /// + /// If no column-specific rule applies, that field independently falls + /// back to the same selector-density rule as the default policy (threshold + /// `32`). A fallback on one field does not force the other fields to use + /// the same strategy. If every decision is identical, execution collapses + /// to the corresponding existing global path. + AutoPerColumn, } impl Default for RowSelectionPolicy { fn default() -> Self { - Self::Auto { threshold: 32 } + Self::Auto { + threshold: DEFAULT_ROW_SELECTION_THRESHOLD, + } } } @@ -63,6 +76,15 @@ pub(crate) enum RowSelectionStrategy { Mask, } +impl RowSelectionStrategy { + pub(crate) fn into_policy(self) -> RowSelectionPolicy { + match self { + Self::Selectors => RowSelectionPolicy::Selectors, + Self::Mask => RowSelectionPolicy::Mask, + } + } +} + /// Cursor for iterating a [`RowSelection`] during execution within a /// [`ReadPlan`](crate::arrow::arrow_reader::ReadPlan). /// @@ -358,7 +380,6 @@ impl LoadedRowRanges { .map(|range| range.end) } - #[cfg(test)] pub(crate) fn ranges(&self) -> &[Range] { &self.0 } diff --git a/parquet/src/arrow/arrow_reader/selection/mod.rs b/parquet/src/arrow/arrow_reader/selection/mod.rs index 7eacf2e52569..d9dd409fe116 100644 --- a/parquet/src/arrow/arrow_reader/selection/mod.rs +++ b/parquet/src/arrow/arrow_reader/selection/mod.rs @@ -42,12 +42,14 @@ mod boolean; mod cursor; mod ranges; mod selector; +mod window; use algebra::{ and_then_mask, and_then_row_selections, and_then_selectors_with_mask, intersect_masks, intersect_row_selections, union_masks, union_row_selections, }; pub use boolean::MaskRunIter; +pub(crate) use boolean::mask_run_count; use boolean::{ MaskSelection, limit_mask, mask_has_at_least_runs, offset_mask, split_off_mask, trim_mask, }; @@ -56,6 +58,9 @@ pub use cursor::{RowSelectionCursor, RowSelectionPolicy}; use ranges::{expand_to_batch_boundaries_from_selectors, scan_ranges_from_selectors}; pub use selector::RowSelector; use selector::{limit_selectors, offset_selectors, split_off_selectors}; +pub(crate) use window::{BatchWindow, ColumnInstruction, RowSelectionExecutionPlan}; + +pub(crate) const DEFAULT_ROW_SELECTION_THRESHOLD: usize = 32; /// [`RowSelection`] represents selecting a subset of rows /// when scanning a parquet file. diff --git a/parquet/src/arrow/arrow_reader/selection/window.rs b/parquet/src/arrow/arrow_reader/selection/window.rs new file mode 100644 index 000000000000..b5db49842e63 --- /dev/null +++ b/parquet/src/arrow/arrow_reader/selection/window.rs @@ -0,0 +1,773 @@ +// 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. + +//! Shared immutable execution windows for per-column row selection. +//! +//! A plan compiles output batch and row-group boundaries once. The logical +//! selection owns lazy, symmetric selector and mask representations shared by +//! every top-level column reader. Columns retain only their ordinal and choose +//! how to replay each row-group chunk. +//! +//! # Alignment invariant +//! +//! Think of the compiled chunks as an instruction stream and each projected +//! column as a lane. A lane interprets the stream through its own strategy, so +//! two lanes issue different reads and skips for the same chunk. What every +//! lane must agree on is how far it advances: the instructions returned by +//! [`RowSelectionExecutionPlan::lower_chunk`] always consume exactly +//! `chunk.physical_rows.len()` physical rows. +//! +//! This is the only thing keeping columns row-aligned once they stop sharing a +//! strategy, so lowering is centralised here rather than spread across the +//! execution loop. A Mask lane in particular must still skip the rows it chose +//! not to decode, even though they contribute nothing to its own output. + +use super::boolean::boolean_mask_from_selectors; +use super::{ + LoadedRowRanges, MaskRunIter, RowSelection, RowSelectionInner, RowSelectionStrategy, + RowSelector, +}; +use crate::errors::{ParquetError, Result}; +use arrow_buffer::BooleanBuffer; +use std::ops::Range; +use std::sync::{Arc, OnceLock}; + +#[derive(Debug)] +struct SelectionRepresentations { + source: SelectionSource, + selectors: OnceLock>, + mask: OnceLock, +} + +#[derive(Debug)] +enum SelectionSource { + Selectors(Arc<[RowSelector]>), + Mask(BooleanBuffer), +} + +impl SelectionRepresentations { + fn new(source: RowSelection) -> Self { + let source = match source.into_inner() { + RowSelectionInner::Selectors(selectors) => SelectionSource::Selectors(selectors.into()), + RowSelectionInner::Mask(mask) => SelectionSource::Mask((*mask).into_mask()), + }; + Self { + source, + selectors: OnceLock::new(), + mask: OnceLock::new(), + } + } + + fn selectors(&self) -> &[RowSelector] { + match &self.source { + SelectionSource::Selectors(selectors) => selectors, + SelectionSource::Mask(mask) => self + .selectors + .get_or_init(|| MaskRunIter::new(mask).collect::>().into()) + .as_ref(), + } + } + + fn mask(&self) -> &BooleanBuffer { + match &self.source { + SelectionSource::Mask(mask) => mask, + SelectionSource::Selectors(selectors) => self + .mask + .get_or_init(|| boolean_mask_from_selectors(selectors)), + } + } +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +struct SelectorPosition { + selector: usize, + offset: usize, +} + +#[derive(Debug)] +pub(crate) struct BatchWindow { + pub(crate) chunks: Range, + pub(crate) selected_rows: usize, +} + +#[derive(Debug)] +pub(crate) struct BatchWindowChunk { + pub(crate) row_group: usize, + pub(crate) physical_rows: Range, + pub(crate) selected_rows: usize, + first_selected_offset: Option, + selector_start: SelectorPosition, + selector_end: SelectorPosition, +} + +#[derive(Debug)] +pub(crate) struct MaskExecutionChunk { + pub(crate) initial_skip: usize, + pub(crate) row_count: usize, + pub(crate) mask: BooleanBuffer, +} + +/// One decode instruction for a single column inside one [`BatchWindowChunk`]. +#[derive(Debug, Eq, PartialEq)] +pub(crate) enum ColumnInstruction { + /// Advance the column reader by `rows` rows without decoding them. + Skip(usize), + /// Decode `rows` rows. + /// + /// `mask` selects which of those rows reach the output. `None` means every + /// decoded row is selected and no filtering is required. + Read { + rows: usize, + mask: Option, + }, +} + +/// The instruction stream for one column inside one chunk. +/// +/// See the [module documentation](self) for the alignment invariant these +/// instructions maintain. +#[derive(Debug)] +pub(crate) enum ChunkInstructions<'a> { + Selectors(SelectorWindowIter<'a>), + Mask { + fragments: std::vec::IntoIter, + /// Fragment whose leading skip was emitted but whose read was not. + pending: Option, + /// Physical rows of the chunk not yet accounted for. + remaining: usize, + }, +} + +impl Iterator for ChunkInstructions<'_> { + type Item = ColumnInstruction; + + fn next(&mut self) -> Option { + match self { + Self::Selectors(selectors) => selectors.next().map(|selector| { + if selector.skip { + ColumnInstruction::Skip(selector.row_count) + } else { + ColumnInstruction::Read { + rows: selector.row_count, + mask: None, + } + } + }), + Self::Mask { + fragments, + pending, + remaining, + } => { + let fragment = match pending.take() { + Some(fragment) => fragment, + None => { + let Some(fragment) = fragments.next() else { + // Rows after the last selected row in this chunk. + // Decoding stops there, but the lane still has to + // advance past them. + return (*remaining != 0) + .then(|| ColumnInstruction::Skip(std::mem::take(remaining))); + }; + *remaining -= fragment.initial_skip; + if fragment.initial_skip != 0 { + let initial_skip = fragment.initial_skip; + *pending = Some(fragment); + return Some(ColumnInstruction::Skip(initial_skip)); + } + fragment + } + }; + *remaining -= fragment.row_count; + Some(ColumnInstruction::Read { + rows: fragment.row_count, + mask: Some(fragment.mask), + }) + } + } + } +} + +#[derive(Debug)] +pub(crate) struct RowSelectionExecutionPlan { + selection: SelectionRepresentations, + batches: Box<[BatchWindow]>, + chunks: Box<[BatchWindowChunk]>, + strategies: Box<[RowSelectionStrategy]>, + loaded_row_ranges: Box<[Option]>, + row_group_offsets: Box<[usize]>, + column_count: usize, +} + +impl RowSelectionExecutionPlan { + pub(crate) fn try_new( + selection: RowSelection, + row_group_rows: &[usize], + column_count: usize, + strategies: Vec, + loaded_row_ranges: Option>>, + batch_size: usize, + ) -> Result { + if batch_size == 0 { + return Err(general_err!( + "Internal Error: per-column row selection requires a non-zero batch size" + )); + } + let matrix_len = row_group_rows.len().saturating_mul(column_count); + if strategies.len() != matrix_len { + return Err(general_err!( + "Internal Error: per-column strategy matrix has {} entries, expected {}", + strategies.len(), + matrix_len + )); + } + let loaded_row_ranges = match loaded_row_ranges { + Some(ranges) if ranges.len() != matrix_len => { + return Err(general_err!( + "Internal Error: per-column loaded-range matrix has {} entries, expected {}", + ranges.len(), + matrix_len + )); + } + Some(ranges) => ranges, + None => std::iter::repeat_with(|| None).take(matrix_len).collect(), + }; + + let mut row_group_offsets = Vec::with_capacity(row_group_rows.len()); + let mut offset = 0usize; + for &rows in row_group_rows { + row_group_offsets.push(offset); + offset = offset + .checked_add(rows) + .ok_or_else(|| general_err!("Internal Error: row-group row count overflow"))?; + } + + let selection = SelectionRepresentations::new(selection.trim()); + let (chunks, batches) = compile_windows(selection.selectors(), row_group_rows, batch_size)?; + Ok(Self { + selection, + batches: batches.into_boxed_slice(), + chunks: chunks.into_boxed_slice(), + strategies: strategies.into_boxed_slice(), + loaded_row_ranges: loaded_row_ranges.into_boxed_slice(), + row_group_offsets: row_group_offsets.into_boxed_slice(), + column_count, + }) + } + + pub(crate) fn batch(&self, batch_idx: usize) -> Option<&BatchWindow> { + self.batches.get(batch_idx) + } + + pub(crate) fn chunks(&self, batch: &BatchWindow) -> &[BatchWindowChunk] { + &self.chunks[batch.chunks.clone()] + } + + pub(crate) fn strategy( + &self, + chunk: &BatchWindowChunk, + column_idx: usize, + ) -> RowSelectionStrategy { + self.strategies[chunk.row_group * self.column_count + column_idx] + } + + /// Lower one chunk into the decode instructions for one column. + /// + /// The returned instructions consume exactly `chunk.physical_rows.len()` + /// physical rows whichever strategy the column uses. See the [module + /// documentation](self) for why that matters. + pub(crate) fn lower_chunk( + &self, + chunk: &BatchWindowChunk, + column_idx: usize, + ) -> Result> { + match self.strategy(chunk, column_idx) { + RowSelectionStrategy::Selectors => Ok(ChunkInstructions::Selectors( + self.selector_instructions(chunk), + )), + RowSelectionStrategy::Mask => Ok(ChunkInstructions::Mask { + fragments: self.mask_execution_chunks(chunk, column_idx)?.into_iter(), + pending: None, + remaining: chunk.physical_rows.len(), + }), + } + } + + fn selector_instructions(&self, chunk: &BatchWindowChunk) -> SelectorWindowIter<'_> { + SelectorWindowIter { + selectors: self.selection.selectors(), + position: chunk.selector_start, + end: chunk.selector_end, + } + } + + /// Lower a Mask window into decode-safe fragments for one column. + /// + /// Sparse page loading can leave physical gaps. Each returned fragment is + /// wholly contained in a loaded range; `initial_skip` is measured from the + /// end of the previous fragment (or the start of `chunk`). Callers must + /// skip any trailing rows not covered by the returned fragments. + /// + /// Fragments are trimmed to the first and last selected row of their range, + /// so unselected rows at either end are skipped rather than decoded and + /// filtered away. + fn mask_execution_chunks( + &self, + chunk: &BatchWindowChunk, + column_idx: usize, + ) -> Result> { + if chunk.first_selected_offset.is_none() { + return Ok(Vec::new()); + } + + let matrix_idx = chunk.row_group * self.column_count + column_idx; + let row_group_offset = self.row_group_offsets[chunk.row_group]; + let mask = self.selection.mask(); + let mut cursor = chunk.physical_rows.start; + let mut selected_rows = 0usize; + let mut fragments = Vec::new(); + + let mut append_range = |range: Range| { + let start = range.start.max(chunk.physical_rows.start); + let end = range.end.min(chunk.physical_rows.end); + if start >= end { + return; + } + + // Trim to the first and last selected row. Leading rows would be + // decoded ahead of the first output row and trailing rows after + // the last one; both are cheaper to skip. Trailing rows matter at + // row-group boundaries, where a chunk can end in a long skip run. + let window = mask.slice(start, end - start); + let mut selected = window.set_slices(); + let Some((first, first_end)) = selected.next() else { + return; + }; + let last_end = selected.last().map_or(first_end, |(_, last_end)| last_end); + + let row_count = last_end - first; + let fragment_mask = window.slice(first, row_count); + selected_rows += fragment_mask.count_set_bits(); + let first_selected = start + first; + fragments.push(MaskExecutionChunk { + initial_skip: first_selected - cursor, + row_count, + mask: fragment_mask, + }); + cursor = first_selected + row_count; + }; + + match &self.loaded_row_ranges[matrix_idx] { + Some(ranges) => { + for range in ranges.ranges() { + append_range(row_group_offset + range.start..row_group_offset + range.end); + } + } + None => append_range(chunk.physical_rows.clone()), + } + + if selected_rows != chunk.selected_rows { + return Err(general_err!( + "Internal Error: loaded pages cover {selected_rows} selected rows, expected {}", + chunk.selected_rows + )); + } + Ok(fragments) + } +} + +fn compile_windows( + selectors: &[RowSelector], + row_group_rows: &[usize], + batch_size: usize, +) -> Result<(Vec, Vec)> { + let mut row_group_ends = Vec::with_capacity(row_group_rows.len()); + let mut total_rows = 0usize; + for &rows in row_group_rows { + total_rows = total_rows + .checked_add(rows) + .ok_or_else(|| general_err!("Internal Error: row-group row count overflow"))?; + row_group_ends.push(total_rows); + } + + let selection_rows = selectors.iter().try_fold(0usize, |rows, selector| { + rows.checked_add(selector.row_count) + }); + if selection_rows.is_none_or(|rows| rows > total_rows) { + return Err(general_err!( + "Internal Error: row selection extends beyond planned row groups" + )); + } + + let mut chunks = Vec::new(); + let mut batches = Vec::new(); + let mut position = SelectorPosition { + selector: 0, + offset: 0, + }; + let mut physical_position = 0usize; + let mut row_group = 0usize; + let mut batch_chunk_start = 0usize; + let mut batch_selected_rows = 0usize; + + while position.selector < selectors.len() { + while row_group < row_group_ends.len() && physical_position == row_group_ends[row_group] { + row_group += 1; + } + if row_group == row_group_ends.len() { + return Err(general_err!( + "Internal Error: row selection extends beyond planned row groups" + )); + } + + let chunk_start = position; + let chunk_physical_start = physical_position; + let mut chunk_selected_rows = 0usize; + let mut first_selected_offset = None; + + while position.selector < selectors.len() + && physical_position < row_group_ends[row_group] + && batch_selected_rows < batch_size + { + let selector = selectors[position.selector]; + let selector_remaining = selector.row_count - position.offset; + let row_group_remaining = row_group_ends[row_group] - physical_position; + let take = if selector.skip { + selector_remaining.min(row_group_remaining) + } else { + selector_remaining + .min(row_group_remaining) + .min(batch_size - batch_selected_rows) + }; + + if !selector.skip { + first_selected_offset.get_or_insert(physical_position - chunk_physical_start); + chunk_selected_rows += take; + batch_selected_rows += take; + } + physical_position += take; + advance_selector_position(&mut position, selector.row_count, take); + } + + if physical_position == chunk_physical_start { + return Err(general_err!( + "Internal Error: per-column window compiler made no progress" + )); + } + chunks.push(BatchWindowChunk { + row_group, + physical_rows: chunk_physical_start..physical_position, + selected_rows: chunk_selected_rows, + first_selected_offset, + selector_start: chunk_start, + selector_end: position, + }); + + if batch_selected_rows == batch_size { + batches.push(BatchWindow { + chunks: batch_chunk_start..chunks.len(), + selected_rows: batch_selected_rows, + }); + batch_chunk_start = chunks.len(); + batch_selected_rows = 0; + } + } + + if batch_selected_rows != 0 { + batches.push(BatchWindow { + chunks: batch_chunk_start..chunks.len(), + selected_rows: batch_selected_rows, + }); + } else if batch_chunk_start != chunks.len() { + return Err(general_err!( + "Internal Error: per-column plan ends in an unselected window" + )); + } + + Ok((chunks, batches)) +} + +fn advance_selector_position(position: &mut SelectorPosition, selector_rows: usize, rows: usize) { + position.offset += rows; + if position.offset == selector_rows { + position.selector += 1; + position.offset = 0; + } +} + +#[derive(Debug)] +pub(crate) struct SelectorWindowIter<'a> { + selectors: &'a [RowSelector], + position: SelectorPosition, + end: SelectorPosition, +} + +impl Iterator for SelectorWindowIter<'_> { + type Item = RowSelector; + + fn next(&mut self) -> Option { + if self.position == self.end { + return None; + } + + let selector = self.selectors[self.position.selector]; + let end_offset = if self.position.selector == self.end.selector { + self.end.offset + } else { + selector.row_count + }; + let row_count = end_offset - self.position.offset; + advance_selector_position(&mut self.position, selector.row_count, row_count); + Some(RowSelector { + row_count, + skip: selector.skip, + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn strategies(row_groups: usize, columns: usize) -> Vec { + vec![RowSelectionStrategy::Selectors; row_groups * columns] + } + + #[test] + fn windows_share_selector_positions_and_cross_row_groups() { + let selection = RowSelection::from(vec![ + RowSelector::skip(2), + RowSelector::select(3), + RowSelector::skip(4), + RowSelector::select(4), + ]); + let plan = + RowSelectionExecutionPlan::try_new(selection, &[5, 8], 2, strategies(2, 2), None, 5) + .unwrap(); + + let first = plan.batch(0).unwrap(); + assert_eq!(first.selected_rows, 5); + let chunks = plan.chunks(first); + assert_eq!(chunks.len(), 2); + assert_eq!(chunks[0].row_group, 0); + assert_eq!(chunks[0].selected_rows, 3); + assert_eq!(chunks[1].row_group, 1); + assert_eq!(chunks[1].selected_rows, 2); + assert_eq!( + chunks + .iter() + .flat_map(|chunk| plan.selector_instructions(chunk)) + .collect::>(), + vec![ + RowSelector::skip(2), + RowSelector::select(3), + RowSelector::skip(4), + RowSelector::select(2), + ] + ); + + let second = plan.batch(1).unwrap(); + assert_eq!(second.selected_rows, 2); + assert_eq!( + plan.selector_instructions(&plan.chunks(second)[0]) + .collect::>(), + vec![RowSelector::select(2)] + ); + assert!(plan.batch(2).is_none()); + } + + #[test] + fn mask_representation_is_shared_and_sliced_to_first_selected_row() { + let selection = RowSelection::from(vec![ + RowSelector::skip(2), + RowSelector::select(2), + RowSelector::skip(3), + RowSelector::select(1), + ]); + let plan = RowSelectionExecutionPlan::try_new( + selection, + &[8], + 1, + vec![RowSelectionStrategy::Mask], + None, + 8, + ) + .unwrap(); + let chunk = &plan.chunks(plan.batch(0).unwrap())[0]; + let mask = plan.mask_execution_chunks(chunk, 0).unwrap(); + assert_eq!(mask.len(), 1); + let mask = &mask[0]; + assert_eq!(mask.initial_skip, 2); + assert_eq!(mask.row_count, 6); + assert_eq!( + mask.mask, + BooleanBuffer::from(vec![true, true, false, false, false, true]) + ); + assert!(std::ptr::eq(plan.selection.mask(), plan.selection.mask())); + } + + #[test] + fn dual_representations_reuse_the_source_and_cache_only_the_other_form() { + let selectors = SelectionRepresentations::new(RowSelection::from(vec![ + RowSelector::skip(1), + RowSelector::select(2), + ])); + let SelectionSource::Selectors(source_selectors) = &selectors.source else { + unreachable!() + }; + assert_eq!(selectors.selectors().as_ptr(), source_selectors.as_ptr()); + assert!(std::ptr::eq(selectors.mask(), selectors.mask())); + + let mask = SelectionRepresentations::new(RowSelection::from(BooleanBuffer::from(vec![ + false, true, true, + ]))); + let SelectionSource::Mask(source_mask) = &mask.source else { + unreachable!() + }; + assert!(std::ptr::eq(mask.mask(), source_mask)); + assert_eq!(mask.selectors().as_ptr(), mask.selectors().as_ptr()); + } + + #[test] + fn mask_fragments_do_not_cross_unloaded_page_ranges() { + let selection = RowSelection::from(vec![ + RowSelector::skip(1), + RowSelector::select(1), + RowSelector::skip(8), + RowSelector::select(1), + ]); + let loaded = LoadedRowRanges::from_selection(RowSelection::from_consecutive_ranges( + [0..4, 9..12].into_iter(), + 12, + )); + let plan = RowSelectionExecutionPlan::try_new( + selection, + &[12], + 1, + vec![RowSelectionStrategy::Mask], + Some(vec![Some(loaded)]), + 8, + ) + .unwrap(); + + let chunk = &plan.chunks(plan.batch(0).unwrap())[0]; + let fragments = plan.mask_execution_chunks(chunk, 0).unwrap(); + assert_eq!(fragments.len(), 2); + // Both fragments stop at their last selected row, so the unselected + // tail of the first loaded range moves into the next initial skip. + assert_eq!(fragments[0].initial_skip, 1); + assert_eq!(fragments[0].row_count, 1); + assert_eq!(fragments[0].mask, BooleanBuffer::from(vec![true])); + assert_eq!(fragments[1].initial_skip, 8); + assert_eq!(fragments[1].row_count, 1); + assert_eq!(fragments[1].mask, BooleanBuffer::from(vec![true])); + } + + #[test] + fn mask_fragments_skip_a_trailing_run_at_a_row_group_boundary() { + // Row group 0 selects two early rows and then skips to its boundary, + // while row group 1 keeps the batch going. Without a trailing trim the + // chunk would decode all 12 rows of row group 0 to emit 2. + let selection = RowSelection::from(vec![ + RowSelector::select(1), + RowSelector::skip(1), + RowSelector::select(1), + RowSelector::skip(9), + RowSelector::select(4), + ]); + let plan = RowSelectionExecutionPlan::try_new( + selection, + &[12, 4], + 1, + vec![RowSelectionStrategy::Mask; 2], + None, + 8, + ) + .unwrap(); + + let chunks = plan.chunks(plan.batch(0).unwrap()); + assert_eq!(chunks.len(), 2); + assert_eq!(chunks[0].physical_rows, 0..12); + + let fragments = plan.mask_execution_chunks(&chunks[0], 0).unwrap(); + assert_eq!(fragments.len(), 1); + assert_eq!(fragments[0].initial_skip, 0); + assert_eq!(fragments[0].row_count, 3); + assert_eq!( + fragments[0].mask, + BooleanBuffer::from(vec![true, false, true]) + ); + + // The lane still advances across the whole chunk. + assert_eq!( + plan.lower_chunk(&chunks[0], 0).unwrap().collect::>(), + vec![ + ColumnInstruction::Read { + rows: 3, + mask: Some(BooleanBuffer::from(vec![true, false, true])), + }, + ColumnInstruction::Skip(9), + ] + ); + } + + #[test] + fn lowered_instructions_consume_every_physical_row_of_a_chunk() { + let selection = RowSelection::from(vec![ + RowSelector::skip(2), + RowSelector::select(2), + RowSelector::skip(3), + RowSelector::select(1), + RowSelector::skip(4), + RowSelector::select(2), + ]); + let plan = RowSelectionExecutionPlan::try_new( + selection, + &[7, 7], + 2, + vec![ + RowSelectionStrategy::Selectors, + RowSelectionStrategy::Mask, + RowSelectionStrategy::Mask, + RowSelectionStrategy::Selectors, + ], + None, + 4, + ) + .unwrap(); + + let mut batch_idx = 0; + while let Some(batch) = plan.batch(batch_idx) { + for chunk in plan.chunks(batch) { + for column_idx in 0..2 { + let instructions = plan.lower_chunk(chunk, column_idx).unwrap(); + let (rows, selected) = + instructions.fold((0, 0), |(rows, selected), item| match item { + ColumnInstruction::Skip(skip) => (rows + skip, selected), + ColumnInstruction::Read { rows: read, mask } => ( + rows + read, + selected + mask.map_or(read, |mask| mask.count_set_bits()), + ), + }); + assert_eq!(rows, chunk.physical_rows.len()); + assert_eq!(selected, chunk.selected_rows); + } + } + batch_idx += 1; + } + assert_ne!(batch_idx, 0); + } +} diff --git a/parquet/src/arrow/async_reader/mod.rs b/parquet/src/arrow/async_reader/mod.rs index d386bec48d0b..d240b0abc711 100644 --- a/parquet/src/arrow/async_reader/mod.rs +++ b/parquet/src/arrow/async_reader/mod.rs @@ -936,6 +936,8 @@ where #[cfg(test)] mod tests { use super::*; + use crate::arrow::arrow_reader::RowSelectionPolicy; + use crate::arrow::arrow_reader::metrics::ArrowReaderMetrics; use crate::arrow::arrow_reader::tests::test_row_numbers_with_multiple_row_groups_helper; use crate::arrow::arrow_reader::{ ArrowPredicateFn, ParquetRecordBatchReaderBuilder, RowFilter, RowSelection, RowSelector, @@ -1261,6 +1263,97 @@ mod tests { assert_eq!(async_batches, sync_batches); } + #[tokio::test] + async fn test_async_auto_per_column_mask_plan_with_sparse_pages() { + let input = RecordBatch::try_from_iter([ + ( + "left", + Arc::new(Int32Array::from_iter_values(0..7300)) as ArrayRef, + ), + ( + "right", + Arc::new(Int32Array::from_iter_values(10_000..17_300)) as ArrayRef, + ), + ]) + .unwrap(); + let props = WriterProperties::builder() + .set_data_page_row_count_limit(32) + .set_write_batch_size(32) + .set_write_page_header_statistics(true) + .build(); + let mut data = Vec::new(); + let mut writer = ArrowWriter::try_new(&mut data, input.schema(), Some(props)).unwrap(); + writer.write(&input).unwrap(); + writer.close().unwrap(); + let data = Bytes::from(data); + + // Enough short runs for the Int32 model to select Mask, plus a large + // middle gap whose pages are not fetched. + let mut selectors = Vec::new(); + for _ in 0..399 { + selectors.push(RowSelector::select(1)); + selectors.push(RowSelector::skip(1)); + } + selectors.push(RowSelector::select(1)); + selectors.push(RowSelector::skip(6500)); + selectors.push(RowSelector::select(1)); + let selection = RowSelection::from(selectors); + + let options = ArrowReaderOptions::new().with_page_index_policy(PageIndexPolicy::Required); + let builder = ParquetRecordBatchStreamBuilder::new_with_options( + TestReader::new(data.clone()), + options, + ) + .await + .unwrap(); + let projection = ProjectionMask::leaves(builder.parquet_schema(), [0, 1]); + let predicate_projection = ProjectionMask::leaves(builder.parquet_schema(), [0]); + let build_filter = || { + let predicate = + ArrowPredicateFn::new(predicate_projection.clone(), |batch: RecordBatch| { + let values = batch.column(0).as_primitive::(); + Ok(BooleanArray::from( + values + .values() + .iter() + .map(|value| value % 3 == 0) + .collect::>(), + )) + }); + RowFilter::new(vec![Box::new(predicate)]) + }; + let metrics = ArrowReaderMetrics::enabled(); + let actual = builder + .with_projection(projection.clone()) + .with_batch_size(64) + .with_metrics(metrics.clone()) + .with_row_selection(selection.clone()) + .with_row_filter(build_filter()) + .with_row_selection_policy(RowSelectionPolicy::AutoPerColumn) + .build() + .unwrap() + .try_collect::>() + .await + .unwrap(); + + let expected = ParquetRecordBatchReaderBuilder::try_new(data) + .unwrap() + .with_projection(projection) + .with_batch_size(64) + .with_row_selection(selection) + .with_row_filter(build_filter()) + .with_row_selection_policy(RowSelectionPolicy::Selectors) + .build() + .unwrap() + .collect::>>() + .unwrap(); + + assert_eq!(actual, expected); + assert_eq!(metrics.row_selection_mask_decisions(), Some(1)); + assert_eq!(metrics.row_selection_selector_decisions(), Some(2)); + assert_eq!(metrics.row_selection_fallback_decisions(), Some(0)); + } + #[tokio::test] async fn test_fuzz_async_reader_selection() { let testdata = arrow::util::test_util::parquet_test_data(); diff --git a/parquet/src/arrow/push_decoder/reader_builder/mod.rs b/parquet/src/arrow/push_decoder/reader_builder/mod.rs index 9444161b8ccf..358bef8001bb 100644 --- a/parquet/src/arrow/push_decoder/reader_builder/mod.rs +++ b/parquet/src/arrow/push_decoder/reader_builder/mod.rs @@ -21,6 +21,9 @@ mod filter; use crate::arrow::ProjectionMask; use crate::arrow::array_reader::{ArrayReaderBuilder, CacheOptions, RowGroupCache}; use crate::arrow::arrow_reader::metrics::ArrowReaderMetrics; +use crate::arrow::arrow_reader::per_column::{ + PerColumnDecision, PerColumnReader, loaded_row_ranges_for_top_level_fields, +}; use crate::arrow::arrow_reader::selection::{LoadedRowRanges, RowSelectionStrategy}; use crate::arrow::arrow_reader::{ ParquetRecordBatchReader, PredicateOptions, ReadPlanBuilder, RowFilter, RowSelection, @@ -542,7 +545,7 @@ impl RowGroupReaderBuilder { } // Make a request for the data needed to evaluate the current predicate - let predicate = filter_info.current(); + let predicate_projection = filter_info.current().projection().clone(); // need to fetch pages the column needs for decoding, figure // that out based on the current selection and projection @@ -551,7 +554,7 @@ impl RowGroupReaderBuilder { row_count, self.batch_size, &self.metadata, - predicate.projection(), // use the predicate's projection + &predicate_projection, // use the predicate's projection ) .with_selection(plan_builder.selection()) // Fetch predicate columns; expand selection only for cached predicate columns @@ -599,36 +602,27 @@ impl RowGroupReaderBuilder { budget, } = row_group_info; - let predicate = filter_info.current(); + let predicate_projection = filter_info.current().projection().clone(); let row_group = data_request.try_into_in_memory_row_group( row_group_idx, row_count, &self.metadata, - predicate.projection(), + &predicate_projection, &mut self.buffers, )?; let cache_options = filter_info.cache_builder().producer(); - let array_reader = ArrayReaderBuilder::new(&row_group, &self.metrics) + let array_reader_builder = ArrayReaderBuilder::new(&row_group, &self.metrics) .with_batch_size(self.batch_size) .with_cache_options(Some(&cache_options)) - .with_parquet_metadata(&self.metadata) - .build_array_reader(self.fields.as_deref(), predicate.projection())?; + .with_parquet_metadata(&self.metadata); // Auto resolution and loaded ranges are projection-specific, so restore the // configured policy before preparing each predicate. plan_builder = plan_builder.with_row_selection_policy(self.row_selection_policy); - // Prepare selection execution for pages pruned during fetch. - plan_builder = prepare_selection_for_page_skipping( - plan_builder, - predicate.projection(), - self.row_group_offset_index(row_group_idx), - row_count, - ); - // When this is the final predicate in the chain and an output // limit is set, tell the filter evaluation to stop once enough // matching rows have been accumulated. @@ -637,15 +631,69 @@ impl RowGroupReaderBuilder { .then(|| budget.selected_row_limit()) .flatten(); - // Evaluate the filter via `with_predicate_options`, opting into - // early termination when this is the final predicate and an - // output limit was set. - let mut predicate_options = - PredicateOptions::new(array_reader, filter_info.current_mut()); - if let Some(limit) = predicate_limit { - predicate_options = predicate_options.with_limit(limit, row_count); + if matches!(self.row_selection_policy, RowSelectionPolicy::AutoPerColumn) { + let loaded_row_ranges = loaded_row_ranges_for_top_level_fields( + self.fields.as_deref(), + &predicate_projection, + plan_builder.selection(), + row_group.offset_index, + row_count, + ); + let predicate_reader = match PerColumnReader::try_new_with_loaded_ranges( + &row_group, + &array_reader_builder, + self.fields.as_deref(), + &predicate_projection, + &plan_builder, + self.batch_size, + &self.metrics, + loaded_row_ranges, + )? { + PerColumnDecision::Engaged(reader) => { + ParquetRecordBatchReader::new_per_column(reader, self.batch_size) + } + PerColumnDecision::Fallback(strategy) => { + let predicate_plan = prepare_selection_for_page_skipping( + plan_builder + .clone() + .with_row_selection_policy(strategy.into_policy()), + &predicate_projection, + row_group.offset_index, + row_count, + ) + .build(); + let array_reader = array_reader_builder.build_array_reader( + self.fields.as_deref(), + &predicate_projection, + )?; + ParquetRecordBatchReader::new(array_reader, predicate_plan) + } + }; + plan_builder = plan_builder.with_predicate_reader_options( + predicate_reader, + filter_info.current_mut(), + predicate_limit, + row_count, + )?; + } else { + // Prepare selection execution for pages pruned during fetch. + plan_builder = prepare_selection_for_page_skipping( + plan_builder, + &predicate_projection, + self.row_group_offset_index(row_group_idx), + row_count, + ); + let array_reader = array_reader_builder + .build_array_reader(self.fields.as_deref(), &predicate_projection)?; + // Evaluate the filter via `with_predicate_options`, opting + // into early termination for the final predicate. + let mut predicate_options = + PredicateOptions::new(array_reader, filter_info.current_mut()); + if let Some(limit) = predicate_limit { + predicate_options = predicate_options.with_limit(limit, row_count); + } + plan_builder = plan_builder.with_predicate_options(predicate_options)?; } - plan_builder = plan_builder.with_predicate_options(predicate_options)?; let row_group_info = RowGroupInfo { row_group_idx, @@ -729,12 +777,17 @@ impl RowGroupReaderBuilder { plan_builder = plan_builder.with_row_selection_policy(self.row_selection_policy); - plan_builder = prepare_selection_for_page_skipping( - plan_builder, - &self.projection, - self.row_group_offset_index(row_group_idx), - row_count, - ); + // AutoPerColumn is lowered only after the top-level fields have + // been planned in WaitingOnData. Existing policies retain the + // projection-wide preparation path exactly. + if !matches!(self.row_selection_policy, RowSelectionPolicy::AutoPerColumn) { + plan_builder = prepare_selection_for_page_skipping( + plan_builder, + &self.projection, + self.row_group_offset_index(row_group_idx), + row_count, + ); + } let row_group_info = RowGroupInfo { row_group_idx, @@ -784,23 +837,68 @@ impl RowGroupReaderBuilder { &mut self.buffers, )?; - let plan = plan_builder.build(); - // if we have any cached results, connect them up - let array_reader_builder = ArrayReaderBuilder::new(&row_group, &self.metrics) + let cache_options: Option> = cache_info + .as_ref() + .map(|cache_info| cache_info.builder().consumer()); + let mut array_reader_builder = ArrayReaderBuilder::new(&row_group, &self.metrics) .with_batch_size(self.batch_size) .with_parquet_metadata(&self.metadata); - let array_reader = if let Some(cache_info) = cache_info.as_ref() { - let cache_options: CacheOptions = cache_info.builder().consumer(); - array_reader_builder - .with_cache_options(Some(&cache_options)) - .build_array_reader(self.fields.as_deref(), &self.projection) - } else { - array_reader_builder - .build_array_reader(self.fields.as_deref(), &self.projection) - }?; + if let Some(cache_options) = cache_options.as_ref() { + array_reader_builder = + array_reader_builder.with_cache_options(Some(cache_options)); + } + + let mut plan_builder = plan_builder; + if matches!( + plan_builder.row_selection_policy(), + RowSelectionPolicy::AutoPerColumn + ) { + let loaded_row_ranges = loaded_row_ranges_for_top_level_fields( + self.fields.as_deref(), + &self.projection, + plan_builder.selection(), + row_group.offset_index, + row_count, + ); + match PerColumnReader::try_new_with_loaded_ranges( + &row_group, + &array_reader_builder, + self.fields.as_deref(), + &self.projection, + &plan_builder, + self.batch_size, + &self.metrics, + loaded_row_ranges, + )? { + PerColumnDecision::Engaged(reader) => { + let reader = + ParquetRecordBatchReader::new_per_column(reader, self.batch_size); + return Ok(NextState::result( + RowGroupDecoderState::Finished, + RowGroupBuildResult::Data { + batch_reader: reader, + remaining_budget: budget, + }, + )); + } + PerColumnDecision::Fallback(strategy) => { + plan_builder = + plan_builder.with_row_selection_policy(strategy.into_policy()); + plan_builder = prepare_selection_for_page_skipping( + plan_builder, + &self.projection, + row_group.offset_index, + row_count, + ); + } + } + } + + let array_reader = array_reader_builder + .build_array_reader(self.fields.as_deref(), &self.projection)?; - let reader = ParquetRecordBatchReader::new(array_reader, plan); + let reader = ParquetRecordBatchReader::new(array_reader, plan_builder.build()); NextState::result( RowGroupDecoderState::Finished, RowGroupBuildResult::Data {