feat(encryption): [11/N] Read encrypted Parquet data files

pyiceberg/io/pyarrow.py

@@ -68,6 +68,7 @@
 from typing_extensions import override
 
 from pyiceberg.conversions import to_bytes
+from pyiceberg.encryption.key_metadata import StandardKeyMetadata
 from pyiceberg.exceptions import ResolveError
 from pyiceberg.expressions import AlwaysTrue, BooleanExpression, BoundIsNaN, BoundIsNull, BoundTerm, Not, Or
 from pyiceberg.expressions.literals import Literal
@@ -1146,9 +1147,10 @@ def _expression_to_complementary_pyarrow(expr: BooleanExpression, schema: Schema
     return expression_to_pyarrow(preserve_expr, schema)
 
 
-@lru_cache
-def _get_file_format(file_format: FileFormat, **kwargs: dict[str, Any]) -> ds.FileFormat:
+def _create_arrow_read_format(file_format: FileFormat, key_metadata: bytes | None, **kwargs: Any) -> ds.FileFormat:
     if file_format == FileFormat.PARQUET:
+        if key_metadata is not None:
+            kwargs["decryption_properties"] = _parquet_decryption_properties(key_metadata)
         return ds.ParquetFileFormat(**kwargs)
     elif file_format == FileFormat.ORC:
         # ORC doesn't support pre_buffer and buffer_size parameters
@@ -1158,11 +1160,31 @@ def _get_file_format(file_format: FileFormat, **kwargs: dict[str, Any]) -> ds.Fi
         raise ValueError(f"Unsupported file format: {file_format}")
 
 
+def _check_encryption_supported(data_file: DataFile) -> None:
+    if data_file.key_metadata is not None and data_file.file_format != FileFormat.PARQUET:
+        raise ValueError(f"Cannot read encrypted {data_file.file_format} file {data_file.file_path}: only Parquet is supported")
+
+
+def _parquet_decryption_properties(key_metadata: bytes) -> pq.FileDecryptionProperties:
+    try:
+        from pyarrow.parquet.encryption import create_decryption_properties
+    except ImportError as e:
+        raise ImportError(f"pyarrow >= 25.0.0 required to read encrypted Parquet files, found {pyarrow.__version__}") from e
+
+    metadata = StandardKeyMetadata.from_bytes(key_metadata)
+    return create_decryption_properties(metadata.encryption_key, aad_prefix=metadata.aad_prefix)
+
+
 def _read_deletes(io: FileIO, data_file: DataFile) -> dict[str, pa.ChunkedArray]:
+    _check_encryption_supported(data_file)
     if data_file.file_format == FileFormat.PARQUET:
         with io.new_input(data_file.file_path).open() as fi:
-            delete_fragment = _get_file_format(
-                data_file.file_format, dictionary_columns=("file_path",), pre_buffer=True, buffer_size=ONE_MEGABYTE
+            delete_fragment = _create_arrow_read_format(
+                data_file.file_format,
+                data_file.key_metadata,
+                dictionary_columns=("file_path",),
+                pre_buffer=True,
+                buffer_size=ONE_MEGABYTE,
             ).make_fragment(fi)
             table = ds.Scanner.from_fragment(fragment=delete_fragment).to_table()
         table = table.unify_dictionaries()
@@ -1172,7 +1194,7 @@ def _read_deletes(io: FileIO, data_file: DataFile) -> dict[str, pa.ChunkedArray]
         }
     elif data_file.file_format == FileFormat.ORC:
         with io.new_input(data_file.file_path).open() as fi:
-            delete_fragment = _get_file_format(data_file.file_format).make_fragment(fi)
+            delete_fragment = _create_arrow_read_format(data_file.file_format, data_file.key_metadata).make_fragment(fi)
             table = ds.Scanner.from_fragment(fragment=delete_fragment).to_table()
             # For ORC, file_path columns are not dictionary-encoded, so we use unique() directly
             return {
@@ -1712,10 +1734,11 @@ def _task_to_record_batches(
     downcast_ns_timestamp_to_us: bool | None = None,
     dictionary_columns: tuple[str, ...] = (),
 ) -> Iterator[pa.RecordBatch]:
+    _check_encryption_supported(task.file)
     format_kwargs: dict[str, Any] = {"pre_buffer": True, "buffer_size": ONE_MEGABYTE * 8}
     if dictionary_columns and task.file.file_format == FileFormat.PARQUET:
         format_kwargs["dictionary_columns"] = dictionary_columns
-    arrow_format = _get_file_format(task.file.file_format, **format_kwargs)
+    arrow_format = _create_arrow_read_format(task.file.file_format, task.file.key_metadata, **format_kwargs)
     with io.new_input(task.file.file_path).open() as fin:
         fragment = arrow_format.make_fragment(fin)
         physical_schema = fragment.physical_schema

tests/io/test_pyarrow.py

@@ -25,6 +25,7 @@
 import warnings
 import zlib
 from collections.abc import Iterator
+from dataclasses import replace
 from datetime import UTC, date, datetime
 from pathlib import Path
 from typing import Any, Literal
@@ -42,6 +43,7 @@
 from pyroaring import BitMap
 
 from pyiceberg.catalog.rest.scan_planning import RESTDataFile, RESTDeleteFile, RESTFileScanTask, RESTPositionDeleteFile
+from pyiceberg.encryption.key_metadata import StandardKeyMetadata
 from pyiceberg.exceptions import ResolveError
 from pyiceberg.expressions import (
     AlwaysFalse,
@@ -80,6 +82,7 @@
     StatsAggregator,
     _check_pyarrow_schema_compatible,
     _ConvertToArrowSchema,
+    _create_arrow_read_format,
     _determine_partitions,
     _primitive_to_physical,
     _read_all_delete_files,
@@ -4132,6 +4135,111 @@ def test_task_to_record_batches_filter_after_positional_deletes_empty_result(tmp
     assert result_batches == []
 
 
+def test_task_to_record_batches_encrypted_by_iceberg_rust() -> None:
+    """The data file was encrypted by iceberg-rust, see tests/table/encryption/README.md."""
+    data_file = DataFile.from_args(
+        content=DataFileContent.DATA,
+        file_path=str(Path(__file__).parent.parent / "table" / "encryption" / "data-v3-encrypted.parquet"),
+        file_format=FileFormat.PARQUET,
+        partition={},
+        record_count=3,
+        file_size_in_bytes=1407,
+        key_metadata=bytes.fromhex("01207d30accaf199ff6f49ae7f3d65cbd8c6022033a904147e53c08d6f71c21c6638852a00"),
+    )
+    table_schema = Schema(NestedField(1, "id", LongType(), required=True), NestedField(2, "data", StringType(), required=False))
+
+    batches = _task_to_record_batches(
+        PyArrowFileIO(),
+        FileScanTask(data_file),
+        bound_row_filter=AlwaysTrue(),
+        projected_schema=table_schema,
+        table_schema=table_schema,
+        projected_field_ids={1, 2},
+        positional_deletes=None,
+        case_sensitive=True,
+    )
+
+    assert pa.Table.from_batches(batches).to_pylist() == [{"id": 1, "data": "a"}, {"id": 2, "data": None}, {"id": 3, "data": "c"}]
+
+
+def test_read_encrypted_parquet_with_position_deletes(tmp_path: Path) -> None:
+    from pyarrow.parquet import encryption
+
+    def write_encrypted(path: str, table: pa.Table, content: DataFileContent) -> DataFile:
+        key_metadata = StandardKeyMetadata.generate()
+        # Java does not store the AAD prefix in the file, so it must come from key metadata
+        encryption_properties = encryption.create_encryption_properties(
+            key_metadata.encryption_key, aad_prefix=key_metadata.aad_prefix, store_aad_prefix=False
+        )
+        pq.write_table(table, path, encryption_properties=encryption_properties)
+        return DataFile.from_args(
+            content=content,
+            file_path=path,
+            file_format=FileFormat.PARQUET,
+            partition={},
+            record_count=len(table),
+            file_size_in_bytes=os.path.getsize(path),
+            key_metadata=key_metadata.to_bytes(),
+        )
+
+    arrow_schema = pa.schema((pa.field("id", pa.int32(), nullable=True, metadata={PYARROW_PARQUET_FIELD_ID_KEY: "1"}),))
+    data_path = str(tmp_path / "data.parquet")
+    data_file = write_encrypted(
+        data_path, pa.table([pa.array([1, 2, 3, 4], type=pa.int32())], schema=arrow_schema), DataFileContent.DATA
+    )
+    delete_file = write_encrypted(
+        str(tmp_path / "deletes.parquet"),
+        pa.table({"file_path": [data_path], "pos": pa.array([2], type=pa.int64())}),
+        DataFileContent.POSITION_DELETES,
+    )
+    table_schema = Schema(NestedField(1, "id", IntegerType(), required=False))
+
+    deletes = _read_deletes(PyArrowFileIO(), delete_file)
+    batches = _task_to_record_batches(
+        PyArrowFileIO(),
+        FileScanTask(data_file),
+        bound_row_filter=AlwaysTrue(),
+        projected_schema=table_schema,
+        table_schema=table_schema,
+        projected_field_ids={1},
+        positional_deletes=[deletes[data_path]],
+        case_sensitive=True,
+    )
+
+    assert pa.Table.from_batches(batches).column("id").to_pylist() == [1, 2, 4]
+
+
+def test_read_encrypted_parquet_fails_closed(tmp_path: Path) -> None:
+    from pyarrow.parquet import encryption
+
+    key_metadata = StandardKeyMetadata.generate()
+    encrypted_path = str(tmp_path / "encrypted.parquet")
+    pq.write_table(
+        pa.table({"id": [1]}),
+        encrypted_path,
+        encryption_properties=encryption.create_encryption_properties(
+            key_metadata.encryption_key, aad_prefix=key_metadata.aad_prefix, store_aad_prefix=False
+        ),
+    )
+    plaintext_path = str(tmp_path / "plaintext.parquet")
+    pq.write_table(pa.table({"id": [1]}), plaintext_path)
+
+    for path, key, error in [
+        (encrypted_path, replace(key_metadata, encryption_key=bytes(16)).to_bytes(), "Failed decryption finalization"),
+        (encrypted_path, replace(key_metadata, aad_prefix=bytes(16)).to_bytes(), "Failed decryption finalization"),
+        (encrypted_path, None, "no decryption found in reader's properties"),
+        (plaintext_path, key_metadata.to_bytes(), "Applying decryption properties on plaintext file"),
+    ]:
+        with pytest.raises(OSError, match=error):
+            with open(path, "rb") as f:
+                _create_arrow_read_format(FileFormat.PARQUET, key).make_fragment(f).to_table()
+
+    for file_format in [FileFormat.ORC, FileFormat.PUFFIN]:
+        delete_file = DataFile.from_args(file_path="deletes", file_format=file_format, key_metadata=key_metadata.to_bytes())
+        with pytest.raises(ValueError, match="file deletes: only Parquet is supported"):
+            _read_deletes(PyArrowFileIO(), delete_file)
+
+
 def test_parse_location_defaults() -> None:
     """Test that parse_location uses defaults."""
 

tests/table/encryption/README.md

@@ -17,7 +17,7 @@
   ~ under the License.
 -->
 
-# Encrypted metadata test data
+# Encrypted test data
 
 `manifest-list-v3-encrypted.avro` is an empty manifest list encrypted by iceberg-rust, copied from its
 `crates/iceberg/testdata/manifests_lists/`. It was written by iceberg-rust's `EncryptionManager`, with a memory KMS
@@ -27,3 +27,7 @@ holding master key `master-1` = bytes `0x00` to `0x0f`. Its keys and snapshot ar
 `manifest-v3-encrypted.avro` is a v3 data manifest with two added entries, written and encrypted by iceberg-rust's
 `ManifestWriter` with key `0123456789abcdef` and AAD prefix `fedcba9876543210`. The key metadata that its manifest list
 entry would record is in `test_read_manifest_encrypted_by_iceberg_rust` in `tests/utils/test_manifest.py`.
+
+`data-v3-encrypted.parquet` is a Parquet data file with three rows, written and encrypted by iceberg-rust's
+`ParquetWriter` with a key from its `EncryptionManager`. The key metadata that its manifest entry would record is in
+`test_task_to_record_batches_encrypted_by_iceberg_rust` in `tests/io/test_pyarrow.py`.