diff --git a/pyiceberg/io/pyarrow.py b/pyiceberg/io/pyarrow.py index c36f1639d9..2b437f25f3 100644 --- a/pyiceberg/io/pyarrow.py +++ b/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(): diff --git a/pyiceberg/table/__init__.py b/pyiceberg/table/__init__.py index 9624eac981..d4f47678bf 100644 --- a/pyiceberg/table/__init__.py +++ b/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 diff --git a/pyiceberg/table/delete_file_index.py b/pyiceberg/table/delete_file_index.py index 3f513aabe5..91365d6ea2 100644 --- a/pyiceberg/table/delete_file_index.py +++ b/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 diff --git a/pyiceberg/table/deletion_vector.py b/pyiceberg/table/deletion_vector.py index 88fb3daf73..9c56af24a2 100644 --- a/pyiceberg/table/deletion_vector.py +++ b/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_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 ] diff --git a/tests/catalog/test_scan_planning_models.py b/tests/catalog/test_scan_planning_models.py index caf571d322..f88e24e92f 100644 --- a/tests/catalog/test_scan_planning_models.py +++ b/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: diff --git a/tests/io/test_pyarrow.py b/tests/io/test_pyarrow.py index b31c18949b..ee2e0be2b4 100644 --- a/tests/io/test_pyarrow.py +++ b/tests/io/test_pyarrow.py @@ -15,16 +15,19 @@ # specific language governing permissions and limitations # under the License. # pylint: disable=protected-access,unused-argument,redefined-outer-name +import json import logging import os +import struct import sys import tempfile import uuid import warnings +import zlib from collections.abc import Iterator from datetime import date, datetime, timezone from pathlib import Path -from typing import Any +from typing import Any, Literal from unittest.mock import MagicMock, patch from uuid import uuid4 @@ -36,7 +39,9 @@ import pytest from packaging import version from pyarrow.fs import AwsDefaultS3RetryStrategy, FileType, LocalFileSystem, S3FileSystem +from pyroaring import BitMap +from pyiceberg.catalog.rest.scan_planning import RESTDataFile, RESTDeleteFile, RESTFileScanTask, RESTPositionDeleteFile from pyiceberg.exceptions import ResolveError from pyiceberg.expressions import ( AlwaysFalse, @@ -75,6 +80,7 @@ _ConvertToArrowSchema, _determine_partitions, _primitive_to_physical, + _read_all_delete_files, _read_deletes, _task_to_record_batches, _to_requested_schema, @@ -87,12 +93,19 @@ schema_to_pyarrow, write_file, ) -from pyiceberg.manifest import DataFile, DataFileContent, FileFormat +from pyiceberg.manifest import DataFile, DataFileContent, FileFormat, ManifestEntry, ManifestEntryStatus from pyiceberg.partitioning import PartitionField, PartitionSpec from pyiceberg.schema import Schema, make_compatible_name, visit from pyiceberg.table import FileScanTask, TableProperties, WriteTask -from pyiceberg.table.metadata import TableMetadataV2 +from pyiceberg.table.delete_file_index import DeleteFileIndex +from pyiceberg.table.deletion_vector import ( + _DV_BLOB_MAGIC_NUMBER, + PROPERTY_REFERENCED_DATA_FILE, + deletion_vectors_from_puffin_file, +) +from pyiceberg.table.metadata import TableMetadataV2, TableMetadataV3 from pyiceberg.table.name_mapping import create_mapping_from_schema +from pyiceberg.table.puffin import MAGIC_BYTES, PuffinFile from pyiceberg.transforms import HourTransform, IdentityTransform from pyiceberg.typedef import UTF8, Properties, Record, TableVersion from pyiceberg.types import ( @@ -1822,6 +1835,319 @@ def test_read_deletes(deletes_file: str, request: pytest.FixtureRequest) -> None assert list(deletes.values())[0] == pa.chunked_array([[1, 3, 5]]) +def _deletion_vector_bitmap_payload(values: list[int] | None = None) -> bytes: + return (1).to_bytes(8, byteorder="little") + (0).to_bytes(4, byteorder="little") + BitMap(values or [1, 3, 5]).serialize() + + +def _deletion_vector_blob(bitmap_payload: bytes) -> bytes: + bitmap_data = struct.pack("I", len(bitmap_data)) + bitmap_data + struct.pack(">I", zlib.crc32(bitmap_data) & 0xFFFFFFFF) + + +def _deletion_vector_puffin_payload(referenced_data_file: str) -> bytes: + dv_blob = _deletion_vector_blob(_deletion_vector_bitmap_payload()) + footer_payload = json.dumps( + { + "blobs": [ + { + "type": "deletion-vector-v1", + "fields": [2147483546], + "snapshot-id": 1, + "sequence-number": 1, + "offset": len(MAGIC_BYTES) + 4, + "length": len(dv_blob), + "properties": {PROPERTY_REFERENCED_DATA_FILE: referenced_data_file}, + } + ], + "properties": {}, + } + ).encode() + + return ( + MAGIC_BYTES + + b"\x00\x00\x00\x00" + + dv_blob + + footer_payload + + len(footer_payload).to_bytes(4, byteorder="little") + + b"\x00\x00\x00\x00" + + MAGIC_BYTES + ) + + +def test_deletion_vectors_from_puffin_file(tmp_path: Path) -> None: + referenced_data_file = f"{tmp_path}/data.parquet" + puffin_payload = _deletion_vector_puffin_payload(referenced_data_file) + delete_file_path = f"{tmp_path}/deletes.puffin" + + with open(delete_file_path, "wb") as f: + f.write(puffin_payload) + + with open(delete_file_path, "rb") as f: + deletion_vectors = deletion_vectors_from_puffin_file(PuffinFile(f.read())) + + assert {dv.referenced_data_file: dv.to_vector() for dv in deletion_vectors} == { + referenced_data_file: pa.chunked_array([[1, 3, 5]]) + } + + +def test_read_deletes_from_whole_puffin_file(tmp_path: Path) -> None: + referenced_data_file = f"{tmp_path}/data.parquet" + delete_file_path = f"{tmp_path}/deletes.puffin" + + with open(delete_file_path, "wb") as f: + f.write(_deletion_vector_puffin_payload(referenced_data_file)) + + deletes = _read_deletes( + PyArrowFileIO(), + DataFile.from_args( + _table_format_version=3, + content=DataFileContent.POSITION_DELETES, + file_path=delete_file_path, + file_format=FileFormat.PUFFIN, + record_count=3, + ), + ) + + assert deletes == {referenced_data_file: pa.chunked_array([[1, 3, 5]])} + + +@pytest.mark.parametrize( + ("missing_field", "error"), + [ + ("content_offset", "content offset is missing"), + ("content_size_in_bytes", "content size is missing"), + ("referenced_data_file", "referenced data file is missing"), + ], +) +def test_read_deletes_rejects_partial_content_reference(tmp_path: Path, missing_field: str, error: str) -> 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] + + with open(delete_file_path, "wb") as f: + f.write(puffin_payload) + + dv = DataFile.from_args( + _table_format_version=3, + content=DataFileContent.POSITION_DELETES, + file_path=delete_file_path, + file_format=FileFormat.PUFFIN, + record_count=3, + content_offset=blob.offset if missing_field != "content_offset" else None, + content_size_in_bytes=blob.length if missing_field != "content_size_in_bytes" else None, + referenced_data_file=referenced_data_file if missing_field != "referenced_data_file" else None, + ) + + with pytest.raises(ValueError, match=error): + _read_deletes(PyArrowFileIO(), dv) + + +def test_read_deletion_vector_blob_from_content_range(tmp_path: Path) -> None: + referenced_data_file = f"{tmp_path}/data.parquet" + dv_blob = _deletion_vector_blob(_deletion_vector_bitmap_payload()) + prefix = b"\x01not-a-puffin-file" + delete_file_path = f"{tmp_path}/deletes.bin" + + with open(delete_file_path, "wb") as f: + f.write(prefix + dv_blob + b"trailing-bytes") + + deletes = _read_deletes( + PyArrowFileIO(), + DataFile.from_args( + _table_format_version=3, + content=DataFileContent.POSITION_DELETES, + file_path=delete_file_path, + file_format=FileFormat.PUFFIN, + record_count=3, + referenced_data_file=referenced_data_file, + content_offset=len(prefix), + content_size_in_bytes=len(dv_blob), + ), + ) + + assert deletes == {referenced_data_file: pa.chunked_array([[1, 3, 5]])} + + +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])) + second_dv_blob = _deletion_vector_blob(_deletion_vector_bitmap_payload([5])) + delete_file_path = f"{tmp_path}/deletes.bin" + + with open(delete_file_path, "wb") as f: + f.write(first_dv_blob + second_dv_blob) + + first_dv = DataFile.from_args( + _table_format_version=3, + content=DataFileContent.POSITION_DELETES, + file_path=delete_file_path, + file_format=FileFormat.PUFFIN, + record_count=2, + referenced_data_file=referenced_data_files[0], + content_offset=0, + content_size_in_bytes=len(first_dv_blob), + ) + second_dv = DataFile.from_args( + _table_format_version=3, + content=DataFileContent.POSITION_DELETES, + file_path=delete_file_path, + file_format=FileFormat.PUFFIN, + record_count=1, + referenced_data_file=referenced_data_files[1], + content_offset=len(first_dv_blob), + content_size_in_bytes=len(second_dv_blob), + ) + data_files = [ + DataFile.from_args( + content=DataFileContent.DATA, + file_path=referenced_data_file, + file_format=FileFormat.PARQUET, + record_count=10, + file_size_in_bytes=100, + ) + for referenced_data_file in referenced_data_files + ] + + deletes = _read_all_delete_files( + PyArrowFileIO(), + [ + FileScanTask(data_file=data_file, delete_files={dv}) + for data_file, dv in zip(data_files, [first_dv, second_dv], strict=True) + ], + ) + + assert deletes == { + referenced_data_files[0]: [pa.chunked_array([[1, 3]])], + referenced_data_files[1]: [pa.chunked_array([[5]])], + } + + +@pytest.mark.parametrize("selected_indices", [(0, 1, 2, 3, 4, 5), (0, 3, 5)]) +@pytest.mark.parametrize("planning_mode", ["local", "rest"]) +def test_read_all_delete_files_from_external_packed_deletion_vectors( + tmp_path: Path, selected_indices: tuple[int, ...], planning_mode: Literal["local", "rest"] +) -> None: + """Apply selected DV ranges from an externally generated .bin file to their data files.""" + fixture_path = Path(__file__).parents[1] / "table" / "deletion_vector" / "v1" / "packed-deletion-vectors.bin" + # (content_offset, content_size_in_bytes, expected deleted row positions) + # Ranges copied from the fixture's V3 manifest; each DV references a different ten-row data file. + content_ranges = [ + (1, 46, [1, 4, 7]), + (47, 48, [0, 3, 6, 9]), + (95, 42, [0]), + (137, 50, [1, 3, 5, 7, 9]), + (187, 52, [0, 2, 4, 6, 8, 9]), + (239, 44, [0, 2]), + ] + schema = Schema(NestedField(1, "id", LongType(), required=True)) + data_files: list[DataFile] = [] + delete_files: list[DataFile] = [] + expected_deletes: dict[str, list[pa.ChunkedArray]] = {} + expected_rows: list[int] = [] + for index, (content_offset, content_size_in_bytes, expected_positions) in enumerate(content_ranges): + data_path = tmp_path / f"data-{index}.parquet" + referenced_data_file = str(data_path) + rows = list(range(index * 10, (index + 1) * 10)) + pq.write_table(pa.table({"id": rows}, schema=schema_to_pyarrow(schema)), data_path) + dv = DataFile.from_args( + _table_format_version=3, + content=DataFileContent.POSITION_DELETES, + file_path=str(fixture_path), + file_format=FileFormat.PUFFIN, + partition=Record(), + record_count=len(expected_positions), + file_size_in_bytes=fixture_path.stat().st_size, + content_offset=content_offset, + content_size_in_bytes=content_size_in_bytes, + referenced_data_file=referenced_data_file, + ) + dv.spec_id = 0 + data_file = DataFile.from_args( + content=DataFileContent.DATA, + file_path=referenced_data_file, + file_format=FileFormat.PARQUET, + partition=Record(), + record_count=10, + file_size_in_bytes=data_path.stat().st_size, + first_row_id=index * 10, + ) + data_file.spec_id = 0 + data_files.append(data_file) + delete_files.append(dv) + if index in selected_indices: + expected_deletes[referenced_data_file] = [pa.chunked_array([expected_positions])] + expected_rows.extend(row for position, row in enumerate(rows) if position not in expected_positions) + + if planning_mode == "local": + delete_index = DeleteFileIndex() + for dv in delete_files: + delete_index.add_delete_file( + ManifestEntry.from_args(status=ManifestEntryStatus.ADDED, sequence_number=2, data_file=dv) + ) + tasks = [ + FileScanTask(data_file=data_files[index], delete_files=delete_index.for_data_file(1, data_files[index])) + for index in selected_indices + ] + else: + rest_deletes: list[RESTDeleteFile] = [ + RESTPositionDeleteFile( + spec_id=0, + file_path=dv.file_path, + file_format=dv.file_format, + file_size_in_bytes=dv.file_size_in_bytes, + record_count=dv.record_count, + content_offset=dv.content_offset, + content_size_in_bytes=dv.content_size_in_bytes, + ) + for dv in delete_files + ] + tasks = [ + FileScanTask.from_rest_response( + RESTFileScanTask( + data_file=RESTDataFile( + spec_id=0, + file_path=data_files[index].file_path, + file_format=FileFormat.PARQUET, + file_size_in_bytes=data_files[index].file_size_in_bytes, + record_count=10, + ), + delete_file_references=[index, index], + ), + rest_deletes, + ) + for index in selected_indices + ] + + for task, index in zip(tasks, selected_indices, strict=True): + assert isinstance(task.delete_files, set) + assert len(task.delete_files) == 1 + dv = next(iter(task.delete_files)) + assert dv.referenced_data_file == task.file.file_path + assert (dv.content_offset, dv.content_size_in_bytes) == content_ranges[index][:2] + + # Repeated tasks must not cause a DV to be loaded more than once. + deletes = _read_all_delete_files(PyArrowFileIO(), iter(tasks + tasks[::-1])) + + assert deletes == expected_deletes + + result = ArrowScan( + table_metadata=TableMetadataV3( + location=str(tmp_path), + last_column_id=1, + schemas=[schema], + current_schema_id=schema.schema_id, + partition_specs=[PartitionSpec()], + next_row_id=60, + ), + io=PyArrowFileIO(), + projected_schema=schema, + row_filter=AlwaysTrue(), + ).to_table(tasks) + + assert sorted(result.column("id").to_pylist()) == expected_rows + + def test_delete(deletes_file: str, request: pytest.FixtureRequest, table_schema_simple: Schema) -> None: # Determine file format from the file extension file_format = FileFormat.PARQUET if deletes_file.endswith(".parquet") else FileFormat.ORC diff --git a/tests/table/deletion_vector/v1/deletion-vector.bin b/tests/table/deletion_vector/v1/deletion-vector.bin new file mode 100644 index 0000000000..1366c99808 Binary files /dev/null and b/tests/table/deletion_vector/v1/deletion-vector.bin differ diff --git a/tests/table/deletion_vector/v1/packed-deletion-vectors.bin b/tests/table/deletion_vector/v1/packed-deletion-vectors.bin new file mode 100644 index 0000000000..628c8d2a66 Binary files /dev/null and b/tests/table/deletion_vector/v1/packed-deletion-vectors.bin differ diff --git a/tests/table/test_delete_file_index.py b/tests/table/test_delete_file_index.py index 09dd9ac81b..6385e091cb 100644 --- a/tests/table/test_delete_file_index.py +++ b/tests/table/test_delete_file_index.py @@ -65,17 +65,24 @@ def _create_partition_delete(sequence_number: int = 1, spec_id: int = 0, partiti def _create_deletion_vector( - sequence_number: int = 1, file_path: str = "s3://bucket/data.parquet", spec_id: int = 0 + sequence_number: int = 1, + file_path: str = "s3://bucket/data.parquet", + spec_id: int = 0, + delete_file_path: str | None = None, + content_offset: int | None = None, + content_size_in_bytes: int | None = None, ) -> ManifestEntry: delete_file = DataFile.from_args( + _table_format_version=3, content=DataFileContent.POSITION_DELETES, - file_path=f"s3://bucket/deletion-vector-{sequence_number}.puffin", + file_path=delete_file_path or f"s3://bucket/deletion-vector-{sequence_number}.puffin", file_format=FileFormat.PUFFIN, partition=Record(), record_count=10, file_size_in_bytes=100, - lower_bounds={PATH_FIELD_ID: file_path.encode()}, - upper_bounds={PATH_FIELD_ID: file_path.encode()}, + referenced_data_file=file_path, + content_offset=content_offset, + content_size_in_bytes=content_size_in_bytes, ) delete_file._spec_id = spec_id return ManifestEntry.from_args(status=ManifestEntryStatus.ADDED, sequence_number=sequence_number, data_file=delete_file) @@ -161,6 +168,45 @@ def test_dvs_treated_as_position_deletes() -> None: assert all(d.content == DataFileContent.POSITION_DELETES for d in result) +def test_dvs_with_same_file_path_are_matched_by_referenced_data_file() -> None: + index = DeleteFileIndex() + shared_file_path = "s3://bucket/deletion-vectors.bin" + + index.add_delete_file( + _create_deletion_vector( + sequence_number=2, + file_path="s3://bucket/a.parquet", + delete_file_path=shared_file_path, + content_offset=4, + content_size_in_bytes=10, + ) + ) + index.add_delete_file( + _create_deletion_vector( + sequence_number=3, + file_path="s3://bucket/b.parquet", + delete_file_path=shared_file_path, + content_offset=40, + content_size_in_bytes=12, + ) + ) + + first_deletes = index.for_data_file(1, _create_data_file("s3://bucket/a.parquet")) + second_deletes = index.for_data_file(1, _create_data_file("s3://bucket/b.parquet")) + assert len(first_deletes) == len(second_deletes) == 1 + assert {dv.referenced_data_file for dv in first_deletes} == {"s3://bucket/a.parquet"} + assert {dv.referenced_data_file for dv in second_deletes} == {"s3://bucket/b.parquet"} + assert not index.for_data_file(1, _create_data_file("s3://bucket/unrelated.parquet")) + assert isinstance(first_deletes, set) + result = [*first_deletes, *second_deletes] + + assert len(result) == 2 + assert {(dv.file_path, dv.content_offset, dv.content_size_in_bytes) for dv in result} == { + (shared_file_path, 4, 10), + (shared_file_path, 40, 12), + } + + def test_cannot_add_after_indexing() -> None: group = PositionDeletes() group.add(_create_positional_delete(sequence_number=1).data_file, 1) diff --git a/tests/table/test_deletion_vector.py b/tests/table/test_deletion_vector.py index 788216f8b3..4b38ea8627 100644 --- a/tests/table/test_deletion_vector.py +++ b/tests/table/test_deletion_vector.py @@ -14,12 +14,23 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. +import struct +import zlib from os import path +from pathlib import Path import pytest from pyroaring import BitMap -from pyiceberg.table.deletion_vector import DeletionVector +from pyiceberg.io.pyarrow import PyArrowFileIO +from pyiceberg.manifest import DataFile, DataFileContent, FileFormat +from pyiceberg.table.deletion_vector import ( + _DV_BLOB_MAGIC_NUMBER, + DeletionVector, + _deserialize_dv_blob, + _has_deletion_vector_content_reference, + read_deletion_vectors, +) def _open_file(file: str) -> bytes: @@ -28,6 +39,93 @@ def _open_file(file: str) -> bytes: return f.read() +def _dv_blob(bitmap_payload: bytes) -> bytes: + bitmap_data = struct.pack("I", len(bitmap_data)) + bitmap_data + struct.pack(">I", zlib.crc32(bitmap_data) & 0xFFFFFFFF) + + +def _bitmap_payload() -> bytes: + return (1).to_bytes(8, byteorder="little") + (0).to_bytes(4, byteorder="little") + BitMap([1, 3, 5]).serialize() + + +def test_deserialize_deletion_vector_blob() -> None: + actual = _deserialize_dv_blob(_dv_blob(_bitmap_payload()), record_count=3) + + assert actual == [BitMap([1, 3, 5])] + + +def test_deserialize_deletion_vector_blob_invalid_length() -> None: + with pytest.raises(ValueError, match="Invalid bitmap data length"): + _deserialize_dv_blob(_dv_blob(_bitmap_payload())[:-1]) + + +def test_deserialize_deletion_vector_blob_invalid_magic() -> None: + bitmap_data = struct.pack("I", len(bitmap_data)) + bitmap_data + struct.pack(">I", zlib.crc32(bitmap_data) & 0xFFFFFFFF) + + with pytest.raises(ValueError, match="Invalid magic number"): + _deserialize_dv_blob(blob) + + +def test_deserialize_deletion_vector_blob_invalid_crc() -> None: + blob = bytearray(_dv_blob(_bitmap_payload())) + blob[-1] ^= 1 + + with pytest.raises(ValueError, match="Invalid CRC"): + _deserialize_dv_blob(bytes(blob)) + + +def test_deserialize_deletion_vector_blob_invalid_cardinality() -> None: + with pytest.raises(ValueError, match="Invalid cardinality"): + _deserialize_dv_blob(_dv_blob(_bitmap_payload()), record_count=4) + + +def _data_file( + content_offset: int | None = None, + content_size_in_bytes: int | None = None, + referenced_data_file: str | None = None, +) -> DataFile: + return DataFile.from_args( + _table_format_version=3, + content=DataFileContent.POSITION_DELETES, + file_path="deletes.puffin", + file_format=FileFormat.PUFFIN, + record_count=1, + content_offset=content_offset, + content_size_in_bytes=content_size_in_bytes, + referenced_data_file=referenced_data_file, + ) + + +def test_has_deletion_vector_content_reference() -> None: + assert not _has_deletion_vector_content_reference(_data_file()) + assert _has_deletion_vector_content_reference(_data_file(content_offset=0)) + assert _has_deletion_vector_content_reference(_data_file(content_size_in_bytes=1)) + assert _has_deletion_vector_content_reference(_data_file(referenced_data_file="data.parquet")) + + +def test_read_external_deletion_vector_content_range() -> None: + """Read an externally generated Iceberg DV blob from a .bin content range.""" + fixture_path = Path(__file__).parent / "deletion_vector" / "v1" / "deletion-vector.bin" + # The V3 manifest points past the one-byte file prefix to the complete DV blob. + dv = DataFile.from_args( + _table_format_version=3, + content=DataFileContent.POSITION_DELETES, + file_path=str(fixture_path), + file_format=FileFormat.PUFFIN, + record_count=3, + content_offset=1, + content_size_in_bytes=46, + referenced_data_file="data.parquet", + ) + + deletion_vectors = read_deletion_vectors(PyArrowFileIO(), dv) + + assert len(deletion_vectors) == 1 + assert deletion_vectors[0].referenced_data_file == "data.parquet" + assert deletion_vectors[0].to_vector().to_pylist() == [1, 4, 5] + + def test_map_empty() -> None: puffin = _open_file("64mapempty.bin")