Merge pull request #76 from wamartin-aml/wamartin/stopInferringJSONtoDataFrame

Stop inferring data types when converting JSON to DataFrame

Author: wamartin-aml

inference_schema/parameter_types/pandas_parameter_type.py

@@ -80,7 +80,7 @@ def deserialize_input(self, input_data):
             raise Exception("Error, unable to convert input of type {} into Pandas Dataframe".format(type(input_data)))
 
         string_stream = StringIO(json.dumps(input_data))
-        data_frame = pd.read_json(string_stream, orient=self.orient)
+        data_frame = pd.read_json(string_stream, orient=self.orient, dtype=False)
 
         if self.apply_column_names:
             data_frame.columns = self.sample_input.columns.copy()

tests/conftest.py

@@ -164,6 +164,29 @@ def pandas_url_func(param):
     return pandas_url_func
 
 
+@pytest.fixture(scope="session")
+def pandas_sample_input_with_categorical():
+    pandas_input_data = {'state': ['characters'], 'cat': ['000']}
+    return pd.DataFrame(data=pandas_input_data)
+
+
+@pytest.fixture(scope="session")
+def decorated_pandas_categorical_func(pandas_sample_input_with_categorical):
+    @input_schema('param', PandasParameterType(pandas_sample_input_with_categorical))
+    def pandas_categorical_func(param):
+        """
+
+        :param param:
+        :type param: pd.DataFrame
+        :return:
+        :rtype: string
+        """
+        assert type(param) is pd.DataFrame
+        return param['cat'][0]
+
+    return pandas_categorical_func
+
+
 @pytest.fixture(scope="session")
 def decorated_spark_func():
     spark_session = SparkSession.builder.config('spark.driver.host', '127.0.0.1').getOrCreate()

tests/test_pandas_parameter_type.py

@@ -67,6 +67,17 @@ def test_pandas_url_handling(self, decorated_pandas_uri_func):
         result = decorated_pandas_uri_func(pandas_input)
         assert website == result
 
+    def test_pandas_categorical_handling(self, decorated_pandas_categorical_func):
+        pandas_input = {'state': ['WA'], 'cat': ['000']}
+        categorical = pandas_input['cat'][0]
+        result = decorated_pandas_categorical_func(pandas_input)
+        assert categorical == result
+
+        pandas_input = {'state': ['WA'], 'cat': ['001']}
+        categorical = pandas_input['cat'][0]
+        result = decorated_pandas_categorical_func(pandas_input)
+        assert categorical == result
+
 
 class TestNestedType(object):