diff --git a/parquet/src/arrow/push_decoder/mod.rs b/parquet/src/arrow/push_decoder/mod.rs index cdb0715edb55..24384471a4ec 100644 --- a/parquet/src/arrow/push_decoder/mod.rs +++ b/parquet/src/arrow/push_decoder/mod.rs @@ -365,6 +365,15 @@ impl ParquetPushDecoder { pub fn buffered_bytes(&self) -> u64 { self.state.buffered_bytes() } + + /// Clear any staged byte ranges currently buffered for future decode work. + /// + /// This clears byte ranges still owned by the decoder's internal + /// `PushBuffers`. It does not affect any data that has already been handed + /// off to an active [`ParquetRecordBatchReader`]. + pub fn clear_all_ranges(&mut self) { + self.state.clear_all_ranges(); + } } /// Internal state machine for the [`ParquetPushDecoder`] @@ -573,6 +582,20 @@ impl ParquetDecoderState { ParquetDecoderState::Finished => 0, } } + + /// Clear any staged ranges currently buffered in the decoder. + fn clear_all_ranges(&mut self) { + match self { + ParquetDecoderState::ReadingRowGroup { + remaining_row_groups, + } => remaining_row_groups.clear_all_ranges(), + ParquetDecoderState::DecodingRowGroup { + record_batch_reader: _, + remaining_row_groups, + } => remaining_row_groups.clear_all_ranges(), + ParquetDecoderState::Finished => {} + } + } } #[cfg(test)] @@ -665,6 +688,55 @@ mod test { assert_eq!(all_output, *TEST_BATCH); } + /// Releasing staged ranges should free speculative buffers without affecting + /// the active row group reader. + #[test] + fn test_decoder_clear_all_ranges() { + let mut decoder = ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata()) + .unwrap() + .with_batch_size(100) + .build() + .unwrap(); + + decoder + .push_range(test_file_range(), TEST_FILE_DATA.clone()) + .unwrap(); + assert_eq!(decoder.buffered_bytes(), test_file_len()); + + // The current row group reader is built from the prefetched bytes, but + // the speculative full-file range remains staged in the decoder. + let batch1 = expect_data(decoder.try_decode()); + assert_eq!(batch1, TEST_BATCH.slice(0, 100)); + assert_eq!(decoder.buffered_bytes(), test_file_len()); + + // All of the buffer is released + decoder.clear_all_ranges(); + assert_eq!(decoder.buffered_bytes(), 0); + + // The active reader still owns the current row group's bytes, so it can + // continue decoding without consulting PushBuffers. + let batch2 = expect_data(decoder.try_decode()); + assert_eq!(batch2, TEST_BATCH.slice(100, 100)); + assert_eq!(decoder.buffered_bytes(), 0); + + // Moving to the next row group now requires the decoder to ask for data + // again because the staged speculative ranges were released. + let ranges = expect_needs_data(decoder.try_decode()); + let num_bytes_requested: u64 = ranges.iter().map(|r| r.end - r.start).sum(); + push_ranges_to_decoder(&mut decoder, ranges); + assert_eq!(decoder.buffered_bytes(), num_bytes_requested); + + let batch3 = expect_data(decoder.try_decode()); + assert_eq!(batch3, TEST_BATCH.slice(200, 100)); + assert_eq!(decoder.buffered_bytes(), 0); + + let batch4 = expect_data(decoder.try_decode()); + assert_eq!(batch4, TEST_BATCH.slice(300, 100)); + assert_eq!(decoder.buffered_bytes(), 0); + + expect_finished(decoder.try_decode()); + } + /// Decode the entire file incrementally, simulating partial reads #[test] fn test_decoder_partial() { diff --git a/parquet/src/arrow/push_decoder/reader_builder/mod.rs b/parquet/src/arrow/push_decoder/reader_builder/mod.rs index d3d78ca7c263..922d8070c064 100644 --- a/parquet/src/arrow/push_decoder/reader_builder/mod.rs +++ b/parquet/src/arrow/push_decoder/reader_builder/mod.rs @@ -212,6 +212,11 @@ impl RowGroupReaderBuilder { self.buffers.buffered_bytes() } + /// Clear any staged ranges currently buffered for future decode work. + pub fn clear_all_ranges(&mut self) { + self.buffers.clear_all_ranges(); + } + /// take the current state, leaving None in its place. /// /// Returns an error if there the state wasn't put back after the previous diff --git a/parquet/src/arrow/push_decoder/remaining.rs b/parquet/src/arrow/push_decoder/remaining.rs index 4613fda08749..2986ca0da8d8 100644 --- a/parquet/src/arrow/push_decoder/remaining.rs +++ b/parquet/src/arrow/push_decoder/remaining.rs @@ -70,6 +70,11 @@ impl RemainingRowGroups { self.row_group_reader_builder.buffered_bytes() } + /// Clear any staged ranges currently buffered for future decode work + pub fn clear_all_ranges(&mut self) { + self.row_group_reader_builder.clear_all_ranges(); + } + /// returns [`ParquetRecordBatchReader`] suitable for reading the next /// group of rows from the Parquet data, or the list of data ranges still /// needed to proceed diff --git a/parquet/src/util/push_buffers.rs b/parquet/src/util/push_buffers.rs index 0c00cf9bd57f..eb4982fb3c6f 100644 --- a/parquet/src/util/push_buffers.rs +++ b/parquet/src/util/push_buffers.rs @@ -154,6 +154,13 @@ impl PushBuffers { self.ranges = new_ranges; self.buffers = new_buffers; } + + /// Clear all buffered ranges and their corresponding data + #[cfg(feature = "arrow")] + pub fn clear_all_ranges(&mut self) { + self.ranges.clear(); + self.buffers.clear(); + } } impl Length for PushBuffers {