Merge pull request #39 from awslabs/normalize-athena-names

Normalize Athena names (tables and columns)

Author: igorborgest

awswrangler/athena.py

@@ -1,6 +1,8 @@
 from time import sleep
 import logging
 import ast
+import re
+import unicodedata
 
 from awswrangler import data_types
 from awswrangler.exceptions import QueryFailed, QueryCancelled
@@ -128,3 +130,33 @@ def repair_table(self, database, table, s3_output=None):
                                   s3_output=s3_output)
         self.wait_query(query_execution_id=query_id)
         return query_id
+
+    @staticmethod
+    def _normalize_name(name):
+        name = "".join(c for c in unicodedata.normalize("NFD", name)
+                       if unicodedata.category(c) != "Mn")
+        name = name.replace(" ", "_")
+        name = name.replace("-", "_")
+        name = name.replace(".", "_")
+        name = re.sub("(.)([A-Z][a-z]+)", r"\1_\2", name)
+        name = re.sub("([a-z0-9])([A-Z])", r"\1_\2", name)
+        name = name.lower()
+        return re.sub(r"(_)\1+", "\\1", name)  # remove repeated underscores
+
+    @staticmethod
+    def normalize_column_name(name):
+        """
+        https://docs.aws.amazon.com/athena/latest/ug/tables-databases-columns-names.html
+        :param name: column name (str)
+        :return: normalized column name (str)
+        """
+        return Athena._normalize_name(name=name)
+
+    @staticmethod
+    def normalize_table_name(name):
+        """
+        https://docs.aws.amazon.com/athena/latest/ug/tables-databases-columns-names.html
+        :param name: table name (str)
+        :return: normalized table name (str)
+        """
+        return Athena._normalize_name(name=name)

awswrangler/glue.py

@@ -3,6 +3,7 @@
 import logging
 
 from awswrangler import data_types
+from awswrangler.athena import Athena
 from awswrangler.exceptions import UnsupportedFileFormat, InvalidSerDe, ApiError
 
 logger = logging.getLogger(__name__)
@@ -64,7 +65,7 @@ def metadata_to_glue(self,
             indexes_position=indexes_position,
             cast_columns=cast_columns)
         table = table if table else Glue.parse_table_name(path)
-        table = table.lower().replace(".", "_")
+        table = Athena.normalize_table_name(name=table)
         if mode == "overwrite":
             self.delete_table_if_exists(database=database, table=table)
         exists = self.does_table_exists(database=database, table=table)
@@ -124,8 +125,13 @@ def create_table(self,
         self._client_glue.create_table(DatabaseName=database,
                                        TableInput=table_input)
 
-    def add_partitions(self, database, table, partition_paths, file_format,
-                       compression, extra_args=None):
+    def add_partitions(self,
+                       database,
+                       table,
+                       partition_paths,
+                       file_format,
+                       compression,
+                       extra_args=None):
         if not partition_paths:
             return None
         partitions = list()
@@ -207,8 +213,12 @@ def parse_table_name(path):
         return path.rpartition("/")[2]
 
     @staticmethod
-    def csv_table_definition(table, partition_cols_schema, schema, path,
-                             compression, extra_args=None):
+    def csv_table_definition(table,
+                             partition_cols_schema,
+                             schema,
+                             path,
+                             compression,
+                             extra_args=None):
         if extra_args is None:
             extra_args = {}
         if partition_cols_schema is None:

awswrangler/pandas.py

@@ -17,7 +17,7 @@
                                     EmptyDataframe, InvalidSerDe,
                                     InvalidCompression)
 from awswrangler.utils import calculate_bounders
-from awswrangler import s3
+from awswrangler import s3, athena
 
 logger = logging.getLogger(__name__)
 
@@ -599,6 +599,7 @@ def to_s3(self,
         :param extra_args: Extra arguments specific for each file formats (E.g. "sep" for CSV)
         :return: List of objects written on S3
         """
+        Pandas.normalize_columns_names_athena(dataframe, inplace=True)
         if compression is not None:
             compression = compression.lower()
         file_format = file_format.lower()
@@ -1024,3 +1025,12 @@ def read_log_query(self,
                 new_row[col_name] = col["value"]
             pre_df.append(new_row)
         return pandas.DataFrame(pre_df)
+
+    @staticmethod
+    def normalize_columns_names_athena(dataframe, inplace=True):
+        if inplace is False:
+            dataframe = dataframe.copy(deep=True)
+        dataframe.columns = [
+            athena.Athena.normalize_column_name(x) for x in dataframe.columns
+        ]
+        return dataframe

testing/test_awswrangler/test_pandas.py

@@ -778,3 +778,49 @@ def test_read_sql_athena_with_time_zone(session, bucket, database):
     assert len(dataframe.columns) == 2
     assert dataframe["type"][0] == "timestamp with time zone"
     assert dataframe["value"][0].year == datetime.utcnow().year
+
+
+def test_normalize_columns_names_athena():
+    dataframe = pandas.DataFrame({
+        "CammelCase": [1, 2, 3],
+        "With Spaces": [4, 5, 6],
+        "With-Dash": [7, 8, 9],
+        "Ãccént": [10, 11, 12],
+    })
+    Pandas.normalize_columns_names_athena(dataframe=dataframe, inplace=True)
+    assert dataframe.columns[0] == "cammel_case"
+    assert dataframe.columns[1] == "with_spaces"
+    assert dataframe.columns[2] == "with_dash"
+    assert dataframe.columns[3] == "accent"
+
+
+def test_to_parquet_with_normalize(
+        session,
+        bucket,
+        database,
+):
+    dataframe = pandas.DataFrame({
+        "CammelCase": [1, 2, 3],
+        "With Spaces": [4, 5, 6],
+        "With-Dash": [7, 8, 9],
+        "Ãccént": [10, 11, 12],
+        "with.dot": [10, 11, 12],
+    })
+    session.pandas.to_parquet(dataframe=dataframe,
+                              database=database,
+                              path=f"s3://{bucket}/TestTable-with.dot/",
+                              mode="overwrite")
+    dataframe2 = None
+    for counter in range(10):
+        dataframe2 = session.pandas.read_sql_athena(
+            sql="select * from test_table_with_dot", database=database)
+        if len(dataframe.index) == len(dataframe2.index):
+            break
+        sleep(2)
+    assert len(dataframe.index) == len(dataframe2.index)
+    assert (len(list(dataframe.columns)) + 1) == len(list(dataframe2.columns))
+    assert dataframe2.columns[0] == "cammel_case"
+    assert dataframe2.columns[1] == "with_spaces"
+    assert dataframe2.columns[2] == "with_dash"
+    assert dataframe2.columns[3] == "accent"
+    assert dataframe2.columns[4] == "with_dot"