feat(parquet): add per-column row selection policy

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"]

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);

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);
+    }
+}

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<i32> {
+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::<Int32Type>(),
+        batch.column(1).as_primitive::<Int32Type>(),
+    ];
+    let strings = [
+        batch.column(2).as_string_view(),
+        batch.column(3).as_string_view(),
+    ];
+    let dictionaries = [
+        batch
+            .column(4)
+            .as_dictionary::<Int32Type>()
+            .downcast_dict::<StringArray>()
+            .unwrap(),
+        batch
+            .column(5)
+            .as_dictionary::<Int32Type>()
+            .downcast_dict::<StringArray>()
+            .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<usize> {
     case.row_groups
         .iter()
         .copied()
@@ -56,10 +200,7 @@ fn expected_selected_global_rows(case: &CaseSpec) -> Vec<i32> {
                 .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()
 }

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,
+    },
+];