From 1b69fd70470f17fc98ecdf0800861ee7e1ba9eca Mon Sep 17 00:00:00 2001 From: Huang Qiwei Date: Thu, 20 Aug 2026 05:26:50 +0800 Subject: [PATCH] fix(parquet): prevent cached Mask reads from crossing unloaded sparse pages (#10735) # Proposed title fix(parquet): prevent cached Mask reads from crossing unloaded sparse pages # Which issue does this PR close? - Closes #10733. # Rationale for this change The async Parquet reader may combine page pruning, predicate caching, and a Mask-backed row selection. `MaskCursor::next_chunk` bounds each chunk by the current loaded row range, but previously retained all trailing skipped rows up to that boundary. For example, if a loaded range starts with one selected row and the rest of the range is skipped, the cursor returned the whole loaded range as `chunk_rows`. A fixed-size cached reader can then extend that read beyond the cached segment and into a sparse page that was never loaded, producing: ```text Invalid offset in sparse column chunk data: ..., no matching page found ``` DataFusion predicate pushdown exposed this in TPC-DS queries 66, 75, and 81. Disabling the predicate cache avoided the failure, but the cache is enabled by default and should be safe with both Auto and explicit Mask selection. # What changes are included in this PR? - Track the position immediately after the last selected row in the current loaded range. - Use that position for both `MaskCursor::position` and `MaskChunk::chunk_rows`, instead of the scan position at the end of the loaded range. - Leave trailing skipped rows to the next cursor step. They become part of `initial_skip`, allowing `ArrayReader::skip_records` to cross unloaded pages. - Update the read-plan unit test to lock down the chunk boundaries. - Add an async regression test with three pages, a cached predicate column, a sparse initial selection, and both Auto and explicit Mask policies. The selected rows and output batch semantics are unchanged. The fix only avoids decoding trailing rows that are not selected. # Are these changes tested? Yes. Focused regression test: ```shell cargo test -p parquet --features async --test arrow_reader 'row_filter::r#async::test_cached_mask_reads_sparse_pages_without_error' -- --exact --nocapture ``` - Test-only commit on top of `bb1e6cd070`: failed with the sparse-column offset error. - This fix: 1 passed, 0 failed. - Mask-focused `arrow_reader` tests: 7 passed, 0 failed. I also ran one-iteration end-to-end correctness smoke tests with DataFusion `52964c2966e855f47b96a15d1aace01baee1f2c1`. For validation only, the Arrow default policy was changed to Mask on top of this PR so the workloads could not fall back to Selectors: - TPC-DS SF10: 99/99 queries succeeded. The 96 queries that had successful pre-fix results had identical row counts; the previously failing Q66, Q75, and Q81 also succeeded. - TPC-H SF10: 22/22 queries succeeded, with identical row counts to the pre-fix baseline. - ClickBench: 43/43 queries succeeded with predicate pushdown explicitly enabled, with identical row counts to the pre-fix baseline. - No `invalid offset`, `no matching page`, panic, or query failure was found in the candidate logs. These workload runs are correctness smoke tests, not performance claims. ## Suggested CI follow-up cc @alamb This bug requires an interaction between predicate pushdown, page pruning, selection representation, and predicate caching, which is difficult to cover with isolated reader tests alone. I suggest adding a CI or scheduled benchmark correctness check that runs representative TPC-DS, TPC-H, and ClickBench queries with **Parquet predicate pushdown enabled**; # Are there any user-facing changes? No API changes. Async Parquet scans using predicate caching and sparse page reads no longer fail when Auto or explicit Mask selection is used. (cherry picked from commit 75406dc4f3e75be2eb6d66a7733a549229414d90) --- .../arrow/array_reader/cached_array_reader.rs | 3 + parquet/src/arrow/arrow_reader/mod.rs | 10 +-- parquet/src/arrow/arrow_reader/read_plan.rs | 10 +-- .../arrow/arrow_reader/selection/cursor.rs | 29 +++++--- .../arrow/push_decoder/reader_builder/mod.rs | 4 +- .../tests/arrow_reader/row_filter/async.rs | 70 ++++++++++++++++++- 6 files changed, 104 insertions(+), 22 deletions(-) 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::>();