Parquet Modular Encryption support

parquet/Cargo.toml

@@ -69,6 +69,7 @@ paste = { version = "1.0" }
 half = { version = "2.1", default-features = false, features = ["num-traits"] }
 sysinfo = { version = "0.32.0", optional = true, default-features = false, features = ["system"] }
 crc32fast = { version = "1.4.2", optional = true, default-features = false }
+ring = { version = "0.17", default-features = false, features = ["std"], optional = true }
 
 [dev-dependencies]
 base64 = { version = "0.22", default-features = false, features = ["std"] }
@@ -97,7 +98,7 @@ zstd-sys = { version = ">=2.0.0, <2.0.14", default-features = false }
 all-features = true
 
 [features]
-default = ["arrow", "snap", "brotli", "flate2", "lz4", "zstd", "base64"]
+default = ["arrow", "snap", "brotli", "flate2", "lz4", "zstd", "base64", "encryption"]
 # Enable lz4
 lz4 = ["lz4_flex"]
 # Enable arrow reader/writer APIs
@@ -120,6 +121,8 @@ zstd = ["dep:zstd", "zstd-sys"]
 sysinfo = ["dep:sysinfo"]
 # Verify 32-bit CRC checksum when decoding parquet pages
 crc = ["dep:crc32fast"]
+# Enable Parquet modular encryption support
+encryption = ["dep:ring"]
 
 
 [[example]]

parquet/src/arrow/arrow_reader/mod.rs

@@ -17,16 +17,15 @@
 
 //! Contains reader which reads parquet data into arrow [`RecordBatch`]
 
-use std::collections::VecDeque;
-use std::sync::Arc;
-
 use arrow_array::cast::AsArray;
 use arrow_array::Array;
 use arrow_array::{RecordBatch, RecordBatchReader};
 use arrow_schema::{ArrowError, DataType as ArrowType, Schema, SchemaRef};
 use arrow_select::filter::prep_null_mask_filter;
 pub use filter::{ArrowPredicate, ArrowPredicateFn, RowFilter};
 pub use selection::{RowSelection, RowSelector};
+use std::collections::VecDeque;
+use std::sync::Arc;
 
 pub use crate::arrow::array_reader::RowGroups;
 use crate::arrow::array_reader::{build_array_reader, ArrayReader};
@@ -42,6 +41,9 @@ mod filter;
 mod selection;
 pub mod statistics;
 
+#[cfg(feature = "encryption")]
+use crate::encryption::ciphers::{CryptoContext, FileDecryptionProperties};
+
 /// Builder for constructing parquet readers into arrow.
 ///
 /// Most users should use one of the following specializations:
@@ -317,7 +319,7 @@ impl ArrowReaderOptions {
     ///
     /// // Create the reader and read the data using the supplied schema.
     /// let mut reader = builder.build().unwrap();
-    /// let _batch = reader.next().unwrap().unwrap();   
+    /// let _batch = reader.next().unwrap().unwrap();
     /// ```
     pub fn with_schema(self, schema: SchemaRef) -> Self {
         Self {
@@ -379,10 +381,20 @@ impl ArrowReaderMetadata {
     /// If `options` has [`ArrowReaderOptions::with_page_index`] true, but
     /// `Self::metadata` is missing the page index, this function will attempt
     /// to load the page index by making an object store request.
-    pub fn load<T: ChunkReader>(reader: &T, options: ArrowReaderOptions) -> Result<Self> {
-        let metadata = ParquetMetaDataReader::new()
-            .with_page_indexes(options.page_index)
+    pub fn load<T: ChunkReader>(
+        reader: &T,
+        options: ArrowReaderOptions,
+        #[cfg(feature = "encryption")] file_decryption_properties: Option<
+            &FileDecryptionProperties,
+        >,
+    ) -> Result<Self> {
+        let metadata = ParquetMetaDataReader::new().with_page_indexes(options.page_index);
+        #[cfg(feature = "encryption")]
+        let metadata = metadata
+            .with_encryption_properties(file_decryption_properties)
             .parse_and_finish(reader)?;
+        #[cfg(not(feature = "encryption"))]
+        let metadata = metadata.parse_and_finish(reader)?;
         Self::try_new(Arc::new(metadata), options)
     }
 
@@ -528,7 +540,23 @@ impl<T: ChunkReader + 'static> ParquetRecordBatchReaderBuilder<T> {
 
     /// Create a new [`ParquetRecordBatchReaderBuilder`] with [`ArrowReaderOptions`]
     pub fn try_new_with_options(reader: T, options: ArrowReaderOptions) -> Result<Self> {
-        let metadata = ArrowReaderMetadata::load(&reader, options)?;
+        let metadata = ArrowReaderMetadata::load(
+            &reader,
+            options,
+            #[cfg(feature = "encryption")]
+            None,
+        )?;
+        Ok(Self::new_with_metadata(reader, metadata))
+    }
+
+    /// Create a new [`ParquetRecordBatchReaderBuilder`] with [`ArrowReaderOptions`] and [`FileDecryptionProperties`]
+    #[cfg(feature = "encryption")]
+    pub fn try_new_with_decryption(
+        reader: T,
+        options: ArrowReaderOptions,
+        file_decryption_properties: Option<&FileDecryptionProperties>,
+    ) -> Result<Self> {
+        let metadata = ArrowReaderMetadata::load(&reader, options, file_decryption_properties)?;
         Ok(Self::new_with_metadata(reader, metadata))
     }
 
@@ -553,6 +581,7 @@ impl<T: ChunkReader + 'static> ParquetRecordBatchReaderBuilder<T> {
     /// # use arrow_schema::{DataType, Field, Schema};
     /// # use parquet::arrow::arrow_reader::{ArrowReaderMetadata, ParquetRecordBatchReader, ParquetRecordBatchReaderBuilder};
     /// # use parquet::arrow::ArrowWriter;
+    /// #
     /// # let mut file: Vec<u8> = Vec::with_capacity(1024);
     /// # let schema = Arc::new(Schema::new(vec![Field::new("i32", DataType::Int32, false)]));
     /// # let mut writer = ArrowWriter::try_new(&mut file, schema.clone(), None).unwrap();
@@ -561,7 +590,7 @@ impl<T: ChunkReader + 'static> ParquetRecordBatchReaderBuilder<T> {
     /// # writer.close().unwrap();
     /// # let file = Bytes::from(file);
     /// #
-    /// let metadata = ArrowReaderMetadata::load(&file, Default::default()).unwrap();
+    /// let metadata = ArrowReaderMetadata::load(&file, Default::default(), None).unwrap();
     /// let mut a = ParquetRecordBatchReaderBuilder::new_with_metadata(file.clone(), metadata.clone()).build().unwrap();
     /// let mut b = ParquetRecordBatchReaderBuilder::new_with_metadata(file, metadata).build().unwrap();
     ///
@@ -677,7 +706,29 @@ impl<T: ChunkReader + 'static> Iterator for ReaderPageIterator<T> {
         let total_rows = rg.num_rows() as usize;
         let reader = self.reader.clone();
 
-        let ret = SerializedPageReader::new(reader, meta, total_rows, page_locations);
+        #[cfg(feature = "encryption")]
+        let crypto_context = if self.metadata.file_decryptor().is_some() {
+            let file_decryptor = Arc::new(self.metadata.file_decryptor().clone().unwrap());
+
+            let crypto_context = CryptoContext::new(
+                rg_idx,
+                self.column_idx,
+                file_decryptor.clone(),
+                file_decryptor,
+            );
+            Some(Arc::new(crypto_context))
+        } else {
+            None
+        };
+
+        let ret = SerializedPageReader::new(
+            reader,
+            meta,
+            total_rows,
+            page_locations,
+            #[cfg(feature = "encryption")]
+            crypto_context,
+        );
         Some(ret.map(|x| Box::new(x) as _))
     }
 }
@@ -788,6 +839,25 @@ impl ParquetRecordBatchReader {
             .build()
     }
 
+    /// Create a new [`ParquetRecordBatchReader`] from the provided chunk reader and [`FileDecryptionProperties`]
+    ///
+    /// Note: this is needed when the parquet file is encrypted
+    // todo: add options or put file_decryption_properties into options
+    #[cfg(feature = "encryption")]
+    pub fn try_new_with_decryption<T: ChunkReader + 'static>(
+        reader: T,
+        batch_size: usize,
+        file_decryption_properties: Option<&FileDecryptionProperties>,
+    ) -> Result<Self> {
+        ParquetRecordBatchReaderBuilder::try_new_with_decryption(
+            reader,
+            Default::default(),
+            file_decryption_properties,
+        )?
+        .with_batch_size(batch_size)
+        .build()
+    }
+
     /// Create a new [`ParquetRecordBatchReader`] from the provided [`RowGroups`]
     ///
     /// Note: this is a low-level interface see [`ParquetRecordBatchReader::try_new`] for a
@@ -943,6 +1013,8 @@ mod tests {
     };
     use arrow_select::concat::concat_batches;
 
+    #[cfg(feature = "encryption")]
+    use crate::arrow::arrow_reader::ArrowReaderMetadata;
     use crate::arrow::arrow_reader::{
         ArrowPredicateFn, ArrowReaderBuilder, ArrowReaderOptions, ParquetRecordBatchReader,
         ParquetRecordBatchReaderBuilder, RowFilter, RowSelection, RowSelector,
@@ -955,6 +1027,8 @@ mod tests {
         BoolType, ByteArray, ByteArrayType, DataType, FixedLenByteArray, FixedLenByteArrayType,
         FloatType, Int32Type, Int64Type, Int96Type,
     };
+    #[cfg(feature = "encryption")]
+    use crate::encryption::ciphers;
     use crate::errors::Result;
     use crate::file::properties::{EnabledStatistics, WriterProperties, WriterVersion};
     use crate::file::writer::SerializedFileWriter;
@@ -1663,6 +1737,190 @@ mod tests {
         assert!(col.value(2).is_nan());
     }
 
+    #[test]
+    #[cfg(feature = "encryption")]
+    fn test_non_uniform_encryption_plaintext_footer() {
+        let testdata = arrow::util::test_util::parquet_test_data();
+        let path = format!("{testdata}/encrypt_columns_plaintext_footer.parquet.encrypted");
+        let file = File::open(path).unwrap();
+
+        let column_1_key = "1234567890123450".as_bytes();
+        let column_2_key = "1234567890123451".as_bytes();
+
+        let decryption_properties = Some(
+            ciphers::FileDecryptionProperties::builder()
+                .with_column_key("float_field".as_bytes().to_vec(), column_1_key.to_vec())
+                .with_column_key("double_field".as_bytes().to_vec(), column_2_key.to_vec())
+                .build(),
+        );
+
+        let metadata =
+            ArrowReaderMetadata::load(&file, Default::default(), decryption_properties.as_ref())
+                .unwrap();
+        let file_metadata = metadata.metadata.file_metadata();
+
+        assert_eq!(file_metadata.num_rows(), 50);
+        assert_eq!(file_metadata.schema_descr().num_columns(), 8);
+        assert_eq!(
+            file_metadata.created_by().unwrap(),
+            "parquet-cpp-arrow version 14.0.0-SNAPSHOT"
+        );
+
+        metadata.metadata.row_groups().iter().for_each(|rg| {
+            assert_eq!(rg.num_columns(), 8);
+            assert_eq!(rg.num_rows(), 50);
+            assert_eq!(rg.total_byte_size(), 3816);
+        });
+
+        let record_reader = ParquetRecordBatchReader::try_new_with_decryption(
+            file,
+            128,
+            decryption_properties.as_ref(),
+        )
+        .unwrap();
+
+        let mut row_count = 0;
+        for batch in record_reader {
+            let batch = batch.unwrap();
+            row_count += batch.num_rows();
+        }
+
+        assert_eq!(row_count, file_metadata.num_rows() as usize);
+    }
+
+    #[test]
+    #[cfg(feature = "encryption")]
+    fn test_non_uniform_encryption() {
+        // Decryption configuration 2: Decrypt using key retriever callback that holds the
+        // keys of two encrypted columns and the footer key. Supply aad_prefix.
+        let testdata = arrow::util::test_util::parquet_test_data();
+        let path = format!("{testdata}/encrypt_columns_and_footer.parquet.encrypted");
+        let file = File::open(path).unwrap();
+
+        let footer_key = "0123456789012345".as_bytes(); // 128bit/16
+        let column_1_key = "1234567890123450".as_bytes();
+        let column_2_key = "1234567890123451".as_bytes();
+
+        let decryption_properties = Some(
+            ciphers::FileDecryptionProperties::builder()
+                .with_footer_key(footer_key.to_vec())
+                .with_column_key("float_field".as_bytes().to_vec(), column_1_key.to_vec())
+                .with_column_key("double_field".as_bytes().to_vec(), column_2_key.to_vec())
+                .with_aad_prefix("tester".as_bytes().to_vec())
+                .build(),
+        );
+
+        let metadata =
+            ArrowReaderMetadata::load(&file, Default::default(), decryption_properties.as_ref())
+                .unwrap();
+        let file_metadata = metadata.metadata.file_metadata();
+
+        assert_eq!(file_metadata.num_rows(), 50);
+        assert_eq!(file_metadata.schema_descr().num_columns(), 8);
+        assert_eq!(
+            file_metadata.created_by().unwrap(),
+            "parquet-cpp-arrow version 19.0.0-SNAPSHOT"
+        );
+    }
+
+    #[test]
+    #[cfg(feature = "encryption")]
+    fn test_uniform_encryption() {
+        let testdata = arrow::util::test_util::parquet_test_data();
+        let path = format!("{testdata}/uniform_encryption.parquet.encrypted");
+        let file = File::open(path).unwrap();
+
+        let key_code: &[u8] = "0123456789012345".as_bytes();
+        let decryption_properties = Some(
+            ciphers::FileDecryptionProperties::builder()
+                .with_footer_key(key_code.to_vec())
+                .build(),
+        );
+
+        let metadata =
+            ArrowReaderMetadata::load(&file, Default::default(), decryption_properties.as_ref())
+                .unwrap();
+        let file_metadata = metadata.metadata.file_metadata();
+
+        assert_eq!(file_metadata.num_rows(), 50);
+        assert_eq!(file_metadata.schema_descr().num_columns(), 8);
+        assert_eq!(
+            file_metadata.created_by().unwrap(),
+            "parquet-cpp-arrow version 19.0.0-SNAPSHOT"
+        );
+
+        metadata.metadata.row_groups().iter().for_each(|rg| {
+            assert_eq!(rg.num_columns(), 8);
+            assert_eq!(rg.num_rows(), 50);
+            assert_eq!(rg.total_byte_size(), 4172);
+        });
+
+        let decryption_properties = Some(
+            ciphers::FileDecryptionProperties::builder()
+                .with_footer_key(key_code.to_vec())
+                .build(),
+        );
+        let record_reader = ParquetRecordBatchReader::try_new_with_decryption(
+            file,
+            128,
+            decryption_properties.as_ref(),
+        )
+        .unwrap();
+
+        let mut row_count = 0;
+        for batch in record_reader {
+            let batch = batch.unwrap();
+            row_count += batch.num_rows();
+
+            let bool_col = batch.column(0).as_boolean();
+            let time_col = batch
+                .column(1)
+                .as_primitive::<types::Time32MillisecondType>();
+            let list_col = batch.column(2).as_list::<i32>();
+            let timestamp_col = batch
+                .column(3)
+                .as_primitive::<types::TimestampNanosecondType>();
+            let f32_col = batch.column(4).as_primitive::<types::Float32Type>();
+            let f64_col = batch.column(5).as_primitive::<types::Float64Type>();
+            let binary_col = batch.column(6).as_binary::<i32>();
+            let fixed_size_binary_col = batch.column(7).as_fixed_size_binary();
+
+            for (i, x) in bool_col.iter().enumerate() {
+                assert_eq!(x.unwrap(), i % 2 == 0);
+            }
+            for (i, x) in time_col.iter().enumerate() {
+                assert_eq!(x.unwrap(), i as i32);
+            }
+            for (i, list_item) in list_col.iter().enumerate() {
+                let list_item = list_item.unwrap();
+                let list_item = list_item.as_primitive::<types::Int64Type>();
+                assert_eq!(list_item.len(), 2);
+                assert_eq!(list_item.value(0), ((i * 2) * 1000000000000) as i64);
+                assert_eq!(list_item.value(1), ((i * 2 + 1) * 1000000000000) as i64);
+            }
+            for x in timestamp_col.iter() {
+                assert!(x.is_some());
+            }
+            for (i, x) in f32_col.iter().enumerate() {
+                assert_eq!(x.unwrap(), i as f32 * 1.1f32);
+            }
+            for (i, x) in f64_col.iter().enumerate() {
+                assert_eq!(x.unwrap(), i as f64 * 1.1111111f64);
+            }
+            for (i, x) in binary_col.iter().enumerate() {
+                assert_eq!(x.is_some(), i % 2 == 0);
+                if let Some(x) = x {
+                    assert_eq!(&x[0..7], b"parquet");
+                }
+            }
+            for (i, x) in fixed_size_binary_col.iter().enumerate() {
+                assert_eq!(x.unwrap(), &[i as u8; 10]);
+            }
+        }
+
+        assert_eq!(row_count, file_metadata.num_rows() as usize);
+    }
+
     #[test]
     fn test_read_float32_float64_byte_stream_split() {
         let path = format!(