alamb commented on code in PR #10288:
URL: https://github.com/apache/arrow-rs/pull/10288#discussion_r3610537564
##########
parquet/src/arrow/arrow_reader/selection.rs:
##########
@@ -824,6 +818,54 @@ impl MaskCursor {
})
}
+ /// Returns the next mask chunk without crossing a row range that was not
loaded.
Review Comment:
This design is really neat -- but it took me quite a while to understand
what it was doing. I was wondering if you might be willing to add some
explanatory comments with the high level design.
I was specifically thinking:
1. An example with data pages, a mask and loaded row ranges
2. An image showing how the concepts are related, simlar to what is on
https://github.com/apache/arrow-rs/issues/8845
For example, I think the `LoadedRowRange` example from this PR description
could be adapted and make this easier to understand
##########
parquet/src/arrow/arrow_reader/mod.rs:
##########
@@ -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 {
Review Comment:
Something else that might be worth considering is using the
https://docs.rs/arrow/latest/arrow/compute/kernels/coalesce/index.html kernel
internally to internally coalesce the results into a large enough record batch
(but I can't remember if DatFusion does that already)
##########
parquet/src/arrow/push_decoder/reader_builder/mod.rs:
##########
@@ -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
Review Comment:
this is a nice cleanup
##########
parquet/src/arrow/arrow_reader/mod.rs:
##########
@@ -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_mask_chunk(batch_size) else {
- return Ok(None);
+ let mut first_filter_mask: Option<BooleanBuffer> = None;
+ let mut combined_filter_mask: Option<BooleanBufferBuilder> =
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 {
Review Comment:
I ran code coverage as follows,
```shell
cargo llvm-cov --html test -p parquet --features=arrow,async
```
And this line seems uncovered -- I am not sure if it is possible to add
coverage
<img width="950" height="530" alt="Image"
src="https://github.com/user-attachments/assets/e488356d-55bb-4f0f-b2eb-9e5330521901";
/>
##########
parquet/src/arrow/arrow_reader/mod.rs:
##########
@@ -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() {
Review Comment:
I spent a while reading this and I think it makes sense, but I found it hard
to follow because:
1. This function is now quite long (several pages)
2. The indent is 4 levels deep
3. it didn't have a high level explanation / example (specifically of how
the first_filter_mask and combined_filter_mask are are related / built up)
Some specific suggestions:
1. Find a way to refactor the code into a function somehow (maybe a
`CombinedFilterMaskBuilder`?)
2. Add a high level description / example of what this is doing (calling
skip/decode to read out a underling batch, and then calling filter record batch
on the result (a refactored function would also add a natural place to add
more documentation)
##########
parquet/src/arrow/arrow_reader/mod.rs:
##########
@@ -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;
Review Comment:
Likewise this one is also uncovered
<img width="980" height="181" alt="Image"
src="https://github.com/user-attachments/assets/7030ceb6-3546-4d3d-99ac-d1e9e5f75d1b";
/>
I know the coverage analysis is somewhat pedantic, but given how much this
code is used in production, I think adding coverage for all these corner cases
is quite valuable
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
