moomindani commented on code in PR #3474:
URL: https://github.com/apache/iceberg-python/pull/3474#discussion_r3488766604
##########
pyiceberg/table/puffin.py:
##########
@@ -75,3 +80,73 @@ def to_vector(self) -> dict[str, "pa.ChunkedArray"]:
from pyiceberg.table.deletion_vector import
deletion_vectors_from_puffin_file # local import avoids the cycle
return {dv.referenced_data_file: dv.to_vector() for dv in
deletion_vectors_from_puffin_file(self)}
+
+
+@dataclass(frozen=True)
+class PuffinBlob:
+ """A blob to write into a Puffin file: its metadata and serialized
payload."""
+
+ metadata: PuffinBlobMetadata
+ payload: bytes
+
+
+class PuffinWriter:
+ """Assembles a Puffin file from blobs and writes it to an output file.
+
+ This writer is format-level and blob-agnostic: callers supply
already-serialized blobs
+ (for example via DeletionVector.to_blob()). Use it as a context manager;
the file is
+ written on exit, after which its size is available via len(output_file).
+ """
+
+ closed: bool
+ _output_file: OutputFile
+ _blobs: list[PuffinBlob]
+ _created_by: str
+
+ def __init__(self, output_file: OutputFile, created_by: str | None = None)
-> None:
+ self.closed = False
+ self._output_file = output_file
+ self._blobs = []
+ self._created_by = created_by if created_by is not None else
f"PyIceberg version {__version__}"
+
+ def __enter__(self) -> "PuffinWriter":
+ """Open the writer."""
+ return self
+
+ def __exit__(
+ self,
+ exc_type: type[BaseException] | None,
+ exc_value: BaseException | None,
+ traceback: TracebackType | None,
+ ) -> None:
+ """Assemble the Puffin file and write it to the output file."""
+ self.closed = True
Review Comment:
Done — `__exit__` now short-circuits when `exc_type is not None`, so a
failed `with` body no longer assembles and writes a half-populated file. Added
`test_puffin_writer_does_not_write_on_exception` to cover it.
##########
pyiceberg/table/deletion_vector.py:
##########
@@ -67,6 +89,21 @@ def _deserialize_bitmap(pl: bytes) -> list[BitMap]:
return bitmaps
+ @staticmethod
+ def _serialize_bitmap(bitmaps: list[BitMap]) -> bytes:
+ # Counterpart of _deserialize_bitmap: number of bitmaps (8 bytes,
little-endian), then for each
+ # non-empty bitmap in ascending key order its key (4 bytes,
little-endian) and serialized payload.
+ non_empty = [(key, bitmap) for key, bitmap in enumerate(bitmaps) if
len(bitmap) > 0]
+
+ with io.BytesIO() as out:
+ out.write(len(non_empty).to_bytes(8, "little"))
+ for key, bitmap in non_empty:
+ 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"))
+ out.write(bitmap.serialize())
Review Comment:
Done — `_serialize_bitmap` now calls `run_optimize()` on each bitmap before
serializing, matching Java's `BitmapPositionDeleteIndex`. Since
`from_positions` builds the bitmaps incrementally with `add()`, a contiguous
100k-position DV drops from ~16 KB to ~25 bytes, and the run-optimized bytes
still round-trip through `_deserialize_bitmap`. Added
`test_serialize_bitmap_run_optimizes_contiguous_run`.
##########
pyiceberg/table/deletion_vector.py:
##########
@@ -37,6 +44,21 @@ def __init__(self, referenced_data_file: str, bitmaps:
list[BitMap]) -> None:
self.referenced_data_file = referenced_data_file
self._bitmaps = bitmaps
+ @classmethod
+ 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")
Review Comment:
Done — `from_positions` now also rejects positions above `MAX_POSITION`,
defined as `((MAX_JAVA_SIGNED - 1) << 32) | 0x80000000` to mirror Java's
`RoaringPositionBitmap.MAX_POSITION` (`toPosition(Integer.MAX_VALUE - 1,
Integer.MIN_VALUE)`). Added `test_from_positions_rejects_out_of_range`.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
