Adding support for cast data types from Pandas to parquet (Nested types not supported yet.)

Author: igorborgest

awswrangler/glue.py

@@ -45,6 +45,32 @@ def get_table_python_types(self, database, table):
         dtypes = self.get_table_athena_types(database=database, table=table)
         return {k: Glue.type_athena2python(v) for k, v in dtypes.items()}
 
+    @staticmethod
+    def type_athena2pyarrow(dtype):
+        dtype = dtype.lower()
+        if dtype == "tinyint":
+            return "int8"
+        if dtype == "smallint":
+            return "int16"
+        elif dtype in ["int", "integer"]:
+            return "int32"
+        elif dtype == "bigint":
+            return "int64"
+        elif dtype == "float":
+            return "float32"
+        elif dtype == "double":
+            return "float64"
+        elif dtype in ["boolean", "bool"]:
+            return "bool"
+        elif dtype in ["string", "char", "varchar", "array", "row", "map"]:
+            return "string"
+        elif dtype == "timestamp":
+            return "timestamp[ns]"
+        elif dtype == "date":
+            return "date32"
+        else:
+            raise UnsupportedType(f"Unsupported Athena type: {dtype}")
+
     @staticmethod
     def type_pyarrow2athena(dtype):
         dtype_str = str(dtype).lower()
@@ -241,21 +267,15 @@ def get_connection_details(self, name):
             Name=name, HidePassword=False)["Connection"]
 
     @staticmethod
-    def _extract_pyarrow_schema(dataframe, preserve_index, cast_columns=None):
+    def _extract_pyarrow_schema(dataframe, preserve_index):
         cols = []
         cols_dtypes = {}
         schema = []
 
-        casted = []
-        if cast_columns is not None:
-            casted = cast_columns.keys()
-
         for name, dtype in dataframe.dtypes.to_dict().items():
             dtype = str(dtype)
             if dtype == "Int64":
                 cols_dtypes[name] = "int64"
-            elif name in casted:
-                cols_dtypes[name] = cast_columns[name]
             else:
                 cols.append(name)
 
@@ -281,18 +301,22 @@ def _build_schema(dataframe,
             partition_cols = []
 
         pyarrow_schema = Glue._extract_pyarrow_schema(
-            dataframe=dataframe,
-            preserve_index=preserve_index,
-            cast_columns=cast_columns)
+            dataframe=dataframe, preserve_index=preserve_index)
 
         schema_built = []
         partition_cols_types = {}
         for name, dtype in pyarrow_schema:
-            athena_type = Glue.type_pyarrow2athena(dtype)
-            if name in partition_cols:
-                partition_cols_types[name] = athena_type
+            if (cast_columns is not None) and (name in cast_columns.keys()):
+                if name in partition_cols:
+                    partition_cols_types[name] = cast_columns[name]
+                else:
+                    schema_built.append((name, cast_columns[name]))
             else:
-                schema_built.append((name, athena_type))
+                athena_type = Glue.type_pyarrow2athena(dtype)
+                if name in partition_cols:
+                    partition_cols_types[name] = athena_type
+                else:
+                    schema_built.append((name, athena_type))
 
         partition_cols_schema_built = [(name, partition_cols_types[name])
                                        for name in partition_cols]

awswrangler/pandas.py

@@ -14,7 +14,7 @@
     AthenaQueryError, EmptyS3Object, LineTerminatorNotFound, EmptyDataframe, \
     InvalidSerDe, InvalidCompression
 from awswrangler.utils import calculate_bounders
-from awswrangler import s3
+from awswrangler import s3, glue
 
 logger = logging.getLogger(__name__)
 
@@ -859,18 +859,21 @@ def write_parquet_dataframe(dataframe,
             if str(dtype) == "Int64":
                 dataframe[name] = dataframe[name].astype("float64")
                 casted_in_pandas.append(name)
-                cast_columns[name] = "int64"
+                cast_columns[name] = "bigint"
                 logger.debug(f"Casting column {name} Int64 to float64")
         table = pyarrow.Table.from_pandas(df=dataframe,
                                           preserve_index=preserve_index,
                                           safe=False)
         if cast_columns:
             for col_name, dtype in cast_columns.items():
                 col_index = table.column_names.index(col_name)
-                table = table.set_column(col_index,
-                                         table.column(col_name).cast(dtype))
+                pyarrow_dtype = glue.Glue.type_athena2pyarrow(dtype)
+                table = table.set_column(
+                    col_index,
+                    table.column(col_name).cast(pyarrow_dtype))
                 logger.debug(
-                    f"Casting column {col_name} ({col_index}) to {dtype}")
+                    f"Casting column {col_name} ({col_index}) to {dtype} ({pyarrow_dtype})"
+                )
         with fs.open(path, "wb") as f:
             parquet.write_table(table,
                                 f,

testing/test_awswrangler/test_cloudwatchlogs.py

@@ -64,7 +64,7 @@ def logstream(cloudformation_outputs, loggroup):
     if token:
         args["sequenceToken"] = token
     client.put_log_events(**args)
-    sleep(180)
+    sleep(300)
     yield logstream
 
 

testing/test_awswrangler/test_pandas.py

@@ -99,7 +99,7 @@ def logstream(cloudformation_outputs, loggroup):
     if token:
         args["sequenceToken"] = token
     client.put_log_events(**args)
-    sleep(180)
+    sleep(300)
     yield logstream
 
 
@@ -243,22 +243,21 @@ def test_to_s3(
         assert len(list(dataframe.columns)) == len(list(dataframe2.columns))
 
 
-def test_to_parquet_with_cast(
+def test_to_parquet_with_cast_int(
         session,
         bucket,
         database,
 ):
     dataframe = pandas.read_csv("data_samples/nano.csv",
                                 dtype={"id": "Int64"},
                                 parse_dates=["date", "time"])
-    print(dataframe.dtypes)
     session.pandas.to_parquet(dataframe=dataframe,
                               database=database,
                               path=f"s3://{bucket}/test/",
                               preserve_index=False,
                               mode="overwrite",
                               procs_cpu_bound=1,
-                              cast_columns={"value": "int64"})
+                              cast_columns={"value": "int"})
     dataframe2 = None
     for counter in range(10):
         dataframe2 = session.pandas.read_sql_athena(sql="select * from test",
@@ -664,3 +663,82 @@ def test_to_parquet_lists(session, bucket, database):
     val = dataframe[dataframe["id"] == 0].iloc[0]["col_list_int"]
     val2 = dataframe2[dataframe2["id"] == 0].iloc[0]["col_list_int"]
     assert val == val2
+
+
+def test_to_parquet_cast(session, bucket, database):
+    dataframe = pandas.DataFrame({
+        "id": [0, 1],
+        "col_int": [[1, 2], [3, 4, 5]],
+        "col_float": [[1.0, 2.0, 3.0], [4.0, 5.0]],
+        "col_string": [["foo"], ["boo", "bar"]],
+        "col_timestamp": [[datetime(2019, 1, 1),
+                           datetime(2019, 1, 2)], [datetime(2019, 1, 3)]],
+        "col_date": [[date(2019, 1, 1), date(2019, 1, 2)], [date(2019, 1, 3)]],
+        "col_list_int": [[[1]], [[2, 3], [4, 5, 6]]],
+        "col_list_list_string": [[[["foo"]]], [[["boo", "bar"]]]],
+    })
+    paths = session.pandas.to_parquet(dataframe=dataframe,
+                                      database=database,
+                                      path=f"s3://{bucket}/test/",
+                                      preserve_index=False,
+                                      mode="overwrite",
+                                      procs_cpu_bound=1)
+    assert len(paths) == 1
+    dataframe2 = None
+    for counter in range(10):
+        dataframe2 = session.pandas.read_sql_athena(
+            sql="select id, col_int, col_float, col_list_int from test",
+            database=database)
+        if len(dataframe.index) == len(dataframe2.index):
+            break
+        sleep(2)
+    assert len(dataframe.index) == len(dataframe2.index)
+    assert 4 == len(list(dataframe2.columns))
+    val = dataframe[dataframe["id"] == 0].iloc[0]["col_list_int"]
+    val2 = dataframe2[dataframe2["id"] == 0].iloc[0]["col_list_int"]
+    assert val == val2
+
+
+def test_to_parquet_with_cast_null(
+        session,
+        bucket,
+        database,
+):
+    dataframe = pandas.DataFrame({
+        "id": [0, 1],
+        "col_null_tinyint": [None, None],
+        "col_null_smallint": [None, None],
+        "col_null_int": [None, None],
+        "col_null_bigint": [None, None],
+        "col_null_float": [None, None],
+        "col_null_double": [None, None],
+        "col_null_string": [None, None],
+        "col_null_date": [None, None],
+        "col_null_timestamp": [None, None],
+    })
+    session.pandas.to_parquet(dataframe=dataframe,
+                              database=database,
+                              path=f"s3://{bucket}/test/",
+                              preserve_index=False,
+                              mode="overwrite",
+                              procs_cpu_bound=1,
+                              cast_columns={
+                                  "col_null_tinyint": "tinyint",
+                                  "col_null_smallint": "smallint",
+                                  "col_null_int": "int",
+                                  "col_null_bigint": "bigint",
+                                  "col_null_float": "float",
+                                  "col_null_double": "double",
+                                  "col_null_string": "string",
+                                  "col_null_date": "date",
+                                  "col_null_timestamp": "timestamp",
+                              })
+    dataframe2 = None
+    for counter in range(10):
+        dataframe2 = session.pandas.read_sql_athena(sql="select * from test",
+                                                    database=database)
+        if len(dataframe.index) == len(dataframe2.index):
+            break
+        sleep(2)
+    assert len(dataframe.index) == len(dataframe2.index)
+    assert len(list(dataframe.columns)) == len(list(dataframe2.columns))