Add new ctas_approach for Pandas.read_sql_athena()

Author: igorborgest

awswrangler/athena.py

@@ -3,6 +3,7 @@
 import logging
 import re
 import unicodedata
+from datetime import datetime, date
 
 from awswrangler.data_types import athena2python
 from awswrangler.exceptions import QueryFailed, QueryCancelled
@@ -162,8 +163,15 @@ def _rows2row(rows: List[Dict[str, List[Dict[str, str]]]],
             vals_varchar: List[Optional[str]] = [x["VarCharValue"] if x else None for x in row["Data"]]
             data: Dict[str, Any] = {}
             for (name, ptype), val in zip(python_types, vals_varchar):
-                if ptype is not None:
-                    data[name] = ptype(val)
+                if val is not None:
+                    if ptype is None:
+                        data[name] = None
+                    elif ptype == date:
+                        data[name] = date(*[int(y) for y in val.split("-")])
+                    elif ptype == datetime:
+                        data[name] = datetime.strptime(val + "000", "%Y-%m-%d %H:%M:%S.%f")
+                    else:
+                        data[name] = ptype(val)
                 else:
                     data[name] = None
             yield data

awswrangler/data_types.py

@@ -1,6 +1,7 @@
 from typing import List, Tuple, Dict, Callable, Optional
 import logging
 from datetime import datetime, date
+from decimal import Decimal
 
 import pyarrow as pa  # type: ignore
 import pandas as pd  # type: ignore
@@ -74,6 +75,8 @@ def athena2python(dtype: str) -> Optional[type]:
         return date
     elif dtype == "unknown":
         return None
+    elif dtype == "decimal":
+        return Decimal
     else:
         raise UnsupportedType(f"Unsupported Athena type: {dtype}")
 

awswrangler/exceptions.py

@@ -80,3 +80,11 @@ class ApiError(Exception):
 
 class InvalidCompression(Exception):
     pass
+
+
+class InvalidTable(Exception):
+    pass
+
+
+class InvalidParameters(Exception):
+    pass

awswrangler/glue.py

@@ -1,10 +1,11 @@
+from typing import Dict, Optional
 from math import ceil
 import re
 import logging
 
 from awswrangler import data_types
 from awswrangler.athena import Athena
-from awswrangler.exceptions import UnsupportedFileFormat, InvalidSerDe, ApiError, UnsupportedType, UndetectedType
+from awswrangler.exceptions import UnsupportedFileFormat, InvalidSerDe, ApiError, UnsupportedType, UndetectedType, InvalidTable, InvalidArguments
 
 logger = logging.getLogger(__name__)
 
@@ -85,7 +86,11 @@ def metadata_to_glue(self,
                                 compression=compression,
                                 extra_args=extra_args)
 
-    def delete_table_if_exists(self, database, table):
+    def delete_table_if_exists(self, table: str = None, database: Optional[str] = None):
+        if database is None and self._session.athena_database is not None:
+            database = self._session.athena_database
+        if database is None:
+            raise InvalidArguments("You must pass a valid database or have one defined in your Session!")
         try:
             self._client_glue.delete_table(DatabaseName=database, Name=table)
         except self._client_glue.exceptions.EntityNotFoundException:
@@ -372,3 +377,16 @@ def _parse_partitions_tuples(objects_paths, partition_cols):
     @staticmethod
     def _parse_partition_values(path, partition_cols):
         return [re.search(f"/{col}=(.*?)/", path).group(1) for col in partition_cols]
+
+    def get_table_location(self, database: str, table: str):
+        """
+        Get table's location on Glue catalog
+
+        :param database: Database name
+        :param table: table name
+        """
+        res: Dict = self._client_glue.get_table(DatabaseName=database, Name=table)
+        try:
+            return res["Table"]["StorageDescriptor"]["Location"]
+        except KeyError:
+            raise InvalidTable(f"{database}.{table}")

awswrangler/pandas.py

@@ -5,7 +5,7 @@
 from math import floor
 import copy
 import csv
-from datetime import datetime
+from datetime import datetime, date
 from decimal import Decimal
 from ast import literal_eval
 
@@ -18,7 +18,8 @@
 
 from awswrangler import data_types
 from awswrangler.exceptions import (UnsupportedWriteMode, UnsupportedFileFormat, AthenaQueryError, EmptyS3Object,
-                                    LineTerminatorNotFound, EmptyDataframe, InvalidSerDe, InvalidCompression)
+                                    LineTerminatorNotFound, EmptyDataframe, InvalidSerDe, InvalidCompression,
+                                    InvalidParameters)
 from awswrangler.utils import calculate_bounders
 from awswrangler import s3
 from awswrangler.athena import Athena
@@ -495,29 +496,100 @@ def read_sql_athena(self,
                         sql: str,
                         database: Optional[str] = None,
                         s3_output: Optional[str] = None,
-                        max_result_size: Optional[int] = None,
                         workgroup: Optional[str] = None,
                         encryption: Optional[str] = None,
-                        kms_key: Optional[str] = None):
+                        kms_key: Optional[str] = None,
+                        ctas_approach: bool = True,
+                        procs_cpu_bound: Optional[int] = None,
+                        max_result_size: Optional[int] = None):
         """
         Executes any SQL query on AWS Athena and return a Dataframe of the result.
-        P.S. If max_result_size is passed, then a iterator of Dataframes is returned.
+        There are two approaches to be defined through ctas_approach parameter:
+        1 - ctas_approach True (Default):
+            Wrap the query with a CTAS and then reads the table data as parquet directly from s3.
+            PROS: Faster and has a better handle of nested types
+            CONS: Can't use max_result_size.
+        2 - ctas_approach False:
+            Does a regular query on Athena and parse the regular CSV result on s3
+            PROS: Accepts max_result_size.
+            CONS: Slower (But stills faster than other libraries that uses the Athena API) and does not handle nested types so well
+
+        P.S. If ctas_approach is False and max_result_size is passed, then a iterator of Dataframes is returned.
         P.S.S. All default values will be inherited from the Session()
 
         :param sql: SQL Query
         :param database: Glue/Athena Database
         :param s3_output: AWS S3 path
-        :param max_result_size: Max number of bytes on each request to S3
         :param workgroup: The name of the workgroup in which the query is being started. (By default uses de Session() workgroup)
         :param encryption: None|'SSE_S3'|'SSE_KMS'|'CSE_KMS'
         :param kms_key: For SSE-KMS and CSE-KMS , this is the KMS key ARN or ID.
-        :return: Pandas Dataframe or Iterator of Pandas Dataframes if max_result_size != None
+        :param ctas_approach: Wraps the query with a CTAS
+        :param procs_cpu_bound: Number of cores used for CPU bound tasks
+        :param max_result_size: Max number of bytes on each request to S3 (VALID ONLY FOR ctas_approach=False)
+        :return: Pandas Dataframe or Iterator of Pandas Dataframes if max_result_size was passed
         """
+        if ctas_approach is True and max_result_size is not None:
+            raise InvalidParameters("ctas_approach can't use max_result_size!")
         if s3_output is None:
             if self._session.athena_s3_output is not None:
                 s3_output = self._session.athena_s3_output
             else:
                 s3_output = self._session.athena.create_athena_bucket()
+        if ctas_approach is False:
+            return self._read_sql_athena_regular(sql=sql,
+                                                 database=database,
+                                                 s3_output=s3_output,
+                                                 workgroup=workgroup,
+                                                 encryption=encryption,
+                                                 kms_key=kms_key,
+                                                 max_result_size=max_result_size)
+        else:
+            return self._read_sql_athena_ctas(sql=sql,
+                                              database=database,
+                                              s3_output=s3_output,
+                                              workgroup=workgroup,
+                                              encryption=encryption,
+                                              kms_key=kms_key,
+                                              procs_cpu_bound=procs_cpu_bound)
+
+    def _read_sql_athena_ctas(self,
+                              sql: str,
+                              s3_output: str,
+                              database: Optional[str] = None,
+                              workgroup: Optional[str] = None,
+                              encryption: Optional[str] = None,
+                              kms_key: Optional[str] = None,
+                              procs_cpu_bound: Optional[int] = None) -> pd.DataFrame:
+        guid: str = pa.compat.guid()
+        name: str = f"temp_table_{guid}"
+        s3_output = s3_output[:-1] if s3_output[-1] == "/" else s3_output
+        path: str = f"{s3_output}/{name}"
+        query: str = f"CREATE TABLE {name}\n" \
+                     f"WITH(\n" \
+                     f"    format = 'Parquet',\n" \
+                     f"    parquet_compression = 'SNAPPY',\n" \
+                     f"    external_location = '{path}'\n" \
+                     f") AS\n" \
+                     f"{sql}"
+        logger.debug(f"query: {query}")
+        query_id: str = self._session.athena.run_query(query=query,
+                                                       database=database,
+                                                       s3_output=s3_output,
+                                                       workgroup=workgroup,
+                                                       encryption=encryption,
+                                                       kms_key=kms_key)
+        self._session.athena.wait_query(query_execution_id=query_id)
+        self._session.glue.delete_table_if_exists(database=database, table=name)
+        return self.read_parquet(path=path, procs_cpu_bound=procs_cpu_bound)
+
+    def _read_sql_athena_regular(self,
+                                 sql: str,
+                                 s3_output: str,
+                                 database: Optional[str] = None,
+                                 workgroup: Optional[str] = None,
+                                 encryption: Optional[str] = None,
+                                 kms_key: Optional[str] = None,
+                                 max_result_size: Optional[int] = None):
         query_execution_id: str = self._session.athena.run_query(query=sql,
                                                                  database=database,
                                                                  s3_output=s3_output,
@@ -542,7 +614,10 @@ def read_sql_athena(self,
             if max_result_size is None:
                 if len(ret.index) > 0:
                     for col in parse_dates:
-                        ret[col] = ret[col].dt.date.replace(to_replace={pd.NaT: None})
+                        if str(ret[col].dtype) == "object":
+                            ret[col] = ret[col].apply(lambda x: date(*[int(y) for y in x.split("-")]))
+                        else:
+                            ret[col] = ret[col].dt.date.replace(to_replace={pd.NaT: None})
                 return ret
             else:
                 return Pandas._apply_dates_to_generator(generator=ret, parse_dates=parse_dates)
@@ -1151,5 +1226,29 @@ def read_parquet(self,
         use_threads: bool = True if procs_cpu_bound > 1 else False
         fs: S3FileSystem = s3.get_fs(session_primitives=self._session.primitives)
         fs = pa.filesystem._ensure_filesystem(fs)
-        return pq.read_table(source=path, columns=columns, filters=filters,
-                             filesystem=fs).to_pandas(use_threads=use_threads)
+        table = pq.read_table(source=path, columns=columns, filters=filters, filesystem=fs, use_threads=use_threads)
+        # Check if we lose some integer during the conversion (Happens when has some null value)
+        integers = [field.name for field in table.schema if str(field.type).startswith("int")]
+        df = table.to_pandas(use_threads=use_threads, integer_object_nulls=True)
+        for c in integers:
+            if not str(df[c].dtype).startswith("int"):
+                df[c] = df[c].astype("Int64")
+        return df
+
+    def read_table(self,
+                   database: str,
+                   table: str,
+                   columns: Optional[List[str]] = None,
+                   filters: Optional[Union[List[Tuple[Any]], List[Tuple[Any]]]] = None,
+                   procs_cpu_bound: Optional[int] = None) -> pd.DataFrame:
+        """
+        Read PARQUET table from S3 using the Glue Catalog location skipping Athena's necessity
+
+        :param database: Database name
+        :param table: table name
+        :param columns: Names of columns to read from the file
+        :param filters: List of filters to apply, like ``[[('x', '=', 0), ...], ...]``.
+        :param procs_cpu_bound: Number of cores used for CPU bound tasks
+        """
+        path: str = self._session.glue.get_table_location(database=database, table=table)
+        return self.read_parquet(path=path, columns=columns, filters=filters, procs_cpu_bound=procs_cpu_bound)

requirements.txt

@@ -1,8 +1,8 @@
 numpy~=1.17.4
 pandas~=0.25.3
 pyarrow~=0.15.1
-botocore~=1.13.34
-boto3~=1.10.34
+botocore~=1.13.35
+boto3~=1.10.35
 s3fs~=0.4.0
 tenacity~=6.0.0
 pg8000~=1.13.2

testing/test_awswrangler/test_athena.py

@@ -1,7 +1,10 @@
 import logging
+from datetime import datetime, date
+from decimal import Decimal
 
 import pytest
 import boto3
+import pandas as pd
 
 from awswrangler import Session
 from awswrangler.exceptions import QueryCancelled, QueryFailed
@@ -193,3 +196,46 @@ def test_query(session, database):
     assert row["_col2"] == 2.0
     assert row["_col3"] is True
     assert row["_col4"] is None
+
+
+def test_query2(session, bucket, database):
+    df = pd.DataFrame({
+        "id": [1, 2, 3],
+        "col_date": [date(194, 1, 12), None, date(2049, 12, 30)],
+        "col_timestamp": [datetime(194, 1, 12, 1, 1, 1, 1000), None,
+                          datetime(2049, 12, 30, 1, 1, 1, 1000)],
+        "col_string": ["foo", None, "boo"],
+        "col_double": [1.1, None, 2.2],
+        "col_decimal": [Decimal((0, (1, 9, 9), -2)), None,
+                        Decimal((0, (1, 9, 0), -2))],
+        "col_int": [1, None, 2]
+    })
+    path = f"s3://{bucket}/test_query2/"
+    session.pandas.to_parquet(dataframe=df,
+                              database=database,
+                              table="test",
+                              path=path,
+                              mode="overwrite",
+                              preserve_index=False)
+    for row in session.athena.query(query="SELECT * FROM test", database=database):
+        if row["id"] == 1:
+            assert row["col_date"] == date(194, 1, 12)
+            assert row["col_timestamp"] == datetime(194, 1, 12, 1, 1, 1, 1000)
+            assert row["col_string"] == "foo"
+            assert row["col_double"] == 1.1
+            assert row["col_decimal"] == Decimal((0, (1, 9, 9), -2))
+            assert row["col_int"] == 1
+        elif row["id"] == 2:
+            assert row["col_date"] is None
+            assert row["col_timestamp"] is None
+            assert row["col_string"] is None
+            assert row["col_double"] is None
+            assert row["col_decimal"] is None
+            assert row["col_int"] is None
+        else:
+            assert row["col_date"] == date(2049, 12, 30)
+            assert row["col_timestamp"] == datetime(2049, 12, 30, 1, 1, 1, 1000)
+            assert row["col_string"] == "boo"
+            assert row["col_double"] == 2.2
+            assert row["col_decimal"] == Decimal((0, (1, 9, 0), -2))
+            assert row["col_int"] == 2