Skip to content

Commit d71d94c

Browse files
committed
Support range-based reads for deletion vectors
1 parent 68898e5 commit d71d94c

11 files changed

Lines changed: 534 additions & 30 deletions

File tree

‎pyiceberg/io/pyarrow.py‎

Lines changed: 4 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -147,11 +147,11 @@
147147
visit_with_partner,
148148
)
149149
from pyiceberg.table import DOWNCAST_NS_TIMESTAMP_TO_US_ON_WRITE, TableProperties
150-
from pyiceberg.table.deletion_vector import deletion_vectors_from_puffin_file
150+
from pyiceberg.table.delete_file import DeleteFileSet
151+
from pyiceberg.table.deletion_vector import read_deletion_vectors
151152
from pyiceberg.table.locations import load_location_provider
152153
from pyiceberg.table.metadata import TableMetadata
153154
from pyiceberg.table.name_mapping import NameMapping, apply_name_mapping
154-
from pyiceberg.table.puffin import PuffinFile
155155
from pyiceberg.transforms import IdentityTransform, TruncateTransform
156156
from pyiceberg.typedef import EMPTY_DICT, Properties, Record, TableVersion
157157
from pyiceberg.types import (
@@ -1149,10 +1149,7 @@ def _read_deletes(io: FileIO, data_file: DataFile) -> dict[str, pa.ChunkedArray]
11491149
for path in table.column("file_path").unique()
11501150
}
11511151
elif data_file.file_format == FileFormat.PUFFIN:
1152-
with io.new_input(data_file.file_path).open() as fi:
1153-
payload = fi.read()
1154-
1155-
return {dv.referenced_data_file: dv.to_vector() for dv in deletion_vectors_from_puffin_file(PuffinFile(payload))}
1152+
return {dv.referenced_data_file: dv.to_vector() for dv in read_deletion_vectors(io, data_file)}
11561153
else:
11571154
raise ValueError(f"Delete file format not supported: {data_file.file_format}")
11581155

@@ -1721,7 +1718,7 @@ def _task_to_record_batches(
17211718

17221719
def _read_all_delete_files(io: FileIO, tasks: Iterable[FileScanTask]) -> dict[str, list[ChunkedArray]]:
17231720
deletes_per_file: dict[str, list[ChunkedArray]] = {}
1724-
unique_deletes = set(itertools.chain.from_iterable([task.delete_files for task in tasks]))
1721+
unique_deletes = DeleteFileSet(itertools.chain.from_iterable(task.delete_files for task in tasks))
17251722
if len(unique_deletes) > 0:
17261723
executor = ExecutorFactory.get_or_create()
17271724
deletes_per_files: Iterator[dict[str, ChunkedArray]] = executor.map(

‎pyiceberg/manifest.py‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -532,6 +532,18 @@ def equality_ids(self) -> list[int] | None:
532532
def sort_order_id(self) -> int | None:
533533
return self._data[15]
534534

535+
@property
536+
def referenced_data_file(self) -> str | None:
537+
return self._data[17] if len(self._data) > 17 else None
538+
539+
@property
540+
def content_offset(self) -> int | None:
541+
return self._data[18] if len(self._data) > 18 else None
542+
543+
@property
544+
def content_size_in_bytes(self) -> int | None:
545+
return self._data[19] if len(self._data) > 19 else None
546+
535547
# Spec ID should not be stored in the file
536548
_spec_id: int
537549

‎pyiceberg/table/__init__.py‎

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,7 @@
4848
from pyiceberg.manifest import DataFile, DataFileContent, ManifestContent, ManifestEntry, ManifestEntryStatus, ManifestFile
4949
from pyiceberg.partitioning import PARTITION_FIELD_ID_START, UNPARTITIONED_PARTITION_SPEC, PartitionKey, PartitionSpec
5050
from pyiceberg.schema import Schema
51+
from pyiceberg.table.delete_file import DeleteFileSet
5152
from pyiceberg.table.delete_file_index import DeleteFileIndex
5253
from pyiceberg.table.inspect import InspectTable
5354
from pyiceberg.table.locations import LocationProvider, load_location_provider
@@ -2232,17 +2233,17 @@ class FileScanTask(ScanTask):
22322233
"""Task representing a data file and its corresponding delete files."""
22332234

22342235
file: DataFile
2235-
delete_files: set[DataFile]
2236+
delete_files: DeleteFileSet
22362237
residual: BooleanExpression
22372238

22382239
def __init__(
22392240
self,
22402241
data_file: DataFile,
2241-
delete_files: set[DataFile] | None = None,
2242+
delete_files: Iterable[DataFile] | None = None,
22422243
residual: BooleanExpression = ALWAYS_TRUE,
22432244
) -> None:
22442245
self.file = data_file
2245-
self.delete_files = delete_files or set()
2246+
self.delete_files = DeleteFileSet(delete_files if delete_files is not None else [])
22462247
self.residual = residual
22472248

22482249
@staticmethod
@@ -2266,7 +2267,7 @@ def from_rest_response(
22662267

22672268
data_file = _rest_file_to_data_file(rest_task.data_file)
22682269

2269-
resolved_deletes: set[DataFile] = set()
2270+
resolved_deletes = DeleteFileSet()
22702271
if rest_task.delete_file_references:
22712272
for idx in rest_task.delete_file_references:
22722273
delete_file = delete_files[idx]

‎pyiceberg/table/delete_file.py‎

Lines changed: 83 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,83 @@
1+
# Licensed to the Apache Software Foundation (ASF) under one
2+
# or more contributor license agreements. See the NOTICE file
3+
# distributed with this work for additional information
4+
# regarding copyright ownership. The ASF licenses this file
5+
# to you under the Apache License, Version 2.0 (the
6+
# "License"); you may not use this file except in compliance
7+
# with the License. You may obtain a copy of the License at
8+
#
9+
# http://www.apache.org/licenses/LICENSE-2.0
10+
#
11+
# Unless required by applicable law or agreed to in writing,
12+
# software distributed under the License is distributed on an
13+
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
# KIND, either express or implied. See the License for the
15+
# specific language governing permissions and limitations
16+
# under the License.
17+
from __future__ import annotations
18+
19+
from collections.abc import Iterable, Iterator, MutableSet
20+
from typing import Any
21+
22+
from pyiceberg.manifest import DataFile
23+
24+
_DeleteFileKey = tuple[str, int | None, int | None]
25+
26+
27+
def _delete_file_key(delete_file: DataFile) -> _DeleteFileKey:
28+
return delete_file.file_path, delete_file.content_offset, delete_file.content_size_in_bytes
29+
30+
31+
class DeleteFileSet(MutableSet[DataFile]):
32+
"""Set-like delete-file collection keyed by location and content range."""
33+
34+
_files: dict[_DeleteFileKey, DataFile]
35+
36+
def __init__(self, delete_files: Iterable[DataFile] = ()) -> None:
37+
self._files = {}
38+
for delete_file in delete_files:
39+
self.add(delete_file)
40+
41+
def __contains__(self, delete_file: object) -> bool:
42+
"""Return whether the delete file is present."""
43+
return isinstance(delete_file, DataFile) and _delete_file_key(delete_file) in self._files
44+
45+
def __iter__(self) -> Iterator[DataFile]:
46+
"""Return an iterator over delete files."""
47+
return iter(self._files.values())
48+
49+
def __len__(self) -> int:
50+
"""Return the number of delete files."""
51+
return len(self._files)
52+
53+
def add(self, delete_file: DataFile) -> None:
54+
self._files[_delete_file_key(delete_file)] = delete_file
55+
56+
def discard(self, delete_file: DataFile) -> None:
57+
self._files.pop(_delete_file_key(delete_file), None)
58+
59+
def update(self, delete_files: Iterable[DataFile]) -> None:
60+
for delete_file in delete_files:
61+
self.add(delete_file)
62+
63+
def __repr__(self) -> str:
64+
"""Return a string representation of the delete file set."""
65+
return f"{type(self).__name__}({list(self)!r})"
66+
67+
def __eq__(self, other: Any) -> bool:
68+
"""Compare delete file sets by delete file identity."""
69+
if isinstance(other, DeleteFileSet):
70+
return self._files.keys() == other._files.keys()
71+
72+
if not isinstance(other, Iterable):
73+
return False
74+
75+
other_keys: set[_DeleteFileKey] = set()
76+
other_count = 0
77+
for delete_file in other:
78+
if not isinstance(delete_file, DataFile):
79+
return False
80+
other_keys.add(_delete_file_key(delete_file))
81+
other_count += 1
82+
83+
return len(other_keys) == other_count and set(self._files) == other_keys

‎pyiceberg/table/delete_file_index.py‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
from pyiceberg.expressions import EqualTo
2222
from pyiceberg.expressions.visitors import _InclusiveMetricsEvaluator
2323
from pyiceberg.manifest import INITIAL_SEQUENCE_NUMBER, POSITIONAL_DELETE_SCHEMA, DataFile, ManifestEntry
24+
from pyiceberg.table.delete_file import DeleteFileSet
2425
from pyiceberg.typedef import Record
2526

2627
PATH_FIELD_ID = 2147483546
@@ -125,11 +126,11 @@ def add_delete_file(self, manifest_entry: ManifestEntry, partition_key: Record |
125126
deletes = self._by_partition.setdefault(key, PositionDeletes())
126127
deletes.add(delete_file, seq)
127128

128-
def for_data_file(self, seq_num: int, data_file: DataFile, partition_key: Record | None = None) -> set[DataFile]:
129+
def for_data_file(self, seq_num: int, data_file: DataFile, partition_key: Record | None = None) -> DeleteFileSet:
129130
if self.is_empty():
130-
return set()
131+
return DeleteFileSet()
131132

132-
deletes: set[DataFile] = set()
133+
deletes = DeleteFileSet()
133134
spec_id = data_file.spec_id or 0
134135

135136
key = _partition_key(spec_id, partition_key)

‎pyiceberg/table/deletion_vector.py‎

Lines changed: 100 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,9 @@
1515
# specific language governing permissions and limitations
1616
# under the License.
1717
import math
18-
from typing import TYPE_CHECKING
18+
import struct
19+
import zlib
20+
from typing import TYPE_CHECKING, cast
1921

2022
from pyroaring import BitMap, FrozenBitMap
2123

@@ -24,9 +26,19 @@
2426
if TYPE_CHECKING:
2527
import pyarrow as pa
2628

29+
from pyiceberg.io import FileIO
30+
from pyiceberg.manifest import DataFile
31+
2732
EMPTY_BITMAP = FrozenBitMap()
2833
MAX_JAVA_SIGNED = int(math.pow(2, 31)) - 1
2934
PROPERTY_REFERENCED_DATA_FILE = "referenced-data-file"
35+
_MAX_DELETION_VECTOR_CONTENT_SIZE = 2**31 - 1
36+
_DV_BLOB_LENGTH = struct.Struct(">I")
37+
_DV_BLOB_MAGIC = struct.Struct("<I")
38+
_DV_BLOB_CRC = struct.Struct(">I")
39+
_DV_BLOB_MAGIC_NUMBER = 1681511377
40+
_ROARING_BITMAP_COUNT_SIZE_BYTES = 8
41+
_DV_BLOB_MIN_SIZE_BYTES = _DV_BLOB_LENGTH.size + _DV_BLOB_MAGIC.size + _ROARING_BITMAP_COUNT_SIZE_BYTES + _DV_BLOB_CRC.size
3042

3143

3244
class DeletionVector:
@@ -77,17 +89,99 @@ def to_vector(self) -> "pa.ChunkedArray":
7789
return self._bitmaps_to_chunked_array(self._bitmaps)
7890

7991

80-
def _extract_vector_payload(blob_payload: bytes) -> bytes:
81-
"""Strip deletion-vector-v1 blob framing: length(4 big-endian) + DV magic(4) ... CRC(4 big-endian)."""
82-
length_prefix = int.from_bytes(blob_payload[0:4], "big")
83-
return blob_payload[8 : 4 + length_prefix]
92+
def _deserialize_dv_blob(blob: bytes, record_count: int | None = None) -> list[BitMap]:
93+
# The DV blob encoding matches Iceberg Java's BitmapPositionDeleteIndex:
94+
# 4-byte big-endian bitmap-data length, 4-byte little-endian magic number,
95+
# portable Roaring bitmap data, and 4-byte big-endian CRC-32.
96+
if len(blob) < _DV_BLOB_MIN_SIZE_BYTES:
97+
raise ValueError(f"Invalid deletion vector blob length: {len(blob)}")
98+
99+
bitmap_data_length = _DV_BLOB_LENGTH.unpack_from(blob)[0]
100+
expected_bitmap_data_length = len(blob) - _DV_BLOB_LENGTH.size - _DV_BLOB_CRC.size
101+
if bitmap_data_length != expected_bitmap_data_length:
102+
raise ValueError(f"Invalid bitmap data length: {bitmap_data_length}, expected {expected_bitmap_data_length}")
103+
104+
bitmap_data_offset = _DV_BLOB_LENGTH.size
105+
crc_offset = bitmap_data_offset + bitmap_data_length
106+
bitmap_data = blob[bitmap_data_offset:crc_offset]
107+
108+
magic_number = _DV_BLOB_MAGIC.unpack_from(bitmap_data)[0]
109+
if magic_number != _DV_BLOB_MAGIC_NUMBER:
110+
raise ValueError(f"Invalid magic number: {magic_number}, expected {_DV_BLOB_MAGIC_NUMBER}")
111+
112+
checksum = zlib.crc32(bitmap_data) & 0xFFFFFFFF
113+
expected_checksum = _DV_BLOB_CRC.unpack_from(blob, crc_offset)[0]
114+
if checksum != expected_checksum:
115+
raise ValueError("Invalid CRC")
116+
117+
bitmaps = DeletionVector._deserialize_bitmap(bitmap_data[_DV_BLOB_MAGIC.size :])
118+
if record_count is not None:
119+
cardinality = sum(len(bitmap) for bitmap in bitmaps)
120+
if cardinality != record_count:
121+
raise ValueError(f"Invalid cardinality: {cardinality}, expected {record_count}")
122+
123+
return bitmaps
124+
125+
126+
def _validate_deletion_vector_content(dv: "DataFile") -> None:
127+
content_offset = dv.content_offset
128+
content_size_in_bytes = dv.content_size_in_bytes
129+
referenced_data_file = dv.referenced_data_file
130+
131+
if content_offset is None:
132+
raise ValueError(f"Invalid deletion vector, content offset is missing: {dv.file_path}")
133+
if content_size_in_bytes is None:
134+
raise ValueError(f"Invalid deletion vector, content size is missing: {dv.file_path}")
135+
if content_offset < 0:
136+
raise ValueError(f"Invalid deletion vector, content offset cannot be negative: {content_offset}")
137+
if content_size_in_bytes < 0:
138+
raise ValueError(f"Invalid deletion vector, content size cannot be negative: {content_size_in_bytes}")
139+
if content_size_in_bytes > _MAX_DELETION_VECTOR_CONTENT_SIZE:
140+
raise ValueError(f"Cannot read deletion vector larger than 2GB: {content_size_in_bytes}")
141+
if referenced_data_file is None:
142+
raise ValueError(f"Invalid deletion vector, referenced data file is missing: {dv.file_path}")
143+
144+
145+
def has_deletion_vector_content_reference(dv: "DataFile") -> bool:
146+
"""Return whether a deletion vector is described by manifest content-range metadata."""
147+
return dv.content_offset is not None or dv.content_size_in_bytes is not None or dv.referenced_data_file is not None
148+
149+
150+
def _read_deletion_vector(io: "FileIO", dv: "DataFile") -> DeletionVector:
151+
_validate_deletion_vector_content(dv)
152+
153+
content_offset = cast(int, dv.content_offset)
154+
content_size_in_bytes = cast(int, dv.content_size_in_bytes)
155+
referenced_data_file = cast(str, dv.referenced_data_file)
156+
157+
with io.new_input(dv.file_path).open() as fi:
158+
fi.seek(content_offset)
159+
payload = fi.read(content_size_in_bytes)
160+
161+
if len(payload) != content_size_in_bytes:
162+
raise ValueError(f"Could not read deletion vector, expected {content_size_in_bytes} bytes, got {len(payload)}")
163+
164+
return DeletionVector(
165+
referenced_data_file=referenced_data_file,
166+
bitmaps=_deserialize_dv_blob(payload, dv.record_count),
167+
)
168+
169+
170+
def read_deletion_vectors(io: "FileIO", dv: "DataFile") -> list[DeletionVector]:
171+
"""Read deletion vectors from a delete file or its manifest content range."""
172+
if has_deletion_vector_content_reference(dv):
173+
return [_read_deletion_vector(io, dv)]
174+
175+
with io.new_input(dv.file_path).open() as fi:
176+
return deletion_vectors_from_puffin_file(PuffinFile(fi.read()))
84177

85178

86179
def deletion_vectors_from_puffin_file(puffin_file: PuffinFile) -> list[DeletionVector]:
180+
"""Read all deletion vectors stored in a Puffin file."""
87181
return [
88182
DeletionVector(
89183
referenced_data_file=blob.properties[PROPERTY_REFERENCED_DATA_FILE],
90-
bitmaps=DeletionVector._deserialize_bitmap(_extract_vector_payload(puffin_file.get_blob_payload(blob))),
184+
bitmaps=_deserialize_dv_blob(puffin_file.get_blob_payload(blob)),
91185
)
92186
for blob in puffin_file.footer.blobs
93187
]

‎tests/avro/test_file.py‎

Lines changed: 14 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -51,6 +51,7 @@
5151
LongType,
5252
NestedField,
5353
StringType,
54+
StructType,
5455
TimestampType,
5556
TimestamptzType,
5657
TimeType,
@@ -89,7 +90,16 @@ def test_missing_schema() -> None:
8990

9091
# helper function to serialize our objects to dicts to enable
9192
# direct comparison with the dicts returned by fastavro
92-
def todict(obj: Any) -> Any:
93+
def todict(obj: Any, struct: StructType | None = None) -> Any:
94+
if struct is not None and isinstance(obj, Record):
95+
return {
96+
field.name: todict(
97+
getattr(obj, field.name),
98+
field.field_type if isinstance(field.field_type, StructType) else None,
99+
)
100+
for field in struct.fields
101+
if hasattr(obj, field.name)
102+
}
93103
if isinstance(obj, dict):
94104
data = []
95105
for k, v in obj.items():
@@ -157,7 +167,7 @@ def test_write_manifest_entry_with_iceberg_read_with_fastavro_v1() -> None:
157167

158168
fa_entry = next(it)
159169

160-
v2_entry = todict(entry)
170+
v2_entry = todict(entry, MANIFEST_ENTRY_SCHEMAS[2].as_struct())
161171

162172
# These are not written in V1
163173
del v2_entry["sequence_number"]
@@ -222,7 +232,7 @@ def test_write_manifest_entry_with_iceberg_read_with_fastavro_v2() -> None:
222232

223233
fa_entry = next(it)
224234

225-
assert todict(entry) == fa_entry
235+
assert todict(entry, MANIFEST_ENTRY_SCHEMAS[2].as_struct()) == fa_entry
226236

227237

228238
@pytest.mark.parametrize("format_version", [1, 2])
@@ -260,7 +270,7 @@ def test_write_manifest_entry_with_fastavro_read_with_iceberg(format_version: Ta
260270
schema = AvroSchemaConversion().iceberg_to_avro(MANIFEST_ENTRY_SCHEMAS[format_version], schema_name="manifest_entry")
261271

262272
with open(tmp_avro_file, "wb") as out:
263-
writer(out, schema, [todict(entry)])
273+
writer(out, schema, [todict(entry, MANIFEST_ENTRY_SCHEMAS[format_version].as_struct())])
264274

265275
# Read as V2
266276
with avro.AvroFile[ManifestEntry](

0 commit comments

Comments
 (0)