alamb commented on code in PR #10288:
URL: https://github.com/apache/arrow-rs/pull/10288#discussion_r3616594726
##########
parquet/src/arrow/arrow_reader/mod.rs:
##########
@@ -1349,6 +1351,123 @@ 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`.
+#[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
Review Comment:
thank you -- this now ready very clearly to me
##########
parquet/src/arrow/arrow_reader/mod.rs:
##########
@@ -1349,6 +1351,123 @@ pub struct ParquetRecordBatchReader {
read_plan: ReadPlan,
}
+/// Accumulates filter masks for decoded chunks in one logical output batch.
Review Comment:
Thank you for this-- I found it much clearer
--
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]
