From 03d882adf392cabd922e4809a7e159fa7d62bc89 Mon Sep 17 00:00:00 2001 From: Charan Rathore <180254320+charan-rathore@users.noreply.github.com> Date: Thu, 8 Oct 2026 10:17:40 +0530 Subject: [PATCH] ENH: serialize scipy sparse arrays (sparray / dok_array) The spmatrix classes are deprecated in favour of sparse arrays, which do not inherit from spmatrix and previously had no serialization handlers registered, so transmitting them failed with "TypeError: Could not serialize object of type ...". Register generic dask serialization for scipy.sparse.sparray and add dok_array handlers mirroring the dok_matrix ones. Parametrize the scipy protocol tests over both legacy matrix and modern array types. Fixes #9353 Signed-off-by: Charan Rathore <180254320+charan-rathore@users.noreply.github.com> --- distributed/protocol/scipy.py | 25 +++++++++++++ distributed/protocol/tests/test_scipy.py | 45 ++++++++++++++++++------ 2 files changed, 59 insertions(+), 11 deletions(-) diff --git a/distributed/protocol/scipy.py b/distributed/protocol/scipy.py index 90a7910693b..6c4d33524f8 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 8bdc2c581b6..13db489b791 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("