Handle matrix subclass serialization (#8480)

Author: fjetter

distributed/protocol/numpy.py

@@ -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):
@@ -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
 
 

distributed/protocol/tests/test_numpy.py

@@ -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):

pyproject.toml

@@ -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 = [