Skip to content

Commit

Permalink
Handle matrix subclass serialization (#8480)
Browse files Browse the repository at this point in the history
  • Loading branch information
fjetter authored Jan 26, 2024
1 parent a170014 commit 22f037f
Show file tree
Hide file tree
Showing 3 changed files with 11 additions and 1 deletion.
10 changes: 9 additions & 1 deletion distributed/protocol/numpy.py
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,13 @@ def itemsize(dt):
return result


@dask_serialize.register(np.matrix)
def serialize_numpy_matrix(x, context=None):
header, frames = serialize_numpy_ndarray(x)
header["matrix"] = True
return header, frames


@dask_serialize.register(np.ndarray)
def serialize_numpy_ndarray(x, context=None):
if x.dtype.hasobject or (x.dtype.flags & np_core.multiarray.LIST_PICKLE):
Expand Down Expand Up @@ -151,7 +158,8 @@ def deserialize_numpy_ndarray(header, frames):
# buffers the decompressed output is deep-copied beforehand into a
# bytearray in order to merge it.
x = np.require(x, requirements=["W"])

if header.get("matrix"):
x = np.asmatrix(x)
return x


Expand Down
1 change: 1 addition & 0 deletions distributed/protocol/tests/test_numpy.py
Original file line number Diff line number Diff line change
Expand Up @@ -77,6 +77,7 @@ def test_serialize():
np.arange(12)[::2], # non-contiguous array
np.ones(shape=(5, 6)).astype(dtype=[("total", "<f8"), ("n", "<f8")]),
np.broadcast_to(np.arange(3), shape=(10, 3)), # zero-strided array
np.matrix([[1, 2], [3, 4]]),
],
)
def test_dumps_serialize_numpy(x):
Expand Down
1 change: 1 addition & 0 deletions pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -151,6 +151,7 @@ filterwarnings = [
'''ignore:datetime\.datetime\.utc(fromtimestamp|now)\(\) is deprecated and scheduled for removal in a future version.*:DeprecationWarning:dateutil''',
# https://github.com/dask/dask/pull/10622
'''ignore:Minimal version of pyarrow will soon be increased to 14.0.1''',
'''ignore:the matrix subclass is not the recommended way''',
]
minversion = "6"
markers = [
Expand Down

0 comments on commit 22f037f

Please sign in to comment.