Enable s3.to_parquet to recieve "zstd" compression type (#1369)

* Add tests for compression type

* feat: add compression type, zstd

* fix: remove debug code

* remove redundant cast to fix mypy error

* style: fix to pass validate.sh

* Revert "remove redundant cast to fix mypy error"

This reverts commit d540098.

* Refactor test

Co-authored-by: Lucas Hanson <[email protected]>
Co-authored-by: Abdel Jaidi <[email protected]>

Author: 3 people

awswrangler/s3/_write.py

@@ -16,6 +16,7 @@
     "bz2": ".bz2",
     "xz": ".xz",
     "zip": ".zip",
+    "zstd": ".zstd",
 }
 
 

awswrangler/s3/_write_parquet.py

@@ -258,7 +258,7 @@ def to_parquet(  # pylint: disable=too-many-arguments,too-many-locals,too-many-b
     index : bool
         True to store the DataFrame index in file, otherwise False to ignore it.
     compression: str, optional
-        Compression style (``None``, ``snappy``, ``gzip``).
+        Compression style (``None``, ``snappy``, ``gzip``, ``zstd``).
     pyarrow_additional_kwargs : Optional[Dict[str, Any]]
         Additional parameters forwarded to pyarrow.
         e.g. pyarrow_additional_kwargs={'coerce_timestamps': 'ns', 'use_deprecated_int96_timestamps': False,
@@ -522,7 +522,7 @@ def to_parquet(  # pylint: disable=too-many-arguments,too-many-locals,too-many-b
 
     # Evaluating compression
     if _COMPRESSION_2_EXT.get(compression, None) is None:
-        raise exceptions.InvalidCompression(f"{compression} is invalid, please use None, 'snappy' or 'gzip'.")
+        raise exceptions.InvalidCompression(f"{compression} is invalid, please use None, 'snappy', 'gzip' or 'zstd'.")
     compression_ext: str = _COMPRESSION_2_EXT[compression]
 
     # Initializing defaults

tests/test_s3_parquet.py

@@ -504,10 +504,11 @@ def test_mixed_types_column(path) -> None:
         wr.s3.to_parquet(df, path, dataset=True, partition_cols=["par"])
 
 
-def test_parquet_plain(path) -> None:
+@pytest.mark.parametrize("compression", [None, "snappy", "gzip", "zstd"])
+def test_parquet_compression(path, compression) -> None:
     df = pd.DataFrame({"id": [1, 2, 3]}, dtype="Int64")
     path_file = f"{path}0.parquet"
-    wr.s3.to_parquet(df=df, path=path_file, compression=None)
+    wr.s3.to_parquet(df=df, path=path_file, compression=compression)
     df2 = wr.s3.read_parquet([path_file])
     assert df.equals(df2)