Adding support for lists to Pandas

Author: igorborgest

awswrangler/athena.py

@@ -1,5 +1,6 @@
 from time import sleep
 import logging
+import ast
 
 from awswrangler.exceptions import UnsupportedType, QueryFailed, QueryCancelled
 
@@ -35,6 +36,8 @@ def _type_athena2pandas(dtype):
             return "datetime64"
         elif dtype == "date":
             return "date"
+        elif dtype == "array":
+            return "literal_eval"
         else:
             raise UnsupportedType(f"Unsupported Athena type: {dtype}")
 
@@ -44,18 +47,21 @@ def get_query_dtype(self, query_execution_id):
         dtype = {}
         parse_timestamps = []
         parse_dates = []
+        converters = {}
         for col_name, col_type in cols_metadata.items():
             ptype = Athena._type_athena2pandas(dtype=col_type)
             if ptype in ["datetime64", "date"]:
                 parse_timestamps.append(col_name)
                 if ptype == "date":
                     parse_dates.append(col_name)
+            elif ptype == "literal_eval":
+                converters[col_name] = ast.literal_eval
             else:
                 dtype[col_name] = ptype
         logger.debug(f"dtype: {dtype}")
         logger.debug(f"parse_timestamps: {parse_timestamps}")
         logger.debug(f"parse_dates: {parse_dates}")
-        return dtype, parse_timestamps, parse_dates
+        return dtype, parse_timestamps, parse_dates, converters
 
     def create_athena_bucket(self):
         """

awswrangler/glue.py

@@ -47,23 +47,25 @@ def get_table_python_types(self, database, table):
 
     @staticmethod
     def type_pyarrow2athena(dtype):
-        dtype = str(dtype).lower()
-        if dtype == "int32":
+        dtype_str = str(dtype).lower()
+        if dtype_str == "int32":
             return "int"
-        elif dtype == "int64":
+        elif dtype_str == "int64":
             return "bigint"
-        elif dtype == "float":
+        elif dtype_str == "float":
             return "float"
-        elif dtype == "double":
+        elif dtype_str == "double":
             return "double"
-        elif dtype == "bool":
+        elif dtype_str == "bool":
             return "boolean"
-        elif dtype == "string":
+        elif dtype_str == "string":
             return "string"
-        elif dtype.startswith("timestamp"):
+        elif dtype_str.startswith("timestamp"):
             return "timestamp"
-        elif dtype.startswith("date"):
+        elif dtype_str.startswith("date"):
             return "date"
+        elif dtype_str.startswith("list"):
+            return f"array<{Glue.type_pyarrow2athena(dtype.value_type)}>"
         else:
             raise UnsupportedType(f"Unsupported Pyarrow type: {dtype}")
 
@@ -260,7 +262,7 @@ def _extract_pyarrow_schema(dataframe, preserve_index, cast_columns=None):
         for field in pyarrow.Schema.from_pandas(df=dataframe[cols],
                                                 preserve_index=preserve_index):
             name = str(field.name)
-            dtype = str(field.type)
+            dtype = field.type
             cols_dtypes[name] = dtype
             if name not in dataframe.columns:
                 schema.append((name, dtype))

awswrangler/pandas.py

@@ -56,6 +56,7 @@ def read_csv(
             parse_dates=False,
             infer_datetime_format=False,
             encoding="utf-8",
+            converters=None,
     ):
         """
         Read CSV file from AWS S3 using optimized strategies.
@@ -76,6 +77,7 @@ def read_csv(
         :param parse_dates: Same as pandas.read_csv()
         :param infer_datetime_format: Same as pandas.read_csv()
         :param encoding: Same as pandas.read_csv()
+        :param converters: Same as pandas.read_csv()
         :return: Pandas Dataframe or Iterator of Pandas Dataframes if max_result_size != None
         """
         bucket_name, key_path = self._parse_path(path)
@@ -99,7 +101,8 @@ def read_csv(
                 escapechar=escapechar,
                 parse_dates=parse_dates,
                 infer_datetime_format=infer_datetime_format,
-                encoding=encoding)
+                encoding=encoding,
+                converters=converters)
         else:
             ret = Pandas._read_csv_once(
                 client_s3=client_s3,
@@ -115,7 +118,8 @@ def read_csv(
                 escapechar=escapechar,
                 parse_dates=parse_dates,
                 infer_datetime_format=infer_datetime_format,
-                encoding=encoding)
+                encoding=encoding,
+                converters=converters)
         return ret
 
     @staticmethod
@@ -135,6 +139,7 @@ def _read_csv_iterator(
             parse_dates=False,
             infer_datetime_format=False,
             encoding="utf-8",
+            converters=None,
     ):
         """
         Read CSV file from AWS S3 using optimized strategies.
@@ -156,6 +161,7 @@ def _read_csv_iterator(
         :param parse_dates: Same as pandas.read_csv()
         :param infer_datetime_format: Same as pandas.read_csv()
         :param encoding: Same as pandas.read_csv()
+        :param converters: Same as pandas.read_csv()
         :return: Pandas Dataframe
         """
         metadata = s3.S3.head_object_with_retry(client=client_s3,
@@ -181,7 +187,8 @@ def _read_csv_iterator(
                 escapechar=escapechar,
                 parse_dates=parse_dates,
                 infer_datetime_format=infer_datetime_format,
-                encoding=encoding)
+                encoding=encoding,
+                converters=converters)
         else:
             bounders = calculate_bounders(num_items=total_size,
                                           max_size=max_result_size)
@@ -234,7 +241,7 @@ def _read_csv_iterator(
                     lineterminator=lineterminator,
                     dtype=dtype,
                     encoding=encoding,
-                )
+                    converters=converters)
                 yield df
                 if count == 1:  # first chunk
                     names = df.columns
@@ -352,6 +359,7 @@ def _read_csv_once(
             parse_dates=False,
             infer_datetime_format=False,
             encoding=None,
+            converters=None,
     ):
         """
         Read CSV file from AWS S3 using optimized strategies.
@@ -372,6 +380,7 @@ def _read_csv_once(
         :param parse_dates: Same as pandas.read_csv()
         :param infer_datetime_format: Same as pandas.read_csv()
         :param encoding: Same as pandas.read_csv()
+        :param converters: Same as pandas.read_csv()
         :return: Pandas Dataframe
         """
         buff = BytesIO()
@@ -392,6 +401,7 @@ def _read_csv_once(
             lineterminator=lineterminator,
             dtype=dtype,
             encoding=encoding,
+            converters=converters,
         )
         buff.close()
         return dataframe
@@ -425,12 +435,13 @@ def read_sql_athena(self,
             message_error = f"Query error: {reason}"
             raise AthenaQueryError(message_error)
         else:
-            dtype, parse_timestamps, parse_dates = self._session.athena.get_query_dtype(
+            dtype, parse_timestamps, parse_dates, converters = self._session.athena.get_query_dtype(
                 query_execution_id=query_execution_id)
             path = f"{s3_output}{query_execution_id}.csv"
             ret = self.read_csv(path=path,
                                 dtype=dtype,
                                 parse_dates=parse_timestamps,
+                                converters=converters,
                                 quoting=csv.QUOTE_ALL,
                                 max_result_size=max_result_size)
             if max_result_size is None:

testing/test_awswrangler/test_pandas.py

@@ -630,3 +630,37 @@ def test_to_parquet_compressed(session, bucket, database, compression):
     assert len(list(dataframe.columns)) == len(list(dataframe2.columns))
     assert dataframe[dataframe["id"] == 1].iloc[0]["name"] == dataframe2[
         dataframe2["id"] == 1].iloc[0]["name"]
+
+
+def test_to_parquet_lists(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