etseidl commented on code in PR #6392: URL: https://github.com/apache/arrow-rs/pull/6392#discussion_r1764498120
########## parquet/src/file/metadata/reader.rs: ########## @@ -0,0 +1,770 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::{io::Read, ops::Range, sync::Arc}; + +use bytes::Bytes; + +use crate::basic::ColumnOrder; +use crate::errors::{ParquetError, Result}; +use crate::file::metadata::{FileMetaData, ParquetMetaData, RowGroupMetaData}; +use crate::file::page_index::index::Index; +use crate::file::page_index::index_reader::{acc_range, decode_column_index, decode_offset_index}; +use crate::file::reader::ChunkReader; +use crate::file::{FOOTER_SIZE, PARQUET_MAGIC}; +use crate::format::{ColumnOrder as TColumnOrder, FileMetaData as TFileMetaData}; +use crate::schema::types; +use crate::schema::types::SchemaDescriptor; +use crate::thrift::{TCompactSliceInputProtocol, TSerializable}; + +#[cfg(feature = "async")] +use crate::arrow::async_reader::MetadataFetch; + +#[cfg(feature = "async")] +use tokio::io::{AsyncRead, AsyncReadExt, AsyncSeek, AsyncSeekExt}; + +#[cfg(feature = "async")] +use crate::arrow::async_reader::AsyncFileReader; + +/// Reads the [`ParquetMetaData`] from the footer of a Parquet file. +/// +/// This function is a wrapper around [`ParquetMetaDataReader`]. The input, which must implement +/// [`ChunkReader`], may be a [`std::fs::File`] or [`Bytes`]. In the latter case, the passed in +/// buffer must contain the contents of the entire file if any of the Parquet [Page Index] +/// structures are to be populated (controlled via the `column_index` and `offset_index` +/// arguments). +/// +/// [Page Index]: https://github.com/apache/parquet-format/blob/master/PageIndex.md +pub fn parquet_metadata_from_file<R: ChunkReader>( + file: &R, + column_index: bool, + offset_index: bool, +) -> Result<ParquetMetaData> { + let mut reader = ParquetMetaDataReader::new() + .with_column_indexes(column_index) + .with_offset_indexes(offset_index); + reader.try_parse(file)?; + reader.finish() +} + +/// Reads the [`ParquetMetaData`] from a byte stream. +/// +/// See [`crate::file::metadata::ParquetMetaDataWriter#output-format`] for a description of +/// the Parquet metadata. +/// +/// # Example +/// ```no_run +/// # use parquet::file::metadata::{ParquetMetaData, ParquetMetaDataReader}; +/// # fn open_parquet_file(path: &str) -> std::fs::File { unimplemented!(); } +/// // read parquet metadata including page indexes +/// let file = open_parquet_file("some_path.parquet"); +/// let mut reader = ParquetMetaDataReader::new() +/// .with_page_indexes(true); +/// reader.try_parse(&file).unwrap(); +/// let metadata = reader.finish().unwrap(); +/// assert!(metadata.column_index().is_some()); +/// assert!(metadata.offset_index().is_some()); +/// ``` +pub struct ParquetMetaDataReader { + metadata: Option<ParquetMetaData>, + column_index: bool, + offset_index: bool, + prefetch_hint: Option<usize>, +} + +impl Default for ParquetMetaDataReader { + fn default() -> Self { + Self::new() + } +} + +impl ParquetMetaDataReader { + /// Create a new [`ParquetMetaDataReader`] + pub fn new() -> Self { + Self { + metadata: None, + column_index: false, + offset_index: false, + prefetch_hint: None, + } + } + + /// Create a new [`ParquetMetaDataReader`] populated with a [`ParquetMetaData`] struct + /// obtained via other means. Primarily intended for use with [`Self::load_page_index()`]. + pub fn new_with_metadata(metadata: ParquetMetaData) -> Self { + Self { + metadata: Some(metadata), + column_index: false, + offset_index: false, + prefetch_hint: None, + } + } + + /// Enable or disable reading the page index structures described in + /// "[Parquet page index]: Layout to Support Page Skipping". Equivalent to: + /// `self.with_column_indexes(val).with_offset_indexes(val)` + /// + /// [Parquet page index]: https://github.com/apache/parquet-format/blob/master/PageIndex.md + pub fn with_page_indexes(self, val: bool) -> Self { + self.with_column_indexes(val).with_offset_indexes(val) + } + + /// Enable or disable reading the Parquet [ColumnIndex] structure. + /// + /// [ColumnIndex]: https://github.com/apache/parquet-format/blob/master/PageIndex.md + pub fn with_column_indexes(mut self, val: bool) -> Self { + self.column_index = val; + self + } + + /// Enable or disable reading the Parquet [OffsetIndex] structure. + /// + /// [OffsetIndex]: https://github.com/apache/parquet-format/blob/master/PageIndex.md + pub fn with_offset_indexes(mut self, val: bool) -> Self { + self.offset_index = val; + self + } + + /// Provide a hint as to the number of bytes needed to fully parse the [`ParquetMetaData`]. + /// Only used for the asynchronous [`Self::try_load()`] and [`Self::try_load_from_tail()`] + /// methods. + /// + /// By default, the reader will first fetch the last 8 bytes of the input file to obtain the + /// size of the footer metadata. A second fetch will be performed to obtain the needed bytes. + /// After parsing the footer metadata, a third fetch will be performed to obtain the bytes + /// needed to decode the page index structures, if they have been requested. To avoid + /// unnecessary fetches, `prefetch` can be set to an estimate of the number of bytes needed + /// to fully decode the [`ParquetMetaData`], which can reduce the number of fetch requests and + /// reduce latency. Setting `prefetch` too small will not trigger an error, but will result + /// in extra fetches being performed. + /// + /// One caveat is that when using [`Self::try_load_from_tail()`], setting `prefetch` to a + /// value larger than the file size will result in an error. + pub fn with_prefetch_hint(mut self, prefetch: Option<usize>) -> Self { + self.prefetch_hint = prefetch; + self + } + + /// Return the parsed [`ParquetMetaData`] struct. + pub fn finish(&mut self) -> Result<ParquetMetaData> { + if self.metadata.is_none() { + return Err(general_err!("could not parse parquet metadata")); + } + Ok(self.metadata.take().unwrap()) + } + + /// Attempts to parse the footer metadata (and optionally page indexes) given a [`ChunkReader`]. + /// If `reader` is [`Bytes`] based, then the buffer must contain sufficient bytes to complete + /// the request. If page indexes are desired, the buffer must contain the entire file, or + /// [`Self::try_parse_range()`] should be used. + pub fn try_parse<R: ChunkReader>(&mut self, reader: &R) -> Result<()> { + self.try_parse_range(reader, 0..reader.len() as usize) + } + + /// Same as [`Self::try_parse()`], but only `file_range` bytes of the original file are + /// available. + // TODO(ets): should this also use IndexOutOfBound when range doesn't include the whole footer? + pub fn try_parse_range<R: ChunkReader>( + &mut self, + reader: &R, + file_range: Range<usize>, + ) -> Result<()> { + self.metadata = Some(Self::parse_metadata(reader)?); + + // we can return if page indexes aren't requested + if !self.column_index && !self.offset_index { + return Ok(()); + } + + // TODO(ets): what is the correct behavior for missing page indexes? MetadataLoader would + // leave them as `None`, while the parser in `index_reader::read_columns_indexes` returns a + // vector of empty vectors. + // I think it's best to leave them as `None`. + + // Get bounds needed for page indexes (if any are present in the file). + let range = self.range_for_page_index(); + let range = match range { + Some(range) => range, + None => return Ok(()), + }; + + // 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. + // TODO(ets): should probably add a new error type...IOOB is a little tortured + if !(file_range.contains(&range.start) && file_range.contains(&range.end)) { + return Err(ParquetError::IndexOutOfBound(range.start, range.end)); + } + + let bytes_needed = range.end - range.start; + let bytes = reader.get_bytes((range.start - file_range.start) as u64, bytes_needed)?; + let offset = range.start; + + self.parse_column_index(&bytes, offset)?; + self.parse_offset_index(&bytes, offset)?; + + Ok(()) + } + + /// Attempts to (asynchronously) parse the footer metadata (and optionally page indexes) + /// given a [`MetadataFetch`]. The file size must be known to use this function. + #[cfg(feature = "async")] + pub async fn try_load<F: MetadataFetch>( + &mut self, + mut fetch: F, + file_size: usize, + ) -> Result<()> { + let (metadata, remainder) = + Self::load_metadata(&mut fetch, file_size, self.get_prefetch_size()).await?; + + self.metadata = Some(metadata); + + // we can return if page indexes aren't requested + if !self.column_index && !self.offset_index { + return Ok(()); + } + + self.load_page_index(fetch, remainder).await + } + + /// Attempts to (asynchronously) parse the footer metadata (and optionally page indexes) + /// given a [`AsyncFileReader`]. The file size need not be known, but this will perform at + /// least two fetches, regardless of the value of `prefetch_hint`, if the page indexes are + /// requested. + #[cfg(feature = "async")] + pub async fn try_load_from_tail<R: AsyncFileReader + AsyncRead + AsyncSeek + Unpin + Send>( Review Comment: I'm echoing the bounds from [here](https://github.com/apache/arrow-rs/blob/5414f1d7c0683c64d69cf721a83c17d677c78a71/parquet/src/arrow/async_reader/mod.rs#L164). `AsyncFileReader` isn't really necessary, though. -- 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]
