feat(encryption): [13/N] Write encrypted Parquet data files

pyiceberg/encryption/manager.py

@@ -74,12 +74,7 @@ def from_table_metadata(cls, kms_client: KeyManagementClient | None, metadata: T
         if kms_client is None:
             raise ValueError(f"Table has {TableProperties.ENCRYPTION_KEY_ID} set but no key management client was provided")
 
-        key_length = property_as_int(metadata.properties, TableProperties.ENCRYPTION_DATA_KEY_LENGTH)
-        key_size = AesKeySize.from_key_length(
-            TableProperties.ENCRYPTION_DATA_KEY_LENGTH_DEFAULT if key_length is None else key_length
-        )
-
-        return cls(kms_client, table_key_id, metadata.encryption_keys, key_size)
+        return cls(kms_client, table_key_id, metadata.encryption_keys, _data_key_size(metadata))
 
     @property
     def encryption_keys(self) -> list[EncryptedKey]:
@@ -170,6 +165,23 @@ def _unwrap_kek(self, kek: EncryptedKey) -> SecureKey:
         return plaintext
 
 
+def _new_data_file_key_metadata(metadata: TableMetadata) -> StandardKeyMetadata | None:
+    """Generate key metadata for a new data file, or return None if the table is not encrypted."""
+    from pyiceberg.table import TableProperties
+
+    if TableProperties.ENCRYPTION_KEY_ID not in metadata.properties:
+        return None
+
+    return StandardKeyMetadata.generate(_data_key_size(metadata))
+
+
+def _data_key_size(metadata: TableMetadata) -> AesKeySize:
+    from pyiceberg.table import TableProperties
+
+    key_length = property_as_int(metadata.properties, TableProperties.ENCRYPTION_DATA_KEY_LENGTH)
+    return AesKeySize.from_key_length(TableProperties.ENCRYPTION_DATA_KEY_LENGTH_DEFAULT if key_length is None else key_length)
+
+
 def _now_ms() -> int:
     return time.time_ns() // 1_000_000
 

pyiceberg/io/fileformat.py

@@ -33,6 +33,7 @@
 if TYPE_CHECKING:
     import pyarrow as pa
 
+    from pyiceberg.encryption.key_metadata import StandardKeyMetadata
     from pyiceberg.io.pyarrow import StatsAggregator
 
 
@@ -159,7 +160,10 @@ def create_writer(
         output_file: OutputFile,
         file_schema: Schema,
         properties: Properties,
-    ) -> FileFormatWriter: ...
+        key_metadata: StandardKeyMetadata | None = None,
+    ) -> FileFormatWriter:
+        """Create a writer for one file, encrypting it with `key_metadata` when it is set."""
+        ...
 
     def add_field_metadata(self, field: NestedField, metadata: dict[bytes, bytes], include_field_ids: bool) -> None:
         """Add format-specific Arrow field metadata."""

pyiceberg/io/pyarrow.py

@@ -69,6 +69,7 @@
 
 from pyiceberg.conversions import to_bytes
 from pyiceberg.encryption.key_metadata import StandardKeyMetadata
+from pyiceberg.encryption.manager import _new_data_file_key_metadata
 from pyiceberg.exceptions import ResolveError
 from pyiceberg.expressions import AlwaysTrue, BooleanExpression, BoundIsNaN, BoundIsNull, BoundTerm, Not, Or
 from pyiceberg.expressions.literals import Literal
@@ -1175,6 +1176,16 @@ def _parquet_decryption_properties(key_metadata: bytes) -> pq.FileDecryptionProp
     return create_decryption_properties(metadata.encryption_key, aad_prefix=metadata.aad_prefix)
 
 
+def _parquet_encryption_properties(key_metadata: StandardKeyMetadata) -> pq.FileEncryptionProperties:
+    try:
+        from pyarrow.parquet.encryption import create_encryption_properties
+    except ImportError as e:
+        raise ImportError(f"pyarrow >= 25.0.0 required to write encrypted Parquet files, found {pyarrow.__version__}") from e
+
+    # Java does not store the AAD prefix in the file, readers take it from key metadata
+    return create_encryption_properties(key_metadata.encryption_key, aad_prefix=key_metadata.aad_prefix, store_aad_prefix=False)
+
+
 def _read_deletes(io: FileIO, data_file: DataFile) -> dict[str, pa.ChunkedArray]:
     _check_encryption_supported(data_file)
     if data_file.file_format == FileFormat.PARQUET:
@@ -2764,13 +2775,21 @@ def data_file_statistics_from_parquet_metadata(
 class ParquetFormatWriter(FileFormatWriter):
     """Writes Arrow tables to a Parquet file."""
 
-    def __init__(self, output_file: OutputFile, file_schema: Schema, properties: Properties) -> None:
+    def __init__(
+        self,
+        output_file: OutputFile,
+        file_schema: Schema,
+        properties: Properties,
+        key_metadata: StandardKeyMetadata | None = None,
+    ) -> None:
         self._output_file = output_file
         self._file_schema = file_schema
         self._properties = properties
         self._writer: pq.ParquetWriter | None = None
         self._fos: OutputStream | None = None
         self._parquet_writer_kwargs = _get_parquet_writer_kwargs(properties)
+        if key_metadata is not None:
+            self._parquet_writer_kwargs["encryption_properties"] = _parquet_encryption_properties(key_metadata)
         self._row_group_size = property_as_int(
             properties=properties,
             property_name=TableProperties.PARQUET_ROW_GROUP_LIMIT,
@@ -2823,8 +2842,9 @@ def create_writer(
         output_file: OutputFile,
         file_schema: Schema,
         properties: Properties,
+        key_metadata: StandardKeyMetadata | None = None,
     ) -> ParquetFormatWriter:
-        return ParquetFormatWriter(output_file, file_schema, properties)
+        return ParquetFormatWriter(output_file, file_schema, properties, key_metadata)
 
     def add_field_metadata(self, field: NestedField, metadata: dict[bytes, bytes], include_field_ids: bool) -> None:
         if include_field_ids:
@@ -2871,7 +2891,10 @@ def write_data_file(task: WriteTask) -> DataFile:
             partition_key=task.partition_key,
         )
         fo = io.new_output(file_path)
-        writer = format_model.create_writer(fo, file_schema, table_metadata.properties)
+        if (key_metadata := _new_data_file_key_metadata(table_metadata)) is None:
+            writer = format_model.create_writer(fo, file_schema, table_metadata.properties)
+        else:
+            writer = format_model.create_writer(fo, file_schema, table_metadata.properties, key_metadata=key_metadata)
         with writer:
             writer.write(arrow_table)
         statistics = writer.result()
@@ -2889,7 +2912,7 @@ def write_data_file(task: WriteTask) -> DataFile:
             # Just copy these from the table for now
             spec_id=table_metadata.default_spec_id,
             equality_ids=None,
-            key_metadata=None,
+            key_metadata=key_metadata.to_bytes() if key_metadata is not None else None,
             **statistics.to_serialized_dict(),
         )
 

tests/io/test_fileformat.py

@@ -48,7 +48,7 @@ def format(self) -> FileFormat:
         def file_extension(self) -> str:
             return "orc"
 
-        def create_writer(self, output_file: Any, file_schema: Any, properties: Any) -> Any:
+        def create_writer(self, output_file: Any, file_schema: Any, properties: Any, key_metadata: Any = None) -> Any:
             raise NotImplementedError
 
         def add_field_metadata(self, field: Any, metadata: Any, include_field_ids: bool) -> None:
@@ -89,7 +89,7 @@ def format(self) -> FileFormat:
     def file_extension(self) -> str:
         return "orc"
 
-    def create_writer(self, output_file: Any, file_schema: Any, properties: Any) -> Any:
+    def create_writer(self, output_file: Any, file_schema: Any, properties: Any, key_metadata: Any = None) -> Any:
         raise NotImplementedError
 
     def add_field_metadata(self, field: Any, metadata: Any, include_field_ids: bool) -> None:

tests/io/test_pyarrow.py

@@ -3678,6 +3678,50 @@ def test_write_file_parquet_round_trip(tmp_path: Path, table_schema_simple: Sche
     assert result.column_names == ["foo", "bar", "baz"]
 
 
+def test_write_file_encrypts_parquet_for_encrypted_table(tmp_path: Path, table_schema_simple: Schema) -> None:
+    from pyarrow.parquet import encryption
+
+    _, task = _simple_write_task_and_metadata(tmp_path, table_schema_simple, {})
+    table_metadata = TableMetadataV3(
+        location=f"file://{tmp_path}",
+        last_column_id=3,
+        current_schema_id=1,
+        schemas=[table_schema_simple],
+        partition_specs=[PartitionSpec()],
+        properties={TableProperties.ENCRYPTION_KEY_ID: "master-1", TableProperties.ENCRYPTION_DATA_KEY_LENGTH: "32"},
+        next_row_id=0,
+    )
+
+    data_file, other_data_file = write_file(
+        io=PyArrowFileIO(), table_metadata=table_metadata, tasks=iter([task, replace(task, task_id=1)])
+    )
+
+    assert data_file.key_metadata is not None and other_data_file.key_metadata is not None
+    key_metadata = StandardKeyMetadata.from_bytes(data_file.key_metadata)
+    other_key_metadata = StandardKeyMetadata.from_bytes(other_data_file.key_metadata)
+    assert len(key_metadata.encryption_key) == 32
+    assert key_metadata.encryption_key != other_key_metadata.encryption_key
+    assert key_metadata.aad_prefix != other_key_metadata.aad_prefix
+    file_path = data_file.file_path.removeprefix("file://")
+    with open(file_path, "rb") as f:
+        assert f.read(4) == b"PARE"
+
+    with pytest.raises(OSError, match="AAD prefix used for file encryption, but not stored in file"):
+        pq.read_table(file_path, decryption_properties=encryption.create_decryption_properties(key_metadata.encryption_key))
+
+    batches = _task_to_record_batches(
+        PyArrowFileIO(),
+        FileScanTask(data_file),
+        bound_row_filter=AlwaysTrue(),
+        projected_schema=table_schema_simple,
+        table_schema=table_schema_simple,
+        projected_field_ids={1, 2, 3},
+        positional_deletes=None,
+        case_sensitive=True,
+    )
+    assert pa.Table.from_batches(batches).to_pydict() == {"foo": ["a", "b", "c"], "bar": [1, 2, 3], "baz": [True, False, True]}
+
+
 def test__to_requested_schema_timestamps(
     arrow_table_schema_with_all_timestamp_precisions: pa.Schema,
     arrow_table_with_all_timestamp_precisions: pa.Table,