From 2b3275474eac5cc8bdfd4fb82d09a0b190f47306 Mon Sep 17 00:00:00 2001 From: Yanzhi Xin Date: Mon, 19 Oct 2026 19:50:35 -0700 Subject: [PATCH] fix: preserve Parquet column key metadata associations --- datafusion/common/src/config.rs | 4 +- .../tests/encryption_metadata_roundtrip.rs | 196 ++++++++++++++++++ 2 files changed, 198 insertions(+), 2 deletions(-) create mode 100644 datafusion/common/tests/encryption_metadata_roundtrip.rs diff --git a/datafusion/common/src/config.rs b/datafusion/common/src/config.rs index fc7202a6c5a45..8685b302aa230 100644 --- a/datafusion/common/src/config.rs +++ b/datafusion/common/src/config.rs @@ -3655,7 +3655,7 @@ impl TryFrom for FileEncryptionProperties { #[cfg(feature = "parquet_encryption")] impl From<&Arc> for ConfigFileEncryptionProperties { fn from(f: &Arc) -> Self { - let (column_names_vec, column_keys_vec, column_metas_vec) = f.column_keys(); + let (column_names_vec, column_keys_vec, _) = f.column_keys(); let mut column_encryption_properties: HashMap< String, @@ -3665,7 +3665,7 @@ impl From<&Arc> for ConfigFileEncryptionProperties { for (i, column_name) in column_names_vec.iter().enumerate() { let column_key_as_hex = hex::encode(&column_keys_vec[i]); let column_metadata_as_hex: Option = - column_metas_vec.get(i).map(hex::encode); + f.column_key_metadata(column_name).map(hex::encode); column_encryption_properties.insert( column_name.clone(), ColumnEncryptionProperties { diff --git a/datafusion/common/tests/encryption_metadata_roundtrip.rs b/datafusion/common/tests/encryption_metadata_roundtrip.rs new file mode 100644 index 0000000000000..33a97247bcdde --- /dev/null +++ b/datafusion/common/tests/encryption_metadata_roundtrip.rs @@ -0,0 +1,196 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Preserve the association between column keys and optional key metadata. + +#![cfg(feature = "parquet_encryption")] + +use std::collections::HashMap; +use std::fs::OpenOptions; +use std::io::Write; +use std::path::PathBuf; +use std::sync::Arc; + +use arrow::array::{ArrayRef, Int32Array}; +use arrow::record_batch::RecordBatch; +use datafusion_common::config::ConfigFileEncryptionProperties; +use datafusion_common::encryption::{FileDecryptionProperties, FileEncryptionProperties}; +use parquet::arrow::ArrowWriter; +use parquet::arrow::arrow_reader::{ArrowReaderOptions, ParquetRecordBatchReaderBuilder}; +use parquet::file::column_crypto_metadata::ColumnCryptoMetaData; +use parquet::file::properties::WriterProperties; +use uuid::Uuid; + +const FOOTER_KEY: &[u8] = b"0123456789abcdef"; +const COLUMN_NAMES: [&str; 4] = ["first", "second", "third", "fourth"]; + +fn properties(metadata: &[Option<&[u8]>]) -> Arc { + let mut builder = FileEncryptionProperties::builder(FOOTER_KEY.to_vec()) + .with_footer_key_metadata(b"footer-key-id".to_vec()); + for (i, metadata) in metadata.iter().enumerate() { + let key = vec![i as u8; 16]; + builder = match metadata { + Some(metadata) => builder.with_column_key_and_metadata( + COLUMN_NAMES[i], + key, + metadata.to_vec(), + ), + None => builder.with_column_key(COLUMN_NAMES[i], key), + }; + } + builder.build().unwrap() +} + +fn assert_roundtrip(metadata: &[Option<&[u8]>]) { + // Each newly built map has its own iteration order. Associations must not + // depend on where a column happens to fall in that order. + for _ in 0..32 { + let original = properties(metadata); + let config = ConfigFileEncryptionProperties::from(&original); + for (i, expected_metadata) in metadata.iter().enumerate() { + let column = &config.column_encryption_properties[COLUMN_NAMES[i]]; + assert_eq!(column.column_key_as_hex, hex::encode(vec![i as u8; 16])); + assert_eq!( + column.column_metadata_as_hex, + expected_metadata.map(hex::encode), + "metadata moved away from column {}", + COLUMN_NAMES[i] + ); + } + + let restored = FileEncryptionProperties::try_from(config).unwrap(); + assert_eq!(restored.footer_key(), FOOTER_KEY); + let (names, keys, _) = restored.column_keys(); + let keys: HashMap<_, _> = names.into_iter().zip(keys).collect(); + for i in 0..metadata.len() { + assert_eq!(keys[COLUMN_NAMES[i]], vec![i as u8; 16]); + } + assert_eq!(original.as_ref(), &restored); + } +} + +#[test] +fn mixed_column_metadata_preserves_each_association() { + assert_roundtrip(&[Some(b"first-key-id"), None, Some(b""), None]); +} + +#[test] +fn all_column_metadata_preserves_each_association() { + assert_roundtrip(&[ + Some(b"first-key-id"), + Some(b"second-key-id"), + Some(b"third-key-id"), + Some(b"fourth-key-id"), + ]); +} + +#[test] +fn absent_column_metadata_remains_absent() { + assert_roundtrip(&[None, None, None, None]); +} + +struct RemoveOnDrop(PathBuf); + +impl Drop for RemoveOnDrop { + fn drop(&mut self) { + let _ = std::fs::remove_file(&self.0); + } +} + +fn assert_encrypted_file_metadata( + properties: Arc, + expected_metadata: &[Option<&[u8]>], +) { + let first: ArrayRef = Arc::new(Int32Array::from(vec![Some(1), None, Some(3)])); + let second: ArrayRef = Arc::new(Int32Array::from(vec![Some(4), Some(5), None])); + let batch = + RecordBatch::try_from_iter(vec![("first", first), ("second", second)]).unwrap(); + + // Write entirely in memory. The temporary file supplies the reader's public + // File-based ChunkReader without adding a byte-buffer dependency to this crate. + let writer_properties = WriterProperties::builder() + .with_file_encryption_properties(properties) + .build(); + let mut writer = + ArrowWriter::try_new(Vec::new(), batch.schema(), Some(writer_properties)) + .unwrap(); + writer.write(&batch).unwrap(); + let bytes = writer.into_inner().unwrap(); + let path = std::env::temp_dir().join(format!( + "datafusion-encryption-metadata-{}.parquet", + Uuid::new_v4() + )); + let mut file = OpenOptions::new() + .read(true) + .write(true) + .create_new(true) + .open(&path) + .unwrap(); + let cleanup = RemoveOnDrop(path); + file.write_all(&bytes).unwrap(); + + let mut decrypt = FileDecryptionProperties::builder(FOOTER_KEY.to_vec()); + for (i, column_name) in COLUMN_NAMES[..2].iter().enumerate() { + decrypt = decrypt.with_column_key(column_name, vec![i as u8; 16]); + } + let options = ArrowReaderOptions::new() + .with_file_decryption_properties(decrypt.build().unwrap()); + let builder = + ParquetRecordBatchReaderBuilder::try_new_with_options(file, options).unwrap(); + for row_group in builder.metadata().row_groups() { + for (i, column) in row_group.columns().iter().enumerate() { + let Some(ColumnCryptoMetaData::ENCRYPTION_WITH_COLUMN_KEY(crypto)) = + column.crypto_metadata() + else { + panic!("expected a column-specific encryption key"); + }; + assert_eq!(crypto.path_in_schema, vec![COLUMN_NAMES[i].to_string()]); + assert_eq!( + crypto.key_metadata.as_deref(), + expected_metadata[i], + "file metadata changed for column {}", + COLUMN_NAMES[i] + ); + } + } + let batches = builder + .build() + .unwrap() + .collect::, _>>() + .unwrap(); + assert_eq!(batches, vec![batch]); + // The reader (and its File) has been consumed before removing the file. + drop(cleanup); +} + +#[test] +fn encrypted_files_preserve_optional_column_metadata_after_conversion() { + for metadata in [ + [Some(b"first-key-id".as_slice()), None], + [None, Some(b"".as_slice())], + ] { + // Repeat with independent maps so metadata placement is not tested with + // only one possible column iteration order. + for _ in 0..16 { + let original = properties(&metadata); + let config = ConfigFileEncryptionProperties::from(&original); + let restored = Arc::new(FileEncryptionProperties::try_from(config).unwrap()); + assert_encrypted_file_metadata(original, &metadata); + assert_encrypted_file_metadata(restored, &metadata); + } + } +}