Merge pull request #36 from awslabs/read-athena-tz

Add support for timestamp with time zone for read from Athena

Author: igorborgest

awswrangler/athena.py

@@ -25,6 +25,7 @@ def get_query_columns_metadata(self, query_execution_id):
     def get_query_dtype(self, query_execution_id):
         cols_metadata = self.get_query_columns_metadata(
             query_execution_id=query_execution_id)
+        logger.debug(f"cols_metadata: {cols_metadata}")
         dtype = {}
         parse_timestamps = []
         parse_dates = []

awswrangler/data_types.py

@@ -18,7 +18,7 @@ def athena2pandas(dtype):
         return "bool"
     elif dtype in ["string", "char", "varchar"]:
         return "str"
-    elif dtype == "timestamp":
+    elif dtype in ["timestamp", "timestamp with time zone"]:
         return "datetime64"
     elif dtype == "date":
         return "date"

testing/test_awswrangler/test_pandas.py

@@ -742,3 +742,14 @@ def test_to_parquet_with_cast_null(
         sleep(2)
     assert len(dataframe.index) == len(dataframe2.index)
     assert len(list(dataframe.columns)) == len(list(dataframe2.columns))
+
+
+def test_read_sql_athena_with_time_zone(session, bucket, database):
+    dataframe = session.pandas.read_sql_athena(
+        sql=
+        "select current_timestamp as value, typeof(current_timestamp) as type",
+        database=database)
+    assert len(dataframe.index) == 1
+    assert len(dataframe.columns) == 2
+    assert dataframe["type"][0] == "timestamp with time zone"
+    assert dataframe["value"][0].year == datetime.utcnow().year