Add more benchmarks for Parquet thrift decoding (#8037)

# Which issue does this PR close?

- Part of #5854.

# Rationale for this change

Before embarking on radical changes to the thrift processing in the
parquet crate, add a few more benchmarks to help evaluate the
performance gains.

# What changes are included in this PR?

Adds a test originally written by @tustvold
(https://github.com/tustvold/arrow-rs/tree/thrift-bench) and exposes a
new public function in the thrift module rather than making the thrift
parser public.

# Are these changes tested?

Benchmark code only, so no tests necessary.

# Are there any user-facing changes?

No, but adds a public function.

Author: etseidl

parquet/benches/metadata.rs

@@ -15,10 +15,141 @@
 // specific language governing permissions and limitations
 // under the License.
 
+use rand::Rng;
+use thrift::protocol::TCompactOutputProtocol;
+
+use arrow::util::test_util::seedable_rng;
 use bytes::Bytes;
 use criterion::*;
 use parquet::file::reader::SerializedFileReader;
 use parquet::file::serialized_reader::ReadOptionsBuilder;
+use parquet::format::{
+    ColumnChunk, ColumnMetaData, CompressionCodec, Encoding, FieldRepetitionType, FileMetaData,
+    RowGroup, SchemaElement, Type,
+};
+use parquet::thrift::TSerializable;
+
+const NUM_COLUMNS: usize = 10_000;
+const NUM_ROW_GROUPS: usize = 10;
+
+fn encoded_meta() -> Vec<u8> {
+    let mut rng = seedable_rng();
+
+    let mut schema = Vec::with_capacity(NUM_COLUMNS + 1);
+    schema.push(SchemaElement {
+        type_: None,
+        type_length: None,
+        repetition_type: None,
+        name: Default::default(),
+        num_children: Some(NUM_COLUMNS as _),
+        converted_type: None,
+        scale: None,
+        precision: None,
+        field_id: None,
+        logical_type: None,
+    });
+    for i in 0..NUM_COLUMNS {
+        schema.push(SchemaElement {
+            type_: Some(Type::FLOAT),
+            type_length: None,
+            repetition_type: Some(FieldRepetitionType::REQUIRED),
+            name: i.to_string(),
+            num_children: None,
+            converted_type: None,
+            scale: None,
+            precision: None,
+            field_id: None,
+            logical_type: None,
+        })
+    }
+
+    let stats = parquet::format::Statistics {
+        min: None,
+        max: None,
+        null_count: Some(0),
+        distinct_count: None,
+        max_value: Some(vec![rng.random(); 8]),
+        min_value: Some(vec![rng.random(); 8]),
+        is_max_value_exact: Some(true),
+        is_min_value_exact: Some(true),
+    };
+
+    let row_groups = (0..NUM_ROW_GROUPS)
+        .map(|i| {
+            let columns = (0..NUM_COLUMNS)
+                .map(|_| ColumnChunk {
+                    file_path: None,
+                    file_offset: 0,
+                    meta_data: Some(ColumnMetaData {
+                        type_: Type::FLOAT,
+                        encodings: vec![Encoding::PLAIN, Encoding::RLE_DICTIONARY],
+                        path_in_schema: vec![],
+                        codec: CompressionCodec::UNCOMPRESSED,
+                        num_values: rng.random(),
+                        total_uncompressed_size: rng.random(),
+                        total_compressed_size: rng.random(),
+                        key_value_metadata: None,
+                        data_page_offset: rng.random(),
+                        index_page_offset: Some(rng.random()),
+                        dictionary_page_offset: Some(rng.random()),
+                        statistics: Some(stats.clone()),
+                        encoding_stats: None,
+                        bloom_filter_offset: None,
+                        bloom_filter_length: None,
+                        size_statistics: None,
+                        geospatial_statistics: None,
+                    }),
+                    offset_index_offset: Some(rng.random()),
+                    offset_index_length: Some(rng.random()),
+                    column_index_offset: Some(rng.random()),
+                    column_index_length: Some(rng.random()),
+                    crypto_metadata: None,
+                    encrypted_column_metadata: None,
+                })
+                .collect();
+
+            RowGroup {
+                columns,
+                total_byte_size: rng.random(),
+                num_rows: rng.random(),
+                sorting_columns: None,
+                file_offset: None,
+                total_compressed_size: Some(rng.random()),
+                ordinal: Some(i as _),
+            }
+        })
+        .collect();
+
+    let file = FileMetaData {
+        schema,
+        row_groups,
+        version: 1,
+        num_rows: rng.random(),
+        key_value_metadata: None,
+        created_by: Some("parquet-rs".into()),
+        column_orders: None,
+        encryption_algorithm: None,
+        footer_signing_key_metadata: None,
+    };
+
+    let mut buf = Vec::with_capacity(1024);
+    {
+        let mut out = TCompactOutputProtocol::new(&mut buf);
+        file.write_to_out_protocol(&mut out).unwrap();
+    }
+    buf
+}
+
+fn get_footer_bytes(data: Bytes) -> Bytes {
+    let footer_bytes = data.slice(data.len() - 8..);
+    let footer_len = footer_bytes[0] as u32
+        | (footer_bytes[1] as u32) << 8
+        | (footer_bytes[2] as u32) << 16
+        | (footer_bytes[3] as u32) << 24;
+    let meta_start = data.len() - footer_len as usize - 8;
+    let meta_end = data.len() - 8;
+    data.slice(meta_start..meta_end)
+}
 
 fn criterion_benchmark(c: &mut Criterion) {
     // Read file into memory to isolate filesystem performance
@@ -36,6 +167,20 @@ fn criterion_benchmark(c: &mut Criterion) {
             SerializedFileReader::new_with_options(data.clone(), options).unwrap()
         })
     });
+
+    let meta_data = get_footer_bytes(data);
+    c.bench_function("decode 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| {
+        b.iter(|| {
+            parquet::thrift::bench_file_metadata(&buf);
+        })
+    });
 }
 
 criterion_group!(benches, criterion_benchmark);

parquet/src/thrift.rs

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