Merge pull request #66 from awslabs/casting-nan

Fixing cast for nan values

Author: igorborgest

awswrangler/pandas.py

@@ -587,20 +587,20 @@ def to_parquet(self,
                           inplace=inplace)
 
     def to_s3(self,
-              dataframe,
-              path,
-              file_format,
-              database=None,
-              table=None,
+              dataframe: pd.DataFrame,
+              path: str,
+              file_format: str,
+              database: Optional[str] = None,
+              table: Optional[str] = None,
               partition_cols=None,
               preserve_index=True,
-              mode="append",
+              mode: str = "append",
               compression=None,
               procs_cpu_bound=None,
               procs_io_bound=None,
               cast_columns=None,
               extra_args=None,
-              inplace=True):
+              inplace: bool = True) -> List[str]:
         """
         Write a Pandas Dataframe on S3
         Optionally writes metadata on AWS Glue.
@@ -621,9 +621,9 @@ def to_s3(self,
         :param inplace: True is cheapest (CPU and Memory) but False leaves your DataFrame intact
         :return: List of objects written on S3
         """
-        if not partition_cols:
+        if partition_cols is None:
             partition_cols = []
-        if not cast_columns:
+        if cast_columns is None:
             cast_columns = {}
         dataframe = Pandas.normalize_columns_names_athena(dataframe, inplace=inplace)
         cast_columns = {Athena.normalize_column_name(k): v for k, v in cast_columns.items()}
@@ -748,20 +748,20 @@ def _data_to_s3_dataset_writer(dataframe,
                                    extra_args=None,
                                    isolated_dataframe=False):
         objects_paths = []
+        dataframe = Pandas._cast_pandas(dataframe=dataframe, cast_columns=cast_columns)
+        cast_columns_materialized = {c: t for c, t in cast_columns.items() if c not in partition_cols}
         if not partition_cols:
             object_path = Pandas._data_to_s3_object_writer(dataframe=dataframe,
                                                            path=path,
                                                            preserve_index=preserve_index,
                                                            compression=compression,
                                                            session_primitives=session_primitives,
                                                            file_format=file_format,
-                                                           cast_columns=cast_columns,
+                                                           cast_columns=cast_columns_materialized,
                                                            extra_args=extra_args,
                                                            isolated_dataframe=isolated_dataframe)
             objects_paths.append(object_path)
         else:
-            dataframe = Pandas._cast_pandas(dataframe=dataframe, cast_columns=cast_columns)
-            cast_columns_materialized = {c: t for c, t in cast_columns.items() if c not in partition_cols}
             dataframe = Pandas._cast_pandas(dataframe=dataframe, cast_columns=cast_columns)
             for keys, subgroup in dataframe.groupby(partition_cols):
                 subgroup = subgroup.drop(partition_cols, axis="columns")
@@ -790,7 +790,7 @@ def _cast_pandas(dataframe: pd.DataFrame, cast_columns: Dict[str, str]) -> pd.Da
             elif pandas_type == "date":
                 dataframe[col] = pd.to_datetime(dataframe[col]).dt.date
             else:
-                dataframe[col] = dataframe[col].astype(pandas_type)
+                dataframe[col] = dataframe[col].astype(pandas_type, skipna=True)
         return dataframe
 
     @staticmethod

data_samples/nan.csv

@@ -0,0 +1,5 @@
+"col1","col2","col3","col4","pt"
+,,,,"1"
+,,,,"2"
+,"foo","bar","baz","1"
+,"foo","bar","baz","2"

testing/test_awswrangler/test_pandas.py

@@ -947,7 +947,7 @@ def test_read_sql_athena_with_nulls(session, bucket, database):
     print(df2.dtypes)
     assert df2.dtypes[0] == "Int64"
     assert df2.dtypes[1] == "bool"
-    assert df2.dtypes[2] == "object"
+    assert df2.dtypes[2] == "bool"
     session.s3.delete_objects(path=path)
 
 
@@ -1149,3 +1149,50 @@ def test_partition_single_row(session, bucket, database, procs):
     assert df2.dtypes[1] == "object"
     assert df2.dtypes[2] == "object"
     session.s3.delete_objects(path=path)
+
+
+@pytest.mark.parametrize("partition_cols", [None, ["pt"]])
+def test_nan_cast(session, bucket, database, partition_cols):
+    dtypes = {"col1": "object", "col2": "object", "col3": "object", "col4": "object", "pt": "object"}
+    df = pd.read_csv("data_samples/nan.csv", dtype=dtypes)
+    print(df)
+    schema = {
+        "col1": "string",
+        "col2": "string",
+        "col3": "string",
+        "col4": "string",
+        "pt": "string",
+    }
+    path = f"s3://{bucket}/test/"
+    session.pandas.to_parquet(dataframe=df,
+                              database=database,
+                              path=path,
+                              partition_cols=partition_cols,
+                              mode="overwrite",
+                              cast_columns=schema)
+    df2 = None
+    for counter in range(10):
+        sleep(1)
+        df2 = session.pandas.read_sql_athena(sql="select * from test", database=database)
+        assert len(list(df.columns)) == len(list(df2.columns)) - 1
+        if len(df.index) == len(df2.index):
+            break
+    print(df2.dtypes)
+    assert len(df.index) == len(df2.index)
+    assert df2.dtypes[0] == "object"
+    assert df2.dtypes[1] == "object"
+    assert df2.dtypes[2] == "object"
+    assert df2.dtypes[3] == "object"
+    assert df2.iloc[:, 0].isna().sum() == 4
+    assert df2.iloc[:, 1].isna().sum() == 2
+    assert df2.iloc[:, 2].isna().sum() == 2
+    assert df2.iloc[:, 3].isna().sum() == 2
+    assert df2.iloc[:, 4].isna().sum() == 0
+    assert df2.iloc[:, 5].isna().sum() == 0
+    if partition_cols is None:
+        assert df2.dtypes[4] == "object"
+        assert df2.dtypes[5] == "Int64"
+    else:
+        assert df2.dtypes[4] == "Int64"
+        assert df2.dtypes[5] == "object"
+    session.s3.delete_objects(path=path)