Remove in memory copy of DataFrame for to_parquet and to_csv.

Author: igorborgest

awswrangler/catalog.py

@@ -886,6 +886,10 @@ def sanitize_table_name(table: str) -> str:
 def drop_duplicated_columns(df: pd.DataFrame) -> pd.DataFrame:
     """Drop all repeated columns (duplicated names).
 
+    Note
+    ----
+    This transformation will run `inplace` and will make changes in the original DataFrame.
+
     Note
     ----
     It is different from Panda's drop_duplicates() function which considers the column values.
@@ -912,11 +916,14 @@ def drop_duplicated_columns(df: pd.DataFrame) -> pd.DataFrame:
     1  2
 
     """
-    duplicated_cols = df.columns.duplicated()
-    duplicated_cols_names: List[str] = list(df.columns[duplicated_cols])
-    if len(duplicated_cols_names) > 0:
-        _logger.warning("Dropping repeated columns: %s", duplicated_cols_names)
-    return df.loc[:, ~duplicated_cols]
+    duplicated = df.columns.duplicated()
+    if duplicated.any():
+        _logger.warning("Dropping duplicated columns...")
+        columns = df.columns.values
+        columns[duplicated] = "AWSDataWranglerDuplicatedMarker"
+        df.columns = columns
+        df.drop(columns="AWSDataWranglerDuplicatedMarker", inplace=True)
+    return df
 
 
 def get_connection(

awswrangler/s3/_write.py

@@ -443,7 +443,7 @@ def to_csv(  # pylint: disable=too-many-arguments,too-many-locals
         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()}
-        df = catalog.drop_duplicated_columns(df=df)
+        catalog.drop_duplicated_columns(df=df)
 
     session: boto3.Session = _utils.ensure_session(session=boto3_session)
     fs: s3fs.S3FileSystem = _utils.get_fs(session=session, s3_additional_kwargs=s3_additional_kwargs)
@@ -829,7 +829,7 @@ def to_parquet(  # pylint: disable=too-many-arguments,too-many-locals
         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()}
-        df = catalog.drop_duplicated_columns(df=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)

testing/test_awswrangler/test_data_lake2.py

@@ -103,11 +103,10 @@ def test_json_chunksize(path):
 def test_parquet_cast_string(path):
     df = pd.DataFrame({"id": [1, 2, 3], "value": ["foo", "boo", "bar"]})
     path_file = f"{path}0.parquet"
-    wr.s3.to_parquet(df, path_file, dtype={"id": "string"})
+    wr.s3.to_parquet(df, path_file, dtype={"id": "string"}, sanitize_columns=False)
     wr.s3.wait_objects_exist([path_file])
     df2 = wr.s3.read_parquet(path_file)
     assert str(df2.id.dtypes) == "string"
-    df2["id"] = df2["id"].astype(int)
     assert df.shape == df2.shape
     for col, row in tuple(itertools.product(df.columns, range(3))):
         assert df[col].iloc[row] == df2[col].iloc[row]
@@ -123,8 +122,6 @@ def test_parquet_cast_string_dataset(path, partition_cols):
     df2 = wr.s3.read_parquet(path, dataset=True).sort_values("id", ignore_index=True)
     assert str(df2.id.dtypes) == "string"
     assert str(df2.c3.dtypes) == "string"
-    df2["id"] = df2["id"].astype(int)
-    df2["c3"] = df2["c3"].astype(float)
     assert df.shape == df2.shape
     for col, row in tuple(itertools.product(df.columns, range(3))):
         assert df[col].iloc[row] == df2[col].iloc[row]
@@ -158,7 +155,7 @@ def test_athena_undefined_column(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.to_parquet(df, path_file, sanitize_columns=True)
     wr.s3.wait_objects_exist([path_file])
     df2 = wr.s3.read_parquet(path_file)
     assert df.shape == df2.shape