diff --git a/parquet/src/arrow/array_reader/cached_array_reader.rs b/parquet/src/arrow/array_reader/cached_array_reader.rs index 73f3ba6c8fb2..7e31ff79203d 100644 --- a/parquet/src/arrow/array_reader/cached_array_reader.rs +++ b/parquet/src/arrow/array_reader/cached_array_reader.rs @@ -135,6 +135,9 @@ impl CachedArrayReader { self.inner_position += skipped; } + // For sparse mask reads, this full-batch fallback relies on `MaskCursor` + // ending every chunk at a selected row. Predicate fetch expands cached + // columns to batch boundaries, so the batch containing that row is loaded. let read = self.inner.read_records(self.batch_size)?; // If there are no remaining records (EOF), return immediately without diff --git a/parquet/src/arrow/arrow_reader/mod.rs b/parquet/src/arrow/arrow_reader/mod.rs index 2517e892fc0c..5341739d98cb 100644 --- a/parquet/src/arrow/arrow_reader/mod.rs +++ b/parquet/src/arrow/arrow_reader/mod.rs @@ -1395,12 +1395,12 @@ 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 `1001` and `1` become `10011`: /// /// ```text -/// append(1000) append(1) +/// append(1001) append(1) /// Empty ───────────────▶ Single ───────────────▶ Combined -/// 1000 10001 +/// 1001 10011 /// (zero copy) (promoted to builder) /// ``` /// @@ -1410,8 +1410,8 @@ pub struct ParquetRecordBatchReader { /// /// ```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 +/// chunk masks: [1 0 0 1] [1] +/// finish(): 1 0 0 1 1 <-- filters the whole batch in one pass /// ``` #[derive(Default)] enum FilterMaskAccumulator { diff --git a/parquet/src/arrow/arrow_reader/read_plan.rs b/parquet/src/arrow/arrow_reader/read_plan.rs index 04f132e1bc1d..5b6cff4e3b58 100644 --- a/parquet/src/arrow/arrow_reader/read_plan.rs +++ b/parquet/src/arrow/arrow_reader/read_plan.rs @@ -645,16 +645,16 @@ mod tests { panic!("expected a Mask cursor"); }; - // The first chunk must end at the loaded range boundary (row 4), not - // continue into the unloaded gap. + // The first chunk stops at its final selected row instead of carrying + // trailing skipped rows to the loaded range boundary. let first = cursor.next_chunk(12).unwrap(); assert_eq!(first.initial_skip, 0); - assert_eq!(first.chunk_rows, 4); + assert_eq!(first.chunk_rows, 1); assert_eq!(first.selected_rows, 1); - // The second chunk skips the gap and decodes only within [10, 12). + // The second chunk skips directly to the next selected row. let second = cursor.next_chunk(12).unwrap(); - assert_eq!(second.initial_skip, 7); + assert_eq!(second.initial_skip, 10); assert_eq!(second.chunk_rows, 1); assert_eq!(second.selected_rows, 1); assert!(cursor.is_empty()); diff --git a/parquet/src/arrow/arrow_reader/selection/cursor.rs b/parquet/src/arrow/arrow_reader/selection/cursor.rs index dcb490746c9e..2b63c098f0f7 100644 --- a/parquet/src/arrow/arrow_reader/selection/cursor.rs +++ b/parquet/src/arrow/arrow_reader/selection/cursor.rs @@ -180,10 +180,11 @@ impl SelectorsCursor { /// 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. +/// The first chunk decodes `[0, 1)` with mask `1`. The next chunk skips to row +/// 11 and decodes `[11, 12)` with mask `1`. When loaded ranges are present, +/// every returned chunk ends at a selected row and never includes trailing +/// unselected rows. [`ParquetRecordBatchReader`] still accumulates both chunks +/// and applies the combined mask `11` once. /// /// [`ParquetRecordBatchReader`]: crate::arrow::arrow_reader::ParquetRecordBatchReader #[derive(Debug)] @@ -258,6 +259,9 @@ impl MaskCursor { } /// Returns the next non-empty mask chunk without crossing an unloaded row range. + /// When loaded ranges are present, every returned chunk ends immediately after a + /// selected row and therefore never contains trailing unselected rows. Those rows + /// remain for the next call's initial skip. /// /// The [`ReadPlan`](crate::arrow::arrow_reader::ReadPlan) removes trailing /// skips before constructing this cursor. Callers therefore only invoke @@ -276,10 +280,13 @@ impl MaskCursor { cursor += 1; } - debug_assert!( - cursor < self.mask.len(), - "ReadPlan must remove trailing skips from Mask selections" - ); + if cursor == self.mask.len() { + return Err(ParquetError::General( + "Internal Error: Mask cursor reached the end without finding a selected row; \ + ReadPlan must remove trailing skips" + .to_string(), + )); + } let loaded_range_end = self .loaded_row_ranges @@ -293,17 +300,19 @@ impl MaskCursor { let mask_start = cursor; let mut selected_rows = 0; + let mut chunk_end = cursor; while cursor < loaded_range_end && cursor < self.mask.len() && selected_rows < batch_size { if self.mask.value(cursor) { selected_rows += 1; + chunk_end = cursor + 1; } cursor += 1; } - self.position = cursor; + self.position = chunk_end; Ok(MaskChunk { initial_skip: mask_start - start_position, - chunk_rows: cursor - mask_start, + chunk_rows: chunk_end - mask_start, selected_rows, mask_start, }) diff --git a/parquet/src/arrow/push_decoder/reader_builder/mod.rs b/parquet/src/arrow/push_decoder/reader_builder/mod.rs index b773f95fd55e..200974af7158 100644 --- a/parquet/src/arrow/push_decoder/reader_builder/mod.rs +++ b/parquet/src/arrow/push_decoder/reader_builder/mod.rs @@ -550,7 +550,9 @@ impl RowGroupReaderBuilder { predicate.projection(), // use the predicate's projection ) .with_selection(plan_builder.selection()) - // Fetch predicate columns; expand selection only for cached predicate columns + // Cached output columns reuse these predicate-stage chunks. Expand their + // selection to cache batch boundaries so a cache miss can safely fetch a + // complete batch from the retained sparse column data. .with_cache_projection(Some(filter_info.cache_projection())) .with_column_chunks(column_chunks) .build(); diff --git a/parquet/tests/arrow_reader/row_filter/async.rs b/parquet/tests/arrow_reader/row_filter/async.rs index fe16e426ea59..3adbb05dae5b 100644 --- a/parquet/tests/arrow_reader/row_filter/async.rs +++ b/parquet/tests/arrow_reader/row_filter/async.rs @@ -35,7 +35,7 @@ use parquet::{ ArrowWriter, ParquetRecordBatchStreamBuilder, ProjectionMask, arrow_reader::{ ArrowPredicateFn, ArrowReaderOptions, RowFilter, RowSelection, RowSelectionPolicy, - RowSelector, + RowSelector, metrics::ArrowReaderMetrics, }, }, file::{ @@ -169,6 +169,74 @@ async fn test_row_filter_full_page_skip_is_handled_async() { } } +#[tokio::test] +async fn test_cached_mask_reads_sparse_pages_without_error() { + let values = (0..60).collect::>(); + let data = make_two_column_i64_file(&values, 20); + + for policy in [ + RowSelectionPolicy::Auto { threshold: 32 }, + RowSelectionPolicy::Mask, + ] { + let metrics = ArrowReaderMetrics::enabled(); + let builder = ParquetRecordBatchStreamBuilder::new_with_options( + TestReader::new(data.clone()), + ArrowReaderOptions::new().with_page_index_policy(PageIndexPolicy::Required), + ) + .await + .unwrap(); + let schema = builder.parquet_schema().clone(); + let projection = ProjectionMask::leaves(&schema, [0]); + let page_first_rows = builder.metadata().offset_index().unwrap()[0][0] + .page_locations() + .iter() + .map(|page| page.first_row_index) + .collect::>(); + assert_eq!(page_first_rows, vec![0, 20, 40]); + + let predicate = ArrowPredicateFn::new(projection.clone(), |batch: RecordBatch| { + Ok(BooleanArray::from(vec![true; batch.num_rows()])) + }); + // Extending the first mask chunk to the 20-row page boundary would make + // the 8-row cache batch at rows 16..24 cross into the unloaded middle page. + let stream = builder + .with_projection(projection) + .with_row_filter(RowFilter::new(vec![Box::new(predicate)])) + .with_row_selection(RowSelection::from(vec![ + RowSelector::select(1), + RowSelector::skip(39), + RowSelector::select(1), + ])) + .with_batch_size(8) + .with_max_predicate_cache_size(1024) + .with_row_selection_policy(policy) + .with_metrics(metrics.clone()) + .build() + .unwrap(); + + let output_schema = stream.schema().clone(); + let batches: Vec = stream.try_collect().await.unwrap(); + let output = concat_batches(&output_schema, &batches).unwrap(); + assert_eq!( + output + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values(), + &[0, 40], + "policy={policy:?}" + ); + assert!( + metrics + .records_read_from_cache() + .expect("metrics are enabled") + > 0, + "predicate cache was not exercised for policy={policy:?}" + ); + } +} + #[tokio::test] async fn test_mask_coalesces_loaded_ranges_to_batch_size() { let values = (0..12).collect::>();