feature: Add sparse matrix serializer (#1739)

Author: bveeramani

src/sagemaker/serializers.py

@@ -20,6 +20,13 @@
 
 import numpy as np
 
+from sagemaker.utils import DeferredError
+
+try:
+    import scipy
+except ImportError as e:
+    scipy = DeferredError(e)
+
 
 class BaseSerializer(abc.ABC):
     """Abstract base class for creation of new serializers.
@@ -183,3 +190,25 @@ def serialize(self, data):
             return json.dumps(data.tolist())
 
         return json.dumps(data)
+
+
+class SparseMatrixSerializer(BaseSerializer):
+    """Serialize a sparse matrix to a buffer using the .npz format."""
+
+    CONTENT_TYPE = "application/x-npz"
+
+    def serialize(self, data):
+        """Serialize a sparse matrix to a buffer using the .npz format.
+
+        Sparse matrices can be in the ``csc``, ``csr``, ``bsr``, ``dia`` or
+        ``coo`` formats.
+
+        Args:
+            data (scipy.sparse.spmatrix): The sparse matrix to serialize.
+
+        Returns:
+            io.BytesIO: A buffer containing the serialized sparse matrix.
+        """
+        buffer = io.BytesIO()
+        scipy.sparse.save_npz(buffer, data)
+        return buffer.getvalue()

tests/unit/sagemaker/test_serializers.py

@@ -18,8 +18,14 @@
 
 import numpy as np
 import pytest
-
-from sagemaker.serializers import CSVSerializer, NumpySerializer, JSONSerializer
+import scipy
+
+from sagemaker.serializers import (
+    CSVSerializer,
+    NumpySerializer,
+    JSONSerializer,
+    SparseMatrixSerializer,
+)
 from tests.unit import DATA_DIR
 
 
@@ -227,3 +233,16 @@ def test_json_serializer_csv_buffer(json_serializer):
         csv_file.seek(0)
         result = json_serializer.serialize(csv_file)
         assert result == validation_value
+
+
+@pytest.fixture
+def sparse_matrix_serializer():
+    return SparseMatrixSerializer()
+
+
+def test_sparse_matrix_serializer(sparse_matrix_serializer):
+    data = scipy.sparse.csc_matrix(np.array([[0, 0, 3], [4, 0, 0]]))
+    stream = io.BytesIO(sparse_matrix_serializer.serialize(data))
+    result = scipy.sparse.load_npz(stream).toarray()
+    expected = data.toarray()
+    assert np.array_equal(result, expected)