fix(parquet): keep mask execution off pages a RowGroups did not load

parquet/src/arrow/array_reader/mod.rs

@@ -53,6 +53,7 @@ pub(crate) mod test_util;
 
 // Note that this crate is public under the `experimental` feature flag.
 use crate::file::metadata::RowGroupMetaData;
+use crate::file::page_index::offset_index::OffsetIndexMetaData;
 pub use builder::{ArrayReaderBuilder, CacheOptions, CacheOptionsBuilder};
 pub use byte_array::make_byte_array_reader;
 pub use byte_array_dictionary::make_byte_array_dictionary_reader;
@@ -166,6 +167,20 @@ pub trait RowGroups {
 
     /// Returns the parquet metadata
     fn metadata(&self) -> &ParquetMetaData;
+
+    /// Returns the [`OffsetIndexMetaData`] for the row group backing this
+    /// collection, if page level information is available.
+    ///
+    /// Implementations that hold only a subset of a column chunk's pages (see
+    /// [`RowSelection::scan_ranges`]) should return the offset index they pruned
+    /// with, so that readers can confine mask based decoding to loaded pages.
+    ///
+    /// Defaults to `None`, meaning no page level information is available.
+    ///
+    /// [`RowSelection::scan_ranges`]: crate::arrow::arrow_reader::RowSelection::scan_ranges
+    fn offset_index(&self) -> Option<&[OffsetIndexMetaData]> {
+        None
+    }
 }
 
 impl RowGroups for Arc<dyn FileReader> {

parquet/src/arrow/arrow_reader/mod.rs

@@ -59,6 +59,8 @@ mod filter;
 pub mod metrics;
 mod read_plan;
 pub(crate) mod selection;
+#[cfg(test)]
+mod sparse_row_group_tests;
 pub mod statistics;
 
 /// Default batch size for reading parquet files
@@ -1615,6 +1617,21 @@ impl ParquetRecordBatchReader {
         batch_size: usize,
         selection: Option<RowSelection>,
     ) -> Result<Self> {
+        // A `RowGroups` may hold only the pages its `RowSelection` touches (see
+        // `RowSelection::scan_ranges`). Mask based decoding must not cross a page
+        // that was never loaded, so when the row group surfaces the offset index it
+        // pruned with, derive the row ranges whose pages are loaded.
+        //
+        // The derivation is conservative: a caller that loaded more pages than the
+        // selection touches still gets a subset of what it actually loaded.
+        let loaded_row_ranges = row_groups.offset_index().and_then(|offset_index| {
+            selection.as_ref()?.loaded_row_ranges(
+                &ProjectionMask::all(),
+                offset_index,
+                row_groups.num_rows(),
+            )
+        });
+
         // note metrics are not supported in this API
         let metrics = ArrowReaderMetrics::disabled();
         let array_reader = ArrayReaderBuilder::new(row_groups, &metrics)
@@ -1624,6 +1641,7 @@ impl ParquetRecordBatchReader {
 
         let read_plan = ReadPlanBuilder::new(batch_size)
             .with_selection(selection)
+            .with_loaded_row_ranges(loaded_row_ranges)
             .build();
 
         Ok(Self {

parquet/src/arrow/arrow_reader/selection/mod.rs

@@ -29,7 +29,8 @@
 //! * `ranges`: mapping a [`RowSelection`] onto page and batch ranges
 //! * `cursor`: iterating a [`RowSelection`] while reading
 
-use crate::file::page_index::offset_index::PageLocation;
+use crate::arrow::ProjectionMask;
+use crate::file::page_index::offset_index::{OffsetIndexMetaData, PageLocation};
 use arrow_array::{Array, BooleanArray};
 use arrow_buffer::{BooleanBuffer, BooleanBufferBuilder};
 use arrow_select::filter::SlicesIterator;
@@ -636,6 +637,40 @@ impl RowSelection {
         }
     }
 
+    /// Row ranges whose backing pages are loaded for every projected column,
+    /// assuming the pages fetched were those returned by [`Self::scan_ranges`].
+    ///
+    /// Reading a column chunk that holds only a subset of pages with
+    /// [`RowSelectionPolicy::Mask`] requires these ranges, so that a decoded
+    /// chunk never crosses a page that was not loaded.
+    ///
+    /// The result is conservative: a caller that loaded *more* pages than
+    /// `scan_ranges` selects (for example after expanding to batch boundaries)
+    /// still gets ranges that are a subset of what it actually loaded.
+    pub(crate) fn loaded_row_ranges(
+        &self,
+        projection: &ProjectionMask,
+        offset_index: &[OffsetIndexMetaData],
+        total_rows: usize,
+    ) -> Option<LoadedRowRanges> {
+        offset_index
+            .iter()
+            .enumerate()
+            .filter_map(|(leaf_idx, column)| {
+                let pages = column.page_locations();
+                (projection.leaf_included(leaf_idx) && !pages.is_empty()).then(|| {
+                    RowSelection::from_consecutive_ranges(
+                        self.row_ranges_for_selected_pages(pages, total_rows)
+                            .into_iter(),
+                        total_rows,
+                    )
+                })
+            })
+            .reduce(|loaded, column| loaded.intersection(&column))
+            .filter(|loaded| loaded.skipped_row_count() != 0)
+            .map(LoadedRowRanges::from_selection)
+    }
+
     /// Expands the selection to align with batch boundaries.
     /// This is needed when using cached array readers to ensure that
     /// the cached data covers full batches.

parquet/src/arrow/arrow_reader/sparse_row_group_tests.rs

@@ -0,0 +1,188 @@
+// 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.
+
+//! Tests for [`ParquetRecordBatchReader::try_new_with_row_groups`] over a
+//! [`RowGroups`](crate::arrow::array_reader::RowGroups) that only holds the data
+//! pages required by the [`RowSelection`] (the pattern
+//! [`RowSelection::scan_ranges`] exists to support).
+
+use std::sync::Arc;
+
+use arrow_array::{ArrayRef, Int64Array, RecordBatch};
+use arrow_select::concat::concat_batches;
+use bytes::Bytes;
+
+use crate::arrow::arrow_reader::selection::RowSelectionStrategy;
+use crate::arrow::arrow_reader::{
+    ArrowReaderMetadata, ArrowReaderOptions, ParquetRecordBatchReader,
+    ParquetRecordBatchReaderBuilder, ReadPlanBuilder, RowSelection, RowSelector,
+};
+use crate::arrow::in_memory_row_group::InMemoryRowGroup;
+use crate::arrow::{ArrowWriter, ProjectionMask, parquet_to_arrow_field_levels};
+use crate::file::metadata::PageIndexPolicy;
+use crate::file::properties::WriterProperties;
+
+pub(super) const TOTAL_ROWS: usize = 1000;
+pub(super) const PAGE_ROWS: usize = 200;
+pub(super) const BATCH_SIZE: usize = 200;
+
+/// A single row group of `TOTAL_ROWS` rows with `PAGE_ROWS` rows per data page.
+pub(super) fn write_test_file() -> Bytes {
+    let values: ArrayRef = Arc::new(Int64Array::from_iter_values(0..TOTAL_ROWS as i64));
+    let batch = RecordBatch::try_from_iter([("value", values)]).unwrap();
+
+    let props = WriterProperties::builder()
+        .set_data_page_row_count_limit(PAGE_ROWS)
+        .set_write_batch_size(PAGE_ROWS)
+        .build();
+
+    let mut buf = Vec::with_capacity(1024);
+    let mut writer = ArrowWriter::try_new(&mut buf, batch.schema(), Some(props)).unwrap();
+    writer.write(&batch).unwrap();
+    writer.close().unwrap();
+    Bytes::from(buf)
+}
+
+/// Many short runs concentrated in the first and fourth data pages, with whole
+/// pages skipped in between.
+///
+/// `RowSelectionPolicy::Auto` prefers Mask when `total_rows < effective_selectors
+/// * threshold`, i.e. when the *average* selector run is shorter than 32 rows —
+/// note that skips count too. Here that is `1000 < 402 * 32`.
+///
+/// The selection still touches only 2 of the 5 data pages, so a caller doing
+/// page pruning fetches 2 pages and leaves 3 unfetched.
+pub(super) fn fragmented_selection() -> RowSelection {
+    let mut selectors = Vec::new();
+    // rows 0..200 (page 0): select every other row
+    for _ in 0..PAGE_ROWS / 2 {
+        selectors.push(RowSelector::select(1));
+        selectors.push(RowSelector::skip(1));
+    }
+    // rows 200..600 (pages 1 and 2): skipped entirely
+    selectors.push(RowSelector::skip(2 * PAGE_ROWS));
+    // rows 600..800 (page 3): select every other row
+    for _ in 0..PAGE_ROWS / 2 {
+        selectors.push(RowSelector::select(1));
+        selectors.push(RowSelector::skip(1));
+    }
+    // rows 800..1000 (page 4): skipped entirely
+    selectors.push(RowSelector::skip(PAGE_ROWS));
+
+    let selection = RowSelection::from(selectors);
+    assert_eq!(selection.row_count(), PAGE_ROWS);
+    selection
+}
+
+pub(super) fn reader_options() -> ArrowReaderOptions {
+    ArrowReaderOptions::new().with_page_index_policy(PageIndexPolicy::Required)
+}
+
+/// Reading a page-pruned `RowGroups` must return exactly the selected rows.
+///
+/// `try_new_with_row_groups` inherits `RowSelectionPolicy::default()` (`Auto`),
+/// which resolves to Mask for this selection, and has no way to be told which
+/// row ranges were actually loaded. Mask execution then decodes across the
+/// unselected gap and touches pages the caller never fetched.
+#[test]
+fn test_try_new_with_row_groups_page_pruned_fragmented_selection() {
+    read_page_pruned();
+}
+
+fn read_page_pruned() {
+    let data = write_test_file();
+    let selection = fragmented_selection();
+    let projection = ProjectionMask::all();
+
+    let metadata = ArrowReaderMetadata::load(&data, reader_options()).unwrap();
+    let metadata = metadata.metadata().clone();
+
+    // Guard against the test going vacuous: it is only meaningful while `Auto`
+    // resolves this selection to Mask.
+    assert_eq!(
+        ReadPlanBuilder::new(BATCH_SIZE)
+            .with_selection(Some(selection.clone()))
+            .resolve_selection_strategy(),
+        RowSelectionStrategy::Mask,
+    );
+    let row_group_meta = metadata.row_group(0);
+    let num_columns = row_group_meta.columns().len();
+    let row_count = row_group_meta.num_rows() as usize;
+
+    let offset_index = metadata
+        .offset_index()
+        .filter(|index| !index.is_empty())
+        .map(|index| index[0].as_slice())
+        .expect("page index required");
+    let num_pages = offset_index[0].page_locations.len();
+    assert!(
+        num_pages > 2,
+        "test needs several data pages per column chunk, got {num_pages}"
+    );
+
+    let mut row_group = InMemoryRowGroup {
+        offset_index: Some(offset_index),
+        column_chunks: vec![None; num_columns],
+        row_count,
+        row_group_idx: 0,
+        metadata: metadata.as_ref(),
+    };
+
+    // Fetch only the pages the selection touches, as an engine doing page
+    // pruning would.
+    let fetch = row_group.fetch_ranges(&projection, Some(&selection), BATCH_SIZE, None);
+    assert!(
+        fetch.ranges.len() < num_pages,
+        "selection should have pruned pages: fetched {} of {num_pages} pages",
+        fetch.ranges.len()
+    );
+    let fetched: Vec<Bytes> = fetch
+        .ranges
+        .iter()
+        .map(|range| data.slice(range.start as usize..range.end as usize))
+        .collect();
+    row_group.fill_column_chunks(&projection, fetch.page_start_offsets, fetched);
+
+    let levels =
+        parquet_to_arrow_field_levels(metadata.file_metadata().schema_descr(), projection, None)
+            .unwrap();
+
+    let reader = ParquetRecordBatchReader::try_new_with_row_groups(
+        &levels,
+        &row_group,
+        BATCH_SIZE,
+        Some(selection.clone()),
+    )
+    .unwrap();
+    let actual = reader.collect::<Result<Vec<_>, _>>().expect("read failed");
+    let actual = concat_batches(&actual[0].schema(), &actual).unwrap();
+
+    // Oracle: the same selection through the high level reader, which is not
+    // affected because it prepares the plan for page skipping.
+    let expected = ParquetRecordBatchReaderBuilder::try_new_with_options(data, reader_options())
+        .unwrap()
+        .with_batch_size(BATCH_SIZE)
+        .with_row_selection(selection)
+        .build()
+        .unwrap()
+        .collect::<Result<Vec<_>, _>>()
+        .unwrap();
+    let expected = concat_batches(&expected[0].schema(), &expected).unwrap();
+
+    assert_eq!(expected.num_rows(), PAGE_ROWS);
+    assert_eq!(actual, expected);
+}

parquet/src/arrow/in_memory_row_group.rs

@@ -234,6 +234,10 @@ impl RowGroups for InMemoryRowGroup<'_> {
     fn metadata(&self) -> &ParquetMetaData {
         self.metadata
     }
+
+    fn offset_index(&self) -> Option<&[OffsetIndexMetaData]> {
+        self.offset_index.filter(|index| !index.is_empty())
+    }
 }
 
 /// An in-memory column chunk.

parquet/src/arrow/push_decoder/reader_builder/mod.rs

@@ -909,26 +909,7 @@ fn loaded_row_ranges_for_projection(
     offset_index: Option<&[OffsetIndexMetaData]>,
     total_rows: usize,
 ) -> Option<LoadedRowRanges> {
-    let selection = selection?;
-    let columns = offset_index?;
-
-    columns
-        .iter()
-        .enumerate()
-        .filter_map(|(leaf_idx, column)| {
-            let pages = column.page_locations();
-            (projection_mask.leaf_included(leaf_idx) && !pages.is_empty()).then(|| {
-                RowSelection::from_consecutive_ranges(
-                    selection
-                        .row_ranges_for_selected_pages(pages, total_rows)
-                        .into_iter(),
-                    total_rows,
-                )
-            })
-        })
-        .reduce(|loaded, column| loaded.intersection(&column))
-        .filter(|loaded| loaded.skipped_row_count() != 0)
-        .map(LoadedRowRanges::from_selection)
+    selection?.loaded_row_ranges(projection_mask, offset_index?, total_rows)
 }
 
 #[cfg(test)]