feat(parquet): RowSelection can be backed by a BooleanBuffer

parquet/Cargo.toml

@@ -276,6 +276,11 @@ name = "row_selector"
 harness = false
 required-features = ["arrow"]
 
+[[bench]]
+name = "row_selector_boolean_buffer"
+harness = false
+required-features = ["arrow"]
+
 [[bench]]
 name = "row_group_index_reader"
 required-features = ["arrow"]

parquet/benches/row_selection_cursor.rs

@@ -20,6 +20,7 @@ use std::sync::Arc;
 
 use arrow_array::builder::StringViewBuilder;
 use arrow_array::{ArrayRef, Float64Array, Int32Array, RecordBatch, StringViewArray};
+use arrow_buffer::BooleanBufferBuilder;
 use arrow_schema::{DataType, Field, Schema};
 use bytes::Bytes;
 use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main};
@@ -36,7 +37,8 @@ const BASE_SEED: u64 = 0xA55AA55A;
 const AVG_SELECTOR_LENGTHS: &[usize] = &[4, 8, 12, 16, 20, 24, 28, 32, 36, 40];
 const COLUMN_WIDTHS: &[usize] = &[2, 4, 8, 16, 32];
 const UTF8VIEW_LENS: &[usize] = &[4, 8, 16, 32, 64, 128, 256];
-const BENCH_MODES: &[BenchMode] = &[BenchMode::ReadSelector, BenchMode::ReadMask];
+const BENCH_MODES: &[BenchMode] = &[BenchMode::Selector, BenchMode::Mask, BenchMode::Auto];
+const BACKINGS: &[Backing] = &[Backing::Selectors, Backing::Mask];
 
 struct DataProfile {
     name: &'static str,
@@ -228,22 +230,25 @@ fn bench_over_lengths(
             (stats.select_ratio * 100.0).round() as u32
         );
 
-        let bench_input = BenchInput {
-            parquet_data: parquet_data.clone(),
-            selection,
-        };
-
-        for &mode in BENCH_MODES {
-            c.bench_with_input(
-                BenchmarkId::new(mode.label(), &suffix),
-                &bench_input,
-                |b, input| {
-                    b.iter(|| {
-                        let total = run_read(&input.parquet_data, &input.selection, mode.policy());
-                        hint::black_box(total);
-                    });
-                },
-            );
+        for &backing in BACKINGS {
+            let bench_input = BenchInput {
+                parquet_data: parquet_data.clone(),
+                selection: backing.convert(&selection),
+            };
+
+            for &mode in BENCH_MODES {
+                c.bench_with_input(
+                    BenchmarkId::new(backing.function_name(mode), &suffix),
+                    &bench_input,
+                    |b, input| {
+                        b.iter(|| {
+                            let total =
+                                run_read(&input.parquet_data, &input.selection, mode.policy());
+                            hint::black_box(total);
+                        });
+                    },
+                );
+            }
         }
     }
 }
@@ -448,22 +453,57 @@ fn sample_length(mean: f64, distribution: &RunDistribution, rng: &mut StdRng) ->
 
 #[derive(Clone, Copy)]
 enum BenchMode {
-    ReadSelector,
-    ReadMask,
+    Selector,
+    Mask,
+    Auto,
 }
 
 impl BenchMode {
     fn label(self) -> &'static str {
         match self {
-            BenchMode::ReadSelector => "read_selector",
-            BenchMode::ReadMask => "read_mask",
+            BenchMode::Selector => "read_selector",
+            BenchMode::Mask => "read_mask",
+            BenchMode::Auto => "read_auto",
         }
     }
 
     fn policy(self) -> RowSelectionPolicy {
         match self {
-            BenchMode::ReadSelector => RowSelectionPolicy::Selectors,
-            BenchMode::ReadMask => RowSelectionPolicy::Mask,
+            BenchMode::Selector => RowSelectionPolicy::Selectors,
+            BenchMode::Mask => RowSelectionPolicy::Mask,
+            BenchMode::Auto => RowSelectionPolicy::default(),
+        }
+    }
+}
+
+/// Which representation backs the input [`RowSelection`] handed to the reader.
+#[derive(Clone, Copy)]
+enum Backing {
+    Selectors,
+    Mask,
+}
+
+impl Backing {
+    /// Benchmark function name for this backing/mode pair. Selector-backed
+    /// inputs keep the historical names so results stay comparable with runs
+    /// from before mask-backed selections existed.
+    fn function_name(self, mode: BenchMode) -> String {
+        match self {
+            Backing::Selectors => mode.label().to_string(),
+            Backing::Mask => format!("{}_mask_backed", mode.label()),
+        }
+    }
+
+    fn convert(self, selection: &RowSelection) -> RowSelection {
+        match self {
+            Backing::Selectors => selection.clone(),
+            Backing::Mask => {
+                let mut builder = BooleanBufferBuilder::new(TOTAL_ROWS);
+                for selector in selection.iter() {
+                    builder.append_n(selector.row_count, !selector.skip);
+                }
+                RowSelection::from_boolean_buffer(builder.finish())
+            }
         }
     }
 }

parquet/benches/row_selector_boolean_buffer.rs

@@ -0,0 +1,233 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use arrow_array::{ArrayRef, BooleanArray, Int32Array, RecordBatch};
+use arrow_buffer::{BooleanBuffer, BooleanBufferBuilder};
+use arrow_schema::{DataType, Field, Schema};
+use bytes::Bytes;
+use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main};
+use parquet::arrow::ArrowWriter;
+use parquet::arrow::arrow_reader::{ParquetRecordBatchReaderBuilder, RowSelection};
+use rand::rngs::StdRng;
+use rand::{Rng, SeedableRng};
+use std::hint;
+use std::sync::Arc;
+
+const TOTAL_ROWS: usize = 3_000_000;
+const SELECTIVITY_CASES: &[Selectivity] = &[
+    Selectivity::new("select01", 1, 100),
+    Selectivity::new("select10", 1, 10),
+    Selectivity::new("select33", 1, 3),
+    Selectivity::new("select80", 4, 5),
+];
+
+/// Generates a deterministic random row mask with the specified selectivity.
+fn generate_random_row_selection(total_rows: usize, selectivity: Selectivity) -> BooleanBuffer {
+    let mut rng = StdRng::seed_from_u64(0x5E1EC7_u64 ^ selectivity.seed());
+    let bools: Vec<bool> = (0..total_rows)
+        .map(|_| rng.random_bool(selectivity.ratio()))
+        .collect();
+    BooleanBuffer::from(bools)
+}
+
+fn generate_fragmented_selection(total_rows: usize, selectivity: Selectivity) -> BooleanBuffer {
+    let mut builder = BooleanBufferBuilder::new(total_rows);
+    for row in 0..total_rows {
+        builder.append(row % selectivity.denominator < selectivity.numerator);
+    }
+    builder.finish()
+}
+
+fn generate_clustered_selection(total_rows: usize, selectivity: Selectivity) -> BooleanBuffer {
+    const RUN: usize = 8 * 1024;
+    const JITTER: usize = RUN / 2;
+    let mut rng = StdRng::seed_from_u64(0xC1057E_u64 ^ selectivity.seed());
+    let mut builder = BooleanBufferBuilder::new(total_rows);
+    let mut rows_remaining = total_rows;
+    let mut run_idx = 0usize;
+
+    while rows_remaining != 0 {
+        let run_len = rng
+            .random_range((RUN - JITTER)..=(RUN + JITTER))
+            .min(rows_remaining);
+        let selected = run_idx % selectivity.denominator < selectivity.numerator;
+        builder.append_n(run_len, selected);
+        rows_remaining -= run_len;
+        run_idx += 1;
+    }
+    builder.finish()
+}
+
+fn boolean_array(mask: &BooleanBuffer) -> BooleanArray {
+    BooleanArray::new(mask.clone(), None)
+}
+
+fn criterion_benchmark(c: &mut Criterion) {
+    let patterns = build_patterns();
+
+    let mut construction = c.benchmark_group("row_selector_boolean_buffer/construction");
+    for case in &patterns {
+        construction.bench_with_input(
+            BenchmarkId::new("from_filters", case.label()),
+            &case.mask,
+            |b, mask| {
+                b.iter(|| {
+                    let array = boolean_array(mask);
+                    let selection = RowSelection::from_filters(&[array]);
+                    hint::black_box(selection);
+                })
+            },
+        );
+
+        construction.bench_with_input(
+            BenchmarkId::new("from_boolean_buffer", case.label()),
+            &case.mask,
+            |b, mask| {
+                b.iter(|| {
+                    let selection = RowSelection::from_boolean_buffer(mask.clone());
+                    hint::black_box(selection);
+                })
+            },
+        );
+    }
+    construction.finish();
+
+    let parquet_data = write_parquet_file(TOTAL_ROWS);
+    let mut reader = c.benchmark_group("row_selector_boolean_buffer/reader");
+    for case in &patterns {
+        reader.bench_with_input(
+            BenchmarkId::new("from_filters", case.label()),
+            &case.mask,
+            |b, mask| {
+                b.iter(|| {
+                    let selection = RowSelection::from_filters(&[boolean_array(mask)]);
+                    let rows = read_rows(&parquet_data, selection);
+                    hint::black_box(rows);
+                })
+            },
+        );
+
+        reader.bench_with_input(
+            BenchmarkId::new("from_boolean_buffer", case.label()),
+            &case.mask,
+            |b, mask| {
+                b.iter(|| {
+                    let selection = RowSelection::from_boolean_buffer(mask.clone());
+                    let rows = read_rows(&parquet_data, selection);
+                    hint::black_box(rows);
+                })
+            },
+        );
+    }
+    reader.finish();
+}
+
+fn build_patterns() -> Vec<BenchCase> {
+    let mut patterns = Vec::new();
+    for selectivity in SELECTIVITY_CASES {
+        patterns.push(BenchCase::new(
+            "fragmented",
+            *selectivity,
+            generate_fragmented_selection(TOTAL_ROWS, *selectivity),
+        ));
+        patterns.push(BenchCase::new(
+            "clustered",
+            *selectivity,
+            generate_clustered_selection(TOTAL_ROWS, *selectivity),
+        ));
+        patterns.push(BenchCase::new(
+            "random",
+            *selectivity,
+            generate_random_row_selection(TOTAL_ROWS, *selectivity),
+        ));
+    }
+    patterns
+}
+
+fn write_parquet_file(total_rows: usize) -> Bytes {
+    let schema = Arc::new(Schema::new(vec![Field::new(
+        "value",
+        DataType::Int32,
+        false,
+    )]));
+    let values: ArrayRef = Arc::new(Int32Array::from_iter_values(
+        (0..total_rows).map(|row| row as i32),
+    ));
+    let batch = RecordBatch::try_new(schema.clone(), vec![values]).unwrap();
+
+    let mut writer = ArrowWriter::try_new(Vec::new(), schema, None).unwrap();
+    writer.write(&batch).unwrap();
+    Bytes::from(writer.into_inner().unwrap())
+}
+
+fn read_rows(parquet_data: &Bytes, selection: RowSelection) -> usize {
+    let reader = ParquetRecordBatchReaderBuilder::try_new(parquet_data.clone())
+        .unwrap()
+        .with_row_selection(selection)
+        .build()
+        .unwrap();
+
+    reader.map(|batch| batch.unwrap().num_rows()).sum::<usize>()
+}
+
+struct BenchCase {
+    pattern: &'static str,
+    selectivity: Selectivity,
+    mask: BooleanBuffer,
+}
+
+impl BenchCase {
+    fn new(pattern: &'static str, selectivity: Selectivity, mask: BooleanBuffer) -> Self {
+        Self {
+            pattern,
+            selectivity,
+            mask,
+        }
+    }
+
+    fn label(&self) -> String {
+        format!("{}/{}", self.pattern, self.selectivity.label)
+    }
+}
+
+#[derive(Clone, Copy)]
+struct Selectivity {
+    label: &'static str,
+    numerator: usize,
+    denominator: usize,
+}
+
+impl Selectivity {
+    const fn new(label: &'static str, numerator: usize, denominator: usize) -> Self {
+        Self {
+            label,
+            numerator,
+            denominator,
+        }
+    }
+
+    fn ratio(self) -> f64 {
+        self.numerator as f64 / self.denominator as f64
+    }
+
+    fn seed(self) -> u64 {
+        ((self.numerator as u64) << 32) ^ self.denominator as u64
+    }
+}
+
+criterion_group!(benches, criterion_benchmark);
+criterion_main!(benches);

parquet/src/arrow/arrow_reader/mod.rs

@@ -24,7 +24,10 @@ use arrow_schema::{ArrowError, DataType as ArrowType, FieldRef, Schema, SchemaRe
 use arrow_select::filter::filter_record_batch;
 pub use filter::{ArrowPredicate, ArrowPredicateFn, RowFilter};
 use selection::MaskCursor;
-pub use selection::{RowSelection, RowSelectionCursor, RowSelectionPolicy, RowSelector};
+pub use selection::{
+    MaskRunIter, RowSelection, RowSelectionCursor, RowSelectionIter, RowSelectionPolicy,
+    RowSelector,
+};
 use std::fmt::{Debug, Formatter};
 use std::sync::Arc;