diff --git a/distributed/protocol/scipy.py b/distributed/protocol/scipy.py index 90a7910693..6c4d33524f 100644 --- a/distributed/protocol/scipy.py +++ b/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 diff --git a/distributed/protocol/tests/test_scipy.py b/distributed/protocol/tests/test_scipy.py index 8bdc2c581b..13db489b79 100644 --- a/distributed/protocol/tests/test_scipy.py +++ b/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("