This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-11438-8c9d69472000e9237683600a5eaf9fc46163ef34 in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
commit 152f0d33c7c7127fca95b065c0ad361d97938e23 Author: Ed Seidl <[email protected]> AuthorDate: Thu Oct 15 13:07:36 2026 +0000 feat: Add separate Parquet `PageIndex` decoder to `ParquetMetaDataPushDecoder` and `ParquetMetaDataReader` (#11438) # Which issue does this PR close? - Part of #7582 # Rationale for this change The current APIs for loading the page indexes all modify the `ParquetMetaData`, installing newly parsed page indexes via `set_page_index`. Users wishing to create a page index cache can currently fetch and decode indexes themselves, but it would be nice to be able to simply parse and return the requested page indexes without installing them in the metadata for return. # What changes are included in this PR? Adds a new `try_decode_page_index` function to `ParquetMetaDataPushDecoder`, as well as helper wrapper functions to `ParquetMetaDataReader`. Note that the required `ParquetMetaData` is consumed in this process, but #11297 should make this more palatable. # Are these changes tested? Yes, includes new tests of the added APIs # Are there any significant user-facing changes? No changes to existing public APIs, adds new functions to the public API. # AI Disclaimer Codex was used in the creation of this PR, but I am responsible for the changes. --------- Co-authored-by: Andrew Lamb <[email protected]> --- parquet/src/file/metadata/parser.rs | 21 +- parquet/src/file/metadata/push_decoder.rs | 327 ++++++++++++++++++++++++++---- parquet/src/file/metadata/reader.rs | 253 ++++++++++++++++++----- 3 files changed, 499 insertions(+), 102 deletions(-) diff --git a/parquet/src/file/metadata/parser.rs b/parquet/src/file/metadata/parser.rs index 22f0564a9c..4599452187 100644 --- a/parquet/src/file/metadata/parser.rs +++ b/parquet/src/file/metadata/parser.rs @@ -20,12 +20,12 @@ //! These functions parse thrift-encoded metadata from a byte slice //! into the corresponding Rust structures -use std::{ops::Range, sync::Arc}; +use std::ops::Range; use bytes::Bytes; use crate::errors::{ParquetError, Result}; -use crate::file::metadata::page_index::{PageIndexBuilder, PageIndexProvider}; +use crate::file::metadata::page_index::{PageIndex, PageIndexBuilder, PageIndexProvider}; use crate::file::metadata::thrift::parquet_metadata_from_bytes; use crate::file::metadata::{ ColumnChunkMask, ColumnChunkMetaData, PageIndexPolicy, ParquetMetaData, ParquetMetaDataOptions, @@ -235,10 +235,10 @@ pub(crate) fn decode_metadata( parquet_metadata_from_bytes(buf, options) } -/// Parses page indexes from the provided bytes and replaces those in the metadata. +/// Parses page indexes from the provided bytes. /// /// Arguments -/// * `metadata` - The ParquetMetaData whose page index will be replaced. +/// * `metadata` - The footer metadata describing the page index locations. /// * `column_index_policy` - The policy for handling column index parsing (e.g., /// Required, Optional, Skip). /// * `offset_index_policy` - The policy for handling offset index parsing (e.g., @@ -250,14 +250,14 @@ pub(crate) fn decode_metadata( /// * `bytes` - The byte slice containing the page index data. /// * `start_offset` - The offset where `bytes` begin in the file. pub(crate) fn parse_page_index( - metadata: &mut ParquetMetaData, + metadata: &ParquetMetaData, column_index_policy: PageIndexPolicy, offset_index_policy: PageIndexPolicy, column_index_mask: &ColumnChunkMask, offset_index_mask: &ColumnChunkMask, bytes: &Bytes, start_offset: u64, -) -> Result<()> { +) -> Result<Option<PageIndex>> { let num_row_groups = metadata.num_row_groups(); let num_columns = metadata.file_metadata().schema_descr().num_columns(); let mut builder = PageIndexBuilder::default(); @@ -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)) } else { - // If no indexes were read, clear the page index - metadata.set_page_index(None); + Ok(None) } - Ok(()) } fn parse_column_index( diff --git a/parquet/src/file/metadata/push_decoder.rs b/parquet/src/file/metadata/push_decoder.rs index 90fd3ee7fa..202fa0d163 100644 --- a/parquet/src/file/metadata/push_decoder.rs +++ b/parquet/src/file/metadata/push_decoder.rs @@ -20,6 +20,7 @@ use crate::DecodeResult; use crate::encryption::decrypt::FileDecryptionProperties; use crate::errors::{ParquetError, Result}; use crate::file::FOOTER_SIZE; +use crate::file::metadata::page_index::PageIndex; use crate::file::metadata::parser::{MetadataParser, parse_page_index}; use crate::file::metadata::{ ColumnChunkMask, FooterTail, PageIndexPolicy, ParquetMetaData, ParquetMetaDataOptions, @@ -274,10 +275,16 @@ impl ParquetMetaDataPushDecoder { /// This can be used to parse and populate the page index structures /// after the metadata has already been decoded. /// - /// When decoding page indexes, any existing page indexes in `metadata` will be - /// replaced based on the configured policies. See - /// [`crate::file::metadata::ParquetMetaDataReader::read_page_indexes`] for details - /// on the replacement behavior. + /// The page indexes can then be decoded in one of two ways: + /// + /// * [`Self::try_decode`] attaches the decoded page indexes to the returned + /// metadata, replacing (or clearing) any existing page indexes + /// See [`ParquetMetaDataReader::read_page_indexes`] for details. + /// * [`Self::try_decode_page_index`] returns the decoded page indexes + /// separately and leaves any page indexes already attached to `metadata` + /// unchanged. + /// + /// [`ParquetMetaDataReader::read_page_indexes`]: [`crate::file::metadata::ParquetMetaDataReader::read_page_indexes`] pub fn try_new_with_metadata(file_len: u64, metadata: ParquetMetaData) -> Result<Self> { let mut new_self = Self::try_new(file_len)?; new_self.state = DecodeState::ReadingPageIndex(Box::new(metadata)); @@ -438,45 +445,23 @@ impl ParquetMetaDataPushDecoder { self.state = DecodeState::ReadingPageIndex(Box::new(metadata)); } - 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, - &self.column_index_mask, - &self.offset_index_mask, - ); - - let Some(page_index_range) = range else { - // There is nothing to decode because both policies are Skip, the masks - // select nothing, or the file has no index ranges. Clear any previously - // loaded indexes while preserving the compatibility behavior that - // PageIndexPolicy::Required accepts files with no page indexes at all. - metadata.set_page_index(None); - self.state = DecodeState::Finished; - return Ok(DecodeResult::Data(*metadata)); - }; - - if !self.buffers.has_range(&page_index_range) { - self.state = DecodeState::ReadingPageIndex(metadata); - return Ok(needs_range(page_index_range)); + DecodeState::ReadingPageIndex(metadata) => { + match self.try_decode_page_index_for_metadata(metadata)? { + DecodeResult::NeedsData(ranges) => { + return Ok(DecodeResult::NeedsData(ranges)); + } + DecodeResult::Data((mut metadata, page_index)) => { + // Install the newly decoded indexes, or clear any existing indexes if + // none were requested or present in the file. + if let Some(page_index) = page_index { + metadata.set_page_index(Some(Arc::new(page_index))); + } else { + metadata.set_page_index(None); + } + return Ok(DecodeResult::Data(*metadata)); + } + DecodeResult::Finished => return Ok(DecodeResult::Finished), } - - let buffer = self.get_bytes(&page_index_range)?; - let offset = page_index_range.start; - parse_page_index( - &mut metadata, - self.column_index_policy, - self.offset_index_policy, - &self.column_index_mask, - &self.offset_index_mask, - &buffer, - offset, - )?; - self.state = DecodeState::Finished; - return Ok(DecodeResult::Data(*metadata)); } DecodeState::Finished => return Ok(DecodeResult::Finished), @@ -489,6 +474,118 @@ impl ParquetMetaDataPushDecoder { } } + /// Try to decode the page indexes from the pushed data, returning the stored + /// [`ParquetMetaData`] together with the decoded [`PageIndex`]. + /// + /// This method can only be called once the footer metadata is known: either + /// after [`Self::try_new_with_metadata`], or after [`Self::try_decode`] has + /// decoded the footer metadata and requested page index data. Calling it + /// earlier returns an error and leaves the decoder unchanged. + /// + /// 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_and_finish`] (or + /// [`ParquetMetaDataReader::load_page_index_and_finish`] for asynchronous + /// I/O), which handles the decoder state machine and I/O requests. See the + /// [`read_page_index_and_finish` example][`ParquetMetaDataReader::read_page_index_and_finish`] + /// for the complete two-stage workflow. + /// + /// # Returns + /// + /// * `Ok(DecodeResult::Data((metadata, Some(page_index))))` - Successfully decoded the + /// requested page indexes based on the configured policies and masks. The newly decoded page + /// index is returned separately; any page index already attached to `metadata` is unchanged. + /// * `Ok(DecodeResult::Data((metadata, None)))` - No page indexes were requested or found + /// according to the configured policies and masks. The stored metadata is still returned. + /// * `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_and_finish`]: crate::file::metadata::ParquetMetaDataReader::read_page_index_and_finish + /// [`ParquetMetaDataReader::load_page_index_and_finish`]: crate::file::metadata::ParquetMetaDataReader::load_page_index_and_finish + pub fn try_decode_page_index( + &mut self, + ) -> Result<DecodeResult<(ParquetMetaData, Option<PageIndex>)>> { + match std::mem::replace(&mut self.state, DecodeState::Intermediate) { + DecodeState::ReadingPageIndex(metadata) => { + Ok(match self.try_decode_page_index_for_metadata(metadata)? { + DecodeResult::NeedsData(ranges) => DecodeResult::NeedsData(ranges), + DecodeResult::Data((metadata, page_index)) => { + DecodeResult::Data((*metadata, page_index)) + } + DecodeResult::Finished => DecodeResult::Finished, + }) + } + DecodeState::Finished => { + self.state = DecodeState::Finished; + Ok(DecodeResult::Finished) + } + other => { + // Restore the original state so the decoder remains usable + self.state = other; + Err(general_err!( + "ParquetMetaDataPushDecoder: try_decode_page_index requires a decoder created with try_new_with_metadata" + )) + } + } + } + + fn try_decode_page_index_for_metadata( + &mut self, + metadata: Box<ParquetMetaData>, + ) -> Result<DecodeResult<(Box<ParquetMetaData>, Option<PageIndex>)>> { + 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(DecodeResult::Data((metadata, None))); + }; + + if !self.buffers.has_range(&page_index_range) { + self.state = DecodeState::ReadingPageIndex(metadata); + return Ok(DecodeResult::NeedsData(vec![page_index_range])); + } + + let buffer = match self.get_bytes(&page_index_range) { + Ok(buffer) => buffer, + Err(error) => { + self.state = DecodeState::ReadingPageIndex(metadata); + return Err(error); + } + }; + let page_index = match 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, + ) { + Ok(page_index) => page_index, + Err(error) => { + self.state = DecodeState::ReadingPageIndex(metadata); + return Err(error); + } + }; + + self.state = DecodeState::Finished; + Ok(DecodeResult::Data((metadata, page_index))) + } + /// Returns the bytes for the given range from the internal buffer fn get_bytes(&self, range: &Range<u64>) -> Result<Bytes> { let start = range.start; @@ -503,7 +600,7 @@ impl ParquetMetaDataPushDecoder { } /// returns a DecodeResults that describes needing the given range -fn needs_range(range: Range<u64>) -> DecodeResult<ParquetMetaData> { +fn needs_range<T: std::fmt::Debug>(range: Range<u64>) -> DecodeResult<T> { DecodeResult::NeedsData(vec![range]) } @@ -568,6 +665,7 @@ pub(crate) fn range_for_page_index( mod tests { use super::*; use crate::arrow::ArrowWriter; + use crate::file::metadata::page_index::PageIndexProvider; use crate::file::properties::WriterProperties; use arrow_array::{ArrayRef, Int64Array, RecordBatch, StringViewArray}; use bytes::Bytes; @@ -619,6 +717,139 @@ mod tests { assert!(masked.end - masked.start < all.end - all.start); } + #[test] + fn test_decode_page_index_after_metadata() { + let metadata = test_metadata_without_page_index(); + assert!(metadata.page_index().is_none()); + + 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 (metadata, page_index) = expect_data(decoder.try_decode_page_index()); + assert!(metadata.page_index().is_none()); + let page_index = 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()).1.is_none()); + expect_finished(decoder.try_decode_page_index()); + + let metadata = test_metadata_without_page_index(); + let mut decoder = + ParquetMetaDataPushDecoder::try_new_with_metadata(test_file_len(), metadata) + .unwrap() + .with_page_index_mask(ColumnChunkMask::none()); + + assert!(expect_data(decoder.try_decode_page_index()).1.is_none()); + expect_finished(decoder.try_decode_page_index()); + } + + #[test] + fn test_page_index_decoder_requires_metadata() { + let mut decoder = ParquetMetaDataPushDecoder::try_new(test_file_len()).unwrap(); + let err = decoder.try_decode_page_index().unwrap_err(); + assert!( + err.to_string() + .contains("requires a decoder created with try_new_with_metadata"), + "unexpected error: {err}" + ); + } + + #[test] + fn test_page_index_decoder_restores_state_after_decode_error() { + 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::Required); + + let ranges = expect_needs_data(decoder.try_decode_page_index()); + let invalid_data = ranges + .iter() + .map(|range| Bytes::from(vec![0; (range.end - range.start) as usize])) + .collect(); + decoder.push_ranges(ranges.clone(), invalid_data).unwrap(); + decoder.try_decode_page_index().unwrap_err(); + + decoder.clear_all_ranges(); + push_ranges_to_metadata_decoder(&mut decoder, ranges); + let (metadata, page_index) = expect_data(decoder.try_decode_page_index()); + assert!(metadata.page_index().is_none()); + assert!(page_index.is_some_and(|index| index.is_complete())); + } + + /// A failed call to `try_decode_page_index` must not leave the decoder in + /// an unusable state: decoding the metadata via `try_decode` should still + /// work afterwards, at every stage of the footer decoding state machine. + #[test] + fn test_decode_after_failed_page_index_decode() { + let file_len = test_file_len(); + let mut decoder = ParquetMetaDataPushDecoder::try_new(file_len).unwrap(); + + // Before any data is pushed, the decoder is waiting for the footer + decoder.try_decode_page_index().unwrap_err(); + let footer_ranges = expect_needs_data(decoder.try_decode()); + assert_eq!(footer_ranges.len(), 1); + push_ranges_to_metadata_decoder(&mut decoder, footer_ranges); + + // The footer is available but not yet decoded + decoder.try_decode_page_index().unwrap_err(); + let metadata_ranges = expect_needs_data(decoder.try_decode()); + assert_eq!(metadata_ranges.len(), 1); + push_ranges_to_metadata_decoder(&mut decoder, metadata_ranges); + + // The metadata is available but not yet decoded + decoder.try_decode_page_index().unwrap_err(); + let index_ranges = expect_needs_data(decoder.try_decode()); + assert_eq!(index_ranges.len(), 1); + + // The metadata is decoded and the decoder is waiting for page index + // data, so try_decode_page_index is now valid and asks for the same range + let index_ranges_again = expect_needs_data(decoder.try_decode_page_index()); + assert_eq!(index_ranges, index_ranges_again); + push_ranges_to_metadata_decoder(&mut decoder, index_ranges); + + // The regular decode path still works and attaches the page index + let metadata = expect_data(decoder.try_decode()); + assert_eq!(metadata.num_row_groups(), 2); + assert!(metadata.page_index().is_some_and(|idx| idx.is_complete())); + + // Once finished, both entry points keep reporting Finished + expect_finished(decoder.try_decode_page_index()); + expect_finished(decoder.try_decode_page_index()); + expect_finished(decoder.try_decode()); + } + /// It is possible to feed some, but not all, of the footer into the metadata decoder /// before asked. This avoids multiple IO requests #[test] @@ -810,6 +1041,16 @@ mod tests { metadata_decoder.push_ranges(ranges, data).unwrap(); } + fn test_metadata_without_page_index() -> 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()]); + let metadata = expect_data(decoder.try_decode()); + assert!(metadata.page_index().is_none()); + metadata + } + /// Expect that the [`DecodeResult`] is a [`DecodeResult::Data`] and return the corresponding element fn expect_data<T: Debug>(result: Result<DecodeResult<T>>) -> T { match result.expect("Expected Ok(DecodeResult::Data(T))") { diff --git a/parquet/src/file/metadata/reader.rs b/parquet/src/file/metadata/reader.rs index 78710600bf..7f1092a700 100644 --- a/parquet/src/file/metadata/reader.rs +++ b/parquet/src/file/metadata/reader.rs @@ -23,6 +23,7 @@ use crate::errors::{ParquetError, Result}; use crate::file::FOOTER_SIZE; #[cfg(feature = "arrow")] use crate::file::metadata::dictionary::decode_dictionary_page; +use crate::file::metadata::page_index::{PageIndex, PageIndexProvider}; use crate::file::metadata::parser::decode_metadata; use crate::file::metadata::thrift::parquet_schema_from_bytes; use crate::file::metadata::{ @@ -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 + /// determine which indexes are returned. The page index is `None` when no indexes were + /// requested or found. + /// + /// This operation consumes `self`, but returns the stored [`ParquetMetaData`] as the first + /// tuple element. Use [`Self::new_with_metadata`] to construct the reader. + /// + /// # Example + /// + /// Read the footer metadata first, then load selected page indexes separately when they are + /// needed: + /// + /// ```no_run + /// # use parquet::errors::Result; + /// use parquet::file::metadata::page_index::PageIndexProvider; + /// use parquet::file::metadata::{ + /// ColumnChunkMask, PageIndexPolicy, ParquetMetaDataReader, + /// }; + /// # fn open_parquet_file() -> std::fs::File { unimplemented!() } + /// # fn example() -> Result<()> { + /// let file = open_parquet_file(); + /// let metadata = ParquetMetaDataReader::new() + /// .with_page_index_policy(PageIndexPolicy::Skip) + /// .parse_and_finish(&file)?; + /// assert!(metadata.page_index().is_none()); + /// + /// let mask = ColumnChunkMask::columns([0]); + /// let (metadata, page_index) = ParquetMetaDataReader::new_with_metadata(metadata) + /// .with_column_index_policy(PageIndexPolicy::Required) + /// .with_offset_index_policy(PageIndexPolicy::Required) + /// .with_column_index_mask(mask.clone()) + /// .with_offset_index_mask(mask) + /// .read_page_index_and_finish(&file)?; + /// + /// // The metadata is returned unchanged and the selected indexes can be cached separately. + /// assert!(metadata.page_index().is_none()); + /// let page_index = page_index.expect("page indexes are required"); + /// assert!(page_index.column_index(0, 0).is_some()); + /// assert!(page_index.offset_index(0, 0).is_some()); + /// # Ok(()) + /// # } + /// ``` + pub fn read_page_index_and_finish<R: ChunkReader>( + self, + reader: &R, + ) -> Result<(ParquetMetaData, Option<PageIndex>)> { + self.read_page_index_sized_and_finish(reader, reader.len()) + } + + /// Reads the page index structures described by the stored [`ParquetMetaData`], returning the + /// metadata and decoded page index, and using `file_size` as the size of the original Parquet + /// file. + /// + /// This variant supports a [`Bytes`] reader containing only a suffix of the file. It returns + /// [`ParquetError::NeedMoreData`] if that suffix does not contain the requested page indexes. + /// See [`Self::read_page_index_and_finish`] for ownership and configuration details. + pub fn read_page_index_sized_and_finish<R: ChunkReader>( + mut self, + reader: &R, + file_size: u64, + ) -> Result<(ParquetMetaData, Option<PageIndex>)> { + self.read_page_index_sized_impl(reader, file_size) + } + + fn read_page_index_sized_impl<R: ChunkReader>( &mut self, reader: &R, file_size: u64, - ) -> Result<()> { + ) -> Result<(ParquetMetaData, Option<PageIndex>)> { let Some(metadata) = self.metadata.take() else { return Err(general_err!( "Tried to read page indexes without ParquetMetaData metadata" @@ -600,34 +662,24 @@ impl ParquetMetaDataReader { .with_metadata_options(self.metadata_options.clone()); let mut push_decoder = self.prepare_push_decoder(push_decoder); - // Get bounds needed for page indexes (if any are present in the file). - let range = match needs_index_data(&mut push_decoder)? { - NeedsIndexData::No(metadata) => { - self.metadata = Some(metadata); - return Ok(()); - } + let range = match needs_page_index_data(&mut push_decoder)? { + NeedsIndexData::No(page_index) => return Ok(page_index), NeedsIndexData::Yes(range) => range, }; - // Check to see if needed range is within `file_range`. Checking `range.end` seems - // redundant, but it guards against `range_for_page_index()` returning garbage. let file_range = file_size.saturating_sub(reader.len())..file_size; if !(file_range.contains(&range.start) && file_range.contains(&range.end)) { - // Requested range starts beyond EOF return if range.end > file_size { Err(eof_err!( "Parquet file too small. Range {range:?} is beyond file bounds {file_size}", )) } else { - // Ask for a larger buffer Err(ParquetError::NeedMoreData( (file_size - range.start).try_into()?, )) }; } - // Perform extra sanity check to make sure `range` and the footer metadata don't - // overlap. if let Some(metadata_size) = self.metadata_size { let metadata_range = file_size.saturating_sub(metadata_size as u64)..file_size; if range.end > metadata_range.start { @@ -637,14 +689,29 @@ impl ParquetMetaDataReader { } } - // add the needed ranges to the decoder let bytes_needed = usize::try_from(range.end - range.start)?; let bytes = reader.get_bytes(range.start - file_range.start, bytes_needed)?; - push_decoder.push_range(range, bytes)?; - let metadata = parse_index_data(&mut push_decoder)?; - self.metadata = Some(metadata); + parse_page_index_data(&mut push_decoder) + } + /// 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`]. + /// + /// See [`Self::read_page_indexes()`] for a description of the replacement behavior on + /// multiple calls. + pub fn read_page_indexes_sized<R: ChunkReader>( + &mut self, + reader: &R, + file_size: u64, + ) -> Result<()> { + let (mut metadata, page_index) = self.read_page_index_sized_impl(reader, file_size)?; + metadata + .set_page_index(page_index.map(|index| Arc::new(index) as Arc<dyn PageIndexProvider>)); + self.metadata = Some(metadata); Ok(()) } @@ -731,6 +798,22 @@ impl ParquetMetaDataReader { self.load_page_index_with_remainder(fetch, None).await } + /// Asynchronously loads the page index structures described by the stored [`ParquetMetaData`], + /// returning the metadata together with the decoded page index. + /// + /// Unlike [`Self::load_page_index`], 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 page index is `None` when no indexes were + /// requested or found. This operation consumes `self`, but returns the stored metadata as the + /// first tuple element. + #[cfg(all(feature = "async", feature = "arrow"))] + pub async fn load_page_index_and_finish<F: MetadataFetch>( + mut self, + fetch: F, + ) -> Result<(ParquetMetaData, Option<PageIndex>)> { + self.load_page_index_with_remainder_impl(fetch, None).await + } + /// Reads and decodes the dictionary page of a column chunk into an Arrow array. /// /// Returns `Ok(None)` if the column chunk has no dictionary page, or if @@ -781,9 +864,24 @@ impl ParquetMetaDataReader { #[cfg(all(feature = "async", feature = "arrow"))] async fn load_page_index_with_remainder<F: MetadataFetch>( &mut self, - mut fetch: F, + fetch: F, remainder: Option<(usize, Bytes)>, ) -> Result<()> { + let (mut metadata, page_index) = self + .load_page_index_with_remainder_impl(fetch, remainder) + .await?; + metadata + .set_page_index(page_index.map(|index| Arc::new(index) as Arc<dyn PageIndexProvider>)); + self.metadata = Some(metadata); + Ok(()) + } + + #[cfg(all(feature = "async", feature = "arrow"))] + async fn load_page_index_with_remainder_impl<F: MetadataFetch>( + &mut self, + mut fetch: F, + remainder: Option<(usize, Bytes)>, + ) -> Result<(ParquetMetaData, Option<PageIndex>)> { let Some(metadata) = self.metadata.take() else { return Err(general_err!("Footer metadata is not present")); }; @@ -800,11 +898,8 @@ impl ParquetMetaDataReader { let mut push_decoder = self.prepare_push_decoder(push_decoder); // Get bounds needed for page indexes (if any are present in the file). - let range = match needs_index_data(&mut push_decoder)? { - NeedsIndexData::No(metadata) => { - self.metadata = Some(metadata); - return Ok(()); - } + let range = match needs_page_index_data(&mut push_decoder)? { + NeedsIndexData::No(page_index) => return Ok(page_index), NeedsIndexData::Yes(range) => range, }; @@ -835,9 +930,7 @@ impl ParquetMetaDataReader { )); } push_decoder.push_range(range.clone(), bytes)?; - let metadata = parse_index_data(&mut push_decoder)?; - self.metadata = Some(metadata); - Ok(()) + parse_page_index_data(&mut push_decoder) } // One-shot parse of footer. @@ -1114,17 +1207,19 @@ impl ParquetMetaDataReader { /// The bounds needed to read page indexes // this is an internal enum, so it is ok to allow differences in enum size -enum NeedsIndexData { +enum NeedsIndexData<T> { /// no additional data is needed (e.g. the indexes weren't requested) - No(ParquetMetaData), + No(T), /// Additional data is needed, with the range that are required Yes(Range<u64>), } -/// Determines a single combined range of bytes needed to read the page indexes, -/// or returns the metadata if no additional data is needed (e.g. if no page indexes are requested) -fn needs_index_data(push_decoder: &mut ParquetMetaDataPushDecoder) -> Result<NeedsIndexData> { - match push_decoder.try_decode()? { +/// Determines a single combined range of bytes needed to read the page indexes, or returns the +/// metadata and decoded page index if no additional data is needed. +fn needs_page_index_data( + push_decoder: &mut ParquetMetaDataPushDecoder, +) -> Result<NeedsIndexData<(ParquetMetaData, Option<PageIndex>)>> { + match push_decoder.try_decode_page_index()? { DecodeResult::NeedsData(ranges) => { let range = ranges .into_iter() @@ -1132,19 +1227,21 @@ fn needs_index_data(push_decoder: &mut ParquetMetaDataPushDecoder) -> Result<Nee .ok_or_else(|| general_err!("Internal error: no ranges provided"))?; Ok(NeedsIndexData::Yes(range)) } - DecodeResult::Data(metadata) => Ok(NeedsIndexData::No(metadata)), + DecodeResult::Data(page_index) => Ok(NeedsIndexData::No(page_index)), DecodeResult::Finished => Err(general_err!("Internal error: decoder was finished")), } } -/// Given a push decoder that has had the needed ranges pushed to it, -/// attempt to decode indexes and return the updated metadata. -fn parse_index_data(push_decoder: &mut ParquetMetaDataPushDecoder) -> Result<ParquetMetaData> { - match push_decoder.try_decode()? { +/// Given a push decoder that has had the needed ranges pushed to it, decode and return the metadata +/// and page indexes. +fn parse_page_index_data( + push_decoder: &mut ParquetMetaDataPushDecoder, +) -> Result<(ParquetMetaData, Option<PageIndex>)> { + match push_decoder.try_decode_page_index()? { DecodeResult::NeedsData(_) => Err(general_err!( "Internal error: decoder still needs data after reading required range" )), - DecodeResult::Data(metadata) => Ok(metadata), + DecodeResult::Data(page_index) => Ok(page_index), DecodeResult::Finished => Err(general_err!("Internal error: decoder was finished")), } } @@ -1331,6 +1428,43 @@ mod tests { ); } + #[test] + #[cfg_attr(miri, ignore)] // Takes too long + fn test_read_page_index_and_finish() { + use crate::file::metadata::page_index::PageIndexProvider; + + let file = get_test_file("alltypes_tiny_pages.parquet"); + let metadata = ParquetMetaDataReader::new() + .with_page_index_policy(PageIndexPolicy::Required) + .parse_and_finish(&file) + .unwrap(); + let original_page_index = Arc::clone(metadata.page_index().unwrap()); + + let reader = ParquetMetaDataReader::new_with_metadata(metadata); + let (metadata, page_index) = reader.read_page_index_and_finish(&file).unwrap(); + assert!(page_index.is_none()); + assert!(Arc::ptr_eq( + metadata.page_index().unwrap(), + &original_page_index + )); + + let reader = ParquetMetaDataReader::new_with_metadata(metadata) + .with_page_index_policy(PageIndexPolicy::Required) + .with_column_index_mask(ColumnChunkMask::columns([0])) + .with_offset_index_mask(ColumnChunkMask::columns([0])); + let (metadata, page_index) = reader.read_page_index_and_finish(&file).unwrap(); + let page_index = page_index.unwrap(); + assert!(Arc::ptr_eq( + metadata.page_index().unwrap(), + &original_page_index + )); + + assert!(page_index.column_index(0, 0).is_some()); + assert!(page_index.offset_index(0, 0).is_some()); + assert!(page_index.column_index(0, 1).is_none()); + assert!(page_index.offset_index(0, 1).is_none()); + } + #[test] fn test_chunk_mask() { let mask = ColumnChunkMask::row_groups_and_columns([0], [1]); @@ -1437,6 +1571,33 @@ mod async_tests { Ok(buf.into()) } + #[tokio::test] + async fn test_load_page_index_and_finish() { + use crate::file::metadata::page_index::PageIndexProvider; + + let mut file = get_test_file("alltypes_tiny_pages.parquet"); + let metadata = ParquetMetaDataReader::new() + .parse_and_finish(&file) + .unwrap(); + let reader = ParquetMetaDataReader::new_with_metadata(metadata) + .with_page_index_policy(PageIndexPolicy::Required) + .with_column_index_mask(ColumnChunkMask::columns([0])) + .with_offset_index_mask(ColumnChunkMask::columns([0])); + let fetch = MetadataFetchFn(|range| futures::future::ready(read_range(&mut file, range))); + + let page_index = reader + .load_page_index_and_finish(fetch) + .await + .unwrap() + .1 + .unwrap(); + + assert!(page_index.column_index(0, 0).is_some()); + assert!(page_index.offset_index(0, 0).is_some()); + assert!(page_index.column_index(0, 1).is_none()); + assert!(page_index.offset_index(0, 1).is_none()); + } + #[tokio::test] async fn test_simple() { let mut file = get_test_file("nulls.snappy.parquet");
