Address review: run-optimize DV bitmaps, cap positions, skip write on error

- DeletionVector._serialize_bitmap now run-optimizes each bitmap before
  serializing, matching Java's BitmapPositionDeleteIndex (a contiguous
  100k-position DV drops from ~16 KB to ~25 bytes; still round-trips).
- from_positions rejects positions above MAX_POSITION, mirroring Java's
  RoaringPositionBitmap.MAX_POSITION, to avoid an OOM in the gap-fill.
- PuffinWriter.__exit__ short-circuits when the with-body raised, so no
  half-populated file is written.

Co-authored-by: Isaac

Author: moomindani

pyiceberg/table/deletion_vector.py

@@ -29,6 +29,10 @@
 
 EMPTY_BITMAP = FrozenBitMap()
 MAX_JAVA_SIGNED = int(math.pow(2, 31)) - 1
+# Largest addressable position, mirroring Java's RoaringPositionBitmap.MAX_POSITION
+# (toPosition(Integer.MAX_VALUE - 1, Integer.MIN_VALUE)): the high 32 bits hold the bitmap
+# key and the low 32 bits the position within that bitmap.
+MAX_POSITION = ((MAX_JAVA_SIGNED - 1) << 32) | 0x80000000
 PROPERTY_REFERENCED_DATA_FILE = "referenced-data-file"
 DELETION_VECTOR_MAGIC = b"\xd1\xd3\x39\x64"
 # Reserved field id of the row position (_pos) metadata column, referenced by
@@ -48,8 +52,8 @@ def __init__(self, referenced_data_file: str, bitmaps: list[BitMap]) -> None:
     def from_positions(cls, referenced_data_file: str, positions: Iterable[int]) -> "DeletionVector":
         bitmaps_by_key: dict[int, BitMap] = {}
         for position in positions:
-            if position < 0:
-                raise ValueError(f"Invalid position: {position}, positions must be non-negative")
+            if position < 0 or position > MAX_POSITION:
+                raise ValueError(f"Invalid position: {position}, must be between 0 and {MAX_POSITION}")
             bitmaps_by_key.setdefault(position >> 32, BitMap()).add(position & 0xFFFFFFFF)
 
         if not bitmaps_by_key:
@@ -101,6 +105,9 @@ def _serialize_bitmap(bitmaps: list[BitMap]) -> bytes:
                 if key > MAX_JAVA_SIGNED:
                     raise ValueError(f"Key {key} is too large, max {MAX_JAVA_SIGNED} to maintain compatibility with Java impl")
                 out.write(key.to_bytes(4, "little"))
+                # Run-length encode before serializing so run-heavy vectors (e.g. contiguous
+                # deletes) stay compact, matching Java's BitmapPositionDeleteIndex.
+                bitmap.run_optimize()
                 out.write(bitmap.serialize())
             return out.getvalue()
 

pyiceberg/table/puffin.py

@@ -121,6 +121,9 @@ def __exit__(
     ) -> None:
         """Assemble the Puffin file and write it to the output file."""
         self.closed = True
+        # If the with-body raised, skip assembling and writing a half-populated file.
+        if exc_type is not None:
+            return
 
         with io.BytesIO() as out:
             out.write(MAGIC_BYTES)

tests/table/test_deletion_vector.py

@@ -21,6 +21,7 @@
 
 from pyiceberg.table.deletion_vector import (
     DELETION_VECTOR_MAGIC,
+    MAX_POSITION,
     PROPERTY_REFERENCED_DATA_FILE,
     ROW_POSITION_FIELD_ID,
     DeletionVector,
@@ -96,10 +97,26 @@ def test_serialize_bitmap_round_trips(positions: list[int]) -> None:
 
 
 def test_from_positions_rejects_negative() -> None:
-    with pytest.raises(ValueError, match="Invalid position: -1, positions must be non-negative"):
+    with pytest.raises(ValueError, match=f"Invalid position: -1, must be between 0 and {MAX_POSITION}"):
         DeletionVector.from_positions("file.parquet", [1, -1, 2])
 
 
+def test_from_positions_rejects_out_of_range() -> None:
+    too_large = MAX_POSITION + 1
+    with pytest.raises(ValueError, match=f"Invalid position: {too_large}, must be between 0 and {MAX_POSITION}"):
+        DeletionVector.from_positions("file.parquet", [1, too_large, 2])
+
+
+def test_serialize_bitmap_run_optimizes_contiguous_run() -> None:
+    # A long contiguous run serializes compactly thanks to run-length encoding and still round-trips.
+    positions = list(range(100_000))
+    dv = DeletionVector.from_positions("file.parquet", positions)
+
+    serialized = DeletionVector._serialize_bitmap(dv._bitmaps)
+    assert len(serialized) < 1_000  # ~16 KB without run optimization
+    assert DeletionVector._deserialize_bitmap(serialized) == dv._bitmaps
+
+
 def test_from_positions_rejects_empty() -> None:
     with pytest.raises(ValueError, match="Deletion vector must contain at least one position"):
         DeletionVector.from_positions("file.parquet", [])

tests/table/test_puffin.py

@@ -143,6 +143,20 @@ def test_puffin_writer_empty(tmp_path: Path) -> None:
     assert deletion_vectors_from_puffin_file(reader) == []
 
 
+def test_puffin_writer_does_not_write_on_exception(tmp_path: Path) -> None:
+    puffin_path = tmp_path / "test.puffin"
+    output_file = PyArrowFileIO().new_output(str(puffin_path))
+
+    with pytest.raises(ValueError, match="boom"):
+        with PuffinWriter(output_file) as writer:
+            writer.add_blob(DeletionVector.from_positions("file.parquet", [1]).to_blob())
+            raise ValueError("boom")
+
+    assert writer.closed
+    # The body raised, so no half-populated file should have been written.
+    assert not output_file.exists()
+
+
 def test_add_blob_to_closed_writer_raises(tmp_path: Path) -> None:
     output_file = PyArrowFileIO().new_output(str(tmp_path / "test.puffin"))
     writer = PuffinWriter(output_file)