alamb commented on code in PR #11438:
URL: https://github.com/apache/arrow-rs/pull/11438#discussion_r4275287028
##########
parquet/src/file/metadata/reader.rs:
##########
@@ -573,19 +574,80 @@ impl ParquetMetaDataReader {
self.read_page_indexes_sized(reader, reader.len())
}
- /// Read the page index structures when a [`ParquetMetaData`] has already
been obtained.
- /// This variant is used when `reader` cannot access the entire Parquet
file (e.g. it is
- /// a [`Bytes`] struct containing the tail of the file).
- /// See [`Self::new_with_metadata()`] and [`Self::has_metadata()`]. Like
- /// [`Self::try_parse_sized()`] this function may return
[`ParquetError::NeedMoreData`].
+ /// Reads the page index structures described by the stored
[`ParquetMetaData`], returning the
+ /// metadata together with the decoded page index.
///
- /// See [`Self::read_page_indexes()`] for a description of the replacement
behavior on
- /// multiple calls.
- pub fn read_page_indexes_sized<R: ChunkReader>(
+ /// Unlike [`Self::read_page_indexes`], this returns the newly decoded
[`PageIndex`] separately;
+ /// any page index already attached to the returned metadata is unchanged.
This is useful for
+ /// populating an external page index cache. The configured page index
policies and masks
Review Comment:
💯
##########
parquet/src/file/metadata/parser.rs:
##########
@@ -285,17 +285,12 @@ pub(crate) fn parse_page_index(
)?;
}
- // Always replace the page index, even if both policies are Skip.
- // This ensures that repeated calls to read_page_indexes replace the
existing index
- // rather than preserving it.
let page_index = builder.build();
if page_index.has_column_indexes() || page_index.has_offset_indexes() {
- metadata.set_page_index(Some(Arc::new(page_index)));
+ Ok(Some(page_index))
Review Comment:
I think this is a much nicer API 👍
##########
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 this is all basically the same as me asking for an example of how to
use this / a demo of incrementally decoding the PageIndex
##########
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 guess I was thinking more of something slightly different:
Read the metadata (but no PageIndex) via
```rust
let mut decoder =
ParquetMetaDataPushDecoder::try_new_with_metadata(test_file_len(), metadata)
.unwrap()
.with_page_index_policy(PageIndexPolicy::Skip);
...
```
And then make a new one like
```rust
let mut decoder =
ParquetMetaDataPushDecoder::try_new_with_metadata(test_file_len(), metadata)
.unwrap()
.with_page_index_mask(ColumnChunkMask::all());
... // decode here and ensure the metadata has the column index
```
--
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]
