From 5fbb0ac1d9e6ba12888e75a9938933359cfa2023 Mon Sep 17 00:00:00 2001 From: "Mads R. B. Kristensen" Date: Wed, 17 Mar 2021 09:40:13 +0100 Subject: [PATCH] pickle_loads(): Handle empty memoryview Fixes https://github.com/dask/distributed/issues/4594 --- distributed/protocol/serialize.py | 8 ++++++-- distributed/protocol/tests/test_pickle.py | 11 +++++++++++ 2 files changed, 17 insertions(+), 2 deletions(-) diff --git a/distributed/protocol/serialize.py b/distributed/protocol/serialize.py index 228f02d78ca..49a2fdab09d 100644 --- a/distributed/protocol/serialize.py +++ b/distributed/protocol/serialize.py @@ -74,9 +74,13 @@ def pickle_loads(header, frames): mv = memoryview(buffers[i]) if writeable[i] == mv.readonly: if mv.readonly: - buffers[i] = memoryview(bytearray(mv)).cast(mv.format, mv.shape) + buf = memoryview(bytearray(mv)) else: - buffers[i] = memoryview(bytes(mv)).cast(mv.format, mv.shape) + buf = memoryview(bytes(mv)) + if buf.nbytes > 0: + buffers[i] = buf.cast(mv.format, mv.shape) + else: + buffers[i] = buf.cast(mv.format) return pickle.loads(x, buffers=buffers) diff --git a/distributed/protocol/tests/test_pickle.py b/distributed/protocol/tests/test_pickle.py index d7a2ad6d1ad..6430ea352b8 100644 --- a/distributed/protocol/tests/test_pickle.py +++ b/distributed/protocol/tests/test_pickle.py @@ -8,6 +8,7 @@ from distributed.protocol import deserialize, serialize from distributed.protocol.pickle import HIGHEST_PROTOCOL, dumps, loads +from distributed.protocol.serialize import pickle_dumps if sys.version_info < (3, 8): try: @@ -71,6 +72,16 @@ def __reduce_ex__(self, protocol): assert isinstance(f[0], bytes) +def test_pickle_empty(): + np = pytest.importorskip("numpy") + x = np.arange(2)[0:0] # Empty view + header, frames = pickle_dumps(x) + header["writeable"] = [False] * len(frames) + y = deserialize(header, frames) + assert memoryview(y).nbytes == 0 + assert memoryview(y).readonly + + def test_pickle_numpy(): np = pytest.importorskip("numpy") x = np.ones(5)