Add Zstandard support to read_pickle/to_pickle

Author: jonashaag

doc/source/getting_started/install.rst

@@ -402,3 +402,13 @@ qtpy                                         Clipboard I/O
 xclip                                        Clipboard I/O on linux
 xsel                                         Clipboard I/O on linux
 ========================= ================== =============================================================
+
+
+Compression
+^^^^^^^^^^^
+
+========================= ================== =============================================================
+Dependency                Minimum Version    Notes
+========================= ================== =============================================================
+Zstandard                                    Zstandard compression
+========================= ================== =============================================================

pandas/compat/_optional.py

@@ -34,6 +34,7 @@
     "xlwt": "1.3.0",
     "xlsxwriter": "1.2.2",
     "numba": "0.50.1",
+    "zstandard": "0.15.2",
 }
 
 # A mapping from import name to package name (on PyPI) for packages where

pandas/io/common.py

@@ -241,7 +241,7 @@ def _get_filepath_or_buffer(
     ----------
     filepath_or_buffer : a url, filepath (str, py.path.local or pathlib.Path),
                          or buffer
-    compression : {{'gzip', 'bz2', 'zip', 'xz', None}}, optional
+    compression : {{'gzip', 'bz2', 'zip', 'xz', 'zstd', None}}, optional
     encoding : the encoding to use to decode bytes, default is 'utf-8'
     mode : str, optional
 
@@ -420,7 +420,7 @@ def file_path_to_url(path: str) -> str:
     return urljoin("file:", pathname2url(path))
 
 
-_compression_to_extension = {"gzip": ".gz", "bz2": ".bz2", "zip": ".zip", "xz": ".xz"}
+_compression_to_extension = {"gzip": ".gz", "bz2": ".bz2", "zip": ".zip", "xz": ".xz", "zstd": ".zst"}
 
 
 def get_compression_method(
@@ -471,10 +471,10 @@ def infer_compression(
     ----------
     filepath_or_buffer : str or file handle
         File path or object.
-    compression : {'infer', 'gzip', 'bz2', 'zip', 'xz', None}
+    compression : {'infer', 'gzip', 'bz2', 'zip', 'xz', 'zstd', None}
         If 'infer' and `filepath_or_buffer` is path-like, then detect
         compression from the following extensions: '.gz', '.bz2', '.zip',
-        or '.xz' (otherwise no compression).
+        '.xz', or '.zst' (otherwise no compression).
 
     Returns
     -------
@@ -556,11 +556,11 @@ def get_handle(
     compression : str or dict, default None
         If string, specifies compression mode. If dict, value at key 'method'
         specifies compression mode. Compression mode must be one of {'infer',
-        'gzip', 'bz2', 'zip', 'xz', None}. If compression mode is 'infer'
-        and `filepath_or_buffer` is path-like, then detect compression from
-        the following extensions: '.gz', '.bz2', '.zip', or '.xz' (otherwise
-        no compression). If dict and compression mode is one of
-        {'zip', 'gzip', 'bz2'}, or inferred as one of the above,
+        'gzip', 'bz2', 'zip', 'xz', 'zstd', None}. If compression mode is
+        'infer' and `filepath_or_buffer` is path-like, then detect compression
+        from the following extensions: '.gz', '.bz2', '.zip', '.xz', or '.zst'
+        (otherwise no compression). If dict and compression mode is one of
+        {'zip', 'gzip', 'bz2', 'zstd'}, or inferred as one of the above,
         other entries passed as additional compression options.
 
         .. versionchanged:: 1.0.0
@@ -572,7 +572,7 @@ def get_handle(
         .. versionchanged:: 1.1.0
 
            Passing compression options as keys in dict is now
-           supported for compression modes 'gzip' and 'bz2' as well as 'zip'.
+           supported for compression modes 'gzip', 'bz2', 'zstd' and 'zip'.
 
     memory_map : bool, default False
         See parsers._parser_params for more information.
@@ -689,6 +689,23 @@ def get_handle(
         elif compression == "xz":
             handle = get_lzma_file(lzma)(handle, ioargs.mode)
 
+        # Zstd Compression
+        elif compression == "zstd":
+            zstd = import_optional_dependency("zstandard")
+            open_args = {
+                arg: compression_args.pop(arg, None)
+                for arg in ["encoding", "errors", "newline"]
+            }
+            if "r" in ioargs.mode:
+                open_args["dctx"] = zstd.ZstdDecompressor(**compression_args)
+            else:
+                open_args["cctx"] = zstd.ZstdCompressor(**compression_args)
+            handle = zstd.open(
+                handle,
+                mode=ioargs.mode,
+                **open_args,
+            )
+
         # Unrecognized Compression
         else:
             msg = f"Unrecognized compression type: {compression}"

pandas/io/pickle.py

@@ -37,11 +37,13 @@ def to_pickle(
         .. versionchanged:: 1.0.0
            Accept URL. URL has to be of S3 or GCS.
 
-    compression : {{'infer', 'gzip', 'bz2', 'zip', 'xz', None}}, default 'infer'
+    compression : {{'infer', 'gzip', 'bz2', 'zip', 'xz', 'zstd', None}},
+        default 'infer'.
         If 'infer' and 'path_or_url' is path-like, then detect compression from
-        the following extensions: '.gz', '.bz2', '.zip', or '.xz' (otherwise no
-        compression) If 'infer' and 'path_or_url' is not path-like, then use
-        None (= no decompression).
+        the following extensions: '.gz', '.bz2', '.zip', '.xz', or '.zst'
+        (otherwise no compression).
+        If 'infer' and 'path_or_url' is not path-like, then use
+        None (= no compression).
     protocol : int
         Int which indicates which protocol should be used by the pickler,
         default HIGHEST_PROTOCOL (see [1], paragraph 12.1.2). The possible
@@ -142,11 +144,13 @@ def read_pickle(
         .. versionchanged:: 1.0.0
            Accept URL. URL is not limited to S3 and GCS.
 
-    compression : {{'infer', 'gzip', 'bz2', 'zip', 'xz', None}}, default 'infer'
+    compression : {{'infer', 'gzip', 'bz2', 'zip', 'xz', 'zstd', None}},
+        default 'infer'.
         If 'infer' and 'path_or_url' is path-like, then detect compression from
-        the following extensions: '.gz', '.bz2', '.zip', or '.xz' (otherwise no
-        compression) If 'infer' and 'path_or_url' is not path-like, then use
-        None (= no decompression).
+        the following extensions: '.gz', '.bz2', '.zip', '.xz', or '.zst'
+        (otherwise no compression).
+        If 'infer' and 'path_or_url' is not path-like, then use
+        None (= no compression).
 
     {storage_options}
 

pandas/tests/io/test_pickle.py

@@ -298,6 +298,7 @@ class TestCompression:
         "bz2": ".bz2",
         "zip": ".zip",
         "xz": ".xz",
+        "zstd": ".zst",
     }
 
     def compress_file(self, src_path, dest_path, compression):
@@ -314,6 +315,9 @@ def compress_file(self, src_path, dest_path, compression):
                 f.write(src_path, os.path.basename(src_path))
         elif compression == "xz":
             f = get_lzma_file(lzma)(dest_path, "w")
+        elif compression == "zstd":
+            zstd = pytest.importorskip("zstandard")
+            f = zstd.open(dest_path, "wb")
         else:
             msg = f"Unrecognized compression type: {compression}"
             raise ValueError(msg)
@@ -350,7 +354,7 @@ def test_write_explicit_bad(self, compression, get_random_path):
                 df = tm.makeDataFrame()
                 df.to_pickle(path, compression=compression)
 
-    @pytest.mark.parametrize("ext", ["", ".gz", ".bz2", ".no_compress", ".xz"])
+    @pytest.mark.parametrize("ext", ["", ".gz", ".bz2", ".no_compress", ".xz", ".zst"])
     def test_write_infer(self, ext, get_random_path):
         base = get_random_path
         path1 = base + ext
@@ -396,7 +400,7 @@ def test_read_explicit(self, compression, get_random_path):
 
             tm.assert_frame_equal(df, df2)
 
-    @pytest.mark.parametrize("ext", ["", ".gz", ".bz2", ".zip", ".no_compress", ".xz"])
+    @pytest.mark.parametrize("ext", ["", ".gz", ".bz2", ".zip", ".no_compress", ".xz", ".zst"])
     def test_read_infer(self, ext, get_random_path):
         base = get_random_path
         path1 = base + ".raw"