moomindani commented on code in PR #3474:
URL: https://github.com/apache/iceberg-python/pull/3474#discussion_r3491972478
##########
pyiceberg/table/puffin.py:
##########
@@ -75,3 +80,76 @@ 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
+ # 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)
+
+ blobs_metadata: list[PuffinBlobMetadata] = []
+ for blob in self._blobs:
+ # offset and length are placeholders on the blob's metadata
until the file is assembled here
+
blobs_metadata.append(blob.metadata.model_copy(update={"offset": out.tell(),
"length": len(blob.payload)}))
+ out.write(blob.payload)
+
+ footer = Footer(blobs=blobs_metadata, properties={"created-by":
self._created_by})
+ footer_payload_bytes = footer.model_dump_json(by_alias=True,
exclude_none=True).encode("utf-8")
+
+ out.write(MAGIC_BYTES)
+ out.write(footer_payload_bytes)
+ out.write(len(footer_payload_bytes).to_bytes(4, "little"))
+ out.write((0).to_bytes(4, "little")) # flags
+ out.write(MAGIC_BYTES)
+
+ puffin_bytes = out.getvalue()
+
+ with self._output_file.create(overwrite=True) as output_stream:
Review Comment:
Yes — it's for idempotency. This matches the other write-once writers in
PyIceberg: the data file writer (`pyarrow.py`) and the manifest writer
(`avro/file.py`) both hardcode `overwrite=True`. Puffin DV files are written to
fresh, UUID-based locations (like data/manifest files), so a collision with a
different file isn't expected; the realistic case where the path already exists
is a retry writing the same file, and `overwrite=True` keeps that idempotent
instead of failing. Only `ToOutputFile.table_metadata` parameterizes
`overwrite`, since metadata writes can legitimately target an existing path.
That said, if you can think of a case where silently overwriting would be a
problem, I'm happy to make it a constructor argument.
--
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]
