Add test_read_subset_of_encrypted_columns_with_plaintext_footer

Author: rok

parquet/src/encryption/decrypt.rs

@@ -614,6 +614,11 @@ impl FileDecryptor {
         Ok(self.footer_decryptor.clone())
     }
 
+    /// Returns true if the column is encrypted with a column key.
+    pub(crate) fn is_column_encrypted(&self, column_name: &str) -> bool {
+        self.decryption_properties.column_keys().0.contains(&column_name.to_string())
+    }
+
     /// Verify the signature of the footer
     pub(crate) fn verify_plaintext_footer_signature(&self, plaintext_footer: &[u8]) -> Result<()> {
         // Plaintext footer format is: [plaintext metadata, nonce, authentication tag]

parquet/src/file/metadata/thrift/encryption.rs

@@ -144,10 +144,16 @@ fn row_group_from_encrypted_thrift(
                 }
                 Some(ColumnCryptoMetaData::ENCRYPTION_WITH_COLUMN_KEY(crypto_metadata)) => {
                     let column_name = crypto_metadata.path_in_schema.join(".");
-                    decryptor.get_column_metadata_decryptor(
-                        column_name.as_str(),
-                        crypto_metadata.key_metadata.as_deref(),
-                    )?
+                    if decryptor.is_column_encrypted(column_name.as_str()) {
+                        decryptor.get_column_metadata_decryptor(
+                            column_name.as_str(),
+                            crypto_metadata.key_metadata.as_deref(),
+                        )?
+                    } else {
+                        // We don't have the key for this column, so we can't decrypt it's metadata.
+                        columns.push(c);
+                        continue;
+                    }
                 }
                 Some(ColumnCryptoMetaData::ENCRYPTION_WITH_FOOTER_KEY) => {
                     decryptor.get_footer_decryptor()?

parquet/tests/encryption/encryption.rs

@@ -720,17 +720,18 @@ fn test_write_uniform_encryption_plaintext_footer() {
 }
 
 #[test]
-pub fn test_row_group_statistics_plaintext_encrypted_write() {
+pub fn test_non_uniform_plaintext_encryption_behaviour() {
     let footer_key = b"0123456789012345".to_vec(); // 128bit/16
     let column_key = b"1234567890123450".to_vec();
 
-    let decryption_properties = FileDecryptionProperties::builder(footer_key.clone())
+    let encryption_properties = FileEncryptionProperties::builder(footer_key.clone())
+        .with_plaintext_footer(true)
         .with_column_key("x", column_key.clone())
+        .with_column_key("y", column_key.clone())
         .build()
         .unwrap();
 
-    let encryption_properties = FileEncryptionProperties::builder(footer_key)
-        .with_plaintext_footer(true)
+    let decryption_properties = FileDecryptionProperties::builder(footer_key.clone())
         .with_column_key("x", column_key.clone())
         .build()
         .unwrap();
@@ -739,16 +740,21 @@ pub fn test_row_group_statistics_plaintext_encrypted_write() {
         .with_file_encryption_properties(encryption_properties)
         .build();
 
-    // Write encrypted data with plaintext footer
+    // Write partly encrypted data with plaintext footer
     let values = Int32Array::from(vec![8, 3, 4, 19, 5]);
+    let values = Arc::new(values);
     let schema = Arc::new(Schema::new(vec![
         Field::new("x", values.data_type().clone(), true),
         Field::new("y", values.data_type().clone(), true),
+        Field::new("z", values.data_type().clone(), true),
     ]));
-
-    let values = Arc::new(values);
-    let record_batches =
-        vec![RecordBatch::try_new(schema.clone(), vec![values.clone(), values.clone()]).unwrap()];
+    let record_batches = vec![
+        RecordBatch::try_new(
+            schema.clone(),
+            vec![values.clone(), values.clone(), values.clone()],
+        )
+        .unwrap(),
+    ];
 
     let temp_file = tempfile::tempfile().unwrap();
     let mut writer = ArrowWriter::try_new(&temp_file, schema, Some(props)).unwrap();
@@ -773,28 +779,32 @@ pub fn test_row_group_statistics_plaintext_encrypted_write() {
         }
     };
 
-    // Check column statistics that are produced on write are complete
+    // Check column statistics produced at write time are available in full
     let row_group = metadata.row_group(0);
     check_column_stats(row_group.column(0), true);
     check_column_stats(row_group.column(1), true);
+    check_column_stats(row_group.column(2), true);
 
     // Check column statistics are read given plaintext footer and available decryption properties
     let options =
         ArrowReaderOptions::default().with_file_decryption_properties(decryption_properties);
-    let reader_metadata = ArrowReaderMetadata::load(&temp_file, options.clone()).unwrap();
+    let reader_metadata = ArrowReaderMetadata::load(&temp_file, options).unwrap();
     let metadata = reader_metadata.metadata();
     let row_group = metadata.row_group(0);
-    // Reader reads encrypted stats since decryption properties are provided
+    // Reader can read plaintext from the unencrypted column
+    // and column x for which the key is provided
     check_column_stats(row_group.column(0), true);
-    check_column_stats(row_group.column(1), true);
+    check_column_stats(row_group.column(1), false);
+    check_column_stats(row_group.column(2), true);
 
     let options = ArrowReaderOptions::default();
-    let reader_metadata = ArrowReaderMetadata::load(&temp_file, options.clone()).unwrap();
+    let reader_metadata = ArrowReaderMetadata::load(&temp_file, options).unwrap();
     let metadata = reader_metadata.metadata();
     let row_group = metadata.row_group(0);
-    // Reader can't read encrypted stats since decryption properties are not provided
+    // Reader can only read plaintext from the unencrypted column if no key is provided
     check_column_stats(row_group.column(0), false);
-    check_column_stats(row_group.column(1), true);
+    check_column_stats(row_group.column(1), false);
+    check_column_stats(row_group.column(2), true);
 }
 
 #[test]