test writer produced metadata

Author: rok

parquet/src/file/metadata/writer.rs

@@ -775,10 +775,12 @@ impl MetadataObjectWriter {
                 )?;
 
                 let encrypt_footer = file_encryptor.properties().encrypt_footer();
-                let statistics = if encrypt_footer {
-                    std::mem::take(&mut column_chunk.statistics)
+                let (statistics, encoding_stats) = if encrypt_footer {
+                    let s = std::mem::take(&mut column_chunk.statistics);
+                    let e = std::mem::take(&mut column_chunk.encoding_stats);
+                    (s, e)
                 } else {
-                    None
+                    (None, None)
                 };
                 // create temp ColumnMetaData that we can encrypt
                 let mut buffer: Vec<u8> = vec![];
@@ -789,6 +791,7 @@ impl MetadataObjectWriter {
                 let ciphertext = column_encryptor.encrypt(&buffer, &aad)?;
                 if encrypt_footer {
                     column_chunk.statistics = statistics;
+                    column_chunk.encoding_stats = encoding_stats;
                 }
                 column_chunk.encrypted_column_metadata = Some(ciphertext);
             }

parquet/tests/encryption/encryption.rs

@@ -755,7 +755,26 @@ pub fn test_row_group_statistics_plaintext_encrypted_write() {
     for batch in record_batches.clone() {
         writer.write(&batch).unwrap();
     }
-    let _file_metadata = writer.close().unwrap();
+    let metadata = writer.close().unwrap();
+
+    // Check column statistics that are produced on write are complete
+    assert_eq!(metadata.num_row_groups(), 1);
+    let row_group = &metadata.row_group(0);
+    assert_eq!(row_group.columns().len(), 2);
+
+    for column in row_group.columns() {
+        assert!(column.page_encoding_stats().is_some());
+        assert!(column.statistics().is_some());
+        let column_stats = column.statistics().unwrap();
+        assert_eq!(
+            column_stats.min_bytes_opt(),
+            Some(3i32.to_le_bytes().as_slice())
+        );
+        assert_eq!(
+            column_stats.max_bytes_opt(),
+            Some(19i32.to_le_bytes().as_slice())
+        );
+    }
 
     // Check column statistics are read given plaintext footer and available decryption properties
     let options =
@@ -765,13 +784,12 @@ pub fn test_row_group_statistics_plaintext_encrypted_write() {
 
     assert_eq!(metadata.num_row_groups(), 1);
 
-    let row_group = &metadata.row_groups()[0];
+    let row_group = &metadata.row_group(0);
     assert_eq!(row_group.columns().len(), 2);
 
     // Statistics should be available from decrypted data
-    assert!(&row_group.column(0).statistics().is_some());
-    assert!(&row_group.column(1).statistics().is_some());
     for column in row_group.columns() {
+        assert!(column.page_encoding_stats().is_some());
         assert!(column.statistics().is_some());
         let column_stats = column.statistics().unwrap();
         assert_eq!(
@@ -791,12 +809,14 @@ pub fn test_row_group_statistics_plaintext_encrypted_write() {
 
     assert_eq!(metadata.num_row_groups(), 1);
 
-    let row_group = &metadata.row_groups()[0];
+    let row_group = &metadata.row_group(0);
     assert_eq!(row_group.columns().len(), 2);
     assert!(&row_group.column(0).statistics().is_none());
     assert!(&row_group.column(1).statistics().is_some());
+    assert!(&row_group.column(0).page_encoding_stats().is_none());
+    assert!(&row_group.column(1).page_encoding_stats().is_some());
 
-    let column_stats = &row_group.columns()[1].statistics().unwrap();
+    let column_stats = &row_group.column(1).statistics().unwrap();
     assert_eq!(
         column_stats.min_bytes_opt(),
         Some(3i32.to_le_bytes().as_slice())