Merge branch 'zwei' into remove-legacy-serde

Author: bveeramani

src/sagemaker/deserializers.py

@@ -22,6 +22,13 @@
 
 import numpy as np
 
+from sagemaker.utils import DeferredError
+
+try:
+    import pandas
+except ImportError as e:
+    pandas = DeferredError(e)
+
 
 class BaseDeserializer(abc.ABC):
     """Abstract base class for creation of new deserializers.
@@ -208,3 +215,31 @@ def deserialize(self, stream, content_type):
             return json.load(codecs.getreader("utf-8")(stream))
         finally:
             stream.close()
+
+
+class PandasDeserializer(BaseDeserializer):
+    """Deserialize CSV or JSON data from an inference endpoint into a pandas dataframe."""
+
+    ACCEPT = "text/csv"
+
+    def deserialize(self, stream, content_type):
+        """Deserialize CSV or JSON data from an inference endpoint into a pandas
+        dataframe.
+
+        If the data is JSON, the data should be formatted in the 'columns' orient.
+        See https://pandas.pydata.org/pandas-docs/stable/reference/api/pandas.read_json.html
+
+        Args:
+            stream (botocore.response.StreamingBody): Data to be deserialized.
+            content_type (str): The MIME type of the data.
+
+        Returns:
+            pandas.DataFrame: The data deserialized into a pandas DataFrame.
+        """
+        if content_type == "text/csv":
+            return pandas.read_csv(stream)
+
+        if content_type == "application/json":
+            return pandas.read_json(stream)
+
+        raise ValueError("%s cannot read content type %s." % (__class__.__name__, content_type))

tests/unit/sagemaker/test_deserializers.py

@@ -13,8 +13,10 @@
 from __future__ import absolute_import
 
 import io
+import json
 
 import numpy as np
+import pandas as pd
 import pytest
 
 from sagemaker.deserializers import (
@@ -24,6 +26,7 @@
     StreamDeserializer,
     NumpyDeserializer,
     JSONDeserializer,
+    PandasDeserializer,
 )
 
 
@@ -171,3 +174,25 @@ def test_json_deserializer_invalid_data(json_deserializer):
     with pytest.raises(ValueError) as error:
         json_deserializer.deserialize(io.BytesIO(b"[[1]"), "application/json")
     assert "column" in str(error)
+
+
+@pytest.fixture
+def pandas_deserializer():
+    return PandasDeserializer()
+
+
+def test_pandas_deserializer_json(pandas_deserializer):
+    data = {"col 1": {"row 1": "a", "row 2": "c"}, "col 2": {"row 1": "b", "row 2": "d"}}
+    stream = io.StringIO(json.dumps(data))
+    result = pandas_deserializer.deserialize(stream, "application/json")
+    expected = pd.DataFrame(
+        [["a", "b"], ["c", "d"]], index=["row 1", "row 2"], columns=["col 1", "col 2"]
+    )
+    assert result.equals(expected)
+
+
+def test_pandas_deserializer_csv(pandas_deserializer):
+    stream = io.StringIO("col 1,col 2\na,b\nc,d")
+    result = pandas_deserializer.deserialize(stream, "text/csv")
+    expected = pd.DataFrame([["a", "b"], ["c", "d"]], columns=["col 1", "col 2"])
+    assert result.equals(expected)