ENH: serialize scipy sparse arrays (sparray / dok_array)

distributed/protocol/scipy.py

@@ -34,3 +34,28 @@ def deserialize_scipy_sparse_dok(header, frames):
     x = x_coo.todok()
 
     return x
+
+
+if hasattr(scipy.sparse, "sparray"):
+    # SciPy >= 1.8: the spmatrix classes are deprecated in favour of sparse
+    # arrays, which do not inherit from spmatrix and need their own handlers.
+    register_generic(scipy.sparse.sparray, "dask", dask_serialize, dask_deserialize)
+
+    @dask_serialize.register(scipy.sparse.dok_array)
+    def serialize_scipy_sparse_dok_array(x):
+        coo_header, coo_frames = dask_serialize(x.tocoo())
+
+        header = {"coo_header": coo_header}
+        frames = coo_frames
+
+        return header, frames
+
+    @dask_deserialize.register(scipy.sparse.dok_array)
+    def deserialize_scipy_sparse_dok_array(header, frames):
+        coo_header = header["coo_header"]
+        coo_frames = frames
+        x_coo = dask_deserialize(coo_header, coo_frames)
+
+        x = x_coo.todok()
+
+        return x

distributed/protocol/tests/test_scipy.py

@@ -13,18 +13,35 @@
 SCIPY_VERSION = Version(scipy.__version__)
 SCIPY_GE_1_15_0 = SCIPY_VERSION.release >= (1, 15, 0)
 
+SPARSE_MATRIX_TYPES = [
+    scipy_sparse.bsr_matrix,
+    scipy_sparse.coo_matrix,
+    scipy_sparse.csc_matrix,
+    scipy_sparse.csr_matrix,
+    scipy_sparse.dia_matrix,
+    scipy_sparse.dok_matrix,
+    scipy_sparse.lil_matrix,
+]
+
+# The spmatrix classes are deprecated in favour of sparse arrays (scipy >= 1.8)
+SPARSE_ARRAY_TYPES = (
+    [
+        scipy_sparse.bsr_array,
+        scipy_sparse.coo_array,
+        scipy_sparse.csc_array,
+        scipy_sparse.csr_array,
+        scipy_sparse.dia_array,
+        scipy_sparse.dok_array,
+        scipy_sparse.lil_array,
+    ]
+    if hasattr(scipy_sparse, "sparray")
+    else []
+)
+
 
 @pytest.mark.parametrize(
     "sparse_type",
-    [
-        scipy_sparse.bsr_matrix,
-        scipy_sparse.coo_matrix,
-        scipy_sparse.csc_matrix,
-        scipy_sparse.csr_matrix,
-        scipy_sparse.dia_matrix,
-        scipy_sparse.dok_matrix,
-        scipy_sparse.lil_matrix,
-    ],
+    SPARSE_MATRIX_TYPES + SPARSE_ARRAY_TYPES,
 )
 @pytest.mark.parametrize(
     "dtype",
@@ -33,16 +50,22 @@
         numpy.dtype("<f8"),
     ],
 )
+# Keep covering the legacy spmatrix classes until SciPy removes them; only the
+# "is being replaced by" deprecation is tolerated, everything else stays an error.
+@pytest.mark.filterwarnings("ignore:.*is being replaced by.*:DeprecationWarning")
 def test_serialize_scipy_sparse(sparse_type, dtype):
     a = numpy.array([[0, 1, 0], [2, 0, 3], [0, 4, 0]], dtype=dtype)
 
     anz = a.nonzero()
-    acoo = scipy_sparse.coo_matrix((a[anz], anz))
+    if sparse_type in SPARSE_MATRIX_TYPES:
+        acoo = scipy_sparse.coo_matrix((a[anz], anz))
+    else:
+        acoo = scipy_sparse.coo_array((a[anz], anz))
     asp = sparse_type(acoo)
 
     header, frames = serialize(asp, serializers=["dask"])
     asp2 = deserialize(header, frames)
 
-    a2 = asp2.todense()
+    a2 = asp2.toarray()
 
     assert (a == a2).all()