feat: Support npz archives in NumpyDeserializer (#3799)

Author: athewsey

src/sagemaker/deserializers.py

@@ -189,7 +189,12 @@ def deserialize(self, stream, content_type):
 
 
 class NumpyDeserializer(SimpleBaseDeserializer):
-    """Deserialize a stream of data in .npy or UTF-8 CSV/JSON format to a numpy array."""
+    """Deserialize a stream of data in .npy, .npz or UTF-8 CSV/JSON format to a numpy array.
+
+    Note that when using application/x-npz archive format, the result will usually be a
+    dictionary-like object containing multiple arrays (as per ``numpy.load()``) - instead of a
+    single array.
+    """
 
     def __init__(self, dtype=None, accept="application/x-npy", allow_pickle=True):
         """Initialize a ``NumpyDeserializer`` instance.
@@ -223,6 +228,11 @@ def deserialize(self, stream, content_type):
                 return np.array(json.load(codecs.getreader("utf-8")(stream)), dtype=self.dtype)
             if content_type == "application/x-npy":
                 return np.load(io.BytesIO(stream.read()), allow_pickle=self.allow_pickle)
+            if content_type == "application/x-npz":
+                try:
+                    return np.load(io.BytesIO(stream.read()), allow_pickle=self.allow_pickle)
+                finally:
+                    stream.close()
         finally:
             stream.close()
 

tests/unit/sagemaker/test_deserializers.py

@@ -164,6 +164,23 @@ def test_numpy_deserializer_from_npy_object_array_with_allow_pickle_false():
         numpy_deserializer.deserialize(stream, "application/x-npy")
 
 
+def test_numpy_deserializer_from_npz(numpy_deserializer):
+    arrays = {
+        "alpha": np.ones((2, 3)),
+        "beta": np.zeros((3, 2)),
+    }
+    stream = io.BytesIO()
+    np.savez_compressed(stream, **arrays)
+    stream.seek(0)
+
+    result = numpy_deserializer.deserialize(stream, "application/x-npz")
+
+    assert isinstance(result, np.lib.npyio.NpzFile)
+    assert set(arrays.keys()) == set(result.keys())
+    for key, arr in arrays.items():
+        assert np.array_equal(arr, result[key])
+
+
 @pytest.fixture
 def json_deserializer():
     return JSONDeserializer()