Feature/python 312 pyarrow strings

earthmover/__init__.py

@@ -7,7 +7,5 @@
 dask.config.set({'dataframe.query-planning': False})
 
 # performance enhancements
-dask.config.set({"dataframe.convert-string": True})
 import pandas as pd
-pd.options.mode.copy_on_write = True
-pd.options.mode.string_storage = "pyarrow"
+pd.options.mode.copy_on_write = True

earthmover/nodes/source.py

@@ -159,7 +159,6 @@ def execute(self):
             if self.optional and not os.path.exists(self.file) or (os.path.isdir(self.file) and not os.listdir(self.file)):
                 self.data = pd.DataFrame(columns=self.columns_list, dtype="string")
             else:
-                dask_config.set({'dataframe.convert-string': False})
                 self.data = self.read_lambda(self.file, self.config)
                 if not self.is_remote:
                     self.size = os.path.getsize(self.file)
@@ -245,7 +244,7 @@ def __get_skiprows(config: 'YamlMapping'):
 
         # We don't want to activate the function inside this helper function.
         read_lambda_mapping = {
-            'csv'       : lambda file, config: dd.read_csv(file, sep=sep, dtype=str, encoding=config.get('encoding', "utf8"), keep_default_na=False, skiprows=__get_skiprows(config)),
+            'csv'       : lambda file, config: dd.read_csv(file, sep=sep, dtype="string", encoding=config.get('encoding', "utf8"), keep_default_na=False, skiprows=__get_skiprows(config)),
             'excel'     : lambda file, config: pd.read_excel(file, sheet_name=config.get("sheet", 0), keep_default_na=False),
             'feather'   : lambda file, _     : pd.read_feather(file),
             'fixedwidth': lambda file, _     : dd.read_fwf(file),
@@ -258,7 +257,7 @@ def __get_skiprows(config: 'YamlMapping'):
             'spss'      : lambda file, _     : pd.read_spss(file),
             'stata'     : lambda file, _     : pd.read_stata(file),
             'xml'       : lambda file, config: pd.read_xml(file, xpath=config.get('xpath', "./*")),
-            'tsv'       : lambda file, config: dd.read_csv(file, sep=sep, dtype=str, encoding=config.get('encoding', "utf8"), keep_default_na=False, skiprows=__get_skiprows(config)),
+            'tsv'       : lambda file, config: dd.read_csv(file, sep=sep, dtype="string", encoding=config.get('encoding', "utf8"), keep_default_na=False, skiprows=__get_skiprows(config)),
         }
         return read_lambda_mapping.get(file_type)
 

earthmover/tests/earthmover.yaml

@@ -280,13 +280,13 @@ transformations:
       - operation: add_columns
         columns:
           {% for month in months %}
-          avg_days_rain_{{month}}: "{%raw%}{{weather.avg_days_of_rain.{%endraw%}{{month}}{%raw%}}}{%endraw%}"
-          avg_high_temp_{{month}}: "{%raw%}{{weather.temperatures.avg_highs.{%endraw%}{{month}}{%raw%}}}{%endraw%}"
-          avg_low_temp_{{month}}: "{%raw%}{{weather.temperatures.avg_lows.{%endraw%}{{month}}{%raw%}}}{%endraw%}"
+          avg_days_rain_{{month}}: "{%raw%}{{fromjson(weather).avg_days_of_rain.{%endraw%}{{month}}{%raw%}}}{%endraw%}"
+          avg_high_temp_{{month}}: "{%raw%}{{fromjson(weather).temperatures.avg_highs.{%endraw%}{{month}}{%raw%}}}{%endraw%}"
+          avg_low_temp_{{month}}: "{%raw%}{{fromjson(weather).temperatures.avg_lows.{%endraw%}{{month}}{%raw%}}}{%endraw%}"
           {% endfor %}
-          avg_days_rain_year: "{%raw%}{{{%endraw%}{%for month in months%}weather.avg_days_of_rain.{{month}}|int{%if not loop.last%} + {%endif%}{%endfor%}{%raw%}}}{%endraw%}"
-          avg_high_temp: "{%raw%}{{(({%endraw%}{%for month in months%}weather.temperatures.avg_highs.{{month}}|int{%if not loop.last%} + {%endif%}{%endfor%}{%raw%})/12)|int}}{%endraw%}"
-          avg_low_temp: "{%raw%}{{(({%endraw%}{%for month in months%}weather.temperatures.avg_lows.{{month}}|int{%if not loop.last%} + {%endif%}{%endfor%}{%raw%})/12)|int}}{%endraw%}"
+          avg_days_rain_year: "{%raw%}{{{%endraw%}{%for month in months%}fromjson(weather).avg_days_of_rain.{{month}}|int{%if not loop.last%} + {%endif%}{%endfor%}{%raw%}}}{%endraw%}"
+          avg_high_temp: "{%raw%}{{(({%endraw%}{%for month in months%}fromjson(weather).temperatures.avg_highs.{{month}}|int{%if not loop.last%} + {%endif%}{%endfor%}{%raw%})/12)|int}}{%endraw%}"
+          avg_low_temp: "{%raw%}{{(({%endraw%}{%for month in months%}fromjson(weather).temperatures.avg_lows.{{month}}|int{%if not loop.last%} + {%endif%}{%endfor%}{%raw%})/12)|int}}{%endraw%}"
       - operation: drop_columns
         columns:
           - weather

earthmover/util.py

@@ -1,11 +1,12 @@
+import ast  # Used for malformed JSON
 import jinja2
 import hashlib
 import json
 import os
 
 from sys import exc_info
 
-from typing import Optional
+from typing import Optional, Union
 from typing import TYPE_CHECKING
 if TYPE_CHECKING:
     from earthmover.error_handler import ErrorHandler
@@ -122,6 +123,17 @@ def jinja2_template_error_lineno():
             return tb.tb_lineno
         tb = tb.tb_next
 
+def fromjson(obj: Union[str, dict]) -> dict:
+    """
+    Helper method to parse malformed JSON with single quotes.
+    """
+    if isinstance(obj, dict):
+        return obj
+
+    try:
+        return json.loads(obj)
+    except json.decoder.JSONDecodeError:
+        return ast.literal_eval(obj)
 
 def build_jinja_template(template_string: str, macros: str = ""):
     """
@@ -132,6 +144,6 @@ def build_jinja_template(template_string: str, macros: str = ""):
     ).from_string(macros.strip() + template_string)
 
     template.globals['md5'] = lambda x: hashlib.md5(x.encode('utf-8')).hexdigest()
-    template.globals['fromjson'] = lambda x: json.loads(x)
+    template.globals['fromjson'] = fromjson
 
     return template