Support range-based reads for deletion vectors

pyiceberg/io/pyarrow.py

@@ -147,11 +147,10 @@
     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.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 (
@@ -1149,10 +1148,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}")
 
@@ -1719,14 +1715,28 @@ def _task_to_record_batches(
             )
 
 
+@dataclass(frozen=True, slots=True)
+class _DeleteFileKey:
+    """Identity of a delete file or a referenced DV range."""
+
+    file_path: str
+    content_offset: int | None
+    content_size_in_bytes: int | None
+
+
 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]))
+    # DataFile equality uses only the path; DVs from different tasks may share a file.
+    unique_deletes: dict[_DeleteFileKey, DataFile] = {}
+    for task in tasks:
+        for delete_file in task.delete_files:
+            key = _DeleteFileKey(delete_file.file_path, delete_file.content_offset, delete_file.content_size_in_bytes)
+            unique_deletes.setdefault(key, delete_file)
     if len(unique_deletes) > 0:
         executor = ExecutorFactory.get_or_create()
         deletes_per_files: Iterator[dict[str, ChunkedArray]] = executor.map(
             lambda args: _read_deletes(*args),
-            [(io, delete_file) for delete_file in unique_deletes],
+            [(io, delete_file) for delete_file in unique_deletes.values()],
         )
         for delete in deletes_per_files:
             for file, arr in delete.items():

pyiceberg/table/__init__.py

@@ -45,7 +45,15 @@
     manifest_evaluator,
 )
 from pyiceberg.io import FileIO, load_file_io
-from pyiceberg.manifest import DataFile, DataFileContent, ManifestContent, ManifestEntry, ManifestEntryStatus, ManifestFile
+from pyiceberg.manifest import (
+    DataFile,
+    DataFileContent,
+    FileFormat,
+    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_index import DeleteFileIndex
@@ -2272,7 +2280,7 @@ def from_rest_response(
                 delete_file = delete_files[idx]
                 if isinstance(delete_file, RESTEqualityDeleteFile):
                     raise NotImplementedError(f"PyIceberg does not yet support equality deletes: {delete_file.file_path}")
-                resolved_deletes.add(_rest_file_to_data_file(delete_file))
+                resolved_deletes.add(_rest_file_to_data_file(delete_file, default_referenced_data_file=data_file.file_path))
 
         return FileScanTask(
             data_file=data_file,
@@ -2281,9 +2289,9 @@ def from_rest_response(
         )
 
 
-def _rest_file_to_data_file(rest_file: RESTContentFile) -> DataFile:
+def _rest_file_to_data_file(rest_file: RESTContentFile, default_referenced_data_file: str | None = None) -> DataFile:
     """Convert a REST content file to a manifest DataFile."""
-    from pyiceberg.catalog.rest.scan_planning import RESTDataFile
+    from pyiceberg.catalog.rest.scan_planning import RESTDataFile, RESTPositionDeleteFile
 
     if isinstance(rest_file, RESTDataFile):
         column_sizes = rest_file.column_sizes.to_dict() if rest_file.column_sizes else None
@@ -2296,6 +2304,17 @@ def _rest_file_to_data_file(rest_file: RESTContentFile) -> DataFile:
         null_value_counts = None
         nan_value_counts = None
 
+    referenced_data_file = None
+    if isinstance(rest_file, RESTPositionDeleteFile):
+        referenced_data_file = rest_file.referenced_data_file
+        # Infer the DV target from its task only for range reads, preserving whole-Puffin fallback.
+        if (
+            referenced_data_file is None
+            and rest_file.file_format == FileFormat.PUFFIN
+            and (rest_file.content_offset is not None or rest_file.content_size_in_bytes is not None)
+        ):
+            referenced_data_file = default_referenced_data_file
+
     data_file = DataFile.from_args(
         content=DataFileContent.from_rest_type(rest_file.content),
         file_path=rest_file.file_path,
@@ -2309,6 +2328,9 @@ def _rest_file_to_data_file(rest_file: RESTContentFile) -> DataFile:
         nan_value_counts=nan_value_counts,
         split_offsets=rest_file.split_offsets,
         sort_order_id=rest_file.sort_order_id,
+        referenced_data_file=referenced_data_file,
+        content_offset=rest_file.content_offset if isinstance(rest_file, RESTPositionDeleteFile) else None,
+        content_size_in_bytes=rest_file.content_size_in_bytes if isinstance(rest_file, RESTPositionDeleteFile) else None,
     )
     data_file.spec_id = rest_file.spec_id
     return data_file

pyiceberg/table/delete_file_index.py

@@ -115,7 +115,7 @@ def is_empty(self) -> bool:
     def add_delete_file(self, manifest_entry: ManifestEntry, partition_key: Record | None = None) -> None:
         delete_file = manifest_entry.data_file
         seq = manifest_entry.sequence_number or INITIAL_SEQUENCE_NUMBER
-        target_path = _referenced_data_file_path(delete_file)
+        target_path = delete_file.referenced_data_file or _referenced_data_file_path(delete_file)
 
         if target_path:
             deletes = self._by_path.setdefault(target_path, PositionDeletes())
@@ -129,6 +129,7 @@ def for_data_file(self, seq_num: int, data_file: DataFile, partition_key: Record
         if self.is_empty():
             return set()
 
+        # DVs are indexed by target path; a valid snapshot has at most one DV per data file.
         deletes: set[DataFile] = set()
         spec_id = data_file.spec_id or 0
 

pyiceberg/table/deletion_vector.py

@@ -15,6 +15,8 @@
 # specific language governing permissions and limitations
 # under the License.
 import math
+import struct
+import zlib
 from typing import TYPE_CHECKING
 
 from pyroaring import BitMap, FrozenBitMap
@@ -24,9 +26,18 @@
 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"
+_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 +88,101 @@ 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") -> tuple[int, int, str]:
+    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_JAVA_SIGNED:
+        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}")
+
+    return content_offset, content_size_in_bytes, referenced_data_file
+
+
+def _has_deletion_vector_content_reference(dv: "DataFile") -> bool:
+    """Return whether any manifest content-reference field is present.
+
+    Partial references must be validated as range reads. Whole-Puffin reads
+    are only used when all content-reference fields are absent.
+    """
+    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:
+    content_offset, content_size_in_bytes, referenced_data_file = _validate_deletion_vector_content(dv)
+
+    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/catalog/test_scan_planning_models.py

@@ -40,6 +40,7 @@
 )
 from pyiceberg.expressions import AlwaysTrue, EqualTo, Reference
 from pyiceberg.manifest import FileFormat
+from pyiceberg.table import FileScanTask
 
 TEST_URI = "https://iceberg-test-catalog/"
 
@@ -526,6 +527,80 @@ def test_plan_scan_with_delete_files(rest_scan_catalog: RestCatalog, requests_mo
     assert len(tasks[0].delete_files) == 1
 
 
+@pytest.mark.parametrize("include_referenced_data_file", [True, False])
+def test_plan_scan_with_shared_dv_file(
+    rest_scan_catalog: RestCatalog, requests_mock: Mocker, include_referenced_data_file: bool
+) -> None:
+    data_files = [_rest_data_file(file_path=f"s3://bucket/data-{index}.parquet") for index in range(2)]
+    delete_files = [
+        {
+            **_rest_position_delete_file(
+                file_path="s3://bucket/deletes.bin", file_format="puffin", file_size_in_bytes=400, content_offset=offset
+            ),
+            **({"referenced-data-file": data_file["file-path"]} if include_referenced_data_file else {}),
+        }
+        for data_file, offset in zip(data_files, [0, 200], strict=True)
+    ]
+    requests_mock.post(
+        f"{TEST_URI}v1/namespaces/db/tables/tbl/plan",
+        json={
+            "status": "completed",
+            "delete-files": delete_files,
+            "file-scan-tasks": [
+                {"data-file": data_file, "delete-file-references": [index, index]} for index, data_file in enumerate(data_files)
+            ],
+        },
+    )
+
+    tasks = rest_scan_catalog.plan_scan(("db", "tbl"), PlanTableScanRequest())
+
+    assert len(tasks) == 2
+    assert all(len(task.delete_files) == 1 for task in tasks)
+    deletes = [dv for task in tasks for dv in task.delete_files]
+    assert len(deletes) == 2
+    assert {(dv.referenced_data_file, dv.content_offset, dv.content_size_in_bytes) for dv in deletes} == {
+        ("s3://bucket/data-0.parquet", 0, 200),
+        ("s3://bucket/data-1.parquet", 200, 200),
+    }
+
+
+@pytest.mark.parametrize("missing_field", ["content-offset", "content-size-in-bytes"])
+@pytest.mark.parametrize("include_referenced_data_file", [True, False])
+def test_rest_task_preserves_partial_dv_reference(missing_field: str, include_referenced_data_file: bool) -> None:
+    fields = {
+        **_rest_position_delete_file(file_format="puffin", content_offset=0),
+        **({"referenced-data-file": "s3://bucket/data.parquet"} if include_referenced_data_file else {}),
+    }
+    fields.pop(missing_field)
+    task = FileScanTask.from_rest_response(
+        RESTFileScanTask.model_validate(
+            {"data-file": _rest_data_file(file_path="s3://bucket/data.parquet"), "delete-file-references": [0]}
+        ),
+        [RESTPositionDeleteFile.model_validate(fields)],
+    )
+    dv = next(iter(task.delete_files))
+
+    assert dv.referenced_data_file == "s3://bucket/data.parquet"
+    assert dv.content_offset == fields.get("content-offset")
+    assert dv.content_size_in_bytes == fields.get("content-size-in-bytes")
+
+
+@pytest.mark.parametrize("file_format", ["puffin", "parquet"])
+def test_rest_task_preserves_whole_delete_file_reference(file_format: str) -> None:
+    fields = _rest_position_delete_file(file_format=file_format)
+    fields.pop("content-offset")
+    fields.pop("content-size-in-bytes")
+    task = FileScanTask.from_rest_response(
+        RESTFileScanTask.model_validate({"data-file": _rest_data_file(), "delete-file-references": [0]}),
+        [RESTPositionDeleteFile.model_validate(fields)],
+    )
+    dv = next(iter(task.delete_files))
+
+    assert dv.referenced_data_file is None
+    assert dv.content_offset is None
+    assert dv.content_size_in_bytes is None
+
+
 def test_plan_scan_async_poll_completes(
     rest_scan_catalog: RestCatalog, requests_mock: Mocker, monkeypatch: pytest.MonkeyPatch
 ) -> None: