Support range-based reads for deletion vectors

pyiceberg/io/pyarrow.py

@@ -146,11 +146,11 @@
     visit_with_partner,
 )
 from pyiceberg.table import DOWNCAST_NS_TIMESTAMP_TO_US_ON_WRITE, TableProperties
-from pyiceberg.table.deletion_vector import deletion_vectors_from_puffin_file
+from pyiceberg.table.delete_file import DeleteFileSet
+from pyiceberg.table.deletion_vector import read_deletion_vectors
 from pyiceberg.table.locations import load_location_provider
 from pyiceberg.table.metadata import TableMetadata
 from pyiceberg.table.name_mapping import NameMapping, apply_name_mapping
-from pyiceberg.table.puffin import PuffinFile
 from pyiceberg.transforms import IdentityTransform, TruncateTransform
 from pyiceberg.typedef import EMPTY_DICT, Properties, Record, TableVersion
 from pyiceberg.types import (
@@ -1141,10 +1141,7 @@ def _read_deletes(io: FileIO, data_file: DataFile) -> dict[str, pa.ChunkedArray]
                 for path in table.column("file_path").unique()
             }
     elif data_file.file_format == FileFormat.PUFFIN:
-        with io.new_input(data_file.file_path).open() as fi:
-            payload = fi.read()
-
-        return {dv.referenced_data_file: dv.to_vector() for dv in deletion_vectors_from_puffin_file(PuffinFile(payload))}
+        return {dv.referenced_data_file: dv.to_vector() for dv in read_deletion_vectors(io, data_file)}
     else:
         raise ValueError(f"Delete file format not supported: {data_file.file_format}")
 
@@ -1713,7 +1710,7 @@ def _task_to_record_batches(
 
 def _read_all_delete_files(io: FileIO, tasks: Iterable[FileScanTask]) -> dict[str, list[ChunkedArray]]:
     deletes_per_file: dict[str, list[ChunkedArray]] = {}
-    unique_deletes = set(itertools.chain.from_iterable([task.delete_files for task in tasks]))
+    unique_deletes = DeleteFileSet(itertools.chain.from_iterable(task.delete_files for task in tasks))
     if len(unique_deletes) > 0:
         executor = ExecutorFactory.get_or_create()
         deletes_per_files: Iterator[dict[str, ChunkedArray]] = executor.map(

pyiceberg/manifest.py

@@ -532,6 +532,18 @@ def equality_ids(self) -> list[int] | None:
     def sort_order_id(self) -> int | None:
         return self._data[15]
 
+    @property
+    def referenced_data_file(self) -> str | None:
+        return self._data[17] if len(self._data) > 17 else None
+
+    @property
+    def content_offset(self) -> int | None:
+        return self._data[18] if len(self._data) > 18 else None
+
+    @property
+    def content_size_in_bytes(self) -> int | None:
+        return self._data[19] if len(self._data) > 19 else None
+
     # Spec ID should not be stored in the file
     _spec_id: int
 

pyiceberg/table/__init__.py

@@ -44,6 +44,7 @@
 from pyiceberg.manifest import DataFile, DataFileContent, ManifestContent, ManifestEntry, ManifestEntryStatus, ManifestFile
 from pyiceberg.partitioning import PARTITION_FIELD_ID_START, UNPARTITIONED_PARTITION_SPEC, PartitionKey, PartitionSpec
 from pyiceberg.schema import Schema
+from pyiceberg.table.delete_file import DeleteFileSet
 from pyiceberg.table.delete_file_index import DeleteFileIndex
 from pyiceberg.table.inspect import InspectTable
 from pyiceberg.table.locations import LocationProvider, load_location_provider
@@ -2058,17 +2059,17 @@ class FileScanTask(ScanTask):
     """Task representing a data file and its corresponding delete files."""
 
     file: DataFile
-    delete_files: set[DataFile]
+    delete_files: DeleteFileSet
     residual: BooleanExpression
 
     def __init__(
         self,
         data_file: DataFile,
-        delete_files: set[DataFile] | None = None,
+        delete_files: Iterable[DataFile] | None = None,
         residual: BooleanExpression = ALWAYS_TRUE,
     ) -> None:
         self.file = data_file
-        self.delete_files = delete_files or set()
+        self.delete_files = DeleteFileSet(delete_files if delete_files is not None else [])
         self.residual = residual
 
     @staticmethod
@@ -2092,7 +2093,7 @@ def from_rest_response(
 
         data_file = _rest_file_to_data_file(rest_task.data_file)
 
-        resolved_deletes: set[DataFile] = set()
+        resolved_deletes = DeleteFileSet()
         if rest_task.delete_file_references:
             for idx in rest_task.delete_file_references:
                 delete_file = delete_files[idx]

pyiceberg/table/delete_file.py

@@ -0,0 +1,83 @@
+# 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.
+from __future__ import annotations
+
+from collections.abc import Iterable, Iterator, MutableSet
+from typing import Any
+
+from pyiceberg.manifest import DataFile
+
+_DeleteFileKey = tuple[str, int | None, int | None]
+
+
+def _delete_file_key(delete_file: DataFile) -> _DeleteFileKey:
+    return delete_file.file_path, delete_file.content_offset, delete_file.content_size_in_bytes
+
+
+class DeleteFileSet(MutableSet[DataFile]):
+    """Set-like delete-file collection keyed by location and content range."""
+
+    _files: dict[_DeleteFileKey, DataFile]
+
+    def __init__(self, delete_files: Iterable[DataFile] = ()) -> None:
+        self._files = {}
+        for delete_file in delete_files:
+            self.add(delete_file)
+
+    def __contains__(self, delete_file: object) -> bool:
+        """Return whether the delete file is present."""
+        return isinstance(delete_file, DataFile) and _delete_file_key(delete_file) in self._files
+
+    def __iter__(self) -> Iterator[DataFile]:
+        """Return an iterator over delete files."""
+        return iter(self._files.values())
+
+    def __len__(self) -> int:
+        """Return the number of delete files."""
+        return len(self._files)
+
+    def add(self, delete_file: DataFile) -> None:
+        self._files[_delete_file_key(delete_file)] = delete_file
+
+    def discard(self, delete_file: DataFile) -> None:
+        self._files.pop(_delete_file_key(delete_file), None)
+
+    def update(self, delete_files: Iterable[DataFile]) -> None:
+        for delete_file in delete_files:
+            self.add(delete_file)
+
+    def __repr__(self) -> str:
+        """Return a string representation of the delete file set."""
+        return f"{type(self).__name__}({list(self)!r})"
+
+    def __eq__(self, other: Any) -> bool:
+        """Compare delete file sets by delete file identity."""
+        if isinstance(other, DeleteFileSet):
+            return self._files.keys() == other._files.keys()
+
+        if not isinstance(other, Iterable):
+            return False
+
+        other_keys: set[_DeleteFileKey] = set()
+        other_count = 0
+        for delete_file in other:
+            if not isinstance(delete_file, DataFile):
+                return False
+            other_keys.add(_delete_file_key(delete_file))
+            other_count += 1
+
+        return len(other_keys) == other_count and set(self._files) == other_keys

pyiceberg/table/delete_file_index.py

@@ -21,6 +21,7 @@
 from pyiceberg.expressions import EqualTo
 from pyiceberg.expressions.visitors import _InclusiveMetricsEvaluator
 from pyiceberg.manifest import INITIAL_SEQUENCE_NUMBER, POSITIONAL_DELETE_SCHEMA, DataFile, ManifestEntry
+from pyiceberg.table.delete_file import DeleteFileSet
 from pyiceberg.typedef import Record
 
 PATH_FIELD_ID = 2147483546
@@ -125,11 +126,11 @@ def add_delete_file(self, manifest_entry: ManifestEntry, partition_key: Record |
             deletes = self._by_partition.setdefault(key, PositionDeletes())
             deletes.add(delete_file, seq)
 
-    def for_data_file(self, seq_num: int, data_file: DataFile, partition_key: Record | None = None) -> set[DataFile]:
+    def for_data_file(self, seq_num: int, data_file: DataFile, partition_key: Record | None = None) -> DeleteFileSet:
         if self.is_empty():
-            return set()
+            return DeleteFileSet()
 
-        deletes: set[DataFile] = set()
+        deletes = DeleteFileSet()
         spec_id = data_file.spec_id or 0
 
         key = _partition_key(spec_id, partition_key)

pyiceberg/table/deletion_vector.py

@@ -15,7 +15,9 @@
 # specific language governing permissions and limitations
 # under the License.
 import math
-from typing import TYPE_CHECKING
+import struct
+import zlib
+from typing import TYPE_CHECKING, cast
 
 from pyroaring import BitMap, FrozenBitMap
 
@@ -24,9 +26,19 @@
 if TYPE_CHECKING:
     import pyarrow as pa
 
+    from pyiceberg.io import FileIO
+    from pyiceberg.manifest import DataFile
+
 EMPTY_BITMAP = FrozenBitMap()
 MAX_JAVA_SIGNED = int(math.pow(2, 31)) - 1
 PROPERTY_REFERENCED_DATA_FILE = "referenced-data-file"
+_MAX_DELETION_VECTOR_CONTENT_SIZE = 2**31 - 1
+_DV_BLOB_LENGTH = struct.Struct(">I")
+_DV_BLOB_MAGIC = struct.Struct("<I")
+_DV_BLOB_CRC = struct.Struct(">I")
+_DV_BLOB_MAGIC_NUMBER = 1681511377
+_ROARING_BITMAP_COUNT_SIZE_BYTES = 8
+_DV_BLOB_MIN_SIZE_BYTES = _DV_BLOB_LENGTH.size + _DV_BLOB_MAGIC.size + _ROARING_BITMAP_COUNT_SIZE_BYTES + _DV_BLOB_CRC.size
 
 
 class DeletionVector:
@@ -77,17 +89,99 @@ def to_vector(self) -> "pa.ChunkedArray":
         return self._bitmaps_to_chunked_array(self._bitmaps)
 
 
-def _extract_vector_payload(blob_payload: bytes) -> bytes:
-    """Strip deletion-vector-v1 blob framing: length(4 big-endian) + DV magic(4) ... CRC(4 big-endian)."""
-    length_prefix = int.from_bytes(blob_payload[0:4], "big")
-    return blob_payload[8 : 4 + length_prefix]
+def _deserialize_dv_blob(blob: bytes, record_count: int | None = None) -> list[BitMap]:
+    # The DV blob encoding matches Iceberg Java's BitmapPositionDeleteIndex:
+    # 4-byte big-endian bitmap-data length, 4-byte little-endian magic number,
+    # portable Roaring bitmap data, and 4-byte big-endian CRC-32.
+    if len(blob) < _DV_BLOB_MIN_SIZE_BYTES:
+        raise ValueError(f"Invalid deletion vector blob length: {len(blob)}")
+
+    bitmap_data_length = _DV_BLOB_LENGTH.unpack_from(blob)[0]
+    expected_bitmap_data_length = len(blob) - _DV_BLOB_LENGTH.size - _DV_BLOB_CRC.size
+    if bitmap_data_length != expected_bitmap_data_length:
+        raise ValueError(f"Invalid bitmap data length: {bitmap_data_length}, expected {expected_bitmap_data_length}")
+
+    bitmap_data_offset = _DV_BLOB_LENGTH.size
+    crc_offset = bitmap_data_offset + bitmap_data_length
+    bitmap_data = blob[bitmap_data_offset:crc_offset]
+
+    magic_number = _DV_BLOB_MAGIC.unpack_from(bitmap_data)[0]
+    if magic_number != _DV_BLOB_MAGIC_NUMBER:
+        raise ValueError(f"Invalid magic number: {magic_number}, expected {_DV_BLOB_MAGIC_NUMBER}")
+
+    checksum = zlib.crc32(bitmap_data) & 0xFFFFFFFF
+    expected_checksum = _DV_BLOB_CRC.unpack_from(blob, crc_offset)[0]
+    if checksum != expected_checksum:
+        raise ValueError("Invalid CRC")
+
+    bitmaps = DeletionVector._deserialize_bitmap(bitmap_data[_DV_BLOB_MAGIC.size :])
+    if record_count is not None:
+        cardinality = sum(len(bitmap) for bitmap in bitmaps)
+        if cardinality != record_count:
+            raise ValueError(f"Invalid cardinality: {cardinality}, expected {record_count}")
+
+    return bitmaps
+
+
+def _validate_deletion_vector_content(dv: "DataFile") -> None:
+    content_offset = dv.content_offset
+    content_size_in_bytes = dv.content_size_in_bytes
+    referenced_data_file = dv.referenced_data_file
+
+    if content_offset is None:
+        raise ValueError(f"Invalid deletion vector, content offset is missing: {dv.file_path}")
+    if content_size_in_bytes is None:
+        raise ValueError(f"Invalid deletion vector, content size is missing: {dv.file_path}")
+    if content_offset < 0:
+        raise ValueError(f"Invalid deletion vector, content offset cannot be negative: {content_offset}")
+    if content_size_in_bytes < 0:
+        raise ValueError(f"Invalid deletion vector, content size cannot be negative: {content_size_in_bytes}")
+    if content_size_in_bytes > _MAX_DELETION_VECTOR_CONTENT_SIZE:
+        raise ValueError(f"Cannot read deletion vector larger than 2GB: {content_size_in_bytes}")
+    if referenced_data_file is None:
+        raise ValueError(f"Invalid deletion vector, referenced data file is missing: {dv.file_path}")
+
+
+def has_deletion_vector_content_reference(dv: "DataFile") -> bool:
+    """Return whether a deletion vector is described by manifest content-range metadata."""
+    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 _read_deletion_vector(io: "FileIO", dv: "DataFile") -> DeletionVector:
+    _validate_deletion_vector_content(dv)
+
+    content_offset = cast(int, dv.content_offset)
+    content_size_in_bytes = cast(int, dv.content_size_in_bytes)
+    referenced_data_file = cast(str, dv.referenced_data_file)
+
+    with io.new_input(dv.file_path).open() as fi:
+        fi.seek(content_offset)
+        payload = fi.read(content_size_in_bytes)
+
+    if len(payload) != content_size_in_bytes:
+        raise ValueError(f"Could not read deletion vector, expected {content_size_in_bytes} bytes, got {len(payload)}")
+
+    return DeletionVector(
+        referenced_data_file=referenced_data_file,
+        bitmaps=_deserialize_dv_blob(payload, dv.record_count),
+    )
+
+
+def read_deletion_vectors(io: "FileIO", dv: "DataFile") -> list[DeletionVector]:
+    """Read deletion vectors from a delete file or its manifest content range."""
+    if has_deletion_vector_content_reference(dv):
+        return [_read_deletion_vector(io, dv)]
+
+    with io.new_input(dv.file_path).open() as fi:
+        return deletion_vectors_from_puffin_file(PuffinFile(fi.read()))
 
 
 def deletion_vectors_from_puffin_file(puffin_file: PuffinFile) -> list[DeletionVector]:
+    """Read all deletion vectors stored in a Puffin file."""
     return [
         DeletionVector(
             referenced_data_file=blob.properties[PROPERTY_REFERENCED_DATA_FILE],
-            bitmaps=DeletionVector._deserialize_bitmap(_extract_vector_payload(puffin_file.get_blob_payload(blob))),
+            bitmaps=_deserialize_dv_blob(puffin_file.get_blob_payload(blob)),
         )
         for blob in puffin_file.footer.blobs
     ]

tests/avro/test_file.py

@@ -51,6 +51,7 @@
     LongType,
     NestedField,
     StringType,
+    StructType,
     TimestampType,
     TimestamptzType,
     TimeType,
@@ -89,7 +90,16 @@ def test_missing_schema() -> None:
 
 # helper function to serialize our objects to dicts to enable
 # direct comparison with the dicts returned by fastavro
-def todict(obj: Any) -> Any:
+def todict(obj: Any, struct: StructType | None = None) -> Any:
+    if struct is not None and isinstance(obj, Record):
+        return {
+            field.name: todict(
+                getattr(obj, field.name),
+                field.field_type if isinstance(field.field_type, StructType) else None,
+            )
+            for field in struct.fields
+            if hasattr(obj, field.name)
+        }
     if isinstance(obj, dict):
         data = []
         for k, v in obj.items():
@@ -157,7 +167,7 @@ def test_write_manifest_entry_with_iceberg_read_with_fastavro_v1() -> None:
 
             fa_entry = next(it)
 
-        v2_entry = todict(entry)
+        v2_entry = todict(entry, MANIFEST_ENTRY_SCHEMAS[2].as_struct())
 
         # These are not written in V1
         del v2_entry["sequence_number"]
@@ -222,7 +232,7 @@ def test_write_manifest_entry_with_iceberg_read_with_fastavro_v2() -> None:
 
             fa_entry = next(it)
 
-        assert todict(entry) == fa_entry
+        assert todict(entry, MANIFEST_ENTRY_SCHEMAS[2].as_struct()) == fa_entry
 
 
 @pytest.mark.parametrize("format_version", [1, 2])
@@ -260,7 +270,7 @@ def test_write_manifest_entry_with_fastavro_read_with_iceberg(format_version: Ta
         schema = AvroSchemaConversion().iceberg_to_avro(MANIFEST_ENTRY_SCHEMAS[format_version], schema_name="manifest_entry")
 
         with open(tmp_avro_file, "wb") as out:
-            writer(out, schema, [todict(entry)])
+            writer(out, schema, [todict(entry, MANIFEST_ENTRY_SCHEMAS[format_version].as_struct())])
 
         # Read as V2
         with avro.AvroFile[ManifestEntry](