Enabling serialization with pydata/sparse.

Author: rok

python/pyarrow/serialization.py

@@ -350,6 +350,35 @@ def _deserialize_scipy_sparse(data):
         pass
 
 
+# ----------------------------------------------------------------------
+# Set up serialization for pydata/sparse tensors.
+
+def _register_pydata_sparse_handlers(serialization_context):
+    try:
+        import sparse
+
+        def _serialize_pydata_sparse(obj):
+            if isinstance(obj, sparse.coo.core.COO):
+                return 'coo', pyarrow.SparseTensorCOO.from_numpy(
+                    obj.data, obj.coords.T, shape=obj.shape)
+
+        def _deserialize_pydata_sparse(data):
+            if data[0] == 'coo':
+                data_array, coords = data[1].to_numpy()
+                return sparse.COO(
+                    data=data_array[:, 0],
+                    coords=coords.T, shape=data[1].shape)
+
+        serialization_context.register_type(
+            sparse.coo.core.COO, 'sparse.coo.core.COO',
+            custom_serializer=_serialize_pydata_sparse,
+            custom_deserializer=_deserialize_pydata_sparse)
+
+    except ImportError:
+        # no pydata/sparse
+        pass
+
+
 def register_default_serialization_handlers(serialization_context):
 
     # ----------------------------------------------------------------------
@@ -403,6 +432,7 @@ def register_default_serialization_handlers(serialization_context):
     _register_collections_serialization_handlers(serialization_context)
     _register_custom_pandas_handlers(serialization_context)
     _register_scipy_handlers(serialization_context)
+    _register_pydata_sparse_handlers(serialization_context)
 
 
 def default_serialization_context():

python/pyarrow/tests/test_serialization.py

@@ -51,6 +51,11 @@
 except ImportError:
     sparse = None
 
+try:
+    import sparse
+except ImportError:
+    sparse = None
+
 
 def assert_equal(obj1, obj2):
     if torch is not None and torch.is_tensor(obj1) and torch.is_tensor(obj2):
@@ -621,6 +626,21 @@ def test_scipy_sparse_tensor_csr_serialization():
     assert np.array_equal(sparse_array.toarray(), result.toarray())
 
 
+@pytest.mark.skipif(not sparse, reason="requires pydata/sparse")
+def test_pydata_sparse__sparse_tensor_coo_serialization():
+    data = np.array([1, 2, 3, 4, 5, 6, 7])
+    row = np.array([0, 0, 2, 3, 1, 3, 0])
+    col = np.array([0, 2, 0, 4, 5, 5, 0])
+    coords = np.vstack([row, col]).T
+    shape = (4, 6)
+
+    sparse_array = sparse.COO(data=data, coords=coords, shape=shape)
+    serialized = pa.serialize(sparse_array)
+    result = serialized.deserialize()
+
+    assert np.array_equal(sparse_array.todense(), result.todense())
+
+
 @pytest.mark.filterwarnings(
     "ignore:the matrix subclass:PendingDeprecationWarning")
 def test_numpy_matrix_serialization(tmpdir):