Skip to content
Merged
Show file tree
Hide file tree
Changes from 4 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 5 additions & 5 deletions parquet/src/arrow/arrow_reader/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1366,12 +1366,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)
/// ```
///
Expand All @@ -1381,8 +1381,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 {
Expand Down
10 changes: 5 additions & 5 deletions parquet/src/arrow/arrow_reader/read_plan.rs
Original file line number Diff line number Diff line change
Expand Up @@ -633,16 +633,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());
Expand Down
31 changes: 21 additions & 10 deletions parquet/src/arrow/arrow_reader/selection/cursor.rs
Original file line number Diff line number Diff line change
Expand Up @@ -176,10 +176,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`. Trimming trailing skipped rows from
Comment thread
hhhizzz marked this conversation as resolved.
Outdated
/// each chunk prevents a fixed-size cached read from extending into an unloaded
/// page. [`ParquetRecordBatchReader`] still accumulates both chunks and applies
/// the combined mask `11` once.
///
/// [`ParquetRecordBatchReader`]: crate::arrow::arrow_reader::ParquetRecordBatchReader
#[derive(Debug)]
Expand Down Expand Up @@ -254,6 +255,11 @@ impl MaskCursor {
}

/// Returns the next non-empty mask chunk without crossing an unloaded row range.
/// When loaded ranges are present, the chunk ends immediately after its last
/// selected row, leaving trailing unselected rows for the next call's initial skip.
/// `CachedArrayReader` can fetch the full cache batch containing that selected
/// row; predicate fetch expands cached columns to batch boundaries so this read
/// stays within fetched data even when the batch crosses the loaded-range end.
///
/// The [`ReadPlan`](crate::arrow::arrow_reader::ReadPlan) removes trailing
/// skips before constructing this cursor. Callers therefore only invoke
Expand All @@ -272,10 +278,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
Expand All @@ -289,17 +298,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,
})
Expand Down
4 changes: 3 additions & 1 deletion parquet/src/arrow/push_decoder/reader_builder/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -554,7 +554,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 because `CachedArrayReader` may fetch
// a full batch around the last selected row of a `MaskCursor` chunk.
.with_cache_projection(Some(filter_info.cache_projection()))
.with_column_chunks(column_chunks)
.build();
Expand Down
70 changes: 69 additions & 1 deletion parquet/tests/arrow_reader/row_filter/async.rs
Original file line number Diff line number Diff line change
Expand Up @@ -35,7 +35,7 @@ use parquet::{
ArrowWriter, ParquetRecordBatchStreamBuilder, ProjectionMask,
arrow_reader::{
ArrowPredicateFn, ArrowReaderOptions, RowFilter, RowSelection, RowSelectionPolicy,
RowSelector,
RowSelector, metrics::ArrowReaderMetrics,
},
},
file::{
Expand Down Expand Up @@ -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() {
Comment thread
hhhizzz marked this conversation as resolved.
let values = (0..60).collect::<Vec<i64>>();
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::<Vec<_>>();
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<RecordBatch> = stream.try_collect().await.unwrap();
let output = concat_batches(&output_schema, &batches).unwrap();
assert_eq!(
output
.column(0)
.as_any()
.downcast_ref::<Int64Array>()
.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::<Vec<i64>>();
Expand Down
Loading