[Minor] Backport changes to metadata benchmark

parquet/benches/metadata.rs

@@ -15,6 +15,7 @@
 // specific language governing permissions and limitations
 // under the License.
 
+use parquet::file::metadata::ParquetMetaDataReader;
 use rand::Rng;
 use thrift::protocol::TCompactOutputProtocol;
 
@@ -25,7 +26,7 @@ use parquet::file::reader::SerializedFileReader;
 use parquet::file::serialized_reader::ReadOptionsBuilder;
 use parquet::format::{
     ColumnChunk, ColumnMetaData, CompressionCodec, Encoding, FieldRepetitionType, FileMetaData,
-    RowGroup, SchemaElement, Type,
+    PageEncodingStats, PageType, RowGroup, SchemaElement, Type,
 };
 use parquet::thrift::TSerializable;
 
@@ -93,7 +94,18 @@ fn encoded_meta() -> Vec<u8> {
                         index_page_offset: Some(rng.random()),
                         dictionary_page_offset: Some(rng.random()),
                         statistics: Some(stats.clone()),
-                        encoding_stats: None,
+                        encoding_stats: Some(vec![
+                            PageEncodingStats {
+                                page_type: PageType::DICTIONARY_PAGE,
+                                encoding: Encoding::PLAIN,
+                                count: 1,
+                            },
+                            PageEncodingStats {
+                                page_type: PageType::DATA_PAGE,
+                                encoding: Encoding::RLE_DICTIONARY,
+                                count: 10,
+                            },
+                        ]),
                         bloom_filter_offset: None,
                         bloom_filter_length: None,
                         size_statistics: None,
@@ -151,6 +163,36 @@ fn get_footer_bytes(data: Bytes) -> Bytes {
     data.slice(meta_start..meta_end)
 }
 
+#[cfg(feature = "arrow")]
+fn rewrite_file(bytes: Bytes) -> (Bytes, FileMetaData) {
+    use arrow::array::RecordBatchReader;
+    use parquet::arrow::{arrow_reader::ParquetRecordBatchReaderBuilder, ArrowWriter};
+    use parquet::file::properties::{EnabledStatistics, WriterProperties};
+
+    let parquet_reader = ParquetRecordBatchReaderBuilder::try_new(bytes)
+        .expect("parquet open")
+        .build()
+        .expect("parquet open");
+    let writer_properties = WriterProperties::builder()
+        .set_statistics_enabled(EnabledStatistics::Page)
+        .set_write_page_header_statistics(true)
+        .build();
+    let mut output = Vec::new();
+    let mut parquet_writer = ArrowWriter::try_new(
+        &mut output,
+        parquet_reader.schema(),
+        Some(writer_properties),
+    )
+    .expect("create arrow writer");
+
+    for maybe_batch in parquet_reader {
+        let batch = maybe_batch.expect("reading batch");
+        parquet_writer.write(&batch).expect("writing data");
+    }
+    let file_meta = parquet_writer.close().expect("finalizing file");
+    (output.into(), file_meta)
+}
+
 fn criterion_benchmark(c: &mut Criterion) {
     // Read file into memory to isolate filesystem performance
     let file = "../parquet-testing/data/alltypes_tiny_pages.parquet";
@@ -168,19 +210,54 @@ fn criterion_benchmark(c: &mut Criterion) {
         })
     });
 
-    let meta_data = get_footer_bytes(data);
-    c.bench_function("decode file metadata", |b| {
+    let meta_data = get_footer_bytes(data.clone());
+    c.bench_function("decode parquet metadata", |b| {
+        b.iter(|| {
+            ParquetMetaDataReader::decode_metadata(&meta_data).unwrap();
+        })
+    });
+
+    c.bench_function("decode thrift file metadata", |b| {
         b.iter(|| {
             parquet::thrift::bench_file_metadata(&meta_data);
         })
     });
 
-    let buf = black_box(encoded_meta()).into();
-    c.bench_function("decode file metadata (wide)", |b| {
+    let buf: Bytes = black_box(encoded_meta()).into();
+    c.bench_function("decode parquet metadata (wide)", |b| {
+        b.iter(|| {
+            ParquetMetaDataReader::decode_metadata(&buf).unwrap();
+        })
+    });
+
+    c.bench_function("decode thrift file metadata (wide)", |b| {
         b.iter(|| {
             parquet::thrift::bench_file_metadata(&buf);
         })
     });
+
+    // rewrite file with page statistics. then read page headers.
+    #[cfg(feature = "arrow")]
+    let (file_bytes, metadata) = rewrite_file(data.clone());
+    #[cfg(feature = "arrow")]
+    c.bench_function("page headers", |b| {
+        b.iter(|| {
+            metadata.row_groups.iter().for_each(|rg| {
+                rg.columns.iter().for_each(|col| {
+                    if let Some(col_meta) = &col.meta_data {
+                        if let Some(dict_offset) = col_meta.dictionary_page_offset {
+                            parquet::thrift::bench_page_header(
+                                &file_bytes.slice(dict_offset as usize..),
+                            );
+                        }
+                        parquet::thrift::bench_page_header(
+                            &file_bytes.slice(col_meta.data_page_offset as usize..),
+                        );
+                    }
+                });
+            });
+        })
+    });
 }
 
 criterion_group!(benches, criterion_benchmark);

parquet/src/thrift.rs

@@ -33,12 +33,20 @@ pub trait TSerializable: Sized {
     fn write_to_out_protocol<T: TOutputProtocol>(&self, o_prot: &mut T) -> thrift::Result<()>;
 }
 
-/// Public function to aid benchmarking.
+// Public function to aid benchmarking. Reads Parquet `FileMetaData` encoded in `bytes`.
+#[doc(hidden)]
 pub fn bench_file_metadata(bytes: &bytes::Bytes) {
     let mut input = TCompactSliceInputProtocol::new(bytes);
     crate::format::FileMetaData::read_from_in_protocol(&mut input).unwrap();
 }
 
+// Public function to aid benchmarking. Reads Parquet `PageHeader` encoded in `bytes`.
+#[doc(hidden)]
+pub fn bench_page_header(bytes: &bytes::Bytes) {
+    let mut prot = TCompactSliceInputProtocol::new(bytes);
+    crate::format::PageHeader::read_from_in_protocol(&mut prot).unwrap();
+}
+
 /// A more performant implementation of [`TCompactInputProtocol`] that reads a slice
 ///
 /// [`TCompactInputProtocol`]: thrift::protocol::TCompactInputProtocol