Skip to content

perf(parquet): avoid repeated page-index range discovery #11276

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
sunchao wants to merge 1 commit into
base: main
Choose a base branch
from
Open

perf(parquet): avoid repeated page-index range discovery #11276

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
100 changes: 93 additions & 7 deletions parquet/src/file/metadata/push_decoder.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -226,6 +226,8 @@ pub struct ParquetMetaDataPushDecoder {
column_index_policy: PageIndexPolicy,
/// policy for loading OffsetIndex (part of the PageIndex)
offset_index_policy: PageIndexPolicy,
/// Index span awaiting bytes, invalidated by policy changes.
pending_page_index_range: Option<Range<u64>>,
/// Underlying buffers
buffers: crate::util::push_buffers::PushBuffers,
/// Encryption API
Expand All @@ -250,6 +252,7 @@ impl ParquetMetaDataPushDecoder {
state: DecodeState::ReadingFooter,
column_index_policy: PageIndexPolicy::Optional,
offset_index_policy: PageIndexPolicy::Optional,
pending_page_index_range: None,
buffers: crate::util::push_buffers::PushBuffers::new(file_len),
metadata_parser: MetadataParser::new(),
})
Expand Down Expand Up @@ -285,18 +288,21 @@ impl ParquetMetaDataPushDecoder {
pub fn with_page_index_policy(mut self, page_index_policy: PageIndexPolicy) -> Self {
self.column_index_policy = page_index_policy;
self.offset_index_policy = page_index_policy;
self.pending_page_index_range = None;
self
}

/// Set the policy for reading the ColumnIndex (part of the PageIndex)
pub fn with_column_index_policy(mut self, column_index_policy: PageIndexPolicy) -> Self {
self.column_index_policy = column_index_policy;
self.pending_page_index_range = None;
self
}

/// Set the policy for reading the OffsetIndex (part of the PageIndex)
pub fn with_offset_index_policy(mut self, offset_index_policy: PageIndexPolicy) -> Self {
self.offset_index_policy = offset_index_policy;
self.pending_page_index_range = None;
self
}

Expand Down Expand Up @@ -406,20 +412,22 @@ impl ParquetMetaDataPushDecoder {
}

DecodeState::ReadingPageIndex(mut metadata) => {
// First determine if any page indexes are needed based on
// the specified policies
let range = range_for_page_index(
&metadata,
self.column_index_policy,
self.offset_index_policy,
);
// Retain the range only while waiting, so decoding or errors clear it.
let range = self.pending_page_index_range.take().or_else(|| {
range_for_page_index(
&metadata,
self.column_index_policy,
self.offset_index_policy,
)
});

let Some(page_index_range) = range else {
self.state = DecodeState::Finished;
return Ok(DecodeResult::Data(*metadata));
};

if !self.buffers.has_range(&page_index_range) {
self.pending_page_index_range = Some(page_index_range.clone());
self.state = DecodeState::ReadingPageIndex(metadata);
return Ok(needs_range(page_index_range));
}
Expand Down Expand Up @@ -490,6 +498,10 @@ pub fn range_for_page_index(
column_index_policy: PageIndexPolicy,
offset_index_policy: PageIndexPolicy,
) -> Option<Range<u64>> {
if column_index_policy == PageIndexPolicy::Skip && offset_index_policy == PageIndexPolicy::Skip
{
return None;
}
let mut range = None;
for c in metadata.row_groups().iter().flat_map(|r| r.columns()) {
if column_index_policy != PageIndexPolicy::Skip {
Expand Down Expand Up @@ -612,10 +624,22 @@ mod tests {

// expect the second request to read the offset indexes
let ranges = expect_needs_data(metadata_decoder.try_decode());
for _ in 0..3 {
assert_eq!(
metadata_decoder.pending_page_index_range,
Some(ranges[0].clone())
);
assert_eq!(expect_needs_data(metadata_decoder.try_decode()), ranges);
}
push_ranges_to_metadata_decoder(&mut metadata_decoder, ranges.clone());
metadata_decoder.clear_all_ranges();
assert_eq!(metadata_decoder.buffers.buffered_bytes(), 0);
assert_eq!(expect_needs_data(metadata_decoder.try_decode()), ranges);
push_ranges_to_metadata_decoder(&mut metadata_decoder, ranges);

// expect the third request to read the actual data
let metadata = expect_data(metadata_decoder.try_decode());
assert!(metadata_decoder.pending_page_index_range.is_none());
expect_finished(metadata_decoder.try_decode());

assert_eq!(metadata.num_row_groups(), 2);
Expand Down Expand Up @@ -651,6 +675,68 @@ mod tests {
assert!(metadata.page_index().is_none()); // of course, we did not read the page index
}

/// Each setter invalidates pending work, including narrower, wider and Skip policies.
#[test]
fn test_metadata_decoder_pending_range_policy_changes() {
let metadata = test_metadata_without_indexes();
let setters: [fn(ParquetMetaDataPushDecoder, PageIndexPolicy) -> ParquetMetaDataPushDecoder;
3] = [
ParquetMetaDataPushDecoder::with_page_index_policy,
ParquetMetaDataPushDecoder::with_column_index_policy,
ParquetMetaDataPushDecoder::with_offset_index_policy,
];
for (column_policy, offset_policy) in [
(PageIndexPolicy::Optional, PageIndexPolicy::Optional),
(PageIndexPolicy::Skip, PageIndexPolicy::Optional),
(PageIndexPolicy::Optional, PageIndexPolicy::Skip),
] {
for setter in setters {
for policy in [
PageIndexPolicy::Skip,
PageIndexPolicy::Optional,
PageIndexPolicy::Required,
] {
let make_decoder = || {
ParquetMetaDataPushDecoder::try_new_with_metadata(
test_file_len(),
metadata.clone(),
)
.unwrap()
.with_column_index_policy(column_policy)
.with_offset_index_policy(offset_policy)
};
let mut decoder = make_decoder();
expect_needs_data(decoder.try_decode());
assert!(decoder.pending_page_index_range.is_some());
decoder = setter(decoder, policy);
assert!(decoder.pending_page_index_range.is_none());
let mut reference = setter(make_decoder(), policy);
push_ranges_to_metadata_decoder(&mut reference, vec![test_file_range()]);
let expected = expect_data(reference.try_decode());

@viirya viirya Sep 29, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Non-blocking: could we also cover an error after a pending request? This fixture has complete indexes, so the Required cases exercise successful decoding but not enforcement when an offset index is missing.

A useful case would remove one chunk’s offset-index metadata, start with Optional, obtain NeedsData, then switch to Required and provide the requested bytes. Assert that decoding reports the missing offset index and leaves pending_page_index_range empty.

A separate corrupt-index-bytes case would cover cache cleanup on a parser error. Both cases passed in scratch tests; adding them here would protect the error-path behavior described by the implementation comment.

if let Some(range) = range_for_page_index(
&metadata,
decoder.column_index_policy,
decoder.offset_index_policy,
) {
assert_eq!(expect_needs_data(decoder.try_decode()), vec![range.clone()]);
push_ranges_to_metadata_decoder(&mut decoder, vec![range]);
}
assert_eq!(expect_data(decoder.try_decode()), expected);
assert!(decoder.pending_page_index_range.is_none());
}
}
}
}

/// Decode only the generated fixture's footer, leaving indexes for an upgrade.
fn test_metadata_without_indexes() -> ParquetMetaData {
let mut decoder = ParquetMetaDataPushDecoder::try_new(test_file_len())
.unwrap()
.with_page_index_policy(PageIndexPolicy::Skip);
push_ranges_to_metadata_decoder(&mut decoder, vec![test_file_range()]);
expect_data(decoder.try_decode())
}

static TEST_BATCH: LazyLock<RecordBatch> = LazyLock::new(|| {
// Input batch has 400 rows, with 3 columns: "a", "b", "c"
// Note c is a different types (so the data page sizes will be different)
Expand Down
Loading