align compression levels

delta-io · Nov 7, 2023 · 96355f0 · 96355f0
1 parent 293586e
commit 96355f0
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 4 deletions.
diff --git a/python/deltalake/_internal.pyi b/python/deltalake/_internal.pyi
@@ -64,7 +64,7 @@ class RawDeltaTable:
         target_size: Optional[int],
         max_concurrent_tasks: Optional[int],
         min_commit_interval: Optional[int],
-        compression: Literal["snappy", "gzip", "brotli", "lz4", "zstd"],
+        compression: str,
     ) -> str: ...
     def z_order_optimize(
         self,
@@ -74,7 +74,7 @@ class RawDeltaTable:
         max_concurrent_tasks: Optional[int],
         max_spill_size: Optional[int],
         min_commit_interval: Optional[int],
-        compression: Literal["snappy", "gzip", "brotli", "lz4", "zstd"],
+        compression: str,
     ) -> str: ...
     def restore(
         self,

diff --git a/python/deltalake/table.py b/python/deltalake/table.py
@@ -109,6 +109,15 @@ class ProtocolVersions(NamedTuple):
 FilterType = Union[FilterConjunctionType, FilterDNFType]
 
 
+def _get_parquet_rs_compression(compression: str):
+    """Returns compression into parquet-rs compression format with pyarrow default compression level"""
+    if compression in ["gzip", "brotli", "zstd"]:
+        compression = (
+            f"{compression}({pyarrow.Codec.default_compression_level(compression)})"
+        )
+    return compression
+
+
 def _check_contains_null(value: Any) -> bool:
     """
     Check if target contains nullish value.
@@ -1377,12 +1386,14 @@ def compact(
         if isinstance(min_commit_interval, timedelta):
             min_commit_interval = int(min_commit_interval.total_seconds())
 
+        compress = _get_parquet_rs_compression(compression)
+
         metrics = self.table._table.compact_optimize(
             partition_filters,
             target_size,
             max_concurrent_tasks,
             min_commit_interval,
-            compression,
+            compress,
         )
         self.table.update_incremental()
         return json.loads(metrics)
@@ -1432,14 +1443,16 @@ def z_order(
         if isinstance(min_commit_interval, timedelta):
             min_commit_interval = int(min_commit_interval.total_seconds())
 
+        compress = _get_parquet_rs_compression(compression)
+
         metrics = self.table._table.z_order_optimize(
             list(columns),
             partition_filters,
             target_size,
             max_concurrent_tasks,
             max_spill_size,
             min_commit_interval,
-            compression,
+            compress,
         )
         self.table.update_incremental()
         return json.loads(metrics)