POC: Add ParquetMetaDataReader

parquet/src/arrow/arrow_reader/mod.rs

@@ -34,9 +34,7 @@ use crate::arrow::schema::{parquet_to_arrow_schema_and_fields, ParquetField};
 use crate::arrow::{parquet_to_arrow_field_levels, FieldLevels, ProjectionMask};
 use crate::column::page::{PageIterator, PageReader};
 use crate::errors::{ParquetError, Result};
-use crate::file::footer;
-use crate::file::metadata::ParquetMetaData;
-use crate::file::page_index::index_reader;
+use crate::file::metadata::{ParquetMetaData, ParquetMetaDataReader};
 use crate::file::reader::{ChunkReader, SerializedPageReader};
 use crate::schema::types::SchemaDescriptor;
 
@@ -382,23 +380,9 @@ impl ArrowReaderMetadata {
     /// `Self::metadata` is missing the page index, this function will attempt
     /// to load the page index by making an object store request.
     pub fn load<T: ChunkReader>(reader: &T, options: ArrowReaderOptions) -> Result<Self> {
-        let mut metadata = footer::parse_metadata(reader)?;
-        if options.page_index {
-            let column_index = metadata
-                .row_groups()
-                .iter()
-                .map(|rg| index_reader::read_columns_indexes(reader, rg.columns()))
-                .collect::<Result<Vec<_>>>()?;
-            metadata.set_column_index(Some(column_index));
-
-            let offset_index = metadata
-                .row_groups()
-                .iter()
-                .map(|rg| index_reader::read_offset_indexes(reader, rg.columns()))
-                .collect::<Result<Vec<_>>>()?;
-
-            metadata.set_offset_index(Some(offset_index))
-        }
+        let metadata = ParquetMetaDataReader::new()
+            .with_page_indexes(options.page_index)
+            .parse(reader)?;
         Self::try_new(Arc::new(metadata), options)
     }
 
@@ -3496,9 +3480,8 @@ mod tests {
                 ArrowReaderOptions::new().with_page_index(true),
             )
             .unwrap();
-            // Although `Vec<Vec<PageLoacation>>` of each row group is empty,
-            // we should read the file successfully.
-            assert!(builder.metadata().offset_index().unwrap()[0].is_empty());
+            // Absent page indexes should not be initialized, and file should still be readable.
+            assert!(builder.metadata().offset_index().is_none());
             let reader = builder.build().unwrap();
             let batches = reader.collect::<Result<Vec<_>, _>>().unwrap();
             assert_eq!(batches.len(), 1);

parquet/src/arrow/async_reader/metadata.rs

@@ -17,8 +17,7 @@
 
 use crate::arrow::async_reader::AsyncFileReader;
 use crate::errors::{ParquetError, Result};
-use crate::file::footer::{decode_footer, decode_metadata};
-use crate::file::metadata::ParquetMetaData;
+use crate::file::metadata::{ParquetMetaData, ParquetMetaDataReader};
 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::FOOTER_SIZE;
@@ -76,27 +75,27 @@ impl<F: MetadataFetch> MetadataLoader<F> {
         let mut footer = [0; FOOTER_SIZE];
         footer.copy_from_slice(&suffix[suffix_len - FOOTER_SIZE..suffix_len]);
 
-        let length = decode_footer(&footer)?;
+        let length = ParquetMetaDataReader::decode_footer(&footer)?;
 
         if file_size < length + FOOTER_SIZE {
             return Err(ParquetError::EOF(format!(
                 "file size of {} is less than footer + metadata {}",
                 file_size,
-                length + 8
+                length + FOOTER_SIZE
             )));
         }
 
         // Did not fetch the entire file metadata in the initial read, need to make a second request
         let (metadata, remainder) = if length > suffix_len - FOOTER_SIZE {
             let metadata_start = file_size - length - FOOTER_SIZE;
             let meta = fetch.fetch(metadata_start..file_size - FOOTER_SIZE).await?;
-            (decode_metadata(&meta)?, None)
+            (ParquetMetaDataReader::decode_metadata(&meta)?, None)
         } else {
             let metadata_start = file_size - length - FOOTER_SIZE - footer_start;
 
             let slice = &suffix[metadata_start..suffix_len - FOOTER_SIZE];
             (
-                decode_metadata(slice)?,
+                ParquetMetaDataReader::decode_metadata(slice)?,
                 Some((footer_start, suffix.slice(..metadata_start))),
             )
         };
@@ -237,8 +236,9 @@ where
     Fut: Future<Output = Result<Bytes>> + Send,
 {
     let fetch = MetadataFetchFn(fetch);
-    let loader = MetadataLoader::load(fetch, file_size, prefetch).await?;
-    Ok(loader.finish())
+    let mut reader = ParquetMetaDataReader::new().with_prefetch_hint(prefetch);
+    reader.try_load(fetch, file_size).await?;
+    reader.finish()
 }
 
 #[cfg(test)]
@@ -332,41 +332,43 @@ mod tests {
         };
 
         let f = MetadataFetchFn(&mut fetch);
-        let mut loader = MetadataLoader::load(f, len, None).await.unwrap();
-        assert_eq!(fetch_count.load(Ordering::SeqCst), 2);
-        loader.load_page_index(true, true).await.unwrap();
+        let mut loader = ParquetMetaDataReader::new().with_page_indexes(true);
+        loader.try_load(f, len).await.unwrap();
         assert_eq!(fetch_count.load(Ordering::SeqCst), 3);
-        let metadata = loader.finish();
+        let metadata = loader.finish().unwrap();
         assert!(metadata.offset_index().is_some() && metadata.column_index().is_some());
 
         // Prefetch just footer exactly
         fetch_count.store(0, Ordering::SeqCst);
         let f = MetadataFetchFn(&mut fetch);
-        let mut loader = MetadataLoader::load(f, len, Some(1729)).await.unwrap();
-        assert_eq!(fetch_count.load(Ordering::SeqCst), 1);
-        loader.load_page_index(true, true).await.unwrap();
+        let mut loader = ParquetMetaDataReader::new()
+            .with_page_indexes(true)
+            .with_prefetch_hint(Some(1729));
+        loader.try_load(f, len).await.unwrap();
         assert_eq!(fetch_count.load(Ordering::SeqCst), 2);
-        let metadata = loader.finish();
+        let metadata = loader.finish().unwrap();
         assert!(metadata.offset_index().is_some() && metadata.column_index().is_some());
 
         // Prefetch more than footer but not enough
         fetch_count.store(0, Ordering::SeqCst);
         let f = MetadataFetchFn(&mut fetch);
-        let mut loader = MetadataLoader::load(f, len, Some(130649)).await.unwrap();
-        assert_eq!(fetch_count.load(Ordering::SeqCst), 1);
-        loader.load_page_index(true, true).await.unwrap();
+        let mut loader = ParquetMetaDataReader::new()
+            .with_page_indexes(true)
+            .with_prefetch_hint(Some(130649));
+        loader.try_load(f, len).await.unwrap();
         assert_eq!(fetch_count.load(Ordering::SeqCst), 2);
-        let metadata = loader.finish();
+        let metadata = loader.finish().unwrap();
         assert!(metadata.offset_index().is_some() && metadata.column_index().is_some());
 
         // Prefetch exactly enough
         fetch_count.store(0, Ordering::SeqCst);
         let f = MetadataFetchFn(&mut fetch);
-        let mut loader = MetadataLoader::load(f, len, Some(130650)).await.unwrap();
-        assert_eq!(fetch_count.load(Ordering::SeqCst), 1);
-        loader.load_page_index(true, true).await.unwrap();
+        let mut loader = ParquetMetaDataReader::new()
+            .with_page_indexes(true)
+            .with_prefetch_hint(Some(130650));
+        loader.try_load(f, len).await.unwrap();
         assert_eq!(fetch_count.load(Ordering::SeqCst), 1);
-        let metadata = loader.finish();
+        let metadata = loader.finish().unwrap();
         assert!(metadata.offset_index().is_some() && metadata.column_index().is_some());
     }
 }

parquet/src/arrow/async_reader/mod.rs

@@ -104,8 +104,7 @@ use crate::bloom_filter::{
 };
 use crate::column::page::{PageIterator, PageReader};
 use crate::errors::{ParquetError, Result};
-use crate::file::footer::{decode_footer, decode_metadata};
-use crate::file::metadata::{ParquetMetaData, RowGroupMetaData};
+use crate::file::metadata::{ParquetMetaData, ParquetMetaDataReader, RowGroupMetaData};
 use crate::file::page_index::offset_index::OffsetIndexMetaData;
 use crate::file::reader::{ChunkReader, Length, SerializedPageReader};
 use crate::file::FOOTER_SIZE;
@@ -186,14 +185,14 @@ impl<T: AsyncRead + AsyncSeek + Unpin + Send> AsyncFileReader for T {
             let mut buf = [0_u8; FOOTER_SIZE];
             self.read_exact(&mut buf).await?;
 
-            let metadata_len = decode_footer(&buf)?;
+            let metadata_len = ParquetMetaDataReader::decode_footer(&buf)?;
             self.seek(SeekFrom::End(-FOOTER_SIZE_I64 - metadata_len as i64))
                 .await?;
 
             let mut buf = Vec::with_capacity(metadata_len);
             self.take(metadata_len as _).read_to_end(&mut buf).await?;
 
-            Ok(Arc::new(decode_metadata(&buf)?))
+            Ok(Arc::new(ParquetMetaDataReader::decode_metadata(&buf)?))
         }
         .boxed()
     }
@@ -220,9 +219,9 @@ impl ArrowReaderMetadata {
             && metadata.offset_index().is_none()
         {
             let m = Arc::try_unwrap(metadata).unwrap_or_else(|e| e.as_ref().clone());
-            let mut loader = MetadataLoader::new(input, m);
-            loader.load_page_index(true, true).await?;
-            metadata = Arc::new(loader.finish())
+            let mut reader = ParquetMetaDataReader::new_with_metadata(m).with_page_indexes(true);
+            reader.load_page_index(input, None).await?;
+            metadata = Arc::new(reader.finish()?)
         }
         Self::try_new(metadata, options)
     }
@@ -909,7 +908,6 @@ mod tests {
     };
     use crate::arrow::schema::parquet_to_arrow_schema_and_fields;
     use crate::arrow::ArrowWriter;
-    use crate::file::footer::parse_metadata;
     use crate::file::page_index::index_reader;
     use crate::file::properties::WriterProperties;
     use arrow::compute::kernels::cmp::eq;
@@ -946,13 +944,17 @@ mod tests {
         }
     }
 
+    fn get_parquet_metadata(data: &Bytes) -> ParquetMetaData {
+        ParquetMetaDataReader::new().parse(data).unwrap()
+    }
+
     #[tokio::test]
     async fn test_async_reader() {
         let testdata = arrow::util::test_util::parquet_test_data();
         let path = format!("{testdata}/alltypes_plain.parquet");
         let data = Bytes::from(std::fs::read(path).unwrap());
 
-        let metadata = parse_metadata(&data).unwrap();
+        let metadata = get_parquet_metadata(&data);
         let metadata = Arc::new(metadata);
 
         assert_eq!(metadata.num_row_groups(), 1);
@@ -1007,7 +1009,7 @@ mod tests {
         let path = format!("{testdata}/alltypes_tiny_pages_plain.parquet");
         let data = Bytes::from(std::fs::read(path).unwrap());
 
-        let metadata = parse_metadata(&data).unwrap();
+        let metadata = get_parquet_metadata(&data);
         let metadata = Arc::new(metadata);
 
         assert_eq!(metadata.num_row_groups(), 1);
@@ -1073,7 +1075,7 @@ mod tests {
         let path = format!("{testdata}/alltypes_tiny_pages_plain.parquet");
         let data = Bytes::from(std::fs::read(path).unwrap());
 
-        let metadata = parse_metadata(&data).unwrap();
+        let metadata = get_parquet_metadata(&data);
         let metadata = Arc::new(metadata);
 
         assert_eq!(metadata.num_row_groups(), 1);
@@ -1117,7 +1119,7 @@ mod tests {
         let path = format!("{testdata}/alltypes_tiny_pages_plain.parquet");
         let data = Bytes::from(std::fs::read(path).unwrap());
 
-        let metadata = parse_metadata(&data).unwrap();
+        let metadata = get_parquet_metadata(&data);
         let metadata = Arc::new(metadata);
 
         assert_eq!(metadata.num_row_groups(), 1);
@@ -1173,7 +1175,7 @@ mod tests {
         let path = format!("{testdata}/alltypes_tiny_pages_plain.parquet");
         let data = Bytes::from(std::fs::read(path).unwrap());
 
-        let metadata = parse_metadata(&data).unwrap();
+        let metadata = get_parquet_metadata(&data);
         let metadata = Arc::new(metadata);
 
         assert_eq!(metadata.num_row_groups(), 1);
@@ -1238,7 +1240,7 @@ mod tests {
         let path = format!("{testdata}/alltypes_tiny_pages_plain.parquet");
         let data = Bytes::from(std::fs::read(path).unwrap());
 
-        let metadata = parse_metadata(&data).unwrap();
+        let metadata = get_parquet_metadata(&data);
         let metadata = Arc::new(metadata);
 
         assert_eq!(metadata.num_row_groups(), 1);
@@ -1317,7 +1319,7 @@ mod tests {
         writer.close().unwrap();
 
         let data: Bytes = buf.into();
-        let metadata = parse_metadata(&data).unwrap();
+        let metadata = get_parquet_metadata(&data);
         let parquet_schema = metadata.file_metadata().schema_descr_ptr();
 
         let test = TestReader {
@@ -1391,7 +1393,7 @@ mod tests {
         writer.close().unwrap();
 
         let data: Bytes = buf.into();
-        let metadata = parse_metadata(&data).unwrap();
+        let metadata = get_parquet_metadata(&data);
 
         assert_eq!(metadata.num_row_groups(), 2);
 
@@ -1479,7 +1481,7 @@ mod tests {
         let path = format!("{testdata}/alltypes_tiny_pages_plain.parquet");
         let data = Bytes::from(std::fs::read(path).unwrap());
 
-        let metadata = parse_metadata(&data).unwrap();
+        let metadata = get_parquet_metadata(&data);
         let parquet_schema = metadata.file_metadata().schema_descr_ptr();
         let metadata = Arc::new(metadata);
 
@@ -1529,7 +1531,7 @@ mod tests {
         let path = format!("{testdata}/alltypes_tiny_pages.parquet");
         let data = Bytes::from(std::fs::read(path).unwrap());
 
-        let metadata = parse_metadata(&data).unwrap();
+        let metadata = get_parquet_metadata(&data);
 
         let offset_index =
             index_reader::read_offset_indexes(&data, metadata.row_group(0).columns())
@@ -1619,7 +1621,7 @@ mod tests {
         let path = format!("{testdata}/alltypes_plain.parquet");
         let data = Bytes::from(std::fs::read(path).unwrap());
 
-        let metadata = parse_metadata(&data).unwrap();
+        let metadata = get_parquet_metadata(&data);
         let file_rows = metadata.file_metadata().num_rows() as usize;
         let metadata = Arc::new(metadata);
 
@@ -1764,7 +1766,7 @@ mod tests {
         let testdata = arrow::util::test_util::parquet_test_data();
         let path = format!("{testdata}/data_index_bloom_encoding_stats.parquet");
         let data = Bytes::from(std::fs::read(path).unwrap());
-        let metadata = parse_metadata(&data).unwrap();
+        let metadata = get_parquet_metadata(&data);
         let metadata = Arc::new(metadata);
         let async_reader = TestReader {
             data: data.clone(),
@@ -1793,7 +1795,7 @@ mod tests {
     }
 
     async fn test_get_row_group_column_bloom_filter(data: Bytes, with_length: bool) {
-        let metadata = parse_metadata(&data).unwrap();
+        let metadata = get_parquet_metadata(&data);
         let metadata = Arc::new(metadata);
 
         assert_eq!(metadata.num_row_groups(), 1);
@@ -1933,7 +1935,7 @@ mod tests {
         writer.close().unwrap();
 
         let data: Bytes = buf.into();
-        let metadata = parse_metadata(&data).unwrap();
+        let metadata = get_parquet_metadata(&data);
         let parquet_schema = metadata.file_metadata().schema_descr_ptr();
 
         let test = TestReader {

parquet/src/arrow/async_reader/store.rs

@@ -24,9 +24,9 @@ use futures::{FutureExt, TryFutureExt};
 
 use object_store::{ObjectMeta, ObjectStore};
 
-use crate::arrow::async_reader::{AsyncFileReader, MetadataLoader};
+use crate::arrow::async_reader::AsyncFileReader;
 use crate::errors::Result;
-use crate::file::metadata::ParquetMetaData;
+use crate::file::metadata::{ParquetMetaData, ParquetMetaDataReader};
 
 /// Reads Parquet files in object storage using [`ObjectStore`].
 ///
@@ -124,15 +124,13 @@ impl AsyncFileReader for ParquetObjectReader {
 
     fn get_metadata(&mut self) -> BoxFuture<'_, Result<Arc<ParquetMetaData>>> {
         Box::pin(async move {
-            let preload_column_index = self.preload_column_index;
-            let preload_offset_index = self.preload_offset_index;
+            let mut reader = ParquetMetaDataReader::new()
+                .with_column_indexes(self.preload_column_index)
+                .with_offset_indexes(self.preload_offset_index)
+                .with_prefetch_hint(self.metadata_size_hint);
             let file_size = self.meta.size;
-            let prefetch = self.metadata_size_hint;
-            let mut loader = MetadataLoader::load(self, file_size, prefetch).await?;
-            loader
-                .load_page_index(preload_column_index, preload_offset_index)
-                .await?;
-            Ok(Arc::new(loader.finish()))
+            reader.try_load(self, file_size).await?;
+            Ok(Arc::new(reader.finish()?))
         })
     }
 }

parquet/src/bin/parquet-concat.rs

@@ -39,6 +39,7 @@
 use clap::Parser;
 use parquet::column::writer::ColumnCloseResult;
 use parquet::errors::{ParquetError, Result};
+use parquet::file::metadata::ParquetMetaDataReader;
 use parquet::file::properties::WriterProperties;
 use parquet::file::writer::SerializedFileWriter;
 use std::fs::File;
@@ -70,7 +71,7 @@ impl Args {
             .iter()
             .map(|x| {
                 let reader = File::open(x)?;
-                let metadata = parquet::file::footer::parse_metadata(&reader)?;
+                let metadata = ParquetMetaDataReader::new().parse(&reader)?;
                 Ok((reader, metadata))
             })
             .collect::<Result<Vec<_>>>()?;