Handle matrix subclass serialization

fjetter · fjetter · commit 3b04a331ae23 · 2024-01-25T15:48:23.000+01:00
diff --git a/distributed/protocol/numpy.py b/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
 
 
diff --git a/distributed/protocol/tests/test_numpy.py b/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):
diff --git a/pyproject.toml b/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 = [

Original file line number	Diff line number	Diff line change
`@@ -77,6 +77,7 @@ def test_serialize():`
`77`	`77`	`np.arange(12)[::2], # non-contiguous array`
`78`	`78`	`np.ones(shape=(5, 6)).astype(dtype=[("total", "<f8"), ("n", "<f8")]),`
`79`	`79`	`np.broadcast_to(np.arange(3), shape=(10, 3)), # zero-strided array`
	`80`	`+ np.matrix([[1, 2], [3, 4]]),`
`80`	`81`	`],`
`81`	`82`	`)`
`82`	`83`	`def test_dumps_serialize_numpy(x):`
Original file line number	Diff line number	Diff line change
`@@ -151,6 +151,7 @@ filterwarnings = [`
`151`	`151`	`'''ignore:datetime\.datetime\.utc(fromtimestamp\|now)\(\) is deprecated and scheduled for removal in a future version.*:DeprecationWarning:dateutil''',`
`152`	`152`	`# https://github.com/dask/dask/pull/10622`
`153`	`153`	`'''ignore:Minimal version of pyarrow will soon be increased to 14.0.1''',`
	`154`	`+ '''ignore:the matrix subclass is not the recommended way''',`
`154`	`155`	`]`
`155`	`156`	`minversion = "6"`
`156`	`157`	`markers = [`