feat(encryption): [12/N] Read encrypted deletion vectors

pyiceberg/io/pyarrow.py

@@ -1161,8 +1161,10 @@ def _create_arrow_read_format(file_format: FileFormat, key_metadata: bytes | Non
 
 
 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")
+    if data_file.key_metadata is not None and data_file.file_format not in (FileFormat.PARQUET, FileFormat.PUFFIN):
+        raise ValueError(
+            f"Cannot read encrypted {data_file.file_format} file {data_file.file_path}: only Parquet and Puffin are supported"
+        )
 
 
 def _parquet_decryption_properties(key_metadata: bytes) -> pq.FileDecryptionProperties:

pyiceberg/table/deletion_vector.py

@@ -21,12 +21,14 @@
 
 from pyroaring import BitMap, FrozenBitMap
 
+from pyiceberg.encryption.io import EncryptedInputFile
+from pyiceberg.encryption.key_metadata import StandardKeyMetadata
 from pyiceberg.table.puffin import PuffinFile
 
 if TYPE_CHECKING:
     import pyarrow as pa
 
-    from pyiceberg.io import FileIO
+    from pyiceberg.io import FileIO, InputFile
     from pyiceberg.manifest import DataFile
 
 EMPTY_BITMAP = FrozenBitMap()
@@ -152,10 +154,18 @@ def _has_deletion_vector_content_reference(dv: "DataFile") -> bool:
     return dv.content_offset is not None or dv.content_size_in_bytes is not None or dv.referenced_data_file is not None
 
 
+def _new_input(io: "FileIO", dv: "DataFile") -> "InputFile":
+    """Return the Puffin file to read, decrypting it when it has key metadata so content offsets refer to the plaintext."""
+    input_file = io.new_input(dv.file_path)
+    if dv.key_metadata is not None:
+        return EncryptedInputFile(input_file, StandardKeyMetadata.from_bytes(dv.key_metadata))
+    return input_file
+
+
 def _read_deletion_vector(io: "FileIO", dv: "DataFile") -> DeletionVector:
     content_offset, content_size_in_bytes, referenced_data_file = _validate_deletion_vector_content(dv)
 
-    with io.new_input(dv.file_path).open() as fi:
+    with _new_input(io, dv).open() as fi:
         fi.seek(content_offset)
         payload = fi.read(content_size_in_bytes)
 
@@ -173,7 +183,7 @@ def read_deletion_vectors(io: "FileIO", dv: "DataFile") -> list[DeletionVector]:
     if _has_deletion_vector_content_reference(dv):
         return [_read_deletion_vector(io, dv)]
 
-    with io.new_input(dv.file_path).open() as fi:
+    with _new_input(io, dv).open() as fi:
         return deletion_vectors_from_puffin_file(PuffinFile(fi.read()))
 
 

tests/io/test_pyarrow.py

@@ -44,6 +44,7 @@
 
 from pyiceberg.catalog.rest.scan_planning import RESTDataFile, RESTDeleteFile, RESTFileScanTask, RESTPositionDeleteFile
 from pyiceberg.encryption.key_metadata import StandardKeyMetadata
+from pyiceberg.encryption.stream import _encrypt_ags1
 from pyiceberg.exceptions import ResolveError
 from pyiceberg.expressions import (
     AlwaysFalse,
@@ -2083,6 +2084,68 @@ def test_read_deletion_vector_blob_from_content_range(tmp_path: Path) -> None:
     assert deletes == {referenced_data_file: pa.chunked_array([[1, 3, 5]])}
 
 
+def _puffin_delete_file(file_path: str, **kwargs: Any) -> DataFile:
+    return DataFile.from_args(
+        _table_format_version=3,
+        content=DataFileContent.POSITION_DELETES,
+        file_path=file_path,
+        file_format=FileFormat.PUFFIN,
+        record_count=3,
+        **kwargs,
+    )
+
+
+def test_read_deletes_from_puffin_file_encrypted_by_iceberg_rust() -> None:
+    """The deletion vector was encrypted by iceberg-rust, see tests/table/encryption/README.md."""
+    file_path = str(Path(__file__).parent.parent / "table" / "encryption" / "dv-v3-encrypted.puffin")
+    referenced_data_file = "s3://bucket/table/data/00000.parquet"
+    key_metadata = bytes.fromhex("01203031323334353637383961626364656602206665646362613938373635343332313002ea04")
+
+    for dv in [
+        _puffin_delete_file(file_path, key_metadata=key_metadata),
+        _puffin_delete_file(
+            file_path,
+            key_metadata=key_metadata,
+            referenced_data_file=referenced_data_file,
+            content_offset=4,
+            content_size_in_bytes=46,
+        ),
+    ]:
+        assert _read_deletes(PyArrowFileIO(), dv) == {referenced_data_file: pa.chunked_array([[1, 3, 5]])}
+
+
+def test_read_deletes_from_encrypted_puffin_file(tmp_path: Path) -> None:
+    referenced_data_file = f"{tmp_path}/data.parquet"
+    delete_file_path = f"{tmp_path}/deletes.puffin"
+    puffin_payload = _deletion_vector_puffin_payload(referenced_data_file)
+    blob = PuffinFile(puffin_payload).footer.blobs[0]
+    key_metadata = StandardKeyMetadata.generate()
+    encrypted = _encrypt_ags1(puffin_payload, key_metadata)
+    key_metadata = replace(key_metadata, file_length=len(encrypted))
+    wrong_key = replace(key_metadata, encryption_key=bytes(16))
+
+    with open(delete_file_path, "wb") as f:
+        f.write(encrypted)
+
+    for content_reference, plaintext_error in [
+        ({}, "Incorrect magic bytes"),
+        (
+            {"referenced_data_file": referenced_data_file, "content_offset": blob.offset, "content_size_in_bytes": blob.length},
+            "Invalid bitmap data length",
+        ),
+    ]:
+        dv = _puffin_delete_file(delete_file_path, key_metadata=key_metadata.to_bytes(), **content_reference)
+        assert _read_deletes(PyArrowFileIO(), dv) == {referenced_data_file: pa.chunked_array([[1, 3, 5]])}
+
+        with pytest.raises(ValueError, match="GCM tag check failed"):
+            _read_deletes(
+                PyArrowFileIO(), _puffin_delete_file(delete_file_path, key_metadata=wrong_key.to_bytes(), **content_reference)
+            )
+
+        with pytest.raises(ValueError, match=plaintext_error):
+            _read_deletes(PyArrowFileIO(), _puffin_delete_file(delete_file_path, **content_reference))
+
+
 def test_read_all_delete_files_keeps_multiple_dv_content_ranges_for_same_path(tmp_path: Path) -> None:
     referenced_data_files = [f"{tmp_path}/data-{index}.parquet" for index in range(2)]
     first_dv_blob = _deletion_vector_blob(_deletion_vector_bitmap_payload([1, 3]))
@@ -4234,10 +4297,9 @@ def test_read_encrypted_parquet_fails_closed(tmp_path: Path) -> None:
             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)
+    delete_file = DataFile.from_args(file_path="deletes", file_format=FileFormat.ORC, key_metadata=key_metadata.to_bytes())
+    with pytest.raises(ValueError, match="file deletes: only Parquet and Puffin are supported"):
+        _read_deletes(PyArrowFileIO(), delete_file)
 
 
 def test_parse_location_defaults() -> None:

tests/table/encryption/README.md

@@ -31,3 +31,8 @@ entry would record is in `test_read_manifest_encrypted_by_iceberg_rust` in `test
 `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`.
+
+`dv-v3-encrypted.puffin` is a Puffin file with one deletion vector, deleting positions 1, 3 and 5 of
+`s3://bucket/table/data/00000.parquet`. It was written and encrypted by iceberg-rust's `PuffinWriter` with key
+`0123456789abcdef` and AAD prefix `fedcba9876543210`. The key metadata that its manifest entry would record is in
+`test_read_deletes_from_puffin_file_encrypted_by_iceberg_rust` in `tests/io/test_pyarrow.py`.