alamb commented on code in PR #11438:
URL: https://github.com/apache/arrow-rs/pull/11438#discussion_r4234232140
##########
parquet/src/file/metadata/push_decoder.rs:
##########
@@ -489,6 +470,97 @@ impl ParquetMetaDataPushDecoder {
}
}
+ /// Try to decode the page indexes from the pushed data, returning the
+ /// decoded [`PageIndex`] or an error if not enough data is available.
+ ///
+ /// This method can only be called after [`Self::try_new_with_metadata`],
Review Comment:
I wonder if we could add an example here of how to use this function -- this
text explains it in words, but I am still not quite sure how it would work
##########
parquet/src/file/metadata/push_decoder.rs:
##########
@@ -489,6 +470,97 @@ impl ParquetMetaDataPushDecoder {
}
}
+ /// Try to decode the page indexes from the pushed data, returning the
+ /// decoded [`PageIndex`] or an error if not enough data is available.
+ ///
+ /// This method can only be called after [`Self::try_new_with_metadata`],
+ /// which initializes the decoder with previously decoded metadata.
+ ///
+ /// This method is useful when you want to:
+ /// 1. First decode the footer metadata without reading the page indexes
+ /// (using [`PageIndexPolicy::Skip`])
+ /// 2. Later selectively decode a subset of the page indexes using
+ /// [`Self::with_column_index_mask`] or [`Self::with_offset_index_mask`]
+ ///
+ /// When the data can be accessed through a [`ChunkReader`], prefer
+ /// [`ParquetMetaDataReader::read_page_index`] (or
+ /// [`ParquetMetaDataReader::read_page_index_async`] for asynchronous
+ /// I/O), which handles the decoder state machine and I/O requests.
+ ///
+ /// # Returns
+ ///
+ /// * `Ok(DecodeResult::Data(Some(page_index)))` - Successfully decoded the
+ /// requested page indexes based on the configured policies and masks
+ /// * `Ok(DecodeResult::Data(None))` - No page indexes were requested
+ /// (both policies set to [`PageIndexPolicy::Skip`] or masks select
nothing)
+ /// * `Ok(DecodeResult::NeedsData(ranges))` - More data is needed to decode
+ /// the page indexes. Push the requested ranges and call this method
again.
+ /// * `Ok(DecodeResult::Finished)` - The decoder has finished and no more
+ /// data can be decoded
+ /// * `Err(_)` - An error occurred during decoding
+ ///
+ /// [`ParquetMetaDataReader::read_page_index`]:
crate::file::metadata::ParquetMetaDataReader::read_page_index
+ /// [`ParquetMetaDataReader::read_page_index_async`]:
crate::file::metadata::ParquetMetaDataReader::read_page_index_async
+ pub fn try_decode_page_index(&mut self) ->
Result<DecodeResult<Option<PageIndex>>> {
+ match std::mem::replace(&mut self.state, DecodeState::Intermediate) {
+ DecodeState::ReadingPageIndex(metadata) => {
+ Ok(match self.try_decode_page_index_for_metadata(metadata)? {
+ PageIndexDecodeResult::NeedsData(ranges) =>
DecodeResult::NeedsData(ranges),
+ PageIndexDecodeResult::Data { page_index, .. } => {
+ DecodeResult::Data(page_index)
+ }
+ })
+ }
+ DecodeState::Finished => Ok(DecodeResult::Finished),
+ _ => Err(general_err!(
+ "ParquetMetaDataPushDecoder: internal error, invalid state"
+ )),
+ }
+ }
+
+ fn try_decode_page_index_for_metadata(
+ &mut self,
+ metadata: Box<ParquetMetaData>,
+ ) -> Result<PageIndexDecodeResult> {
+ let range = range_for_page_index(
+ &metadata,
+ self.column_index_policy,
+ self.offset_index_policy,
+ &self.column_index_mask,
+ &self.offset_index_mask,
+ );
+
+ let Some(page_index_range) = range else {
+ self.state = DecodeState::Finished;
+ return Ok(PageIndexDecodeResult::Data {
+ metadata,
+ page_index: None,
+ });
+ };
+
+ if !self.buffers.has_range(&page_index_range) {
+ self.state = DecodeState::ReadingPageIndex(metadata);
+ return
Ok(PageIndexDecodeResult::NeedsData(vec![page_index_range]));
+ }
+
+ let buffer = self.get_bytes(&page_index_range)?;
+ let page_index = parse_page_index(
+ &metadata,
+ self.column_index_policy,
+ self.offset_index_policy,
+ &self.column_index_mask,
+ &self.offset_index_mask,
+ &buffer,
+ page_index_range.start,
+ )?;
+
+ self.state = DecodeState::Finished;
Review Comment:
Does setting the state to `Finished` mean that we can only try and decode
page metadata once and that subsequent calls to `try_decode_page_index` will
error? I suppose that makes sense
But I do think an end to end example showing how to use these APIs to add
PageIndex to an existing ParquetMetadata would be super helpful
Like is it read the metadata without a page index and then create a new
reader with that metadata configured to just read the relevant parts of the
page index 🤔
##########
parquet/src/file/metadata/push_decoder.rs:
##########
@@ -523,6 +595,14 @@ enum DecodeState {
Intermediate,
}
+enum PageIndexDecodeResult {
Review Comment:
This looks pretty similar to
`DecodeResult<(Box<ParquetMetaData>,Option<PageIndex>)` 🤔
##########
parquet/src/file/metadata/reader.rs:
##########
@@ -573,6 +574,78 @@ impl ParquetMetaDataReader {
self.read_page_indexes_sized(reader, reader.len())
}
+ /// Reads and returns the page index structures described by the stored
+ /// [`ParquetMetaData`].
+ ///
+ /// Unlike [`Self::read_page_indexes`], this returns the decoded
[`PageIndex`] separately
+ /// instead of attaching it to the metadata. This is useful for populating
an external page
+ /// index cache. The configured page index policies and masks determine
which indexes are
+ /// returned.
+ ///
+ /// This operation consumes `self` and the stored [`ParquetMetaData`]. Use
+ /// [`Self::new_with_metadata`] to construct the reader.
+ pub fn read_page_index<R: ChunkReader>(self, reader: &R) ->
Result<Option<PageIndex>> {
Review Comment:
How about `read_page_index_and_finish` and `load_page_index_and_finish` to
mirror existing APIs (suggested by Claude)?
It is also kind of strange that this consumes `self` -- which has an
embedded `ParquetMetaData` -- maybe it could return the `(ParquetMetaData,
Option<PageIndex>)` instead 🤔
##########
parquet/src/file/metadata/push_decoder.rs:
##########
@@ -619,6 +700,66 @@ mod tests {
assert!(masked.end - masked.start < all.end - all.start);
}
+ #[test]
+ fn test_page_index_decoder_incremental() {
+ let metadata = test_metadata_without_page_index();
+ let mut decoder =
+ ParquetMetaDataPushDecoder::try_new_with_metadata(test_file_len(),
metadata)
+ .unwrap()
+ .with_column_index_policy(PageIndexPolicy::Required)
+ .with_offset_index_policy(PageIndexPolicy::Required)
+
.with_column_index_mask(ColumnChunkMask::row_groups_and_columns([1], [0, 2]))
+
.with_offset_index_mask(ColumnChunkMask::row_groups_and_columns([0], [1]));
+
+ let ranges = expect_needs_data(decoder.try_decode_page_index());
+ assert_eq!(ranges.len(), 1);
+ push_ranges_to_metadata_decoder(&mut decoder, ranges);
+
+ let page_index = expect_data(decoder.try_decode_page_index()).unwrap();
+ expect_finished(decoder.try_decode_page_index());
+
+ for row_group_idx in 0..2 {
+ for column_idx in 0..3 {
+ assert_eq!(
+ page_index.column_index(row_group_idx,
column_idx).is_some(),
+ row_group_idx == 1 && matches!(column_idx, 0 | 2)
+ );
+ assert_eq!(
+ page_index.offset_index(row_group_idx,
column_idx).is_some(),
+ row_group_idx == 0 && column_idx == 1
+ );
+ }
+ }
+ }
+
+ #[test]
+ fn test_page_index_decoder_no_indexes_requested() {
+ let metadata = test_metadata_without_page_index();
+ let mut decoder =
+ ParquetMetaDataPushDecoder::try_new_with_metadata(test_file_len(),
metadata)
+ .unwrap()
+ .with_page_index_policy(PageIndexPolicy::Skip);
+
+ assert!(expect_data(decoder.try_decode_page_index()).is_none());
+ expect_finished(decoder.try_decode_page_index());
+
+ let metadata = test_metadata_without_page_index();
Review Comment:
I think another interesting test would be to make the second
ParquetMetaDataPushDecoder with a page_index_mask to read it as a subsequent
step
--
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]
