Fix DataFrame sanitize for single files in to_parquet(). #240

Author: igorborgest

awswrangler/catalog.py

@@ -640,7 +640,7 @@ def table(
 
 def _sanitize_name(name: str) -> str:
     name = "".join(c for c in unicodedata.normalize("NFD", name) if unicodedata.category(c) != "Mn")  # strip accents
-    name = re.sub("[^A-Za-z0-9_]+", "_", name)  # Removing non alphanumeric characters
+    name = re.sub("[^A-Za-z0-9_]+", "_", name)  # Replacing non alphanumeric characters by underscore
     return re.sub("([a-z0-9])([A-Z])", r"\1_\2", name).lower()  # Converting CamelCase to snake_case
 
 

awswrangler/s3.py

@@ -1043,11 +1043,19 @@ def to_parquet(  # pylint: disable=too-many-arguments
         )
     if df.empty is True:
         raise exceptions.EmptyDataFrame()
-    session: boto3.Session = _utils.ensure_session(session=boto3_session)
+
+    # Sanitize table to respect Athena's standards
     partition_cols = partition_cols if partition_cols else []
     dtype = dtype if dtype else {}
     columns_comments = columns_comments if columns_comments else {}
     partitions_values: Dict[str, List[str]] = {}
+    df = catalog.sanitize_dataframe_columns_names(df=df)
+    partition_cols = [catalog.sanitize_column_name(p) for p in partition_cols]
+    dtype = {catalog.sanitize_column_name(k): v.lower() for k, v in dtype.items()}
+    columns_comments = {catalog.sanitize_column_name(k): v for k, v in columns_comments.items()}
+    df = catalog.drop_duplicated_columns(df=df)
+
+    session: boto3.Session = _utils.ensure_session(session=boto3_session)
     cpus: int = _utils.ensure_cpu_count(use_threads=use_threads)
     fs: s3fs.S3FileSystem = _utils.get_fs(session=session, s3_additional_kwargs=s3_additional_kwargs)
     compression_ext: Optional[str] = _COMPRESSION_2_EXT.get(compression, None)
@@ -1075,16 +1083,11 @@ def to_parquet(  # pylint: disable=too-many-arguments
         ]
     else:
         mode = "append" if mode is None else mode
-        if (database is not None) and (table is not None):  # Normalize table to respect Athena's standards
-            df = catalog.sanitize_dataframe_columns_names(df=df)
-            partition_cols = [catalog.sanitize_column_name(p) for p in partition_cols]
-            dtype = {catalog.sanitize_column_name(k): v.lower() for k, v in dtype.items()}
-            columns_comments = {catalog.sanitize_column_name(k): v for k, v in columns_comments.items()}
+        if (database is not None) and (table is not None):
             exist: bool = catalog.does_table_exist(database=database, table=table, boto3_session=session)
             if (exist is True) and (mode in ("append", "overwrite_partitions")):
                 for k, v in catalog.get_table_types(database=database, table=table, boto3_session=session).items():
                     dtype[k] = v
-        df = catalog.drop_duplicated_columns(df=df)
         paths, partitions_values = _to_parquet_dataset(
             df=df,
             path=path,

testing/test_awswrangler/test_data_lake.py

@@ -1681,4 +1681,17 @@ def test_athena_undefined_column(database):
     with pytest.raises(wr.exceptions.InvalidArgumentValue):
         wr.athena.read_sql_query("SELECT 1", database)
     with pytest.raises(wr.exceptions.InvalidArgumentValue):
-        wr.athena.read_sql_query("SELECT NULL", database)
+        wr.athena.read_sql_query("SELECT NULL AS my_null", database)
+
+
+def test_to_parquet_file_sanitize(path):
+    df = pd.DataFrame({"C0": [0, 1], "camelCase": [2, 3], "c**--2": [4, 5]})
+    path_file = f"{path}0.parquet"
+    wr.s3.to_parquet(df, path_file)
+    wr.s3.wait_objects_exist([path_file])
+    df2 = wr.s3.read_parquet(path_file)
+    assert df.shape == df2.shape
+    assert list(df2.columns) == ["c0", "camel_case", "c_2"]
+    assert df2.c0.sum() == 1
+    assert df2.camel_case.sum() == 5
+    assert df2.c_2.sum() == 9