geruh commented on code in PR #2822:
URL: https://github.com/apache/iceberg-python/pull/2822#discussion_r2605149847
##########
pyiceberg/table/puffin.py:
##########
@@ -114,3 +147,97 @@ def __init__(self, puffin: bytes) -> None:
def to_vector(self) -> dict[str, "pa.ChunkedArray"]:
return {path: _bitmaps_to_chunked_array(bitmaps) for path, bitmaps in
self._deletion_vectors.items()}
+
+
+class PuffinWriter:
+ _blobs: list[PuffinBlobMetadata]
+ _blob_payloads: list[bytes]
+
+ def __init__(self) -> None:
+ self._blobs = []
+ self._blob_payloads = []
+
+ def add(
+ self,
+ positions: Iterable[int],
+ referenced_data_file: str,
+ ) -> None:
+ # 1. Create bitmaps from positions
+ bitmaps: dict[int, BitMap] = {}
+ cardinality = 0
+ for pos in positions:
+ cardinality += 1
+ key = pos >> 32
+ low_bits = pos & 0xFFFFFFFF
+ if key not in bitmaps:
+ bitmaps[key] = BitMap()
+ bitmaps[key].add(low_bits)
+
+ # 2. Serialize bitmaps for the vector payload
+ vector_payload = _serialize_bitmaps(bitmaps)
+
+ # 3. Construct the full blob payload for deletion-vector-v1
+ with io.BytesIO() as blob_payload_buffer:
+ # Magic bytes for DV
+ blob_payload_buffer.write(DELETION_VECTOR_MAGIC)
+ # The vector itself
+ blob_payload_buffer.write(vector_payload)
+
+ # The content for CRC calculation
+ crc_content = blob_payload_buffer.getvalue()
+ crc32 = zlib.crc32(crc_content)
+
+ # The full blob to be stored in the Puffin file
+ with io.BytesIO() as full_blob_buffer:
+ # Combined length of the vector and magic bytes stored as 4
bytes, big-endian
+ full_blob_buffer.write(len(crc_content).to_bytes(4, "big"))
+ # The content (magic + vector)
+ full_blob_buffer.write(crc_content)
+ # A CRC-32 checksum of the magic bytes and serialized vector
as 4 bytes, big-endian
+ full_blob_buffer.write(crc32.to_bytes(4, "big"))
+
+ self._blob_payloads.append(full_blob_buffer.getvalue())
+
+ # 4. Create blob metadata
+ properties = {PROPERTY_REFERENCED_DATA_FILE: referenced_data_file,
"cardinality": str(cardinality)}
+
+ self._blobs.append(
+ PuffinBlobMetadata(
+ type="deletion-vector-v1",
+ fields=[],
+ snapshot_id=-1,
+ sequence_number=-1,
+ offset=0, # Will be set later
+ length=0, # Will be set later
+ properties=properties,
+ compression_codec=None, # Explicitly None
+ )
+ )
+
+ def finish(self) -> bytes:
+ with io.BytesIO() as out:
+ payload_buffer = io.BytesIO()
+ for blob_payload in self._blob_payloads:
+ payload_buffer.write(blob_payload)
+
+ updated_blobs_metadata: list[PuffinBlobMetadata] = []
+ current_offset = 4 # Start after file magic (4 bytes)
+ for i, blob_payload in enumerate(self._blob_payloads):
+ original_metadata_dict =
self._blobs[i].model_dump(by_alias=True, exclude_none=True)
+ original_metadata_dict["offset"] = current_offset
+ original_metadata_dict["length"] = len(blob_payload)
+
updated_blobs_metadata.append(PuffinBlobMetadata(**original_metadata_dict))
+ current_offset += len(blob_payload)
+
+ footer = Footer(blobs=updated_blobs_metadata)
+ footer_payload_bytes = footer.model_dump_json(by_alias=True,
exclude_none=True).encode("utf-8")
+
+ # Final assembly
+ out.write(MAGIC_BYTES)
+ out.write(payload_buffer.getvalue())
+ out.write(footer_payload_bytes)
Review Comment:
We need to write `MAGIC_BYTES` before the payload as well
https://github.com/apache/iceberg/blob/main/core/src/main/java/org/apache/iceberg/puffin/PuffinWriter.java#L152
--
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]
