store encodings as bitmask

Author: etseidl

parquet/src/basic.rs

@@ -724,6 +724,58 @@ impl FromStr for Encoding {
     }
 }
 
+const MAX_ENCODING: i32 = Encoding::BYTE_STREAM_SPLIT as i32;
+
+/// Given a Thrift input stream, read a vector of [`Encoding`] enums and convert to a bitmask.
+pub(super) fn thrift_encodings_to_mask<'a, R: ThriftCompactInputProtocol<'a>>(
+    prot: &mut R,
+) -> Result<i32> {
+    let mut mask = 0;
+
+    let list_ident = prot.read_list_begin()?;
+    for _ in 0..list_ident.size {
+        let val = i32::read_thrift(prot)?;
+        if (0..=MAX_ENCODING).contains(&val) {
+            mask |= 1 << val;
+        }
+    }
+    Ok(mask)
+}
+
+#[allow(deprecated)]
+fn i32_to_encoding(val: i32) -> Encoding {
+    match val {
+        0 => Encoding::PLAIN,
+        2 => Encoding::PLAIN_DICTIONARY,
+        3 => Encoding::RLE,
+        4 => Encoding::BIT_PACKED,
+        5 => Encoding::DELTA_BINARY_PACKED,
+        6 => Encoding::DELTA_LENGTH_BYTE_ARRAY,
+        7 => Encoding::DELTA_BYTE_ARRAY,
+        8 => Encoding::RLE_DICTIONARY,
+        9 => Encoding::BYTE_STREAM_SPLIT,
+        _ => panic!("Impossible encoding {val}"),
+    }
+}
+
+pub(super) fn encodings_to_mask<'a, I>(encodings: I) -> i32
+where
+    I: Iterator<Item = &'a Encoding>,
+{
+    let mut mask = 0;
+    for e in encodings {
+        mask |= 1 << (*e as i32);
+    }
+    mask
+}
+
+pub(super) fn mask_to_encodings_vec(mask: i32) -> Vec<Encoding> {
+    (0..=MAX_ENCODING)
+        .filter(|i| mask & (1 << i) != 0)
+        .map(i32_to_encoding)
+        .collect()
+}
+
 // ----------------------------------------------------------------------
 // Mirrors thrift enum `CompressionCodec`
 

parquet/src/column/writer/mod.rs

@@ -28,6 +28,7 @@ use std::str;
 
 use crate::basic::{
     BoundaryOrder, Compression, ConvertedType, Encoding, LogicalType, PageType, Type,
+    encodings_to_mask,
 };
 use crate::column::page::{CompressedPage, Page, PageWriteSpec, PageWriter};
 use crate::column::writer::encoder::{ColumnValueEncoder, ColumnValueEncoderImpl, ColumnValues};
@@ -1190,7 +1191,7 @@ impl<'a, E: ColumnValueEncoder> GenericColumnWriter<'a, E> {
 
         let mut builder = ColumnChunkMetaData::builder(self.descr.clone())
             .set_compression(self.codec)
-            .set_encodings(self.encodings.iter().cloned().collect())
+            .set_encodings_mask(encodings_to_mask(self.encodings.iter()))
             .set_page_encoding_stats(self.encoding_stats.clone())
             .set_total_compressed_size(total_compressed_size)
             .set_total_uncompressed_size(total_uncompressed_size)
@@ -1734,7 +1735,7 @@ mod tests {
         assert_eq!(r.rows_written, 4);
 
         let metadata = r.metadata;
-        assert_eq!(metadata.encodings(), &vec![Encoding::PLAIN, Encoding::RLE]);
+        assert_eq!(metadata.encodings(), vec![Encoding::PLAIN, Encoding::RLE]);
         assert_eq!(metadata.num_values(), 4); // just values
         assert_eq!(metadata.dictionary_page_offset(), None);
     }
@@ -2096,7 +2097,7 @@ mod tests {
         let metadata = r.metadata;
         assert_eq!(
             metadata.encodings(),
-            &vec![Encoding::PLAIN, Encoding::RLE, Encoding::RLE_DICTIONARY]
+            vec![Encoding::PLAIN, Encoding::RLE, Encoding::RLE_DICTIONARY]
         );
         assert_eq!(metadata.num_values(), 4);
         assert_eq!(metadata.compressed_size(), 20);
@@ -2222,7 +2223,7 @@ mod tests {
         let metadata = r.metadata;
         assert_eq!(
             metadata.encodings(),
-            &vec![Encoding::PLAIN, Encoding::RLE, Encoding::RLE_DICTIONARY]
+            vec![Encoding::PLAIN, Encoding::RLE, Encoding::RLE_DICTIONARY]
         );
         assert_eq!(metadata.num_values(), 4);
         assert_eq!(metadata.compressed_size(), 20);

parquet/src/file/metadata/encryption.rs

@@ -20,7 +20,7 @@
 use std::io::Write;
 
 use crate::{
-    basic::{Compression, Encoding, Type},
+    basic::{Compression, thrift_encodings_to_mask},
     encryption::decrypt::{FileDecryptionProperties, FileDecryptor},
     errors::{ParquetError, Result},
     file::{
@@ -115,32 +115,6 @@ pub(crate) struct FileCryptoMetaData<'a> {
 }
 );
 
-type CompressionCodec = Compression;
-
-thrift_struct!(
-struct ColumnMetaData<'a> {
-  1: required Type r#type
-  2: required list<Encoding> encodings
-  // we don't expose path_in_schema so skip
-  //3: required list<string> path_in_schema
-  4: required CompressionCodec codec
-  5: required i64 num_values
-  6: required i64 total_uncompressed_size
-  7: required i64 total_compressed_size
-  // we don't expose key_value_metadata so skip
-  //8: optional list<KeyValue> key_value_metadata
-  9: required i64 data_page_offset
-  10: optional i64 index_page_offset
-  11: optional i64 dictionary_page_offset
-  12: optional Statistics<'a> statistics
-  13: optional list<PageEncodingStats> encoding_stats;
-  14: optional i64 bloom_filter_offset;
-  15: optional i32 bloom_filter_length;
-  16: optional SizeStatistics size_statistics;
-  17: optional GeospatialStatistics geospatial_statistics;
-}
-);
-
 fn row_group_from_encrypted_thrift(
     mut rg: RowGroupMetaData,
     decryptor: Option<&FileDecryptor>,
@@ -206,8 +180,7 @@ fn row_group_from_encrypted_thrift(
                 })?;
 
             // parse decrypted buffer and then replace fields in 'c'
-            let mut prot = ThriftSliceInputProtocol::new(decrypted_cc_buf.as_slice());
-            let col_meta = ColumnMetaData::read_thrift(&mut prot)?;
+            let col_meta = read_column_metadata(decrypted_cc_buf.as_slice())?;
 
             let (
                 unencoded_byte_array_data_bytes,
@@ -234,7 +207,7 @@ fn row_group_from_encrypted_thrift(
             c.data_page_offset = col_meta.data_page_offset;
             c.index_page_offset = col_meta.index_page_offset;
             c.dictionary_page_offset = col_meta.dictionary_page_offset;
-            c.statistics = convert_stats(col_meta.r#type, col_meta.statistics)?;
+            c.statistics = convert_stats(d.physical_type(), col_meta.statistics)?;
             c.encoding_stats = col_meta.encoding_stats;
             c.bloom_filter_offset = col_meta.bloom_filter_offset;
             c.bloom_filter_length = col_meta.bloom_filter_length;
@@ -400,3 +373,172 @@ fn get_file_decryptor(
         )),
     }
 }
+
+#[derive(Clone, Debug, Eq, PartialEq)]
+struct ColumnMetaData<'a> {
+    encodings: i32,
+    codec: Compression,
+    num_values: i64,
+    total_uncompressed_size: i64,
+    total_compressed_size: i64,
+    data_page_offset: i64,
+    index_page_offset: Option<i64>,
+    dictionary_page_offset: Option<i64>,
+    statistics: Option<Statistics<'a>>,
+    encoding_stats: Option<Vec<PageEncodingStats>>,
+    bloom_filter_offset: Option<i64>,
+    bloom_filter_length: Option<i32>,
+    size_statistics: Option<SizeStatistics>,
+    geospatial_statistics: Option<GeospatialStatistics>,
+}
+
+fn read_column_metadata<'a>(buf: &'a [u8]) -> Result<ColumnMetaData<'a>> {
+    let mut prot = ThriftSliceInputProtocol::new(buf);
+
+    let mut encodings: Option<i32> = None;
+    let mut codec: Option<Compression> = None;
+    let mut num_values: Option<i64> = None;
+    let mut total_uncompressed_size: Option<i64> = None;
+    let mut total_compressed_size: Option<i64> = None;
+    let mut data_page_offset: Option<i64> = None;
+    let mut index_page_offset: Option<i64> = None;
+    let mut dictionary_page_offset: Option<i64> = None;
+    let mut statistics: Option<Statistics> = None;
+    let mut encoding_stats: Option<Vec<PageEncodingStats>> = None;
+    let mut bloom_filter_offset: Option<i64> = None;
+    let mut bloom_filter_length: Option<i32> = None;
+    let mut size_statistics: Option<SizeStatistics> = None;
+    let mut geospatial_statistics: Option<GeospatialStatistics> = None;
+
+    // `ColumnMetaData`. Read inline for performance sake.
+    // struct ColumnMetaData {
+    //   1: required Type type
+    //   2: required list<Encoding> encodings
+    //   3: required list<string> path_in_schema
+    //   4: required CompressionCodec codec
+    //   5: required i64 num_values
+    //   6: required i64 total_uncompressed_size
+    //   7: required i64 total_compressed_size
+    //   8: optional list<KeyValue> key_value_metadata
+    //   9: required i64 data_page_offset
+    //   10: optional i64 index_page_offset
+    //   11: optional i64 dictionary_page_offset
+    //   12: optional Statistics statistics;
+    //   13: optional list<PageEncodingStats> encoding_stats;
+    //   14: optional i64 bloom_filter_offset;
+    //   15: optional i32 bloom_filter_length;
+    //   16: optional SizeStatistics size_statistics;
+    //   17: optional GeospatialStatistics geospatial_statistics;
+    // }
+    let mut last_field_id = 0i16;
+    loop {
+        let field_ident = prot.read_field_begin(last_field_id)?;
+        if field_ident.field_type == FieldType::Stop {
+            break;
+        }
+        match field_ident.id {
+            // 1: type is never used, we can use the column descriptor
+            2 => {
+                let val = thrift_encodings_to_mask(&mut prot)?;
+                encodings = Some(val);
+            }
+            // 3: path_in_schema is redundant
+            4 => {
+                codec = Some(Compression::read_thrift(&mut prot)?);
+            }
+            5 => {
+                num_values = Some(i64::read_thrift(&mut prot)?);
+            }
+            6 => {
+                total_uncompressed_size = Some(i64::read_thrift(&mut prot)?);
+            }
+            7 => {
+                total_compressed_size = Some(i64::read_thrift(&mut prot)?);
+            }
+            // 8: we don't expose this key value
+            9 => {
+                data_page_offset = Some(i64::read_thrift(&mut prot)?);
+            }
+            10 => {
+                index_page_offset = Some(i64::read_thrift(&mut prot)?);
+            }
+            11 => {
+                dictionary_page_offset = Some(i64::read_thrift(&mut prot)?);
+            }
+            12 => {
+                statistics = Some(Statistics::read_thrift(&mut prot)?);
+            }
+            13 => {
+                let val =
+                    read_thrift_vec::<PageEncodingStats, ThriftSliceInputProtocol>(&mut prot)?;
+                encoding_stats = Some(val);
+            }
+            14 => {
+                bloom_filter_offset = Some(i64::read_thrift(&mut prot)?);
+            }
+            15 => {
+                bloom_filter_length = Some(i32::read_thrift(&mut prot)?);
+            }
+            16 => {
+                let val = SizeStatistics::read_thrift(&mut prot)?;
+                size_statistics = Some(val);
+            }
+            17 => {
+                let val = GeospatialStatistics::read_thrift(&mut prot)?;
+                geospatial_statistics = Some(val);
+            }
+            _ => {
+                prot.skip(field_ident.field_type)?;
+            }
+        };
+        last_field_id = field_ident.id;
+    }
+
+    let Some(encodings) = encodings else {
+        return Err(ParquetError::General(
+            "Required field encodings is missing".to_owned(),
+        ));
+    };
+    let Some(codec) = codec else {
+        return Err(ParquetError::General(
+            "Required field codec is missing".to_owned(),
+        ));
+    };
+    let Some(num_values) = num_values else {
+        return Err(ParquetError::General(
+            "Required field num_values is missing".to_owned(),
+        ));
+    };
+    let Some(total_uncompressed_size) = total_uncompressed_size else {
+        return Err(ParquetError::General(
+            "Required field total_uncompressed_size is missing".to_owned(),
+        ));
+    };
+    let Some(total_compressed_size) = total_compressed_size else {
+        return Err(ParquetError::General(
+            "Required field total_compressed_size is missing".to_owned(),
+        ));
+    };
+    let Some(data_page_offset) = data_page_offset else {
+        return Err(ParquetError::General(
+            "Required field data_page_offset is missing".to_owned(),
+        ));
+    };
+
+    Ok(ColumnMetaData {
+        encodings,
+        num_values,
+        codec,
+        total_uncompressed_size,
+        total_compressed_size,
+        data_page_offset,
+        index_page_offset,
+        dictionary_page_offset,
+        statistics,
+        encoding_stats,
+        bloom_filter_offset,
+        bloom_filter_length,
+        size_statistics,
+        geospatial_statistics,
+    })
+}