[thrift-remodel] Add thrift write support (#8237)

# Which issue does this PR close?
**Note: this targets a feature branch, not main**

- Part of #5854.

# Rationale for this change

Begins adding custom thrift write support.

# What changes are included in this PR?

Adds traits to aid in writing of thrift and modifies thrift macros to
generate writing code.

# Are these changes tested?

Yes, adds some roundtrip tests to validate encoded data can be decoded
to the same state.

# Are there any user-facing changes?

No

Author: etseidl

parquet/src/basic.rs

@@ -17,13 +17,18 @@
 
 //! Column chunk encryption metadata
 
+use std::io::Write;
+
 use crate::errors::{ParquetError, Result};
 use crate::format::{
     ColumnCryptoMetaData as TColumnCryptoMetaData,
     EncryptionWithColumnKey as TEncryptionWithColumnKey,
     EncryptionWithFooterKey as TEncryptionWithFooterKey,
 };
-use crate::parquet_thrift::{FieldType, ThriftCompactInputProtocol};
+use crate::parquet_thrift::{
+    ElementType, FieldType, ThriftCompactInputProtocol, ThriftCompactOutputProtocol, WriteThrift,
+    WriteThriftField,
+};
 use crate::{thrift_struct, thrift_union};
 
 // define this and ColumnCryptoMetadata here so they're only defined when
@@ -84,6 +89,7 @@ pub fn to_thrift(column_crypto_metadata: &ColumnCryptoMetaData) -> TColumnCrypto
 #[cfg(test)]
 mod tests {
     use super::*;
+    use crate::parquet_thrift::tests::test_roundtrip;
 
     #[test]
     fn test_encryption_with_footer_key_from_thrift() {
@@ -101,4 +107,24 @@ mod tests {
 
         assert_eq!(try_from_thrift(&to_thrift(&metadata)).unwrap(), metadata);
     }
+
+    #[test]
+    fn test_column_crypto_roundtrip() {
+        test_roundtrip(ColumnCryptoMetaData::ENCRYPTION_WITH_FOOTER_KEY);
+
+        let path_in_schema = vec!["foo".to_owned(), "bar".to_owned(), "really".to_owned()];
+        let key_metadata = vec![1u8; 32];
+        test_roundtrip(ColumnCryptoMetaData::ENCRYPTION_WITH_COLUMN_KEY(
+            EncryptionWithColumnKey {
+                path_in_schema: path_in_schema.clone(),
+                key_metadata: None,
+            },
+        ));
+        test_roundtrip(ColumnCryptoMetaData::ENCRYPTION_WITH_COLUMN_KEY(
+            EncryptionWithColumnKey {
+                path_in_schema,
+                key_metadata: Some(key_metadata),
+            },
+        ));
+    }
 }

parquet/src/file/column_crypto_metadata.rs

@@ -125,7 +125,10 @@ use crate::{
 };
 use crate::{
     basic::{ColumnOrder, Compression, Encoding, Type},
-    parquet_thrift::{FieldType, ThriftCompactInputProtocol},
+    parquet_thrift::{
+        ElementType, FieldType, ThriftCompactInputProtocol, ThriftCompactOutputProtocol,
+        WriteThrift, WriteThriftField,
+    },
 };
 use crate::{
     data_type::private::ParquetValueType, file::page_index::offset_index::OffsetIndexMetaData,
@@ -135,6 +138,7 @@ use crate::{
     thrift_struct,
 };
 pub use reader::{FooterTail, PageIndexPolicy, ParquetMetaDataReader};
+use std::io::Write;
 use std::ops::Range;
 use std::sync::Arc;
 pub use writer::ParquetMetaDataWriter;

parquet/src/file/metadata/mod.rs

@@ -17,6 +17,7 @@
 
 // a collection of generated structs used to parse thrift metadata
 
+use std::io::Write;
 use std::sync::Arc;
 
 #[cfg(feature = "encryption")]
@@ -33,7 +34,10 @@ use crate::{
         page_encoding_stats::PageEncodingStats,
         statistics::ValueStatistics,
     },
-    parquet_thrift::{FieldType, ThriftCompactInputProtocol},
+    parquet_thrift::{
+        ElementType, FieldType, ThriftCompactInputProtocol, ThriftCompactOutputProtocol,
+        WriteThrift, WriteThriftField,
+    },
     schema::types::{parquet_schema_from_array, ColumnDescriptor, SchemaDescriptor},
     thrift_struct,
     util::bit_util::FromBytes,
@@ -507,3 +511,45 @@ impl<'a> TryFrom<&mut ThriftCompactInputProtocol<'a>> for ParquetMetaData {
         Ok(ParquetMetaData::new(fmd, row_groups))
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use crate::file::metadata::thrift_gen::BoundingBox;
+    use crate::parquet_thrift::tests::test_roundtrip;
+
+    #[test]
+    fn test_bounding_box_roundtrip() {
+        test_roundtrip(BoundingBox {
+            xmin: 0.1.into(),
+            xmax: 10.3.into(),
+            ymin: 0.001.into(),
+            ymax: 128.5.into(),
+            zmin: None,
+            zmax: None,
+            mmin: None,
+            mmax: None,
+        });
+
+        test_roundtrip(BoundingBox {
+            xmin: 0.1.into(),
+            xmax: 10.3.into(),
+            ymin: 0.001.into(),
+            ymax: 128.5.into(),
+            zmin: Some(11.0.into()),
+            zmax: Some(1300.0.into()),
+            mmin: None,
+            mmax: None,
+        });
+
+        test_roundtrip(BoundingBox {
+            xmin: 0.1.into(),
+            xmax: 10.3.into(),
+            ymin: 0.001.into(),
+            ymax: 128.5.into(),
+            zmin: Some(11.0.into()),
+            zmax: Some(1300.0.into()),
+            mmin: Some(3.7.into()),
+            mmax: Some(42.0.into()),
+        });
+    }
+}

parquet/src/file/metadata/thrift_gen.rs

@@ -17,9 +17,14 @@
 
 //! Per-page encoding information.
 
+use std::io::Write;
+
 use crate::basic::{Encoding, PageType};
 use crate::errors::{ParquetError, Result};
-use crate::parquet_thrift::{FieldType, ThriftCompactInputProtocol};
+use crate::parquet_thrift::{
+    ElementType, FieldType, ThriftCompactInputProtocol, ThriftCompactOutputProtocol, WriteThrift,
+    WriteThriftField,
+};
 use crate::thrift_struct;
 
 // TODO: This should probably all be moved to thrift_gen

parquet/src/file/page_encoding_stats.rs

@@ -26,8 +26,12 @@ use crate::file::page_index::column_index::{
 };
 use crate::file::page_index::offset_index::OffsetIndexMetaData;
 use crate::file::reader::ChunkReader;
-use crate::parquet_thrift::{FieldType, ThriftCompactInputProtocol};
+use crate::parquet_thrift::{
+    ElementType, FieldType, ThriftCompactInputProtocol, ThriftCompactOutputProtocol, WriteThrift,
+    WriteThriftField,
+};
 use crate::thrift_struct;
+use std::io::Write;
 use std::ops::Range;
 
 /// Computes the covering range of two optional ranges

parquet/src/file/page_index/index_reader.rs

@@ -19,7 +19,12 @@
 //!
 //! [`OffsetIndex`]: https://github.com/apache/parquet-format/blob/master/PageIndex.md
 
-use crate::parquet_thrift::{FieldType, ThriftCompactInputProtocol};
+use std::io::Write;
+
+use crate::parquet_thrift::{
+    ElementType, FieldType, ThriftCompactInputProtocol, ThriftCompactOutputProtocol, WriteThrift,
+    WriteThriftField,
+};
 use crate::{
     errors::{ParquetError, Result},
     thrift_struct,
@@ -193,3 +198,36 @@ fn read_page_location<'a>(prot: &mut ThriftCompactInputProtocol<'a>) -> Result<P
         first_row_index,
     })
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::parquet_thrift::tests::test_roundtrip;
+
+    #[test]
+    fn test_offset_idx_roundtrip() {
+        let page_locations = [
+            PageLocation {
+                offset: 0,
+                compressed_page_size: 10,
+                first_row_index: 0,
+            },
+            PageLocation {
+                offset: 10,
+                compressed_page_size: 20,
+                first_row_index: 100,
+            },
+        ]
+        .to_vec();
+        let unenc = [0i64, 100i64].to_vec();
+
+        test_roundtrip(OffsetIndexMetaData {
+            page_locations: page_locations.clone(),
+            unencoded_byte_array_data_bytes: Some(unenc),
+        });
+        test_roundtrip(OffsetIndexMetaData {
+            page_locations,
+            unencoded_byte_array_data_bytes: None,
+        });
+    }
+}