Fix issue with column rename for column with special chars (#136)

Author: george-zubrienko

spark_utils/dataframes/functions.py

@@ -23,7 +23,7 @@
 """
   Helper functions for Spark Dataframes
 """
-
+import re
 from typing import List, Iterator, Tuple
 
 from datetime import datetime
@@ -130,25 +130,15 @@ def rename_column(name: str) -> str:
     :return:
     """
 
-    illegals = [
-        " ",
-        ",",
-        ";",
-        "{",
-        "}",
-        "(",
-        ")",
-        "\t",
-        "=",
-        "/",
-        "\\",
-        ".",
-    ]
+    return re.sub(r"\W+", "", name)
+
 
-    for illegal in illegals:
-        name = name.replace(illegal, "")
+def safe_encode(column_name: str) -> str:
+    """
+    Adds `` around the column name so columns with unsupported chars are resolved
+    """
 
-    return name
+    return f"`{column_name}`"
 
 
 def rename_columns(dataframe: DataFrame) -> DataFrame:
@@ -158,7 +148,7 @@ def rename_columns(dataframe: DataFrame) -> DataFrame:
     :param dataframe: Source dataframe
     :return: Dataframe with renamed columns
     """
-    return dataframe.select([col(c).alias(rename_column(c)) for c in dataframe.columns])
+    return dataframe.select([col(safe_encode(c)).alias(rename_column(c)) for c in dataframe.columns])
 
 
 def _max_timestamp(dataframe: DataFrame, timestamp_column: str, timestamp_column_format: str) -> datetime:

test/test_common_functions.py

@@ -7,6 +7,7 @@
 from pyspark.sql import DataFrame
 from pyspark.sql import SparkSession
 
+from spark_utils.dataframes.functions import rename_column
 from spark_utils.models.job_socket import JobSocket
 from spark_utils.common.functions import read_from_socket, write_to_socket
 
@@ -97,3 +98,18 @@ def test_job_socket_serialize(sep: str, test_base_path: str):
     )
 
     assert socket.serialize(separator=sep) == f"{socket.alias}{sep}{socket.data_path}{sep}{socket.data_format}"
+
+
+@pytest.mark.parametrize(
+    "funky_name, expected_name",
+    [
+        ("a--bc", "abc"),
+        (".abc", "abc"),
+        ("a bc", "abc"),
+        ("a\\bc", "abc"),
+        ("a/bc", "abc"),
+        ("a\t{};,bc", "abc"),
+    ],
+)
+def test_column_rename(funky_name: str, expected_name: str):
+    assert expected_name == rename_column(funky_name)