This is an automated email from the ASF dual-hosted git repository.
alamb pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
The following commit(s) were added to refs/heads/main by this push:
new 0cc07133c6 fix(parquet): support mask filtering across skipped pages
(#10288)
0cc07133c6 is described below
commit 0cc07133c659489c2c03a2ad08994b46910e152f
Author: Huang Qiwei <[email protected]>
AuthorDate: Tue Jul 21 04:19:08 2026 +0800
fix(parquet): support mask filtering across skipped pages (#10288)
# Which issue does this PR close?
- Closes #8845.
- Related: #9239, #7456, #8846, and #9591.
# Rationale for this change
Mask-backed row selection decodes every row covered by a mask chunk
before applying Arrow's filter kernel. Page pruning can leave sparse
in-memory column chunks where complete pages were never fetched. If a
mask chunk crosses one of those gaps, decoding requests data that is not
present and fails with an invalid sparse-page offset.
Previously, `Auto` avoided this failure by falling back to selectors
when page pruning was detected, while an explicit
`RowSelectionPolicy::Mask` could still fail. This PR makes skipped pages
safe for Mask execution, allowing explicit Mask to work and `Auto` to
retain its density-based strategy choice.
# What changes are included in this PR?
The reader now computes `LoadedRowRanges`: absolute row-group intervals
whose backing pages are loaded for every Parquet leaf in the current
projection.
For each predicate projection and the final output projection, it:
1. Maps the pages selected for each projected leaf to absolute row
ranges.
2. Intersects those ranges across the projected leaves.
3. Gives the resulting ranges to Mask execution.
4. Ends each decoded mask chunk at a loaded-range boundary.
5. Uses `skip_records()` to cross gaps without decoding missing pages.
These ranges are decoder safety boundaries, not output batch boundaries.
`ParquetRecordBatchReader` accumulates multiple safe chunks and their
mask bits until it reaches `batch_size`, then consumes and filters once.
Sparse pages therefore do not fragment one logical batch into one batch
per loaded range.
When no projected pages were skipped, `LoadedRowRanges` is omitted and
the existing Mask fast path is unchanged. No masks are intersected or
rewritten; the original selection mask is still applied by Arrow's
filter kernel.
## Loaded row range example
Consider two projected columns with different page boundaries and the
selection mask `100000000001`:
<table>
<thead>
<tr><th>Row partition</th><th><code>[0, 4)</code></th><th><code>[4,
6)</code></th><th><code>[6, 8)</code></th><th><code>[8,
10)</code></th><th><code>[10, 12)</code></th></tr>
</thead>
<tbody>
<tr><th>Selection
mask</th><td><code>1000</code></td><td><code>00</code></td><td><code>00</code></td><td><code>00</code></td><td><code>01</code></td></tr>
<tr><th>Column A pages</th><td>loaded A0</td><td colspan="2">not loaded
A1 <code>[4, 8)</code></td><td colspan="2">loaded A2 <code>[8,
12)</code></td></tr>
<tr><th>Column B pages</th><td colspan="2">loaded B0 <code>[0,
6)</code></td><td colspan="2">not loaded B1 <code>[6,
10)</code></td><td>loaded B2</td></tr>
<tr><th>LoadedRowRanges</th><td><strong>loaded</strong></td><td
colspan="3">unloaded gap</td><td><strong>loaded</strong></td></tr>
</tbody>
</table>
```text
Column A: [0, 4) U [8, 12)
Column B: [0, 6) U [10, 12)
LoadedRowRanges: [0, 4) U [10, 12)
```
The intersection does not need to match any one column's page
boundaries. It is a row-domain guarantee: every row in the result has
backing page data in every projected leaf. For example, `[0, 4)` is a
valid subset of Column B's `[0, 6)` page.
## Mask execution and batch behavior
With `batch_size = 12`, the reader processes the example as follows:
<table>
<thead>
<tr><th>Cursor position</th><th>Decoder action</th><th>Selected rows
accumulated</th></tr>
</thead>
<tbody>
<tr><td>0</td><td>Read <code>[0, 4)</code> with mask
<code>1000</code></td><td>1</td></tr>
<tr><td>4</td><td><code>skip_records(7)</code> to row 11, then read
<code>[11, 12)</code> with mask <code>1</code></td><td>2</td></tr>
<tr><td>End</td><td>Consume once and apply the combined mask</td><td>one
2-row batch</td></tr>
</tbody>
</table>
The unloaded gap is never decoded, while the loaded-range boundary
remains invisible to callers as a batch boundary.
# Are these changes tested?
Coverage includes:
- page-to-row-range conversion and intersections across different column
page boundaries;
- explicit Mask and `Auto`, with predicate caching enabled and disabled;
- reads with and without an initial row selection;
- multiple predicates with projection-specific loaded ranges;
- sparse loaded ranges coalescing to the requested logical batch size;
- nested `List<Struct<...>>` projections whose leaves have different
page boundaries;
- masks trimmed before the end of their enclosing loaded range.
Validated with:
```bash
cargo test -p parquet --lib --features arrow
cargo test -p parquet --test arrow_reader --features async
```
The current GitHub CI run is green, including Parquet tests, Clippy,
MSRV, macOS, Windows, wasm32, and PySpark integration.
# Are there any user-facing changes?
There are no public API or source compatibility changes.
Explicit Mask execution no longer fails when page pruning creates sparse
column data. `Auto` no longer falls back to selectors solely because
pages were pruned, and sparse loaded ranges no longer fragment logical
output batches. These changes can affect execution performance
characteristics, but not output semantics.
---------
Co-authored-by: Andrew Lamb <[email protected]>
---
parquet/src/arrow/arrow_reader/mod.rs | 243 +++++++++-----
parquet/src/arrow/arrow_reader/read_plan.rs | 39 ++-
parquet/src/arrow/arrow_reader/selection.rs | 250 ++++++++++++---
.../src/arrow/push_decoder/reader_builder/mod.rs | 176 +++++++---
parquet/tests/arrow_reader/row_filter/async.rs | 355 ++++++++++++++++++---
parquet/tests/arrow_reader/row_filter/sync.rs | 2 +-
6 files changed, 845 insertions(+), 220 deletions(-)
diff --git a/parquet/src/arrow/arrow_reader/mod.rs
b/parquet/src/arrow/arrow_reader/mod.rs
index 814b825050..621c1c812a 100644
--- a/parquet/src/arrow/arrow_reader/mod.rs
+++ b/parquet/src/arrow/arrow_reader/mod.rs
@@ -18,10 +18,12 @@
//! Contains reader which reads parquet data into arrow [`RecordBatch`]
use arrow_array::cast::AsArray;
-use arrow_array::{Array, RecordBatch, RecordBatchReader};
+use arrow_array::{BooleanArray, RecordBatch, RecordBatchReader};
+use arrow_buffer::{BooleanBuffer, BooleanBufferBuilder};
use arrow_schema::{ArrowError, DataType as ArrowType, FieldRef, Schema,
SchemaRef};
use arrow_select::filter::filter_record_batch;
pub use filter::{ArrowPredicate, ArrowPredicateFn, RowFilter};
+use selection::MaskCursor;
pub use selection::{RowSelection, RowSelectionCursor, RowSelectionPolicy,
RowSelector};
use std::fmt::{Debug, Formatter};
use std::sync::Arc;
@@ -1349,6 +1351,140 @@ pub struct ParquetRecordBatchReader {
read_plan: ReadPlan,
}
+/// Accumulates filter masks for decoded chunks in one logical output batch.
+///
+/// The first chunk keeps its [`BooleanBuffer`] without copying. A second chunk
+/// promotes the accumulator to a [`BooleanBufferBuilder`], and later chunks
are
+/// appended to it. For example, chunks `1000` and `1` become `10001`:
+///
+/// ```text
+/// append(1000) append(1)
+/// Empty ───────────────▶ Single ───────────────▶ Combined
+/// 1000 10001
+/// (zero copy) (promoted to builder)
+/// ```
+///
+/// The combined mask lines up with the rows that [`ArrayReader::read_records`]
+/// buffered across the decoded chunks (see [`read_mask_batch`]). Consuming the
+/// buffered batch and filtering it with the accumulated mask yields the
output.
+///
+/// ```text
+/// decoded rows: 0 1 2 3 11 <-- buffered by the array reader
+/// chunk masks: [1 0 0 0] [1]
+/// finish(): 1 0 0 0 1 <-- filters the whole batch in one pass
+/// ```
+#[derive(Default)]
+enum FilterMaskAccumulator {
+ #[default]
+ Empty,
+ Single(BooleanBuffer),
+ Combined(BooleanBufferBuilder),
+}
+
+impl FilterMaskAccumulator {
+ fn append(&mut self, mask: BooleanBuffer) {
+ *self = match std::mem::take(self) {
+ Self::Empty => Self::Single(mask),
+ Self::Single(first) => {
+ let mut combined = BooleanBufferBuilder::new(first.len() +
mask.len());
+ combined.append_buffer(&first);
+ combined.append_buffer(&mask);
+ Self::Combined(combined)
+ }
+ Self::Combined(mut combined) => {
+ combined.append_buffer(&mask);
+ Self::Combined(combined)
+ }
+ };
+ }
+
+ fn finish(self) -> Option<BooleanBuffer> {
+ match self {
+ Self::Empty => None,
+ Self::Single(mask) => Some(mask),
+ Self::Combined(combined) => Some(combined.build()),
+ }
+ }
+}
+
+/// Converts the projection buffered by `array_reader` into a record batch.
+fn consume_record_batch(array_reader: &mut dyn ArrayReader) ->
Result<RecordBatch> {
+ let array = array_reader.consume_batch()?;
+ let struct_array = array.as_struct_opt().ok_or_else(|| {
+ ArrowError::ParquetError("Struct array reader should return struct
array".to_string())
+ })?;
+ Ok(RecordBatch::from(struct_array))
+}
+
+/// Reads one logical Mask batch, potentially spanning multiple loaded ranges.
+///
+/// Each [`MaskCursor`] chunk is safe to decode because it stays within loaded
+/// pages. Gaps are crossed with [`ArrayReader::skip_records`], while decoded
+/// arrays and their mask fragments remain buffered. Once `batch_size` selected
+/// rows have accumulated, this consumes the underlying batch and filters it
+/// once with the combined mask.
+fn read_mask_batch(
+ array_reader: &mut dyn ArrayReader,
+ mask_cursor: &mut MaskCursor,
+ batch_size: usize,
+) -> Result<Option<RecordBatch>> {
+ let mut selected_rows = 0;
+ let mut filter_mask = FilterMaskAccumulator::default();
+
+ while selected_rows < batch_size && !mask_cursor.is_empty() {
+ let mask_chunk = mask_cursor.next_chunk(batch_size - selected_rows)?;
+
+ if mask_chunk.initial_skip > 0 {
+ let skipped = array_reader.skip_records(mask_chunk.initial_skip)?;
+ if skipped != mask_chunk.initial_skip {
+ return Err(general_err!(
+ "failed to skip rows, expected {}, got {}",
+ mask_chunk.initial_skip,
+ skipped
+ ));
+ }
+ }
+
+ let mask = mask_cursor.mask_values_for(&mask_chunk)?;
+ let read = array_reader.read_records(mask_chunk.chunk_rows)?;
+ if read == 0 {
+ return Err(general_err!(
+ "reached end of column while expecting {} rows",
+ mask_chunk.chunk_rows
+ ));
+ }
+ if read != mask_chunk.chunk_rows {
+ return Err(general_err!(
+ "insufficient rows read from array reader - expected {}, got
{}",
+ mask_chunk.chunk_rows,
+ read
+ ));
+ }
+
+ filter_mask.append(mask.values().clone());
+ selected_rows += mask_chunk.selected_rows;
+ }
+
+ if selected_rows == 0 {
+ return Ok(None);
+ }
+
+ let filter_mask = filter_mask
+ .finish()
+ .ok_or_else(|| general_err!("Internal Error: decoded Mask batch has no
filter values"))?;
+ let batch = consume_record_batch(array_reader)?;
+ let filtered_batch = filter_record_batch(&batch,
&BooleanArray::from(filter_mask))?;
+ if filtered_batch.num_rows() != selected_rows {
+ return Err(general_err!(
+ "filtered rows mismatch selection - expected {}, got {}",
+ selected_rows,
+ filtered_batch.num_rows()
+ ));
+ }
+
+ Ok(Some(filtered_batch))
+}
+
impl Debug for ParquetRecordBatchReader {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
f.debug_struct("ParquetRecordBatchReader")
@@ -1383,74 +1519,7 @@ impl ParquetRecordBatchReader {
}
match self.read_plan.row_selection_cursor_mut() {
RowSelectionCursor::Mask(mask_cursor) => {
- // Stream the record batch reader using contiguous segments of
the selection
- // mask, avoiding the need to materialize intermediate
`RowSelector` ranges.
- while !mask_cursor.is_empty() {
- let Some(mask_chunk) =
mask_cursor.next_mask_chunk(batch_size) else {
- return Ok(None);
- };
-
- if mask_chunk.initial_skip > 0 {
- let skipped =
self.array_reader.skip_records(mask_chunk.initial_skip)?;
- if skipped != mask_chunk.initial_skip {
- return Err(general_err!(
- "failed to skip rows, expected {}, got {}",
- mask_chunk.initial_skip,
- skipped
- ));
- }
- }
-
- if mask_chunk.chunk_rows == 0 {
- if mask_cursor.is_empty() && mask_chunk.selected_rows
== 0 {
- return Ok(None);
- }
- continue;
- }
-
- let mask = mask_cursor.mask_values_for(&mask_chunk)?;
-
- let read =
self.array_reader.read_records(mask_chunk.chunk_rows)?;
- if read == 0 {
- return Err(general_err!(
- "reached end of column while expecting {} rows",
- mask_chunk.chunk_rows
- ));
- }
- if read != mask_chunk.chunk_rows {
- return Err(general_err!(
- "insufficient rows read from array reader -
expected {}, got {}",
- mask_chunk.chunk_rows,
- read
- ));
- }
-
- let array = self.array_reader.consume_batch()?;
- // The column reader exposes the projection as a struct
array; convert this
- // into a record batch before applying the boolean filter
mask.
- let struct_array = array.as_struct_opt().ok_or_else(|| {
- ArrowError::ParquetError(
- "Struct array reader should return struct
array".to_string(),
- )
- })?;
-
- let filtered_batch =
- filter_record_batch(&RecordBatch::from(struct_array),
&mask)?;
-
- if filtered_batch.num_rows() != mask_chunk.selected_rows {
- return Err(general_err!(
- "filtered rows mismatch selection - expected {},
got {}",
- mask_chunk.selected_rows,
- filtered_batch.num_rows()
- ));
- }
-
- if filtered_batch.num_rows() == 0 {
- continue;
- }
-
- return Ok(Some(filtered_batch));
- }
+ return read_mask_batch(self.array_reader.as_mut(),
mask_cursor, batch_size);
}
RowSelectionCursor::Selectors(selectors_cursor) => {
while read_records < batch_size &&
!selectors_cursor.is_empty() {
@@ -1494,15 +1563,11 @@ impl ParquetRecordBatchReader {
RowSelectionCursor::All => {
self.array_reader.read_records(batch_size)?;
}
- };
-
- let array = self.array_reader.consume_batch()?;
- let struct_array = array.as_struct_opt().ok_or_else(|| {
- ArrowError::ParquetError("Struct array reader should return struct
array".to_string())
- })?;
+ }
- Ok(if struct_array.len() > 0 {
- Some(RecordBatch::from(struct_array))
+ let batch = consume_record_batch(self.array_reader.as_mut())?;
+ Ok(if batch.num_rows() > 0 {
+ Some(batch)
} else {
None
})
@@ -1623,7 +1688,7 @@ pub(crate) mod tests {
Time64MicrosecondType,
};
use arrow_array::*;
- use arrow_buffer::{ArrowNativeType, Buffer, IntervalDayTime, NullBuffer,
i256};
+ use arrow_buffer::{ArrowNativeType, BooleanBuffer, Buffer,
IntervalDayTime, NullBuffer, i256};
use arrow_data::{ArrayData, ArrayDataBuilder};
use arrow_schema::{DataType as ArrowDataType, Field, Fields, Schema,
SchemaRef, TimeUnit};
use arrow_select::concat::concat_batches;
@@ -1631,6 +1696,28 @@ pub(crate) mod tests {
use half::f16;
use num_traits::PrimInt;
+ #[test]
+ fn filter_mask_accumulator_handles_empty_single_and_multiple_chunks() {
+ let first = BooleanBuffer::from(vec![true, false, false, false]);
+ let second = BooleanBuffer::from(vec![true]);
+ let third = BooleanBuffer::from(vec![false, true]);
+
+ assert!(super::FilterMaskAccumulator::default().finish().is_none());
+
+ let mut single = super::FilterMaskAccumulator::default();
+ single.append(first.clone());
+ assert_eq!(single.finish().unwrap(), first);
+
+ let mut combined = super::FilterMaskAccumulator::default();
+ combined.append(BooleanBuffer::from(vec![true, false, false, false]));
+ combined.append(second);
+ combined.append(third);
+ assert_eq!(
+ combined.finish().unwrap(),
+ BooleanBuffer::from(vec![true, false, false, false, true, false,
true])
+ );
+ }
+
#[test]
fn test_arrow_reader_all_columns() {
let file =
get_test_file("parquet/generated_simple_numerics/blogs.parquet");
diff --git a/parquet/src/arrow/arrow_reader/read_plan.rs
b/parquet/src/arrow/arrow_reader/read_plan.rs
index 3d80526c82..09381d8d8c 100644
--- a/parquet/src/arrow/arrow_reader/read_plan.rs
+++ b/parquet/src/arrow/arrow_reader/read_plan.rs
@@ -19,8 +19,9 @@
//! from a Parquet file
use crate::arrow::array_reader::ArrayReader;
-use crate::arrow::arrow_reader::selection::RowSelectionPolicy;
-use crate::arrow::arrow_reader::selection::RowSelectionStrategy;
+use crate::arrow::arrow_reader::selection::{
+ LoadedRowRanges, RowSelectionPolicy, RowSelectionStrategy,
+};
use crate::arrow::arrow_reader::{
ArrowPredicate, ParquetRecordBatchReader, RowSelection,
RowSelectionCursor, RowSelector,
};
@@ -29,6 +30,7 @@ use arrow_array::{Array, BooleanArray};
use arrow_buffer::BooleanBuffer;
use arrow_select::filter::prep_null_mask_filter;
use std::collections::VecDeque;
+use std::sync::Arc;
/// Options for [`ReadPlanBuilder::with_predicate_options`].
pub struct PredicateOptions<'a> {
@@ -84,6 +86,8 @@ pub struct ReadPlanBuilder {
selection: Option<RowSelection>,
/// Policy to use when materializing the row selection
row_selection_policy: RowSelectionPolicy,
+ /// Row ranges with page data loaded for the current projection.
+ loaded_row_ranges: Option<Arc<LoadedRowRanges>>,
}
impl ReadPlanBuilder {
@@ -93,6 +97,7 @@ impl ReadPlanBuilder {
batch_size,
selection: None,
row_selection_policy: RowSelectionPolicy::default(),
+ loaded_row_ranges: None,
}
}
@@ -110,6 +115,11 @@ impl ReadPlanBuilder {
self
}
+ pub(crate) fn with_loaded_row_ranges(mut self, ranges:
Option<LoadedRowRanges>) -> Self {
+ self.loaded_row_ranges = ranges.map(Arc::new);
+ self
+ }
+
/// Returns the current row selection policy
pub fn row_selection_policy(&self) -> &RowSelectionPolicy {
&self.row_selection_policy
@@ -304,17 +314,17 @@ impl ReadPlanBuilder {
batch_size,
selection,
row_selection_policy: _,
+ loaded_row_ranges,
} = self;
let selection = selection.map(|s| s.trim());
let row_selection_cursor = selection
.map(|s| {
- let trimmed = s.trim();
- let selectors: Vec<RowSelector> = trimmed.into();
+ let selectors: Vec<RowSelector> = s.into();
match selection_strategy {
RowSelectionStrategy::Mask => {
- RowSelectionCursor::new_mask_from_selectors(selectors)
+ RowSelectionCursor::new_mask_from_selectors(selectors,
loaded_row_ranges)
}
RowSelectionStrategy::Selectors =>
RowSelectionCursor::new_selectors(selectors),
}
@@ -476,6 +486,25 @@ mod tests {
);
}
+ #[test]
+ fn mask_plan_trims_trailing_skips_before_chunking() {
+ let mut plan = ReadPlanBuilder::new(8)
+ .with_selection(Some(RowSelection::from(vec![
+ RowSelector::select(1),
+ RowSelector::skip(7),
+ ])))
+ .with_row_selection_policy(RowSelectionPolicy::Mask)
+ .build();
+ let RowSelectionCursor::Mask(cursor) = plan.row_selection_cursor_mut()
else {
+ panic!("expected a Mask cursor");
+ };
+
+ let chunk = cursor.next_chunk(8).unwrap();
+ assert_eq!(chunk.chunk_rows, 1);
+ assert_eq!(chunk.selected_rows, 1);
+ assert!(cursor.is_empty());
+ }
+
#[test]
fn with_predicate_options_limit_pads_tail_when_no_prior_selection() {
use crate::arrow::ProjectionMask;
diff --git a/parquet/src/arrow/arrow_reader/selection.rs
b/parquet/src/arrow/arrow_reader/selection.rs
index 2ddf812f9c..a60feabfc2 100644
--- a/parquet/src/arrow/arrow_reader/selection.rs
+++ b/parquet/src/arrow/arrow_reader/selection.rs
@@ -15,20 +15,17 @@
// specific language governing permissions and limitations
// under the License.
-use crate::arrow::ProjectionMask;
use crate::errors::ParquetError;
-use crate::file::page_index::offset_index::{OffsetIndexMetaData, PageLocation};
+use crate::file::page_index::offset_index::PageLocation;
use arrow_array::{Array, BooleanArray};
use arrow_buffer::{BooleanBuffer, BooleanBufferBuilder};
use arrow_select::filter::SlicesIterator;
use std::cmp::Ordering;
use std::collections::VecDeque;
use std::ops::Range;
+use std::sync::Arc;
/// Policy for picking a strategy to materialise [`RowSelection`] during
execution.
-///
-/// Note that this is a user-provided preference, and the actual strategy used
-/// may differ based on safety considerations (e.g. page skipping).
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum RowSelectionPolicy {
/// Use a queue of [`RowSelector`] values
@@ -50,8 +47,8 @@ impl Default for RowSelectionPolicy {
/// Fully resolved strategy for materializing [`RowSelection`] during
execution.
///
-/// This is determined from a combination of user preference (via
[`RowSelectionPolicy`])
-/// and safety considerations (e.g. page skipping).
+/// This is determined by [`RowSelectionPolicy`], including selector density
for
+/// [`RowSelectionPolicy::Auto`].
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub(crate) enum RowSelectionStrategy {
/// Use a queue of [`RowSelector`] values
@@ -250,37 +247,32 @@ impl RowSelection {
ranges
}
- /// Returns true if this selection would skip any data pages within the
provided columns
- fn selection_skips_any_page(
+ /// Returns the complete row ranges of the pages selected by
[`Self::scan_ranges`].
+ pub(crate) fn row_ranges_for_selected_pages(
&self,
- projection: &ProjectionMask,
- columns: &[OffsetIndexMetaData],
- ) -> bool {
- columns.iter().enumerate().any(|(leaf_idx, column)| {
- if !projection.leaf_included(leaf_idx) {
- return false;
- }
-
- let locations = column.page_locations();
- if locations.is_empty() {
- return false;
+ page_locations: &[PageLocation],
+ total_rows: usize,
+ ) -> Vec<Range<usize>> {
+ let mut selected_pages =
self.scan_ranges(page_locations).into_iter().peekable();
+ let mut row_ranges = Vec::new();
+
+ for (idx, page) in page_locations.iter().enumerate() {
+ let Some(selected_page) = selected_pages.peek() else {
+ break;
+ };
+ if selected_page.start != page.offset as u64 {
+ continue;
}
+ selected_pages.next();
- let ranges = self.scan_ranges(locations);
- !ranges.is_empty() && ranges.len() < locations.len()
- })
- }
-
- /// Returns true if selectors should be forced, preventing mask
materialisation
- pub(crate) fn should_force_selectors(
- &self,
- projection: &ProjectionMask,
- offset_index: Option<&[OffsetIndexMetaData]>,
- ) -> bool {
- match offset_index {
- Some(columns) => self.selection_skips_any_page(projection,
columns),
- None => false,
+ let end = page_locations
+ .get(idx + 1)
+ .map(|next| next.first_row_index as usize)
+ .unwrap_or(total_rows);
+ row_ranges.push(page.first_row_index as usize..end);
}
+
+ row_ranges
}
/// Splits off the first `row_count` from this [`RowSelection`]
@@ -765,11 +757,32 @@ fn union_row_selections(left: &[RowSelector], right:
&[RowSelector]) -> RowSelec
///
/// This is best for dense selections where there are many small skips
/// or selections. For example, selecting every other row.
+///
+/// When page pruning produces sparse column data, `loaded_row_ranges` limits
+/// each decoded chunk to rows whose pages are loaded for every projected leaf.
+/// For example, two projected columns can have different page boundaries:
+///
+/// ```text
+/// Row ranges: [0, 4) [4, 6) [6, 8) [8, 10) [10, 12)
+/// Selection mask: 1000 00 00 00 01
+/// Column A pages: loaded | missing [4, 8) | loaded [8, 12)
+/// Column B pages: loaded [0, 6) | missing [6, 10) | loaded
+/// LoadedRowRanges: [0, 4) [10, 12)
+/// ```
+///
+/// The first chunk decodes `[0, 4)` with mask `1000`. The next chunk skips to
+/// row 11 and decodes `[11, 12)` with mask `1`. The loaded ranges are decode
+/// boundaries, not output batch boundaries: [`ParquetRecordBatchReader`]
+/// accumulates both chunks and applies the combined mask `10001` once.
+///
+/// [`ParquetRecordBatchReader`]: super::ParquetRecordBatchReader
#[derive(Debug)]
pub struct MaskCursor {
mask: BooleanBuffer,
/// Current absolute offset into the selection
position: usize,
+ /// Row ranges whose backing pages are loaded for every projected column.
+ loaded_row_ranges: Option<Arc<LoadedRowRanges>>,
}
impl MaskCursor {
@@ -780,13 +793,19 @@ impl MaskCursor {
/// Advance through the mask representation, producing the next chunk
summary
pub fn next_mask_chunk(&mut self, batch_size: usize) -> Option<MaskChunk> {
- let (initial_skip, chunk_rows, selected_rows, mask_start,
end_position) = {
- let mask = &self.mask;
+ if self.is_empty() {
+ return None;
+ }
- if self.position >= mask.len() {
- return None;
- }
+ Some(self.next_mask_chunk_non_empty(batch_size))
+ }
+ /// Produces the next chunk for a non-empty, trailing-skip-free mask.
+ fn next_mask_chunk_non_empty(&mut self, batch_size: usize) -> MaskChunk {
+ debug_assert!(!self.is_empty());
+
+ let (initial_skip, chunk_rows, selected_rows, mask_start,
end_position) = {
+ let mask = &self.mask;
let start_position = self.position;
let mut cursor = start_position;
let mut initial_skip = 0;
@@ -795,6 +814,10 @@ impl MaskCursor {
initial_skip += 1;
cursor += 1;
}
+ debug_assert!(
+ cursor < mask.len(),
+ "ReadPlan must remove trailing skips from Mask selections"
+ );
let mask_start = cursor;
let mut chunk_rows = 0;
@@ -816,11 +839,63 @@ impl MaskCursor {
self.position = end_position;
- Some(MaskChunk {
+ MaskChunk {
initial_skip,
chunk_rows,
selected_rows,
mask_start,
+ }
+ }
+
+ /// Returns the next non-empty mask chunk without crossing an unloaded row
range.
+ ///
+ /// The [`ReadPlan`](crate::arrow::arrow_reader::ReadPlan) removes trailing
+ /// skips before constructing this cursor. Callers therefore only invoke
+ /// this method for a non-empty mask that has another selected row.
+ pub(crate) fn next_chunk(&mut self, batch_size: usize) ->
Result<MaskChunk, ParquetError> {
+ debug_assert!(batch_size > 0);
+ debug_assert!(!self.is_empty());
+
+ if self.loaded_row_ranges.is_none() {
+ return Ok(self.next_mask_chunk_non_empty(batch_size));
+ }
+
+ let start_position = self.position;
+ let mut cursor = start_position;
+ while cursor < self.mask.len() && !self.mask.value(cursor) {
+ cursor += 1;
+ }
+
+ debug_assert!(
+ cursor < self.mask.len(),
+ "ReadPlan must remove trailing skips from Mask selections"
+ );
+
+ let loaded_range_end = self
+ .loaded_row_ranges
+ .as_ref()
+ .and_then(|ranges| ranges.end_containing(cursor))
+ .ok_or_else(|| {
+ ParquetError::General(format!(
+ "Internal Error: selected row {cursor} has no loaded page
range"
+ ))
+ })?;
+
+ let mask_start = cursor;
+ let mut selected_rows = 0;
+ while cursor < loaded_range_end && cursor < self.mask.len() &&
selected_rows < batch_size {
+ if self.mask.value(cursor) {
+ selected_rows += 1;
+ }
+ cursor += 1;
+ }
+
+ self.position = cursor;
+ Ok(MaskChunk {
+ initial_skip: mask_start - start_position,
+ chunk_rows: cursor - mask_start,
+ selected_rows,
+ mask_start,
})
}
@@ -885,6 +960,39 @@ pub struct MaskChunk {
pub mask_start: usize,
}
+/// Row ranges whose backing pages are loaded for every projected column.
+#[derive(Clone, Debug)]
+pub(crate) struct LoadedRowRanges(Vec<Range<usize>>);
+
+impl LoadedRowRanges {
+ pub(crate) fn from_selection(selection: RowSelection) -> Self {
+ let selectors: Vec<RowSelector> = selection.into();
+ let mut position = 0;
+ let ranges = selectors
+ .into_iter()
+ .filter_map(|selector| {
+ let start = position;
+ position += selector.row_count;
+ (!selector.skip).then_some(start..position)
+ })
+ .collect();
+ Self(ranges)
+ }
+
+ fn end_containing(&self, row: usize) -> Option<usize> {
+ let idx = self.0.partition_point(|range| range.end <= row);
+ self.0
+ .get(idx)
+ .filter(|range| range.start <= row)
+ .map(|range| range.end)
+ }
+
+ #[cfg(test)]
+ pub(crate) fn ranges(&self) -> &[Range<usize>] {
+ &self.0
+ }
+}
+
/// Cursor for iterating a [`RowSelection`] during execution within a
/// [`ReadPlan`](crate::arrow::arrow_reader::ReadPlan).
///
@@ -902,10 +1010,21 @@ pub enum RowSelectionCursor {
impl RowSelectionCursor {
/// Create a [`MaskCursor`] cursor backed by a bitmask, from an existing
set of selectors
- pub(crate) fn new_mask_from_selectors(selectors: Vec<RowSelector>) -> Self
{
+ pub(crate) fn new_mask_from_selectors(
+ selectors: Vec<RowSelector>,
+ loaded_row_ranges: Option<Arc<LoadedRowRanges>>,
+ ) -> Self {
+ debug_assert!(
+ selectors
+ .last()
+ .map(|selector| !selector.skip)
+ .unwrap_or(true),
+ "Mask selectors must not end with a skip"
+ );
Self::Mask(MaskCursor {
mask: boolean_mask_from_selectors(&selectors),
position: 0,
+ loaded_row_ranges,
})
}
@@ -1466,6 +1585,10 @@ mod tests {
// assert_eq!(mask, vec![false, true, true, false, true, true, false]);
assert_eq!(ranges, vec![10..20, 20..30, 40..50, 50..60]);
+ assert_eq!(
+ selection.row_ranges_for_selected_pages(&index, 70),
+ vec![10..20, 20..30, 40..50, 50..60]
+ );
let selection = RowSelection::from(vec![
// Skip first page
@@ -1537,6 +1660,51 @@ mod tests {
assert_eq!(ranges, vec![10..20, 20..30, 30..40]);
}
+ #[test]
+ fn test_loaded_mask_chunk_stops_at_trimmed_mask_end() {
+ let loaded =
LoadedRowRanges::from_selection(RowSelection::from_consecutive_ranges(
+ std::iter::once(0..5),
+ 10,
+ ));
+ let RowSelectionCursor::Mask(mut cursor) =
RowSelectionCursor::new_mask_from_selectors(
+ vec![RowSelector::select(1)],
+ Some(loaded.into()),
+ ) else {
+ unreachable!()
+ };
+
+ let chunk = cursor.next_chunk(10).unwrap();
+ assert_eq!(chunk.chunk_rows, 1);
+ assert!(cursor.is_empty());
+ }
+
+ #[test]
+ fn test_next_mask_chunk_until_cursor_is_empty() {
+ let RowSelectionCursor::Mask(mut cursor) =
RowSelectionCursor::new_mask_from_selectors(
+ vec![
+ RowSelector::skip(2),
+ RowSelector::select(2),
+ RowSelector::skip(1),
+ RowSelector::select(1),
+ ],
+ None,
+ ) else {
+ unreachable!()
+ };
+
+ let first = cursor.next_mask_chunk(2).unwrap();
+ assert_eq!(first.initial_skip, 2);
+ assert_eq!(first.chunk_rows, 2);
+ assert_eq!(first.selected_rows, 2);
+
+ let second = cursor.next_mask_chunk(2).unwrap();
+ assert_eq!(second.initial_skip, 1);
+ assert_eq!(second.chunk_rows, 1);
+ assert_eq!(second.selected_rows, 1);
+
+ assert!(cursor.next_mask_chunk(2).is_none());
+ }
+
#[test]
fn test_from_ranges() {
let ranges = [1..3, 4..6, 6..6, 8..8, 9..10];
diff --git a/parquet/src/arrow/push_decoder/reader_builder/mod.rs
b/parquet/src/arrow/push_decoder/reader_builder/mod.rs
index dacf1a2caa..b773f95fd5 100644
--- a/parquet/src/arrow/push_decoder/reader_builder/mod.rs
+++ b/parquet/src/arrow/push_decoder/reader_builder/mod.rs
@@ -21,7 +21,7 @@ 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::selection::RowSelectionStrategy;
+use crate::arrow::arrow_reader::selection::{LoadedRowRanges,
RowSelectionStrategy};
use crate::arrow::arrow_reader::{
ParquetRecordBatchReader, PredicateOptions, ReadPlanBuilder, RowFilter,
RowSelection,
RowSelectionPolicy,
@@ -613,20 +613,16 @@ impl RowGroupReaderBuilder {
.with_parquet_metadata(&self.metadata)
.build_array_reader(self.fields.as_deref(),
predicate.projection())?;
- // Reset to original policy before each predicate so the
override
- // can detect page skipping for THIS predicate's columns.
- // Without this reset, a prior predicate's override (e.g. Mask)
- // carries forward and the check returns early, missing
unfetched
- // pages for subsequent predicates.
+ // 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 to evaluate the filter.
- // Note: first update the selection strategy to properly
handle any pages
- // pruned during fetch
- plan_builder = override_selector_strategy_if_needed(
+ // 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
@@ -729,10 +725,11 @@ impl RowGroupReaderBuilder {
plan_builder =
plan_builder.with_row_selection_policy(self.row_selection_policy);
- plan_builder = override_selector_strategy_if_needed(
+ 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 {
@@ -857,7 +854,7 @@ impl RowGroupReaderBuilder {
}
}
-/// Override the selection strategy if needed.
+/// Prepare row selection execution when page pruning produced sparse column
data.
///
/// Some pages can be skipped during row-group construction if they are not
read
/// by the selections. This means that the data pages for those rows are never
@@ -865,58 +862,147 @@ impl RowGroupReaderBuilder {
/// `RowSelections` selection works because `skip_records()` handles this
/// case and skips the page accordingly.
///
-/// However, with the current mask design, all values must be read and decoded
-/// and then a mask filter is applied. Thus if any pages are skipped during
-/// row-group construction, the data pages are missing and cannot be decoded.
+/// However, with the current mask design, all values covered by a mask chunk
+/// are decoded before the mask filter is applied. Thus a chunk cannot cross a
+/// page that was skipped during row-group construction.
///
/// A simple example:
/// * the page size is 2, the mask is 100001, row selection should be read(1)
skip(4) read(1)
/// * the `ColumnChunkData` would be page1(10), page2(skipped), page3(01)
///
-/// Using the row selection to skip(4), page2 won't be read at all, so in this
-/// case we can't decode all the rows and apply a mask. To correctly apply the
-/// bit mask, we need all 6 values be read, but page2 is not in memory.
-fn override_selector_strategy_if_needed(
+/// Mask execution records the row ranges loaded for every projected column, so
+/// each mask chunk stays within loaded data and `skip_records()` crosses the
+/// gaps. This applies both to an explicit mask policy and to Auto when it
+/// resolves to mask execution.
+fn prepare_selection_for_page_skipping(
plan_builder: ReadPlanBuilder,
projection_mask: &ProjectionMask,
offset_index: Option<&[OffsetIndexMetaData]>,
+ total_rows: usize,
) -> ReadPlanBuilder {
- // override only applies to Auto policy, If the policy is already Mask or
Selectors, respect that
- let RowSelectionPolicy::Auto { .. } = plan_builder.row_selection_policy()
else {
- return plan_builder;
- };
-
- let preferred_strategy = plan_builder.resolve_selection_strategy();
-
- let force_selectors = matches!(preferred_strategy,
RowSelectionStrategy::Mask)
- && plan_builder.selection().is_some_and(|selection| {
- selection.should_force_selectors(projection_mask, offset_index)
- });
-
- let resolved_strategy = if force_selectors {
- RowSelectionStrategy::Selectors
- } else {
- preferred_strategy
- };
-
- // override the plan builder strategy with the resolved one
- let new_policy = match resolved_strategy {
- RowSelectionStrategy::Mask => RowSelectionPolicy::Mask,
- RowSelectionStrategy::Selectors => RowSelectionPolicy::Selectors,
- };
-
- plan_builder.with_row_selection_policy(new_policy)
+ match plan_builder.resolve_selection_strategy() {
+ RowSelectionStrategy::Mask => {
+ let loaded = loaded_row_ranges_for_projection(
+ plan_builder.selection(),
+ projection_mask,
+ offset_index,
+ total_rows,
+ );
+ plan_builder
+ .with_row_selection_policy(RowSelectionPolicy::Mask)
+ .with_loaded_row_ranges(loaded)
+ }
+ RowSelectionStrategy::Selectors => {
+
plan_builder.with_row_selection_policy(RowSelectionPolicy::Selectors)
+ }
+ }
+}
+
+/// Computes row ranges for which every projected column has page data loaded.
+fn loaded_row_ranges_for_projection(
+ selection: Option<&RowSelection>,
+ projection_mask: &ProjectionMask,
+ 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)
}
#[cfg(test)]
mod tests {
use super::*;
use crate::arrow::arrow_reader::{RowSelection, RowSelector};
+ use crate::file::page_index::offset_index::PageLocation;
#[test]
// Verify that the size of RowGroupDecoderState does not grow too large
fn test_structure_size() {
- assert_eq!(std::mem::size_of::<RowGroupDecoderState>(), 232);
+ assert_eq!(std::mem::size_of::<RowGroupDecoderState>(), 240);
+ }
+
+ #[test]
+ fn test_loaded_row_ranges_intersect_column_page_boundaries() {
+ let column = |first_rows: &[i64]| OffsetIndexMetaData {
+ page_locations: first_rows
+ .iter()
+ .enumerate()
+ .map(|(idx, first_row_index)| PageLocation {
+ offset: (idx * 10) as i64,
+ compressed_page_size: 10,
+ first_row_index: *first_row_index,
+ })
+ .collect(),
+ unencoded_byte_array_data_bytes: None,
+ };
+ let columns = vec![column(&[0, 4, 8]), column(&[0, 6, 10])];
+ let selection = RowSelection::from(vec![
+ RowSelector::skip(1),
+ RowSelector::select(1),
+ RowSelector::skip(9),
+ RowSelector::select(1),
+ ]);
+
+ let loaded = loaded_row_ranges_for_projection(
+ Some(&selection),
+ &ProjectionMask::all(),
+ Some(&columns),
+ 12,
+ )
+ .unwrap();
+
+ assert_eq!(loaded.ranges(), &[0..4, 10..12]);
+ }
+
+ #[test]
+ fn test_auto_keeps_mask_when_page_pruning_skips_pages() {
+ let columns = vec![OffsetIndexMetaData {
+ page_locations: [0, 2, 4, 6, 8, 10]
+ .into_iter()
+ .enumerate()
+ .map(|(idx, first_row_index)| PageLocation {
+ offset: (idx * 10) as i64,
+ compressed_page_size: 10,
+ first_row_index,
+ })
+ .collect(),
+ unencoded_byte_array_data_bytes: None,
+ }];
+ let selection = RowSelection::from(vec![
+ RowSelector::select(1),
+ RowSelector::skip(10),
+ RowSelector::select(1),
+ ]);
+ let plan_builder = ReadPlanBuilder::new(12)
+ .with_selection(Some(selection))
+ .with_row_selection_policy(RowSelectionPolicy::Auto { threshold:
32 });
+
+ let prepared = prepare_selection_for_page_skipping(
+ plan_builder,
+ &ProjectionMask::all(),
+ Some(&columns),
+ 12,
+ );
+
+ assert_eq!(prepared.row_selection_policy(), &RowSelectionPolicy::Mask);
}
#[test]
diff --git a/parquet/tests/arrow_reader/row_filter/async.rs
b/parquet/tests/arrow_reader/row_filter/async.rs
index 66840bb814..885326687a 100644
--- a/parquet/tests/arrow_reader/row_filter/async.rs
+++ b/parquet/tests/arrow_reader/row_filter/async.rs
@@ -16,7 +16,7 @@
// under the License.
use crate::io::TestReader;
-use std::sync::Arc;
+use std::sync::{Arc, Mutex};
use arrow::{
array::AsArray,
@@ -44,25 +44,13 @@ use parquet::{
},
};
-#[tokio::test]
-async fn test_row_filter_full_page_skip_is_handled_async() {
- let first_value: i64 = 1111;
- let last_value: i64 = 9999;
- let num_rows: usize = 12;
-
- // build data with row selection average length 4
- // The result would be (1111 XXXX) ... (4 page in the middle)... (XXXX
9999)
- // The Row Selection would be [1111, (skip 10), 9999]
+fn make_two_column_i64_file(values: &[i64], rows_per_page: usize) -> Bytes {
let schema = Arc::new(Schema::new(vec![
Field::new("key", DataType::Int64, false),
Field::new("value", DataType::Int64, false),
]));
-
- let mut int_values: Vec<i64> = (0..num_rows as i64).collect();
- int_values[0] = first_value;
- int_values[num_rows - 1] = last_value;
- let keys = Int64Array::from(int_values.clone());
- let values = Int64Array::from(int_values.clone());
+ let keys = Int64Array::from(values.to_vec());
+ let values = Int64Array::from(values.to_vec());
let batch = RecordBatch::try_new(
Arc::clone(&schema),
vec![Arc::new(keys) as ArrayRef, Arc::new(values) as ArrayRef],
@@ -70,15 +58,29 @@ async fn test_row_filter_full_page_skip_is_handled_async() {
.unwrap();
let props = WriterProperties::builder()
- .set_write_batch_size(2)
- .set_data_page_row_count_limit(2)
+ .set_write_batch_size(rows_per_page)
+ .set_data_page_row_count_limit(rows_per_page)
.build();
-
let mut buffer = Vec::new();
let mut writer = ArrowWriter::try_new(&mut buffer, schema,
Some(props)).unwrap();
writer.write(&batch).unwrap();
writer.close().unwrap();
- let data = Bytes::from(buffer);
+ Bytes::from(buffer)
+}
+
+#[tokio::test]
+async fn test_row_filter_full_page_skip_is_handled_async() {
+ let first_value: i64 = 1111;
+ let last_value: i64 = 9999;
+ let num_rows: usize = 12;
+
+ // build data with row selection average length 4
+ // The result would be (1111 XXXX) ... (4 page in the middle)... (XXXX
9999)
+ // The Row Selection would be [1111, (skip 10), 9999]
+ let mut int_values: Vec<i64> = (0..num_rows as i64).collect();
+ int_values[0] = first_value;
+ int_values[num_rows - 1] = last_value;
+ let data = make_two_column_i64_file(&int_values, 2);
let builder = ParquetRecordBatchStreamBuilder::new_with_options(
TestReader::new(data.clone()),
@@ -89,35 +91,282 @@ async fn test_row_filter_full_page_skip_is_handled_async()
{
let schema = builder.parquet_schema().clone();
let filter_mask = ProjectionMask::leaves(&schema, [0]);
- let make_predicate = |mask: ProjectionMask| {
- ArrowPredicateFn::new(mask, move |batch: RecordBatch| {
- let column = batch.column(0);
- let match_first = eq(column,
&Int64Array::new_scalar(first_value))?;
- let match_second = eq(column,
&Int64Array::new_scalar(last_value))?;
- or(&match_first, &match_second)
- })
- };
+ for policy in [
+ RowSelectionPolicy::Auto { threshold: 32 },
+ RowSelectionPolicy::Mask,
+ ] {
+ for max_predicate_cache_size in [None, Some(0)] {
+ for with_initial_selection in [false, true] {
+ let predicate_batch_sizes = Arc::new(Mutex::new(Vec::new()));
+ let observed_batch_sizes = Arc::clone(&predicate_batch_sizes);
+ let predicate =
+ ArrowPredicateFn::new(filter_mask.clone(), move |batch:
RecordBatch| {
+
observed_batch_sizes.lock().unwrap().push(batch.num_rows());
+ let column = batch.column(0);
+ let match_first = eq(column,
&Int64Array::new_scalar(first_value))?;
+ let match_second = eq(column,
&Int64Array::new_scalar(last_value))?;
+ or(&match_first, &match_second)
+ });
+
+ let stream = ParquetRecordBatchStreamBuilder::new_with_options(
+ TestReader::new(data.clone()),
+
ArrowReaderOptions::new().with_page_index_policy(PageIndexPolicy::Required),
+ )
+ .await
+ .unwrap()
+ .with_row_filter(RowFilter::new(vec![Box::new(predicate)]))
+ .with_batch_size(12)
+ .with_row_selection_policy(policy);
+ let stream = if with_initial_selection {
+ stream.with_row_selection(RowSelection::from(vec![
+ RowSelector::select(1),
+ RowSelector::skip(10),
+ RowSelector::select(1),
+ ]))
+ } else {
+ stream
+ };
+ let stream = match max_predicate_cache_size {
+ Some(size) => stream.with_max_predicate_cache_size(size),
+ None => stream,
+ }
+ .build()
+ .unwrap();
- let predicate = make_predicate(filter_mask.clone());
+ let schema = stream.schema().clone();
+ let batches: Vec<_> = stream.try_collect().await.unwrap();
+ let batch_sizes = batches
+ .iter()
+ .map(RecordBatch::num_rows)
+ .collect::<Vec<_>>();
+ assert_eq!(
+ batch_sizes,
+ vec![2],
+ "policy={policy:?}, cache={max_predicate_cache_size:?},
initial={with_initial_selection}"
+ );
+ assert_eq!(
+ *predicate_batch_sizes.lock().unwrap(),
+ if with_initial_selection {
+ vec![2]
+ } else {
+ vec![12]
+ },
+ "policy={policy:?}, cache={max_predicate_cache_size:?},
initial={with_initial_selection}"
+ );
+
+ let result = concat_batches(&schema, &batches).unwrap();
+ assert_eq!(result.num_rows(), 2);
+ assert_eq!(
+ result
+ .column(0)
+ .as_any()
+ .downcast_ref::<Int64Array>()
+ .unwrap(),
+ &Int64Array::from(vec![first_value, last_value])
+ );
+ }
+ }
+ }
+}
- // The batch size is set to 12 to read all rows in one go after filtering
- // If the Reader chooses mask to handle filter, it might cause panic
because the mid 4 pages may not be decoded.
+#[tokio::test]
+async fn test_mask_coalesces_loaded_ranges_to_batch_size() {
+ let values = (0..12).collect::<Vec<i64>>();
+ let data = make_two_column_i64_file(&values, 2);
let stream = ParquetRecordBatchStreamBuilder::new_with_options(
- TestReader::new(data.clone()),
+ TestReader::new(data),
ArrowReaderOptions::new().with_page_index_policy(PageIndexPolicy::Required),
)
.await
.unwrap()
- .with_row_filter(RowFilter::new(vec![Box::new(predicate)]))
- .with_batch_size(12)
- .with_row_selection_policy(RowSelectionPolicy::Auto { threshold: 32 })
+ .with_row_selection(RowSelection::from(vec![
+ RowSelector::select(1),
+ RowSelector::skip(4),
+ RowSelector::select(1),
+ RowSelector::skip(5),
+ RowSelector::select(1),
+ ]))
+ .with_batch_size(2)
+ .with_row_selection_policy(RowSelectionPolicy::Mask)
.build()
.unwrap();
let schema = stream.schema().clone();
- let batches: Vec<_> = stream.try_collect().await.unwrap();
- let result = concat_batches(&schema, &batches).unwrap();
- assert_eq!(result.num_rows(), 2);
+ let batches: Vec<RecordBatch> = stream.try_collect().await.unwrap();
+ assert_eq!(
+ batches
+ .iter()
+ .map(RecordBatch::num_rows)
+ .collect::<Vec<_>>(),
+ vec![2, 1]
+ );
+
+ let batch = concat_batches(&schema, &batches).unwrap();
+ assert_eq!(
+ batch
+ .column(0)
+ .as_any()
+ .downcast_ref::<Int64Array>()
+ .unwrap()
+ .values(),
+ &[0, 5, 11]
+ );
+}
+
+#[tokio::test]
+async fn test_mask_stops_after_trailing_skipped_pages() {
+ let values = (0..8).collect::<Vec<i64>>();
+ let data = make_two_column_i64_file(&values, 2);
+ let stream = ParquetRecordBatchStreamBuilder::new_with_options(
+ TestReader::new(data),
+
ArrowReaderOptions::new().with_page_index_policy(PageIndexPolicy::Required),
+ )
+ .await
+ .unwrap()
+ .with_row_selection(RowSelection::from(vec![
+ RowSelector::select(1),
+ RowSelector::skip(4),
+ RowSelector::select(1),
+ RowSelector::skip(2),
+ ]))
+ .with_batch_size(8)
+ .with_row_selection_policy(RowSelectionPolicy::Mask)
+ .build()
+ .unwrap();
+
+ let batches: Vec<RecordBatch> = stream.try_collect().await.unwrap();
+ assert_eq!(
+ batches
+ .iter()
+ .map(RecordBatch::num_rows)
+ .collect::<Vec<_>>(),
+ vec![2]
+ );
+ assert_eq!(
+ batches[0]
+ .column(0)
+ .as_any()
+ .downcast_ref::<Int64Array>()
+ .unwrap()
+ .values(),
+ &[0, 5]
+ );
+}
+
+#[tokio::test]
+async fn test_mask_nested_projection_with_different_page_boundaries() {
+ use arrow_array::{
+ Array, UInt32Array,
+ builder::{Int32Builder, ListBuilder, StringBuilder, StructBuilder},
+ };
+ use arrow_schema::Fields;
+ use parquet::file::properties::WriterVersion;
+
+ let num_rows = 14usize;
+ let long_prefix = "x".repeat(500);
+ let struct_fields = Fields::from(vec![
+ Field::new("x", DataType::Int32, false),
+ Field::new("y", DataType::Utf8, false),
+ ]);
+ let mut list_builder =
ListBuilder::new(StructBuilder::from_fields(struct_fields, 0));
+ for row in 0..num_rows {
+ let values = list_builder.values();
+ for item in 0..if row % 2 == 0 { 2 } else { 3 } {
+ values
+ .field_builder::<Int32Builder>(0)
+ .unwrap()
+ .append_value(row as i32 * 10 + item);
+ values
+ .field_builder::<StringBuilder>(1)
+ .unwrap()
+ .append_value(format!("{long_prefix}_{row}_{item}"));
+ values.append(true);
+ }
+ list_builder.append(true);
+ }
+ let values = list_builder.finish();
+ let schema = Arc::new(Schema::new(vec![Field::new(
+ "values",
+ values.data_type().clone(),
+ false,
+ )]));
+ let input = RecordBatch::try_new(schema.clone(), vec![Arc::new(values) as
ArrayRef]).unwrap();
+
+ // The fixed-width leaf fits in one page, while the long string leaf
flushes
+ // frequently and therefore has different page boundaries.
+ let props = WriterProperties::builder()
+ .set_writer_version(WriterVersion::PARQUET_2_0)
+ .set_write_batch_size(2)
+ .set_dictionary_enabled(false)
+ .set_data_page_size_limit(512)
+ .build();
+ let mut buffer = Vec::new();
+ let mut writer = ArrowWriter::try_new(&mut buffer, schema,
Some(props)).unwrap();
+ writer.write(&input).unwrap();
+ writer.close().unwrap();
+
+ let builder = ParquetRecordBatchStreamBuilder::new_with_options(
+ TestReader::new(Bytes::from(buffer)),
+
ArrowReaderOptions::new().with_page_index_policy(PageIndexPolicy::Required),
+ )
+ .await
+ .unwrap();
+ let page_first_rows = builder.metadata().offset_index().unwrap()[0]
+ .iter()
+ .map(|column| {
+ column
+ .page_locations()
+ .iter()
+ .map(|page| page.first_row_index)
+ .collect::<Vec<_>>()
+ })
+ .collect::<Vec<_>>();
+ assert_eq!(page_first_rows.len(), 2);
+ assert_eq!(page_first_rows[0], vec![0]);
+ assert_ne!(page_first_rows[0], page_first_rows[1]);
+
+ let variable_width_pages = &page_first_rows[1];
+ assert_eq!(variable_width_pages.first(), Some(&0));
+ let selected_pages = [0, 6, 13]
+ .map(|row| variable_width_pages.partition_point(|first_row| *first_row
<= row) - 1);
+ assert!(
+ selected_pages
+ .windows(2)
+ .all(|pages| pages[1] > pages[0] + 1),
+ "selected rows must leave unloaded pages between them:
{page_first_rows:?}"
+ );
+
+ let stream = builder
+ .with_row_selection(RowSelection::from(vec![
+ RowSelector::select(1),
+ RowSelector::skip(5),
+ RowSelector::select(1),
+ RowSelector::skip(6),
+ RowSelector::select(1),
+ ]))
+ .with_batch_size(2)
+ .with_row_selection_policy(RowSelectionPolicy::Mask)
+ .build()
+ .unwrap();
+
+ let output_schema = stream.schema().clone();
+ let batches: Vec<RecordBatch> = stream.try_collect().await.unwrap();
+ assert_eq!(
+ batches
+ .iter()
+ .map(RecordBatch::num_rows)
+ .collect::<Vec<_>>(),
+ vec![2, 1]
+ );
+
+ let output = concat_batches(&output_schema, &batches).unwrap();
+ let expected = arrow::compute::take(
+ input.column(0).as_ref(),
+ &UInt32Array::from(vec![0, 6, 13]),
+ None,
+ )
+ .unwrap();
+ assert_eq!(output.column(0).to_data(), expected.to_data());
}
#[tokio::test]
@@ -526,19 +775,16 @@ async fn test_predicate_pushdown_with_skipped_pages() {
}
}
-/// Regression test: when multiple predicates are used, the first predicate's
-/// override of the selection strategy (to Mask) must NOT carry forward to
-/// subsequent predicates. Each predicate must get a fresh Auto policy so the
-/// override can detect page skipping for that predicate's specific columns.
+/// Regression test: Auto resolution and loaded row ranges must be refreshed
for
+/// each predicate's current selection and projection.
///
/// Scenario:
/// - Dense initial RowSelection (alternating select/skip) covers all pages →
Auto resolves to Mask
/// - Predicate 1 evaluates on column A, narrows selection to skip middle pages
/// - Predicate 2's column B is fetched sparsely with the narrowed selection
(missing middle pages)
-/// - Without the fix, the override for predicate 2 returns early
(policy=Mask, not Auto),
-/// so Mask is used and tries to read missing pages → "Invalid offset" error
+/// - Predicate 2 remains safe with Mask because its loaded row ranges are
recomputed for column B
#[tokio::test]
-async fn test_multi_predicate_mask_policy_carryover() {
+async fn test_multi_predicate_auto_mask_with_sparse_pages() {
// 300 rows, 1 row group, 100 rows per page (3 pages)
let num_rows = 300usize;
let rows_per_page = 100;
@@ -604,7 +850,10 @@ async fn test_multi_predicate_mask_policy_carryover() {
// Predicate 2 on value_col: keeps rows where value_col < 250
// This column is fetched AFTER predicate 1 narrows the selection.
// Its sparse data will be missing the middle page.
- let pred2 = ArrowPredicateFn::new(ProjectionMask::roots(&schema_descr,
[1]), |batch| {
+ let predicate_batch_sizes = Arc::new(Mutex::new(Vec::new()));
+ let observed_batch_sizes = Arc::clone(&predicate_batch_sizes);
+ let pred2 = ArrowPredicateFn::new(ProjectionMask::roots(&schema_descr,
[1]), move |batch| {
+ observed_batch_sizes.lock().unwrap().push(batch.num_rows());
let col = batch.column(0).as_primitive::<Int32Type>();
Ok(BooleanArray::from_iter(
col.iter().map(|v| v.map(|val| val < 250)),
@@ -612,8 +861,6 @@ async fn test_multi_predicate_mask_policy_carryover() {
});
let row_filter = RowFilter::new(vec![Box::new(pred1), Box::new(pred2)]);
-
- // Output projection: both columns
let projection = ProjectionMask::roots(&schema_descr, [0, 1]);
let stream = builder
@@ -624,9 +871,17 @@ async fn test_multi_predicate_mask_policy_carryover() {
.build()
.unwrap();
- // Without the fix, this panics with:
- // "Invalid offset in sparse column chunk data: ..., no matching page
found."
+ // Without loaded-range-aware Mask reads, this panics with an invalid
sparse-page offset.
let batches: Vec<RecordBatch> = stream.try_collect().await.unwrap();
+ assert_eq!(*predicate_batch_sizes.lock().unwrap(), vec![100]);
+ assert_eq!(
+ batches
+ .iter()
+ .map(RecordBatch::num_rows)
+ .collect::<Vec<_>>(),
+ vec![75]
+ );
+
let batch = concat_batches(&batches[0].schema(), &batches).unwrap();
// Verify results: rows where filter_col==0 AND value_col<250 AND original
alternating selection
diff --git a/parquet/tests/arrow_reader/row_filter/sync.rs
b/parquet/tests/arrow_reader/row_filter/sync.rs
index 77a75220dc..a224c95c8f 100644
--- a/parquet/tests/arrow_reader/row_filter/sync.rs
+++ b/parquet/tests/arrow_reader/row_filter/sync.rs
@@ -191,7 +191,7 @@ fn test_row_filter_full_page_skip_is_handled() {
.unwrap();
// Predicate pruning used to panic once mask-backed plans removed whole
pages.
- // Collecting into batches validates the plan now downgrades to selectors
instead.
+ // Collecting into batches validates that Mask can read around unloaded
pages.
let schema = reader.schema().clone();
let batches = reader.collect::<Result<Vec<_>, _>>().unwrap();
let result = concat_batches(&schema, &batches).unwrap();
