fix(parquet): support mask filtering across skipped pages

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,14 +1688,36 @@ 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;
     use bytes::Bytes;
     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");

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;