[thrift-remodel] Rework thrift reader API

parquet/src/basic.rs

@@ -26,8 +26,8 @@ use std::{fmt, str};
 
 pub use crate::compression::{BrotliLevel, GzipLevel, ZstdLevel};
 use crate::parquet_thrift::{
-    ElementType, FieldType, ThriftCompactInputProtocol, ThriftCompactOutputProtocol, WriteThrift,
-    WriteThriftField,
+    ElementType, FieldType, ReadThrift, ThriftCompactInputProtocol, ThriftCompactOutputProtocol,
+    WriteThrift, WriteThriftField,
 };
 use crate::{thrift_enum, thrift_struct, thrift_union_all_empty};
 
@@ -165,9 +165,8 @@ pub enum ConvertedType {
     INTERVAL,
 }
 
-impl<'a> TryFrom<&mut ThriftCompactInputProtocol<'a>> for ConvertedType {
-    type Error = ParquetError;
-    fn try_from(prot: &mut ThriftCompactInputProtocol<'a>) -> Result<Self> {
+impl<'a, R: ThriftCompactInputProtocol<'a>> ReadThrift<'a, R> for ConvertedType {
+    fn read_thrift(prot: &mut R) -> Result<Self> {
         let val = prot.read_i32()?;
         Ok(match val {
             0 => Self::UTF8,
@@ -361,12 +360,9 @@ pub enum LogicalType {
     },
 }
 
-impl<'a> TryFrom<&mut ThriftCompactInputProtocol<'a>> for LogicalType {
-    type Error = ParquetError;
-    fn try_from(prot: &mut ThriftCompactInputProtocol<'a>) -> Result<Self> {
-        prot.read_struct_begin()?;
-
-        let field_ident = prot.read_field_begin()?;
+impl<'a, R: ThriftCompactInputProtocol<'a>> ReadThrift<'a, R> for LogicalType {
+    fn read_thrift(prot: &mut R) -> Result<Self> {
+        let field_ident = prot.read_field_begin(0)?;
         if field_ident.field_type == FieldType::Stop {
             return Err(general_err!("received empty union from remote LogicalType"));
         }
@@ -388,7 +384,7 @@ impl<'a> TryFrom<&mut ThriftCompactInputProtocol<'a>> for LogicalType {
                 Self::Enum
             }
             5 => {
-                let val = DecimalType::try_from(&mut *prot)?;
+                let val = DecimalType::read_thrift(&mut *prot)?;
                 Self::Decimal {
                     scale: val.scale,
                     precision: val.precision,
@@ -399,21 +395,21 @@ impl<'a> TryFrom<&mut ThriftCompactInputProtocol<'a>> for LogicalType {
                 Self::Date
             }
             7 => {
-                let val = TimeType::try_from(&mut *prot)?;
+                let val = TimeType::read_thrift(&mut *prot)?;
                 Self::Time {
                     is_adjusted_to_u_t_c: val.is_adjusted_to_u_t_c,
                     unit: val.unit,
                 }
             }
             8 => {
-                let val = TimestampType::try_from(&mut *prot)?;
+                let val = TimestampType::read_thrift(&mut *prot)?;
                 Self::Timestamp {
                     is_adjusted_to_u_t_c: val.is_adjusted_to_u_t_c,
                     unit: val.unit,
                 }
             }
             10 => {
-                let val = IntType::try_from(&mut *prot)?;
+                let val = IntType::read_thrift(&mut *prot)?;
                 Self::Integer {
                     is_signed: val.is_signed,
                     bit_width: val.bit_width,
@@ -440,19 +436,19 @@ impl<'a> TryFrom<&mut ThriftCompactInputProtocol<'a>> for LogicalType {
                 Self::Float16
             }
             16 => {
-                let val = VariantType::try_from(&mut *prot)?;
+                let val = VariantType::read_thrift(&mut *prot)?;
                 Self::Variant {
                     specification_version: val.specification_version,
                 }
             }
             17 => {
-                let val = GeometryType::try_from(&mut *prot)?;
+                let val = GeometryType::read_thrift(&mut *prot)?;
                 Self::Geometry {
                     crs: val.crs.map(|s| s.to_owned()),
                 }
             }
             18 => {
-                let val = GeographyType::try_from(&mut *prot)?;
+                let val = GeographyType::read_thrift(&mut *prot)?;
                 Self::Geography {
                     crs: val.crs.map(|s| s.to_owned()),
                     algorithm: val.algorithm,
@@ -465,13 +461,12 @@ impl<'a> TryFrom<&mut ThriftCompactInputProtocol<'a>> for LogicalType {
                 }
             }
         };
-        let field_ident = prot.read_field_begin()?;
+        let field_ident = prot.read_field_begin(field_ident.id)?;
         if field_ident.field_type != FieldType::Stop {
             return Err(general_err!(
                 "Received multiple fields for union from remote LogicalType"
             ));
         }
-        prot.read_struct_end()?;
         Ok(ret)
     }
 }
@@ -756,9 +751,8 @@ pub enum Compression {
     LZ4_RAW,
 }
 
-impl<'a> TryFrom<&mut ThriftCompactInputProtocol<'a>> for Compression {
-    type Error = ParquetError;
-    fn try_from(prot: &mut ThriftCompactInputProtocol<'a>) -> Result<Self> {
+impl<'a, R: ThriftCompactInputProtocol<'a>> ReadThrift<'a, R> for Compression {
+    fn read_thrift(prot: &mut R) -> Result<Self> {
         let val = prot.read_i32()?;
         Ok(match val {
             0 => Self::UNCOMPRESSED,
@@ -1124,12 +1118,9 @@ impl ColumnOrder {
     }
 }
 
-impl<'a> TryFrom<&mut ThriftCompactInputProtocol<'a>> for ColumnOrder {
-    type Error = ParquetError;
-
-    fn try_from(prot: &mut ThriftCompactInputProtocol<'a>) -> Result<Self> {
-        prot.read_struct_begin()?;
-        let field_ident = prot.read_field_begin()?;
+impl<'a, R: ThriftCompactInputProtocol<'a>> ReadThrift<'a, R> for ColumnOrder {
+    fn read_thrift(prot: &mut R) -> Result<Self> {
+        let field_ident = prot.read_field_begin(0)?;
         if field_ident.field_type == FieldType::Stop {
             return Err(general_err!("Received empty union from remote ColumnOrder"));
         }
@@ -1144,13 +1135,12 @@ impl<'a> TryFrom<&mut ThriftCompactInputProtocol<'a>> for ColumnOrder {
                 Self::UNKNOWN
             }
         };
-        let field_ident = prot.read_field_begin()?;
+        let field_ident = prot.read_field_begin(field_ident.id)?;
         if field_ident.field_type != FieldType::Stop {
             return Err(general_err!(
                 "Received multiple fields for union from remote ColumnOrder"
             ));
         }
-        prot.read_struct_end()?;
         Ok(ret)
     }
 }

parquet/src/file/column_crypto_metadata.rs

@@ -26,8 +26,8 @@ use crate::format::{
     EncryptionWithFooterKey as TEncryptionWithFooterKey,
 };
 use crate::parquet_thrift::{
-    ElementType, FieldType, ThriftCompactInputProtocol, ThriftCompactOutputProtocol, WriteThrift,
-    WriteThriftField,
+    read_thrift_vec, ElementType, FieldType, ReadThrift, ThriftCompactInputProtocol,
+    ThriftCompactOutputProtocol, WriteThrift, WriteThriftField,
 };
 use crate::{thrift_struct, thrift_union};
 

parquet/src/file/metadata/mod.rs

@@ -121,8 +121,8 @@ use crate::{
 use crate::{
     basic::{ColumnOrder, Compression, Encoding, Type},
     parquet_thrift::{
-        ElementType, FieldType, ThriftCompactInputProtocol, ThriftCompactOutputProtocol,
-        WriteThrift, WriteThriftField,
+        ElementType, FieldType, ReadThrift, ThriftCompactInputProtocol,
+        ThriftCompactOutputProtocol, WriteThrift, WriteThriftField,
     },
 };
 use crate::{

parquet/src/file/metadata/reader.rs

@@ -19,7 +19,7 @@ use std::{io::Read, ops::Range};
 
 #[cfg(feature = "encryption")]
 use crate::encryption::decrypt::{CryptoContext, FileDecryptionProperties};
-use crate::parquet_thrift::ThriftCompactInputProtocol;
+use crate::parquet_thrift::{ReadThrift, ThriftSliceInputProtocol};
 use bytes::Bytes;
 
 use crate::errors::{ParquetError, Result};
@@ -962,8 +962,8 @@ impl ParquetMetaDataReader {
     ///
     /// [Parquet Spec]: https://github.com/apache/parquet-format#metadata
     pub fn decode_metadata(buf: &[u8]) -> Result<ParquetMetaData> {
-        let mut prot = ThriftCompactInputProtocol::new(buf);
-        ParquetMetaData::try_from(&mut prot)
+        let mut prot = ThriftSliceInputProtocol::new(buf);
+        ParquetMetaData::read_thrift(&mut prot)
     }
 }
 

parquet/src/file/metadata/thrift_gen.rs

@@ -35,8 +35,8 @@ use crate::{
         statistics::ValueStatistics,
     },
     parquet_thrift::{
-        ElementType, FieldType, ThriftCompactInputProtocol, ThriftCompactOutputProtocol,
-        WriteThrift, WriteThriftField,
+        read_thrift_vec, ElementType, FieldType, ReadThrift, ThriftCompactInputProtocol,
+        ThriftCompactOutputProtocol, WriteThrift, WriteThriftField,
     },
     schema::types::{parquet_schema_from_array, ColumnDescriptor, SchemaDescriptor},
     thrift_struct, thrift_union,
@@ -46,6 +46,7 @@ use crate::{
 use crate::{
     encryption::decrypt::{FileDecryptionProperties, FileDecryptor},
     file::column_crypto_metadata::ColumnCryptoMetaData,
+    parquet_thrift::ThriftSliceInputProtocol,
     schema::types::SchemaDescPtr,
 };
 
@@ -669,8 +670,8 @@ fn row_group_from_encrypted_thrift(
                         )
                     })?;
 
-            let mut prot = ThriftCompactInputProtocol::new(decrypted_cc_buf.as_slice());
-            let col_meta = ColumnMetaData::try_from(&mut prot)?;
+            let mut prot = ThriftSliceInputProtocol::new(decrypted_cc_buf.as_slice());
+            let col_meta = ColumnMetaData::read_thrift(&mut prot)?;
             c.meta_data = Some(col_meta);
             columns.push(convert_column(c, d.clone())?);
         } else {
@@ -699,14 +700,14 @@ pub(crate) fn parquet_metadata_with_encryption(
     encrypted_footer: bool,
     buf: &[u8],
 ) -> Result<ParquetMetaData> {
-    let mut prot = ThriftCompactInputProtocol::new(buf);
+    let mut prot = ThriftSliceInputProtocol::new(buf);
     let mut file_decryptor = None;
     let decrypted_fmd_buf;
 
     if encrypted_footer {
         if let Some(file_decryption_properties) = file_decryption_properties {
             let t_file_crypto_metadata: FileCryptoMetaData =
-                FileCryptoMetaData::try_from(&mut prot)
+                FileCryptoMetaData::read_thrift(&mut prot)
                     .map_err(|e| general_err!("Could not parse crypto metadata: {}", e))?;
             let supply_aad_prefix = match &t_file_crypto_metadata.encryption_algorithm {
                 EncryptionAlgorithm::AES_GCM_V1(algo) => algo.supply_aad_prefix,
@@ -734,7 +735,7 @@ pub(crate) fn parquet_metadata_with_encryption(
                         "Provided footer key and AAD were unable to decrypt parquet footer"
                     )
                 })?;
-            prot = ThriftCompactInputProtocol::new(decrypted_fmd_buf.as_ref());
+            prot = ThriftSliceInputProtocol::new(decrypted_fmd_buf.as_ref());
 
             file_decryptor = Some(decryptor);
         } else {
@@ -744,7 +745,7 @@ pub(crate) fn parquet_metadata_with_encryption(
         }
     }
 
-    let file_meta = super::thrift_gen::FileMetaData::try_from(&mut prot)
+    let file_meta = super::thrift_gen::FileMetaData::read_thrift(&mut prot)
         .map_err(|e| general_err!("Could not parse metadata: {}", e))?;
 
     let version = file_meta.version;
@@ -852,10 +853,9 @@ pub(super) fn get_file_decryptor(
 
 /// Create ParquetMetaData from thrift input. Note that this only decodes the file metadata in
 /// the Parquet footer. Page indexes will need to be added later.
-impl<'a> TryFrom<&mut ThriftCompactInputProtocol<'a>> for ParquetMetaData {
-    type Error = ParquetError;
-    fn try_from(prot: &mut ThriftCompactInputProtocol<'a>) -> Result<Self> {
-        let file_meta = super::thrift_gen::FileMetaData::try_from(prot)?;
+impl<'a, R: ThriftCompactInputProtocol<'a>> ReadThrift<'a, R> for ParquetMetaData {
+    fn read_thrift(prot: &mut R) -> Result<Self> {
+        let file_meta = super::thrift_gen::FileMetaData::read_thrift(prot)?;
 
         let version = file_meta.version;
         let num_rows = file_meta.num_rows;

parquet/src/file/page_encoding_stats.rs

@@ -20,10 +20,10 @@
 use std::io::Write;
 
 use crate::basic::{Encoding, PageType};
-use crate::errors::{ParquetError, Result};
+use crate::errors::Result;
 use crate::parquet_thrift::{
-    ElementType, FieldType, ThriftCompactInputProtocol, ThriftCompactOutputProtocol, WriteThrift,
-    WriteThriftField,
+    ElementType, FieldType, ReadThrift, ThriftCompactInputProtocol, ThriftCompactOutputProtocol,
+    WriteThrift, WriteThriftField,
 };
 use crate::thrift_struct;
 

parquet/src/file/page_index/index_reader.rs

@@ -27,8 +27,8 @@ use crate::file::page_index::column_index::{
 use crate::file::page_index::offset_index::OffsetIndexMetaData;
 use crate::file::reader::ChunkReader;
 use crate::parquet_thrift::{
-    ElementType, FieldType, ThriftCompactInputProtocol, ThriftCompactOutputProtocol, WriteThrift,
-    WriteThriftField,
+    read_thrift_vec, ElementType, FieldType, ReadThrift, ThriftCompactInputProtocol,
+    ThriftCompactOutputProtocol, ThriftSliceInputProtocol, WriteThrift, WriteThriftField,
 };
 use crate::thrift_struct;
 use std::io::Write;
@@ -136,15 +136,15 @@ pub fn read_offset_indexes<R: ChunkReader>(
 }
 
 pub(crate) fn decode_offset_index(data: &[u8]) -> Result<OffsetIndexMetaData, ParquetError> {
-    let mut prot = ThriftCompactInputProtocol::new(data);
+    let mut prot = ThriftSliceInputProtocol::new(data);
 
     // Try to read fast-path first. If that fails, fall back to slower but more robust
     // decoder.
     match OffsetIndexMetaData::try_from_fast(&mut prot) {
         Ok(offset_index) => Ok(offset_index),
         Err(_) => {
-            prot = ThriftCompactInputProtocol::new(data);
-            OffsetIndexMetaData::try_from(&mut prot)
+            prot = ThriftSliceInputProtocol::new(data);
+            OffsetIndexMetaData::read_thrift(&mut prot)
         }
     }
 }
@@ -166,8 +166,8 @@ pub(crate) fn decode_column_index(
     data: &[u8],
     column_type: Type,
 ) -> Result<ColumnIndexMetaData, ParquetError> {
-    let mut prot = ThriftCompactInputProtocol::new(data);
-    let index = ThriftColumnIndex::try_from(&mut prot)?;
+    let mut prot = ThriftSliceInputProtocol::new(data);
+    let index = ThriftColumnIndex::read_thrift(&mut prot)?;
 
     let index = match column_type {
         Type::BOOLEAN => {

parquet/src/file/page_index/offset_index.rs

@@ -22,8 +22,8 @@
 use std::io::Write;
 
 use crate::parquet_thrift::{
-    ElementType, FieldType, ThriftCompactInputProtocol, ThriftCompactOutputProtocol, WriteThrift,
-    WriteThriftField,
+    read_thrift_vec, ElementType, FieldType, ReadThrift, ThriftCompactInputProtocol,
+    ThriftCompactOutputProtocol, WriteThrift, WriteThriftField,
 };
 use crate::{
     errors::{ParquetError, Result},
@@ -113,7 +113,9 @@ impl OffsetIndexMetaData {
     // Fast-path read of offset index. This works because we expect all field deltas to be 1,
     // and there's no nesting beyond PageLocation, so no need to save the last field id. Like
     // read_page_locations(), this will fail if absolute field id's are used.
-    pub(super) fn try_from_fast<'a>(prot: &mut ThriftCompactInputProtocol<'a>) -> Result<Self> {
+    pub(super) fn try_from_fast<'a, R: ThriftCompactInputProtocol<'a>>(
+        prot: &mut R,
+    ) -> Result<Self> {
         // Offset index is a struct with 2 fields. First field is an array of PageLocations,
         // the second an optional array of i64.
 
@@ -140,7 +142,7 @@ impl OffsetIndexMetaData {
                     "encountered unknown field while reading OffsetIndex"
                 ));
             }
-            let vec = Vec::<i64>::try_from(&mut *prot)?;
+            let vec = read_thrift_vec::<i64, R>(&mut *prot)?;
             unencoded_byte_array_data_bytes = Some(vec);
 
             // this one should be Stop
@@ -164,7 +166,7 @@ impl OffsetIndexMetaData {
 
 // Note: this will fail if the fields are either out of order, or if a suboptimal
 // encoder doesn't use field deltas.
-fn read_page_location<'a>(prot: &mut ThriftCompactInputProtocol<'a>) -> Result<PageLocation> {
+fn read_page_location<'a, R: ThriftCompactInputProtocol<'a>>(prot: &mut R) -> Result<PageLocation> {
     // there are 3 fields, all mandatory, so all field deltas should be 1
     let (field_type, delta) = prot.read_field_header()?;
     if delta != 1 || field_type != FieldType::I64 as u8 {