From 4fef17851c2008a0ebdbb743c725fc4abb494304 Mon Sep 17 00:00:00 2001 From: Qiwei Huang Date: Thu, 16 Jul 2026 22:22:49 +0800 Subject: [PATCH 1/4] Preserve mask execution across skipped parquet pages --- parquet/src/arrow/arrow_reader/mod.rs | 2 +- parquet/src/arrow/arrow_reader/read_plan.rs | 17 +- parquet/src/arrow/arrow_reader/selection.rs | 171 +++++++++++++---- .../arrow/push_decoder/reader_builder/mod.rs | 176 +++++++++++++----- .../tests/arrow_reader/row_filter/async.rs | 65 ++++--- 5 files changed, 322 insertions(+), 109 deletions(-) diff --git a/parquet/src/arrow/arrow_reader/mod.rs b/parquet/src/arrow/arrow_reader/mod.rs index 814b8250508e..b9fc8eda9ca1 100644 --- a/parquet/src/arrow/arrow_reader/mod.rs +++ b/parquet/src/arrow/arrow_reader/mod.rs @@ -1386,7 +1386,7 @@ impl ParquetRecordBatchReader { // 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 { + let Some(mask_chunk) = mask_cursor.next_chunk(batch_size)? else { return Ok(None); }; diff --git a/parquet/src/arrow/arrow_reader/read_plan.rs b/parquet/src/arrow/arrow_reader/read_plan.rs index 3d80526c8282..dd95ef899b36 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, /// 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>, } 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) -> 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,6 +314,7 @@ impl ReadPlanBuilder { batch_size, selection, row_selection_policy: _, + loaded_row_ranges, } = self; let selection = selection.map(|s| s.trim()); @@ -314,7 +325,7 @@ impl ReadPlanBuilder { let selectors: Vec = trimmed.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), } diff --git a/parquet/src/arrow/arrow_reader/selection.rs b/parquet/src/arrow/arrow_reader/selection.rs index 2ddf812f9c39..141ca329b1c4 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> { + 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`] @@ -770,6 +762,8 @@ 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>, } impl MaskCursor { @@ -824,6 +818,54 @@ impl MaskCursor { }) } + /// Returns the next mask chunk without crossing a row range that was not loaded. + pub(crate) fn next_chunk( + &mut self, + batch_size: usize, + ) -> Result, ParquetError> { + if self.loaded_row_ranges.is_none() { + return Ok(self.next_mask_chunk(batch_size)); + } + + let start_position = self.position; + let mut cursor = start_position; + while cursor < self.mask.len() && !self.mask.value(cursor) { + cursor += 1; + } + + if cursor == self.mask.len() { + self.position = cursor; + return Ok(None); + } + + 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(Some(MaskChunk { + initial_skip: mask_start - start_position, + chunk_rows: cursor - mask_start, + selected_rows, + mask_start, + })) + } + /// Materialise the boolean values for a mask-backed chunk pub fn mask_values_for(&self, chunk: &MaskChunk) -> Result { if chunk.mask_start.saturating_add(chunk.chunk_rows) > self.mask.len() { @@ -885,6 +927,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>); + +impl LoadedRowRanges { + pub(crate) fn from_selection(selection: RowSelection) -> Self { + let selectors: Vec = 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 { + 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] { + &self.0 + } +} + /// Cursor for iterating a [`RowSelection`] during execution within a /// [`ReadPlan`](crate::arrow::arrow_reader::ReadPlan). /// @@ -902,10 +977,14 @@ 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) -> Self { + pub(crate) fn new_mask_from_selectors( + selectors: Vec, + loaded_row_ranges: Option>, + ) -> Self { Self::Mask(MaskCursor { mask: boolean_mask_from_selectors(&selectors), position: 0, + loaded_row_ranges, }) } @@ -1466,6 +1545,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 +1620,24 @@ 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().unwrap(); + assert_eq!(chunk.chunk_rows, 1); + assert!(cursor.is_empty()); + } + #[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 dacf1a2caad9..b773f95fd55e 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 { + 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::(), 232); + assert_eq!(std::mem::size_of::(), 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 66840bb8147b..b12108cc3b45 100644 --- a/parquet/tests/arrow_reader/row_filter/async.rs +++ b/parquet/tests/arrow_reader/row_filter/async.rs @@ -98,26 +98,44 @@ async fn test_row_filter_full_page_skip_is_handled_async() { }) }; - let predicate = make_predicate(filter_mask.clone()); + for policy in [ + RowSelectionPolicy::Auto { threshold: 32 }, + RowSelectionPolicy::Mask, + ] { + let predicate = make_predicate(filter_mask.clone()); - // 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. - 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(RowSelectionPolicy::Auto { threshold: 32 }) - .build() - .unwrap(); + // A batch larger than the selection forces Mask to cross the skipped pages. + let stream = ParquetRecordBatchStreamBuilder::new_with_options( + TestReader::new(data.clone()), + ArrowReaderOptions::new().with_page_index_policy(PageIndexPolicy::Required), + ) + .await + .unwrap() + .with_row_selection(RowSelection::from(vec![ + RowSelector::select(1), + RowSelector::skip(10), + RowSelector::select(1), + ])) + .with_row_filter(RowFilter::new(vec![Box::new(predicate)])) + .with_batch_size(12) + .with_row_selection_policy(policy) + .with_max_predicate_cache_size(0) + .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 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); + assert_eq!( + result + .column(0) + .as_any() + .downcast_ref::() + .unwrap(), + &Int64Array::from(vec![first_value, last_value]) + ); + } } #[tokio::test] @@ -526,19 +544,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; From 152d58167f4e626832b15f7b48dc5732c0d2728f Mon Sep 17 00:00:00 2001 From: Qiwei Huang Date: Fri, 17 Jul 2026 09:19:16 +0800 Subject: [PATCH 2/4] Coalesce mask reads across sparse parquet pages Accumulate loaded-range-safe mask chunks until batch_size so sparse page pruning does not fragment logical batches or predicate evaluation. Cover cache policies, initial selections, multiple predicates, and nested columns with different page boundaries. --- parquet/src/arrow/arrow_reader/mod.rs | 91 +++-- .../tests/arrow_reader/row_filter/async.rs | 320 ++++++++++++++---- parquet/tests/arrow_reader/row_filter/sync.rs | 2 +- 3 files changed, 318 insertions(+), 95 deletions(-) diff --git a/parquet/src/arrow/arrow_reader/mod.rs b/parquet/src/arrow/arrow_reader/mod.rs index b9fc8eda9ca1..4f9d2229c339 100644 --- a/parquet/src/arrow/arrow_reader/mod.rs +++ b/parquet/src/arrow/arrow_reader/mod.rs @@ -18,7 +18,8 @@ //! 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}; @@ -1381,13 +1382,17 @@ impl ParquetRecordBatchReader { if batch_size == 0 { return Ok(None); } - match self.read_plan.row_selection_cursor_mut() { + let filter_mask = 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_chunk(batch_size)? else { - return Ok(None); + let mut first_filter_mask: Option = None; + let mut combined_filter_mask: Option = None; + + // Each chunk stays within loaded pages, while the in-progress array and + // filter mask accumulate chunks until the logical batch is full. + while read_records < batch_size && !mask_cursor.is_empty() { + let remaining = batch_size - read_records; + let Some(mask_chunk) = mask_cursor.next_chunk(remaining)? else { + break; }; if mask_chunk.initial_skip > 0 { @@ -1403,7 +1408,7 @@ impl ParquetRecordBatchReader { if mask_chunk.chunk_rows == 0 { if mask_cursor.is_empty() && mask_chunk.selected_rows == 0 { - return Ok(None); + break; } continue; } @@ -1424,33 +1429,33 @@ impl ParquetRecordBatchReader { 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; + match combined_filter_mask.as_mut() { + Some(combined) => combined.append_buffer(mask.values()), + None => match first_filter_mask.take() { + Some(first) => { + let mut combined = + BooleanBufferBuilder::new(first.len() + mask.len()); + combined.append_buffer(&first); + combined.append_buffer(mask.values()); + combined_filter_mask = Some(combined); + } + None => first_filter_mask = Some(mask.values().clone()), + }, } + read_records += mask_chunk.selected_rows; + } - return Ok(Some(filtered_batch)); + if read_records == 0 { + return Ok(None); } + + let filter_mask = match combined_filter_mask { + Some(combined) => combined.build(), + None => first_filter_mask.ok_or_else(|| { + general_err!("Internal Error: decoded Mask batch has no filter values") + })?, + }; + Some(BooleanArray::from(filter_mask)) } RowSelectionCursor::Selectors(selectors_cursor) => { while read_records < batch_size && !selectors_cursor.is_empty() { @@ -1490,9 +1495,11 @@ impl ParquetRecordBatchReader { rec => read_records += rec, }; } + None } RowSelectionCursor::All => { self.array_reader.read_records(batch_size)?; + None } }; @@ -1501,8 +1508,24 @@ impl ParquetRecordBatchReader { ArrowError::ParquetError("Struct array reader should return struct array".to_string()) })?; - Ok(if struct_array.len() > 0 { - Some(RecordBatch::from(struct_array)) + let batch = RecordBatch::from(struct_array); + let batch = match filter_mask { + Some(mask) => { + let filtered_batch = filter_record_batch(&batch, &mask)?; + if filtered_batch.num_rows() != read_records { + return Err(general_err!( + "filtered rows mismatch selection - expected {}, got {}", + read_records, + filtered_batch.num_rows() + )); + } + filtered_batch + } + None => batch, + }; + + Ok(if batch.num_rows() > 0 { + Some(batch) } else { None }) diff --git a/parquet/tests/arrow_reader/row_filter/async.rs b/parquet/tests/arrow_reader/row_filter/async.rs index b12108cc3b45..7a2a323c6e4e 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 = (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 = (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,53 +91,242 @@ 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, ] { - let predicate = make_predicate(filter_mask.clone()); + 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(); - // A batch larger than the selection forces Mask to cross the skipped pages. - let stream = ParquetRecordBatchStreamBuilder::new_with_options( - TestReader::new(data.clone()), - ArrowReaderOptions::new().with_page_index_policy(PageIndexPolicy::Required), - ) - .await - .unwrap() + let schema = stream.schema().clone(); + let batches: Vec<_> = stream.try_collect().await.unwrap(); + let batch_sizes = batches + .iter() + .map(RecordBatch::num_rows) + .collect::>(); + 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::() + .unwrap(), + &Int64Array::from(vec![first_value, last_value]) + ); + } + } + } +} + +#[tokio::test] +async fn test_mask_coalesces_loaded_ranges_to_batch_size() { + let values = (0..12).collect::>(); + 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(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(); + assert_eq!( + batches + .iter() + .map(RecordBatch::num_rows) + .collect::>(), + vec![2, 1] + ); + + let batch = concat_batches(&schema, &batches).unwrap(); + assert_eq!( + batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values(), + &[0, 5, 11] + ); +} + +#[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::(0) + .unwrap() + .append_value(row as i32 * 10 + item); + values + .field_builder::(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::>() + }) + .collect::>(); + 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(10), + RowSelector::skip(5), + RowSelector::select(1), + RowSelector::skip(6), RowSelector::select(1), ])) - .with_row_filter(RowFilter::new(vec![Box::new(predicate)])) - .with_batch_size(12) - .with_row_selection_policy(policy) - .with_max_predicate_cache_size(0) + .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); - assert_eq!( - result - .column(0) - .as_any() - .downcast_ref::() - .unwrap(), - &Int64Array::from(vec![first_value, last_value]) - ); - } + let output_schema = stream.schema().clone(); + let batches: Vec = stream.try_collect().await.unwrap(); + assert_eq!( + batches + .iter() + .map(RecordBatch::num_rows) + .collect::>(), + 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] @@ -619,7 +810,10 @@ async fn test_multi_predicate_auto_mask_with_sparse_pages() { // 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::(); Ok(BooleanArray::from_iter( col.iter().map(|v| v.map(|val| val < 250)), @@ -627,8 +821,6 @@ async fn test_multi_predicate_auto_mask_with_sparse_pages() { }); 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 @@ -639,9 +831,17 @@ async fn test_multi_predicate_auto_mask_with_sparse_pages() { .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 = stream.try_collect().await.unwrap(); + assert_eq!(*predicate_batch_sizes.lock().unwrap(), vec![100]); + assert_eq!( + batches + .iter() + .map(RecordBatch::num_rows) + .collect::>(), + 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 77a75220dc2e..a224c95c8f3b 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::, _>>().unwrap(); let result = concat_batches(&schema, &batches).unwrap(); From 240565e241cefee8c2030904cf9dd92727a46b72 Mon Sep 17 00:00:00 2001 From: Qiwei Huang Date: Sun, 19 Jul 2026 21:34:01 +0800 Subject: [PATCH 3/4] Clarify sparse mask batch execution Make the trailing-skip invariant explicit and remove unreachable empty-chunk handling. Extract mask accumulation from next_inner and document how loaded ranges coalesce into one logical batch. --- parquet/src/arrow/arrow_reader/mod.rs | 243 +++++++++++------- parquet/src/arrow/arrow_reader/read_plan.rs | 22 +- parquet/src/arrow/arrow_reader/selection.rs | 107 ++++++-- .../tests/arrow_reader/row_filter/async.rs | 40 +++ 4 files changed, 292 insertions(+), 120 deletions(-) diff --git a/parquet/src/arrow/arrow_reader/mod.rs b/parquet/src/arrow/arrow_reader/mod.rs index 4f9d2229c339..1d7da0598a5a 100644 --- a/parquet/src/arrow/arrow_reader/mod.rs +++ b/parquet/src/arrow/arrow_reader/mod.rs @@ -23,6 +23,7 @@ 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; @@ -1350,6 +1351,123 @@ 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`. +#[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 { + 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 { + 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> { + 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") @@ -1382,80 +1500,9 @@ impl ParquetRecordBatchReader { if batch_size == 0 { return Ok(None); } - let filter_mask = match self.read_plan.row_selection_cursor_mut() { + match self.read_plan.row_selection_cursor_mut() { RowSelectionCursor::Mask(mask_cursor) => { - let mut first_filter_mask: Option = None; - let mut combined_filter_mask: Option = None; - - // Each chunk stays within loaded pages, while the in-progress array and - // filter mask accumulate chunks until the logical batch is full. - while read_records < batch_size && !mask_cursor.is_empty() { - let remaining = batch_size - read_records; - let Some(mask_chunk) = mask_cursor.next_chunk(remaining)? else { - break; - }; - - 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 { - break; - } - 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 - )); - } - match combined_filter_mask.as_mut() { - Some(combined) => combined.append_buffer(mask.values()), - None => match first_filter_mask.take() { - Some(first) => { - let mut combined = - BooleanBufferBuilder::new(first.len() + mask.len()); - combined.append_buffer(&first); - combined.append_buffer(mask.values()); - combined_filter_mask = Some(combined); - } - None => first_filter_mask = Some(mask.values().clone()), - }, - } - read_records += mask_chunk.selected_rows; - } - - if read_records == 0 { - return Ok(None); - } - - let filter_mask = match combined_filter_mask { - Some(combined) => combined.build(), - None => first_filter_mask.ok_or_else(|| { - general_err!("Internal Error: decoded Mask batch has no filter values") - })?, - }; - Some(BooleanArray::from(filter_mask)) + 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() { @@ -1495,35 +1542,13 @@ impl ParquetRecordBatchReader { rec => read_records += rec, }; } - None } RowSelectionCursor::All => { self.array_reader.read_records(batch_size)?; - None - } - }; - - 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()) - })?; - - let batch = RecordBatch::from(struct_array); - let batch = match filter_mask { - Some(mask) => { - let filtered_batch = filter_record_batch(&batch, &mask)?; - if filtered_batch.num_rows() != read_records { - return Err(general_err!( - "filtered rows mismatch selection - expected {}, got {}", - read_records, - filtered_batch.num_rows() - )); - } - filtered_batch } - None => batch, - }; + } + let batch = consume_record_batch(self.array_reader.as_mut())?; Ok(if batch.num_rows() > 0 { Some(batch) } else { @@ -1646,7 +1671,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; @@ -1654,6 +1679,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 dd95ef899b36..09381d8d8c9e 100644 --- a/parquet/src/arrow/arrow_reader/read_plan.rs +++ b/parquet/src/arrow/arrow_reader/read_plan.rs @@ -321,8 +321,7 @@ impl ReadPlanBuilder { let row_selection_cursor = selection .map(|s| { - let trimmed = s.trim(); - let selectors: Vec = trimmed.into(); + let selectors: Vec = s.into(); match selection_strategy { RowSelectionStrategy::Mask => { RowSelectionCursor::new_mask_from_selectors(selectors, loaded_row_ranges) @@ -487,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 141ca329b1c4..a60feabfc200 100644 --- a/parquet/src/arrow/arrow_reader/selection.rs +++ b/parquet/src/arrow/arrow_reader/selection.rs @@ -757,6 +757,25 @@ 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, @@ -774,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 { - 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; @@ -789,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; @@ -810,21 +839,25 @@ impl MaskCursor { self.position = end_position; - Some(MaskChunk { + MaskChunk { initial_skip, chunk_rows, selected_rows, mask_start, - }) + } } - /// Returns the next mask chunk without crossing a row range that was not loaded. - pub(crate) fn next_chunk( - &mut self, - batch_size: usize, - ) -> Result, ParquetError> { + /// 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 { + debug_assert!(batch_size > 0); + debug_assert!(!self.is_empty()); + if self.loaded_row_ranges.is_none() { - return Ok(self.next_mask_chunk(batch_size)); + return Ok(self.next_mask_chunk_non_empty(batch_size)); } let start_position = self.position; @@ -833,10 +866,10 @@ impl MaskCursor { cursor += 1; } - if cursor == self.mask.len() { - self.position = cursor; - return Ok(None); - } + debug_assert!( + cursor < self.mask.len(), + "ReadPlan must remove trailing skips from Mask selections" + ); let loaded_range_end = self .loaded_row_ranges @@ -858,12 +891,12 @@ impl MaskCursor { } self.position = cursor; - Ok(Some(MaskChunk { + Ok(MaskChunk { initial_skip: mask_start - start_position, chunk_rows: cursor - mask_start, selected_rows, mask_start, - })) + }) } /// Materialise the boolean values for a mask-backed chunk @@ -981,6 +1014,13 @@ impl RowSelectionCursor { selectors: Vec, loaded_row_ranges: Option>, ) -> 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, @@ -1633,11 +1673,38 @@ mod tests { unreachable!() }; - let chunk = cursor.next_chunk(10).unwrap().unwrap(); + 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/tests/arrow_reader/row_filter/async.rs b/parquet/tests/arrow_reader/row_filter/async.rs index 7a2a323c6e4e..885326687a19 100644 --- a/parquet/tests/arrow_reader/row_filter/async.rs +++ b/parquet/tests/arrow_reader/row_filter/async.rs @@ -213,6 +213,46 @@ async fn test_mask_coalesces_loaded_ranges_to_batch_size() { ); } +#[tokio::test] +async fn test_mask_stops_after_trailing_skipped_pages() { + let values = (0..8).collect::>(); + 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 = stream.try_collect().await.unwrap(); + assert_eq!( + batches + .iter() + .map(RecordBatch::num_rows) + .collect::>(), + vec![2] + ); + assert_eq!( + batches[0] + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values(), + &[0, 5] + ); +} + #[tokio::test] async fn test_mask_nested_projection_with_different_page_boundaries() { use arrow_array::{ From c9af8dffc93a73a9486bf9cdfb5282dcaeab41c4 Mon Sep 17 00:00:00 2001 From: Andrew Lamb Date: Mon, 20 Jul 2026 14:37:38 -0400 Subject: [PATCH 4/4] Add example diagram to FilterMaskAccumulator --- parquet/src/arrow/arrow_reader/mod.rs | 19 ++++++++++++++++++- 1 file changed, 18 insertions(+), 1 deletion(-) diff --git a/parquet/src/arrow/arrow_reader/mod.rs b/parquet/src/arrow/arrow_reader/mod.rs index 1d7da0598a5a..621c1c812afd 100644 --- a/parquet/src/arrow/arrow_reader/mod.rs +++ b/parquet/src/arrow/arrow_reader/mod.rs @@ -1355,7 +1355,24 @@ pub struct ParquetRecordBatchReader { /// /// 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`. +/// 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]