Keep DV range identity internal to scan aggregation

Restore native scan task sets by routing deletion vectors to their referenced data files. Preserve REST range references and cover local and REST scans of packed DV fixtures.

Author: KaiqiJinWow

pyiceberg/io/pyarrow.py

@@ -147,7 +147,6 @@
     visit_with_partner,
 )
 from pyiceberg.table import DOWNCAST_NS_TIMESTAMP_TO_US_ON_WRITE, TableProperties
-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
@@ -1716,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 = DeleteFileSet(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,10 +45,17 @@
     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 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
@@ -2233,17 +2240,17 @@ class FileScanTask(ScanTask):
     """Task representing a data file and its corresponding delete files."""
 
     file: DataFile
-    delete_files: DeleteFileSet
+    delete_files: set[DataFile]
     residual: BooleanExpression
 
     def __init__(
         self,
         data_file: DataFile,
-        delete_files: Iterable[DataFile] | None = None,
+        delete_files: set[DataFile] | None = None,
         residual: BooleanExpression = ALWAYS_TRUE,
     ) -> None:
         self.file = data_file
-        self.delete_files = DeleteFileSet(delete_files if delete_files is not None else [])
+        self.delete_files = delete_files or set()
         self.residual = residual
 
     @staticmethod
@@ -2267,13 +2274,13 @@ def from_rest_response(
 
         data_file = _rest_file_to_data_file(rest_task.data_file)
 
-        resolved_deletes: list[DataFile] = []
+        resolved_deletes: set[DataFile] = set()
         if rest_task.delete_file_references:
             for idx in rest_task.delete_file_references:
                 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.append(_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,
@@ -2282,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
@@ -2297,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,
@@ -2310,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.py

@@ -21,7 +21,6 @@
 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
@@ -116,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())
@@ -126,25 +125,26 @@ 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) -> DeleteFileSet:
+    def for_data_file(self, seq_num: int, data_file: DataFile, partition_key: Record | None = None) -> set[DataFile]:
         if self.is_empty():
-            return DeleteFileSet()
+            return set()
 
-        deletes: list[DataFile] = []
+        # 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
 
         key = _partition_key(spec_id, partition_key)
         partition_deletes = self._by_partition.get(key)
         if partition_deletes:
             for delete_file in partition_deletes.filter_by_seq(seq_num):
                 if _applies_to_data_file(delete_file, data_file):
-                    deletes.append(delete_file)
+                    deletes.add(delete_file)
 
         path_deletes = self._by_path.get(data_file.file_path)
         if path_deletes:
-            deletes.extend(path_deletes.filter_by_seq(seq_num))
+            deletes.update(path_deletes.filter_by_seq(seq_num))
 
-        return DeleteFileSet(deletes)
+        return deletes
 
     def referenced_delete_files(self) -> list[DataFile]:
         data_files: list[DataFile] = []

pyiceberg/table/delete_file_index.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: