mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 21:14:14 +08:00
26358fa2d8
This PR adds a new util class that enables efficient appending of sequential byte data with custom serialization and deserialization. Pull Request resolved: https://github.com/pytorch/pytorch/pull/148226 Approved by: https://github.com/aorenste
110 lines
2.9 KiB
Python
110 lines
2.9 KiB
Python
from collections.abc import Iterable
|
|
from typing import Callable, Generic, TypeVar
|
|
|
|
|
|
T = TypeVar("T")
|
|
|
|
_ENCODING_VERSION: int = 1
|
|
|
|
__all__ = ["AppendingByteSerializer"]
|
|
|
|
|
|
#######################################
|
|
# Helper classes
|
|
#######################################
|
|
|
|
|
|
class BytesWriter:
|
|
def __init__(self, preallocate_size: int) -> None:
|
|
self._data = bytearray(preallocate_size)
|
|
|
|
def write_uint64(self, i: int) -> None:
|
|
self._data.extend(i.to_bytes(8, byteorder="big", signed=False))
|
|
|
|
def write_str(self, s: str) -> None:
|
|
payload = s.encode("utf-8")
|
|
self.write_bytes(payload)
|
|
|
|
def write_bytes(self, b: bytes) -> None:
|
|
self.write_uint64(len(b))
|
|
self._data.extend(b)
|
|
|
|
def to_bytes(self) -> bytes:
|
|
return bytes(self._data)
|
|
|
|
|
|
class BytesReader:
|
|
def __init__(self, data: bytes) -> None:
|
|
self._data = data
|
|
self._i = 0
|
|
|
|
def is_finished(self) -> bool:
|
|
return len(self._data) == self._i
|
|
|
|
def read_uint64(self) -> int:
|
|
result = int.from_bytes(
|
|
self._data[self._i : self._i + 8], byteorder="big", signed=False
|
|
)
|
|
self._i += 8
|
|
return result
|
|
|
|
def read_str(self) -> str:
|
|
return self.read_bytes().decode("utf-8")
|
|
|
|
def read_bytes(self) -> bytes:
|
|
size = self.read_uint64()
|
|
result = self._data[self._i : self._i + size]
|
|
self._i += size
|
|
return result
|
|
|
|
|
|
#######################################
|
|
# AppendingByteSerializer
|
|
#######################################
|
|
|
|
|
|
class AppendingByteSerializer(Generic[T]):
|
|
"""
|
|
Provides efficient serialization and deserialization of list of bytes
|
|
Note that this does not provide any guarantees around byte order
|
|
"""
|
|
|
|
_serialize_fn: Callable[[BytesWriter, T], None]
|
|
_writer: BytesWriter
|
|
_preallocate_size: int
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
serialize_fn: Callable[[BytesWriter, T], None],
|
|
preallocate_size: int = 0,
|
|
) -> None:
|
|
self._serialize_fn = serialize_fn
|
|
self._preallocate_size = preallocate_size
|
|
self.clear()
|
|
|
|
def clear(self) -> None:
|
|
self._writer = BytesWriter(preallocate_size=self._preallocate_size)
|
|
# First 8-bytes are for version
|
|
self._writer.write_uint64(_ENCODING_VERSION)
|
|
|
|
def append(self, data: T) -> None:
|
|
self._serialize_fn(self._writer, data)
|
|
|
|
def extend(self, elems: Iterable[T]) -> None:
|
|
for elem in elems:
|
|
self.append(elem)
|
|
|
|
def to_bytes(self) -> bytes:
|
|
return self._writer.to_bytes()
|
|
|
|
@staticmethod
|
|
def to_list(data: bytes, *, deserialize_fn: Callable[[BytesReader], T]) -> list[T]:
|
|
reader = BytesReader(data)
|
|
assert reader.read_uint64() == _ENCODING_VERSION
|
|
|
|
result: list[T] = []
|
|
while not reader.is_finished():
|
|
result.append(deserialize_fn(reader))
|
|
return result
|