From 12bccecce68ffef6eabb8c8039680b77e09df8d0 Mon Sep 17 00:00:00 2001 From: Andrew Lamb Date: Fri, 12 Sep 2025 13:50:18 -0400 Subject: [PATCH] Move state machine into ParquetMetadataDecoder --- parquet/src/file/metadata/parser.rs | 82 ++++++- parquet/src/file/metadata/push_decoder.rs | 277 ++++++++++++++++------ parquet/src/file/metadata/reader.rs | 260 ++++++++++++-------- 3 files changed, 457 insertions(+), 162 deletions(-) diff --git a/parquet/src/file/metadata/parser.rs b/parquet/src/file/metadata/parser.rs index a68f14d4d7aa..c5dd031c0031 100644 --- a/parquet/src/file/metadata/parser.rs +++ b/parquet/src/file/metadata/parser.rs @@ -43,6 +43,86 @@ use crate::encryption::{ #[cfg(feature = "encryption")] use crate::format::EncryptionAlgorithm; +/// Helper struct for metadata parsing +/// +/// This structure parses thrift-encoded bytes into the correct Rust structs, +/// such as [`ParquetMetaData`], handling decryption if necessary. +// +// Note this structure is used to minimize the number of +// places need to add `#[cfg(feature = "encryption")]` checks. +pub(crate) use inner::MetadataParser; + +#[cfg(feature = "encryption")] +mod inner { + use super::*; + use crate::encryption::decrypt::FileDecryptionProperties; + use crate::errors::Result; + + /// API for decoding metadata that may be encrypted + #[derive(Debug, Default)] + pub(crate) struct MetadataParser { + // the credentials and keys needed to decrypt metadata + file_decryption_properties: Option>, + } + + impl MetadataParser { + pub fn new() -> Self { + MetadataParser::default() + } + + pub(crate) fn with_file_decryption_properties( + mut self, + file_decryption_properties: Option>, + ) -> Self { + self.file_decryption_properties = file_decryption_properties; + self + } + + pub(crate) fn decode_metadata( + &self, + buf: &[u8], + encrypted_footer: bool, + ) -> Result { + decode_metadata_with_encryption( + buf, + encrypted_footer, + self.file_decryption_properties.as_deref(), + ) + } + } +} + +#[cfg(not(feature = "encryption"))] +mod inner { + use super::*; + use crate::errors::Result; + /// parallel implementation when encryption feature is not enabled + /// + /// This has the same API as the encryption-enabled version + #[derive(Debug, Default)] + pub(crate) struct MetadataParser; + + impl MetadataParser { + pub(crate) fn new() -> Self { + MetadataParser + } + + pub(crate) fn decode_metadata( + &self, + buf: &[u8], + encrypted_footer: bool, + ) -> Result { + if encrypted_footer { + Err(general_err!( + "Parquet file has an encrypted footer but the encryption feature is disabled" + )) + } else { + decode_metadata(buf) + } + } + } +} + /// Decodes [`ParquetMetaData`] from the provided bytes. /// /// Typically this is used to decode the metadata from the end of a parquet @@ -288,7 +368,7 @@ fn parse_single_offset_index( /// [Parquet Spec]: https://github.com/apache/parquet-format#metadata /// [Parquet Encryption Spec]: https://parquet.apache.org/docs/file-format/data-pages/encryption/ #[cfg(feature = "encryption")] -pub(crate) fn decode_metadata_with_encryption( +fn decode_metadata_with_encryption( buf: &[u8], encrypted_footer: bool, file_decryption_properties: Option<&FileDecryptionProperties>, diff --git a/parquet/src/file/metadata/push_decoder.rs b/parquet/src/file/metadata/push_decoder.rs index 811caf4fd46c..7203d08f97a6 100644 --- a/parquet/src/file/metadata/push_decoder.rs +++ b/parquet/src/file/metadata/push_decoder.rs @@ -15,23 +15,37 @@ // specific language governing permissions and limitations // under the License. +#[cfg(feature = "encryption")] +use crate::encryption::decrypt::FileDecryptionProperties; use crate::errors::ParquetError; -use crate::file::metadata::{PageIndexPolicy, ParquetMetaData, ParquetMetaDataReader}; +use crate::file::metadata::parser::{parse_column_index, parse_offset_index, MetadataParser}; +use crate::file::metadata::{FooterTail, PageIndexPolicy, ParquetMetaData}; +use crate::file::page_index::index_reader::acc_range; +use crate::file::reader::ChunkReader; +use crate::file::FOOTER_SIZE; use crate::DecodeResult; +use bytes::Bytes; +use std::ops::Range; /// A push decoder for [`ParquetMetaData`]. /// -/// This structure implements a push API based version of the [`ParquetMetaDataReader`], which -/// decouples the IO from the metadata decoding logic. +/// This structure implements a push API for decoding Parquet metadata, which +/// decouples IO from the metadata decoding logic (sometimes referred to as +/// [Sans-IO]). /// -/// You can use this decoder to customize your IO operations, as shown in the -/// examples below for minimizing bytes read, prefetching data, or -/// using async IO. +/// See [`ParquetMetaDataReader`] for a pull-based API that incorporates IO and +/// is simpler to use for basic use cases. This decoder is best for customizing +/// your IO operations to minimize bytes read, prefetch data, or use async IO. +/// +/// [Sans-IO]: https://sans-io.readthedocs.io +/// [`ParquetMetaDataReader`]: crate::file::metadata::ParquetMetaDataReader /// /// # Example /// /// The most basic usage is to feed the decoder with the necessary byte ranges -/// as requested as shown below. +/// as requested as shown below. This minimizes the number of bytes read, but +/// requires the most IO operations - one to read the footer and then one +/// to read the metadata, and possibly more if page indexes are requested. /// /// ```rust /// # use std::ops::Range; @@ -192,9 +206,16 @@ use crate::DecodeResult; /// [`AsyncRead`]: tokio::io::AsyncRead #[derive(Debug)] pub struct ParquetMetaDataPushDecoder { - done: bool, - metadata_reader: ParquetMetaDataReader, + /// Decoding state + state: DecodeState, + /// policy for loading ColumnIndex (part of the PageIndex) + column_index_policy: PageIndexPolicy, + /// policy for loading OffsetIndex (part of the PageIndex) + offset_index_policy: PageIndexPolicy, + /// Underlying buffers buffers: crate::util::push_buffers::PushBuffers, + /// Encryption API + metadata_parser: MetadataParser, } impl ParquetMetaDataPushDecoder { @@ -211,16 +232,38 @@ impl ParquetMetaDataPushDecoder { ))); }; - let metadata_reader = - ParquetMetaDataReader::new().with_page_index_policy(PageIndexPolicy::Optional); - Ok(Self { - done: false, - metadata_reader, + state: DecodeState::ReadingFooter, + column_index_policy: PageIndexPolicy::Optional, + offset_index_policy: PageIndexPolicy::Optional, buffers: crate::util::push_buffers::PushBuffers::new(file_len), + metadata_parser: MetadataParser::new(), }) } + /// Begin decoding from the given footer tail. + pub(crate) fn try_new_with_footer_tail( + file_len: u64, + footer_tail: FooterTail, + ) -> Result { + let mut new_self = Self::try_new(file_len)?; + new_self.state = DecodeState::ReadingMetadata(footer_tail); + Ok(new_self) + } + + /// Create a decoder with the given `ParquetMetaData` already known. + /// + /// This can be used to parse and populate the page index structures + /// after the metadata has already been decoded. + pub fn try_new_with_metadata( + file_len: u64, + metadata: ParquetMetaData, + ) -> Result { + let mut new_self = Self::try_new(file_len)?; + new_self.state = DecodeState::ReadingPageIndex(Box::new(metadata)); + Ok(new_self) + } + /// Enable or disable reading the page index structures described in /// "[Parquet page index] Layout to Support Page Skipping". /// @@ -232,9 +275,32 @@ impl ParquetMetaDataPushDecoder { /// /// [Parquet page index]: https://github.com/apache/parquet-format/blob/master/PageIndex.md pub fn with_page_index_policy(mut self, page_index_policy: PageIndexPolicy) -> Self { - self.metadata_reader = self - .metadata_reader - .with_page_index_policy(page_index_policy); + self.column_index_policy = page_index_policy; + self.offset_index_policy = page_index_policy; + 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 + } + + /// 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 + } + + #[cfg(feature = "encryption")] + /// Provide decryption properties for decoding encrypted Parquet files + pub(crate) fn with_file_decryption_properties( + mut self, + file_decryption_properties: Option>, + ) -> Self { + self.metadata_parser = self + .metadata_parser + .with_file_decryption_properties(file_decryption_properties); self } @@ -259,14 +325,13 @@ impl ParquetMetaDataPushDecoder { /// example on [`Self`] pub fn push_ranges( &mut self, - ranges: Vec>, - buffers: Vec, - ) -> std::result::Result<(), String> { - if self.done { - return Err( + ranges: Vec>, + buffers: Vec, + ) -> Result<(), ParquetError> { + if matches!(&self.state, DecodeState::Finished) { + return Err(general_err!( "ParquetMetaDataPushDecoder: cannot push data after decoding is finished" - .to_string(), - ); + )); } self.buffers.push_ranges(ranges, buffers); Ok(()) @@ -274,55 +339,135 @@ impl ParquetMetaDataPushDecoder { /// Try to decode the metadata from the pushed data, returning the /// decoded metadata or an error if not enough data is available. - pub fn try_decode( - &mut self, - ) -> std::result::Result, ParquetError> { - if self.done { - return Ok(DecodeResult::Finished); - } - - // need to have the last 8 bytes of the file to decode the metadata + pub fn try_decode(&mut self) -> Result, ParquetError> { let file_len = self.buffers.file_len(); - if !self.buffers.has_range(&(file_len - 8..file_len)) { - #[expect(clippy::single_range_in_vec_init)] - return Ok(DecodeResult::NeedsData(vec![file_len - 8..file_len])); + let footer_len = FOOTER_SIZE as u64; + loop { + match std::mem::replace(&mut self.state, DecodeState::Intermediate) { + DecodeState::ReadingFooter => { + // need to have the last 8 bytes of the file to decode the metadata + let footer_start = file_len.saturating_sub(footer_len); + let footer_range = footer_start..file_len; + + if !self.buffers.has_range(&footer_range) { + self.state = DecodeState::ReadingFooter; + return Ok(needs_range(footer_range)); + } + let footer_bytes = self.get_bytes(&footer_range)?; + let footer_tail = FooterTail::try_from(footer_bytes.as_ref())?; + + self.state = DecodeState::ReadingMetadata(footer_tail); + continue; + } + + DecodeState::ReadingMetadata(footer_tail) => { + let metadata_len: u64 = footer_tail.metadata_length() as u64; + let metadata_start = file_len - footer_len - metadata_len; + let metadata_end = metadata_start + metadata_len; + let metadata_range = metadata_start..metadata_end; + + if !self.buffers.has_range(&metadata_range) { + self.state = DecodeState::ReadingMetadata(footer_tail); + return Ok(needs_range(metadata_range)); + } + + let metadata = self.metadata_parser.decode_metadata( + &self.get_bytes(&metadata_range)?, + footer_tail.is_encrypted_footer(), + )?; + self.state = DecodeState::ReadingPageIndex(Box::new(metadata)); + continue; + } + + DecodeState::ReadingPageIndex(mut metadata) => { + let range = range_for_page_index( + &metadata, + self.column_index_policy, + self.offset_index_policy, + ); + + let Some(page_index_range) = range else { + // no ranges means no page indexes are needed + 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)); + } + + let buffer = self.get_bytes(&page_index_range)?; + let offset = page_index_range.start; + parse_column_index(&mut metadata, self.column_index_policy, &buffer, offset)?; + parse_offset_index(&mut metadata, self.offset_index_policy, &buffer, offset)?; + self.state = DecodeState::Finished; + return Ok(DecodeResult::Data(*metadata)); + } + + DecodeState::Finished => return Ok(DecodeResult::Finished), + DecodeState::Intermediate => { + return Err(general_err!( + "ParquetMetaDataPushDecoder: internal error, invalid state" + )); + } + } } + } - // Try to parse the metadata from the buffers we have. - // - // If we don't have enough data, returns a `ParquetError::NeedMoreData` - // with the number of bytes needed to complete the metadata parsing. - // - // If we have enough data, returns `Ok(())` and we can complete - // the metadata parsing. - let maybe_metadata = self - .metadata_reader - .try_parse_sized(&self.buffers, self.buffers.file_len()); - - match maybe_metadata { - Ok(()) => { - // Metadata successfully parsed, proceed to decode the row groups - let metadata = self.metadata_reader.finish()?; - self.done = true; - Ok(DecodeResult::Data(metadata)) - } + /// Returns the bytes for the given range from the internal buffer + fn get_bytes(&self, range: &Range) -> Result { + let start = range.start; + let raw_len = range.end - range.start; + let len: usize = raw_len.try_into().map_err(|_| { + ParquetError::General(format!( + "ParquetMetaDataPushDecoder: Range length too large to fit in usize: {raw_len}", + )) + })?; + self.buffers.get_bytes(start, len) + } +} - Err(ParquetError::NeedMoreData(needed)) => { - let needed = needed as u64; - let Some(start_offset) = file_len.checked_sub(needed) else { - return Err(ParquetError::General(format!( - "Parquet metadata reader needs at least {needed} bytes, but file length is only {file_len}" - ))); - }; - let needed_range = start_offset..start_offset + needed; - // needs `needed_range` bytes at the end of the file - Ok(DecodeResult::NeedsData(vec![needed_range])) - } - Err(ParquetError::NeedMoreDataRange(range)) => Ok(DecodeResult::NeedsData(vec![range])), +/// returns a DecodeResults that describes needing the given range +fn needs_range(range: Range) -> DecodeResult { + DecodeResult::NeedsData(vec![range]) +} + +/// Decoding state machine +#[derive(Debug)] +enum DecodeState { + /// Reading the last 8 bytes of the file + ReadingFooter, + /// Reading the metadata thrift structure + ReadingMetadata(FooterTail), + // Actively reading the page index + ReadingPageIndex(Box), + // Decoding is complete + Finished, + /// State left during the `try_decode` method so something valid is present. + /// This state should never be observed. + Intermediate, +} - Err(e) => Err(e), // some other error, pass back +/// Returns the byte range needed to read the offset/page indexes, based on the +/// specified policies +/// +/// Returns None if no page indexes are needed +pub fn range_for_page_index( + metadata: &ParquetMetaData, + column_index_policy: PageIndexPolicy, + offset_index_policy: PageIndexPolicy, +) -> Option> { + let mut range = None; + for c in metadata.row_groups().iter().flat_map(|r| r.columns()) { + if column_index_policy != PageIndexPolicy::Skip { + range = acc_range(range, c.column_index_range()); + } + if offset_index_policy != PageIndexPolicy::Skip { + range = acc_range(range, c.offset_index_range()); } } + range } // These tests use the arrow writer to create a parquet file in memory diff --git a/parquet/src/file/metadata/reader.rs b/parquet/src/file/metadata/reader.rs index 4b8c57175d4e..6b2302787a9a 100644 --- a/parquet/src/file/metadata/reader.rs +++ b/parquet/src/file/metadata/reader.rs @@ -15,21 +15,19 @@ // specific language governing permissions and limitations // under the License. -use std::{io::Read, ops::Range}; - #[cfg(feature = "encryption")] use crate::encryption::decrypt::FileDecryptionProperties; use crate::errors::{ParquetError, Result}; -use crate::file::metadata::{FooterTail, ParquetMetaData}; -use crate::file::page_index::index_reader::acc_range; +use crate::file::metadata::{FooterTail, ParquetMetaData, ParquetMetaDataPushDecoder}; use crate::file::reader::ChunkReader; use crate::file::FOOTER_SIZE; +use bytes::Bytes; +use std::{io::Read, ops::Range}; #[cfg(all(feature = "async", feature = "arrow"))] use crate::arrow::async_reader::{MetadataFetch, MetadataSuffixFetch}; -#[cfg(feature = "encryption")] -use crate::file::metadata::parser::decode_metadata_with_encryption; -use crate::file::metadata::parser::{decode_metadata, parse_column_index, parse_offset_index}; +use crate::file::metadata::parser::decode_metadata; +use crate::DecodeResult; /// Reads [`ParquetMetaData`] from a byte stream, with either synchronous or /// asynchronous I/O. @@ -40,8 +38,6 @@ use crate::file::metadata::parser::{decode_metadata, parse_column_index, parse_o /// /// See the [`ParquetMetaDataPushDecoder`] for an API that does not require I/O. /// -/// [`ParquetMetaDataPushDecoder`]: crate::file::metadata::push_decoder::ParquetMetaDataPushDecoder -/// /// # Format Notes /// /// Parquet metadata is not necessarily contiguous in a Parquet file: a portion is stored @@ -76,7 +72,7 @@ pub struct ParquetMetaDataReader { // `self.parse_metadata` is called. metadata_size: Option, #[cfg(feature = "encryption")] - file_decryption_properties: Option, + file_decryption_properties: Option>, } /// Describes the policy for reading page indexes @@ -186,7 +182,7 @@ impl ParquetMetaDataReader { mut self, properties: Option<&FileDecryptionProperties>, ) -> Self { - self.file_decryption_properties = properties.cloned(); + self.file_decryption_properties = properties.cloned().map(std::sync::Arc::new); self } @@ -220,8 +216,6 @@ impl ParquetMetaDataReader { /// .with_page_indexes(true) /// .parse_and_finish(&file).unwrap(); /// ``` - /// - /// [`Bytes`]: bytes::Bytes pub fn parse_and_finish(mut self, reader: &R) -> Result { self.try_parse(reader)?; self.finish() @@ -232,8 +226,6 @@ impl ParquetMetaDataReader { /// If `reader` is [`Bytes`] based, then the buffer must contain sufficient bytes to complete /// the request, and must include the Parquet footer. If page indexes are desired, the buffer /// must contain the entire file, or [`Self::try_parse_sized()`] should be used. - /// - /// [`Bytes`]: bytes::Bytes pub fn try_parse(&mut self, reader: &R) -> Result<()> { self.try_parse_sized(reader, reader.len()) } @@ -310,8 +302,6 @@ impl ParquetMetaDataReader { /// } /// let metadata = reader.finish().unwrap(); /// ``` - /// - /// [`Bytes`]: bytes::Bytes pub fn try_parse_sized(&mut self, reader: &R, file_size: u64) -> Result<()> { self.metadata = match self.parse_metadata(reader) { Ok(metadata) => Some(metadata), @@ -352,24 +342,31 @@ impl ParquetMetaDataReader { /// 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`]. - /// - /// [`Bytes`]: bytes::Bytes pub fn read_page_indexes_sized( &mut self, reader: &R, file_size: u64, ) -> Result<()> { - // Get bounds needed for page indexes (if any are present in the file). - let Some(range) = self.range_for_page_index() else { - return Ok(()); - }; - - let Some(metadata) = self.metadata.as_mut() else { + let Some(metadata) = self.metadata.take() else { return Err(general_err!( "Tried to read page indexes without ParquetMetaData metadata" )); }; + let push_decoder = ParquetMetaDataPushDecoder::try_new_with_metadata(file_size, metadata)? + .with_offset_index_policy(self.offset_index) + .with_column_index_policy(self.column_index); + 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(()); + } + 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; @@ -398,12 +395,13 @@ 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)?; - let offset = range.start; - parse_column_index(metadata, self.column_index, &bytes, offset)?; - parse_offset_index(metadata, self.offset_index, &bytes, offset)?; + push_decoder.push_ranges(vec![range], vec![bytes])?; + let metadata = parse_index_data(&mut push_decoder)?; + self.metadata = Some(metadata); Ok(()) } @@ -492,15 +490,27 @@ impl ParquetMetaDataReader { async fn load_page_index_with_remainder( &mut self, mut fetch: F, - remainder: Option<(usize, bytes::Bytes)>, + remainder: Option<(usize, Bytes)>, ) -> Result<()> { - // Get bounds needed for page indexes (if any are present in the file). - let Some(range) = self.range_for_page_index() else { - return Ok(()); + let Some(metadata) = self.metadata.take() else { + return Err(general_err!("Footer metadata is not present")); }; - let Some(metadata) = self.metadata.as_mut() else { - return Err(general_err!("Footer metadata is not present")); + // in this case we don't actually know what the file size is, so just use u64::MAX + // this is ok since the offsets in the metadata are always valid + let file_size = u64::MAX; + let push_decoder = ParquetMetaDataPushDecoder::try_new_with_metadata(file_size, metadata)? + .with_offset_index_policy(self.offset_index) + .with_column_index_policy(self.column_index); + 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(()); + } + NeedsIndexData::Yes(range) => range, }; let bytes = match &remainder { @@ -517,31 +527,12 @@ impl ParquetMetaDataReader { // Sanity check assert_eq!(bytes.len() as u64, range.end - range.start); - - parse_column_index(metadata, self.column_index, &bytes, range.start)?; - parse_offset_index(metadata, self.offset_index, &bytes, range.start)?; - + push_decoder.push_ranges(vec![range.clone()], vec![bytes])?; + let metadata = parse_index_data(&mut push_decoder)?; + self.metadata = Some(metadata); Ok(()) } - fn range_for_page_index(&self) -> Option> { - // sanity check - self.metadata.as_ref()?; - - // Get bounds needed for page indexes (if any are present in the file). - let mut range = None; - let metadata = self.metadata.as_ref().unwrap(); - for c in metadata.row_groups().iter().flat_map(|r| r.columns()) { - if self.column_index != PageIndexPolicy::Skip { - range = acc_range(range, c.column_index_range()); - } - if self.offset_index != PageIndexPolicy::Skip { - range = acc_range(range, c.offset_index_range()); - } - } - range - } - // One-shot parse of footer. // Side effect: this will set `self.metadata_size` fn parse_metadata(&mut self, chunk_reader: &R) -> Result { @@ -566,10 +557,8 @@ impl ParquetMetaDataReader { } let start = file_size - footer_metadata_len as u64; - self.decode_footer_metadata( - chunk_reader.get_bytes(start, metadata_len)?.as_ref(), - &footer, - ) + let bytes = chunk_reader.get_bytes(start, metadata_len)?; + self.decode_footer_metadata(bytes, file_size, footer) } /// Return the number of bytes to read in the initial pass. If `prefetch_size` has @@ -587,10 +576,10 @@ impl ParquetMetaDataReader { #[cfg(all(feature = "async", feature = "arrow"))] async fn load_metadata( - &self, + &mut self, fetch: &mut F, file_size: u64, - ) -> Result<(ParquetMetaData, Option<(usize, bytes::Bytes)>)> { + ) -> Result<(ParquetMetaData, Option<(usize, Bytes)>)> { let prefetch = self.get_prefetch_size() as u64; if file_size < FOOTER_SIZE as u64 { @@ -635,14 +624,14 @@ impl ParquetMetaDataReader { let meta = fetch .fetch(metadata_start..(file_size - FOOTER_SIZE as u64)) .await?; - Ok((self.decode_footer_metadata(&meta, &footer)?, None)) + Ok((self.decode_footer_metadata(meta, file_size, footer)?, None)) } else { let metadata_start = (file_size - (length + FOOTER_SIZE) as u64 - footer_start) .try_into() .expect("metadata length should never be larger than u32"); - let slice = &suffix[metadata_start..suffix_len - FOOTER_SIZE]; + let slice = suffix.slice(metadata_start..suffix_len - FOOTER_SIZE); Ok(( - self.decode_footer_metadata(slice, &footer)?, + self.decode_footer_metadata(slice, file_size, footer)?, Some((footer_start as usize, suffix.slice(..metadata_start))), )) } @@ -650,9 +639,9 @@ impl ParquetMetaDataReader { #[cfg(all(feature = "async", feature = "arrow"))] async fn load_metadata_via_suffix( - &self, + &mut self, fetch: &mut F, - ) -> Result<(ParquetMetaData, Option<(usize, bytes::Bytes)>)> { + ) -> Result<(ParquetMetaData, Option<(usize, Bytes)>)> { let prefetch = self.get_prefetch_size(); let suffix = fetch.fetch_suffix(prefetch as _).await?; @@ -671,6 +660,9 @@ impl ParquetMetaDataReader { let footer = Self::decode_footer_tail(&footer)?; let length = footer.metadata_length(); + // fake file size as we are only parsing the footer metadata here + // (cant be parsing page indexes without the full file size) + let file_size = (length + FOOTER_SIZE) as u64; // Did not fetch the entire file metadata in the initial read, need to make a second request let metadata_offset = length + FOOTER_SIZE; @@ -685,22 +677,21 @@ impl ParquetMetaDataReader { )); } - Ok(( - // need to slice off the footer or decryption fails - self.decode_footer_metadata(&meta.slice(0..length), &footer)?, - None, - )) + // need to slice off the footer or decryption fails + let meta = meta.slice(0..length); + Ok((self.decode_footer_metadata(meta, file_size, footer)?, None)) } else { let metadata_start = suffix_len - metadata_offset; - let slice = &suffix[metadata_start..suffix_len - FOOTER_SIZE]; + let slice = suffix.slice(metadata_start..suffix_len - FOOTER_SIZE); Ok(( - self.decode_footer_metadata(slice, &footer)?, + self.decode_footer_metadata(slice, file_size, footer)?, Some((0, suffix.slice(..metadata_start))), )) } } /// Decodes a [`FooterTail`] from the provided 8-byte slice. + // todo deprecate pub fn decode_footer_tail(slice: &[u8; FOOTER_SIZE]) -> Result { FooterTail::try_new(slice) } @@ -725,27 +716,69 @@ impl ParquetMetaDataReader { /// /// [Parquet Spec]: https://github.com/apache/parquet-format#metadata pub(crate) fn decode_footer_metadata( - &self, - buf: &[u8], - footer_tail: &FooterTail, + &mut self, + buf: Bytes, + file_size: u64, + footer_tail: FooterTail, ) -> Result { - #[cfg(feature = "encryption")] - let result = decode_metadata_with_encryption( - buf, - footer_tail.is_encrypted_footer(), - self.file_decryption_properties.as_ref(), - ); - #[cfg(not(feature = "encryption"))] - let result = { - if footer_tail.is_encrypted_footer() { - Err(general_err!( - "Parquet file has an encrypted footer but the encryption feature is disabled" - )) - } else { - Self::decode_metadata(buf) - } - }; - result + // The push decoder expects the metadata to be at the end of the file + // (... data ...) + (metadata) + (footer) + // so we need to provide the starting offset of the metadata + // within the file. + let ending_offset = file_size.checked_sub(FOOTER_SIZE as u64).ok_or_else(|| { + general_err!( + "file size {file_size} is smaller than footer size {}", + FOOTER_SIZE + ) + })?; + + let starting_offset = ending_offset.checked_sub(buf.len() as u64).ok_or_else(|| { + general_err!( + "file size {file_size} is smaller than buffer size {} + footer size {}", + buf.len(), + FOOTER_SIZE + ) + })?; + + let range = starting_offset..ending_offset; + + let push_decoder = + ParquetMetaDataPushDecoder::try_new_with_footer_tail(file_size, footer_tail)? + // NOTE: DO NOT enable page indexes here, they are handled separately + .with_page_index_policy(PageIndexPolicy::Skip); + + let mut push_decoder = self.prepare_push_decoder(push_decoder); + push_decoder.push_ranges(vec![range], vec![buf])?; + match push_decoder.try_decode()? { + DecodeResult::Data(metadata) => Ok(metadata), + DecodeResult::Finished => Err(general_err!( + "could not parse parquet metadata -- previously finished" + )), + DecodeResult::NeedsData(ranges) => Err(general_err!( + "could not parse parquet metadata, needs ranges {:?}", + ranges + )), + } + } + + /// Prepares a push decoder and runs it to decode the metadata. + #[cfg(feature = "encryption")] + fn prepare_push_decoder( + &self, + push_decoder: ParquetMetaDataPushDecoder, + ) -> ParquetMetaDataPushDecoder { + push_decoder.with_file_decryption_properties( + self.file_decryption_properties + .as_ref() + .map(std::sync::Arc::clone), + ) + } + #[cfg(not(feature = "encryption"))] + fn prepare_push_decoder( + &self, + push_decoder: ParquetMetaDataPushDecoder, + ) -> ParquetMetaDataPushDecoder { + push_decoder } /// Decodes [`ParquetMetaData`] from the provided bytes. @@ -761,12 +794,49 @@ impl ParquetMetaDataReader { } } +/// The bounds needed to read page indexes +// this is an internal enum, so it is ok to allow differences in enum size +#[allow(clippy::large_enum_variant)] +enum NeedsIndexData { + /// no additional data is needed (e.g. the indexes weren't requested) + No(ParquetMetaData), + /// Additional data is needed, with the range that are required + Yes(Range), +} + +/// 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 { + match push_decoder.try_decode()? { + DecodeResult::NeedsData(ranges) => { + let range = ranges + .into_iter() + .reduce(|a, b| a.start.min(b.start)..a.end.max(b.end)) + .ok_or_else(|| general_err!("Internal error: no ranges provided"))?; + Ok(NeedsIndexData::Yes(range)) + } + DecodeResult::Data(metadata) => Ok(NeedsIndexData::No(metadata)), + 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 { + match push_decoder.try_decode()? { + DecodeResult::NeedsData(_) => Err(general_err!( + "Internal error: decoder still needs data after reading required range" + )), + DecodeResult::Data(metadata) => Ok(metadata), + DecodeResult::Finished => Err(general_err!("Internal error: decoder was finished")), + } +} + #[cfg(test)] mod tests { use super::*; use crate::file::reader::Length; use crate::util::test_common::file_util::get_test_file; - use bytes::Bytes; use std::ops::Range; #[test]