def _deserialize_output_memory(obj_id: plasma.ObjectID,
client: plasma.PlasmaClient) -> Any:
"""Gets data from memory.
Args:
obj_id: The ID of the object to retrieve from the plasma store.
client: A PlasmaClient to interface with the in-memory object
store.
Returns:
The unserialized data from the store corresponding to the
`obj_id`.
Raises:
ObjectNotFoundError: If the specified `obj_id` is not in the
store.
ValueError: If the serialization type in the metadata is not
valid.
"""
obj_ids = [obj_id]
# TODO: the get_buffers allows for batch, which we want to use in
# the future.
buffers = client.get_buffers(obj_ids, with_meta=True, timeout_ms=1000)
# Since we currently know that we are only restrieving one buffer,
# we can instantly get its metadata and buffer.
metadata, buffer = buffers[0]
# Getting the buffer timed out. We conclude that the object has not
# yet been written to the store and maybe never will.
if metadata is None and buffer is None:
raise ObjectNotFoundError(
f'Object with ObjectID "{obj_id}" does not exist in store.')
metadata = metadata.decode("utf-8").split(Config.__METADATA_SEPARATOR__)
_, _, serialization, _ = metadata
if serialization == Serialization.ARROW_TABLE.name:
# Read all batches as a table.
stream = pa.ipc.open_stream(buffer)
return stream.read_all()
elif serialization == Serialization.ARROW_BATCH.name:
# Return the first batch (the only one).
stream = pa.ipc.open_stream(buffer)
return [b for b in stream][0]
elif serialization == Serialization.PICKLE.name:
# Can load the buffer directly because its a bytes-like-object:
# https://docs.python.org/3/library/pickle.html#pickle.loads
return pickle.loads(buffer)
else:
raise ValueError(
"Object was serialized with an unsupported serialization")
def _get_output_memory(obj_id: plasma.ObjectID,
client: plasma.PlasmaClient) -> Any:
"""Gets data from memory.
Args:
obj_id: The ID of the object to retrieve from the plasma store.
client: A PlasmaClient to interface with the in-memory object
store.
Returns:
The unserialized data from the store corresponding to the
`obj_id`.
Raises:
ObjectNotFoundError: If the specified `obj_id` is not in the
store.
"""
obj_ids = [obj_id]
# TODO: the get_buffers allows for batch, which we want to use in
# the future.
buffers = client.get_buffers(obj_ids, with_meta=True, timeout_ms=1000)
# Since we currently know that we are only restrieving one buffer,
# we can instantly get its metadata and buffer.
metadata, buffer = buffers[0]
# Getting the buffer timed out. We conclude that the object has not
# yet been written to the store and maybe never will.
if metadata is None and buffer is None:
raise ObjectNotFoundError(
f'Object with ObjectID "{obj_id}" does not exist in store.')
buffers_bytes = buffer.to_pybytes()
obj = pa.deserialize(buffers_bytes)
# If the metadata stated that the object was pickled, then we need
# to additionally unpickle the obj.
if metadata == bytes(f"{Config.IDENTIFIER_SERIALIZATION};arrowpickle",
"utf-8"):
obj = pickle.loads(obj)
return obj