Add Zstandard support to read_pickle/to_pickle

Author: jonashaag

MANIFEST.in

@@ -36,6 +36,7 @@ global-exclude *.xpt
 global-exclude *.cpt
 global-exclude *.xz
 global-exclude *.zip
+global-exclude *.zst
 global-exclude *~
 global-exclude .DS_Store
 global-exclude .git*

ci/deps/actions-38-slow.yaml

@@ -35,3 +35,4 @@ dependencies:
   - xlsxwriter
   - xlwt
   - numba
+  - zstandard

ci/deps/actions-39-slow.yaml

@@ -38,6 +38,7 @@ dependencies:
   - xlsxwriter
   - xlwt
   - pyreadstat
+  - zstandard
   - pip
   - pip:
     - pyxlsb

ci/deps/actions-39.yaml

@@ -37,6 +37,7 @@ dependencies:
   - xlsxwriter
   - xlwt
   - pyreadstat
+  - zstandard
   - pip
   - pip:
     - pyxlsb

ci/deps/azure-macos-38.yaml

@@ -31,6 +31,7 @@ dependencies:
   - xlrd
   - xlsxwriter
   - xlwt
+  - zstandard
   - pip
   - pip:
     - cython>=0.29.24

ci/deps/azure-windows-38.yaml

@@ -33,3 +33,4 @@ dependencies:
   - xlrd
   - xlsxwriter
   - xlwt
+  - zstandard

ci/deps/azure-windows-39.yaml

@@ -37,6 +37,7 @@ dependencies:
   - xlsxwriter
   - xlwt
   - pyreadstat
+  - zstandard
   - pip
   - pip:
     - pyxlsb

ci/deps/circle-38-arm64.yaml

@@ -16,6 +16,7 @@ dependencies:
   - numpy
   - python-dateutil
   - pytz
+  - zstandard
   - pip
   - flask
   - pip:

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
+========================= ================== =============================================================

doc/source/user_guide/io.rst

@@ -316,14 +316,14 @@ chunksize : int, default ``None``
 Quoting, compression, and file format
 +++++++++++++++++++++++++++++++++++++
 
-compression : {``'infer'``, ``'gzip'``, ``'bz2'``, ``'zip'``, ``'xz'``, ``None``, ``dict``}, default ``'infer'``
+compression : {``'infer'``, ``'gzip'``, ``'bz2'``, ``'zip'``, ``'xz'``, ``'zstd'``, ``None``, ``dict``}, default ``'infer'``
   For on-the-fly decompression of on-disk data. If 'infer', then use gzip,
-  bz2, zip, or xz if ``filepath_or_buffer`` is path-like ending in '.gz', '.bz2',
-  '.zip', or '.xz', respectively, and no decompression otherwise. If using 'zip',
+  bz2, zip, xz, or zstandard if ``filepath_or_buffer`` is path-like ending in '.gz', '.bz2',
+  '.zip', '.xz', '.zst', respectively, and no decompression otherwise. If using 'zip',
   the ZIP file must contain only one data file to be read in.
   Set to ``None`` for no decompression. Can also be a dict with key ``'method'``
-  set to one of {``'zip'``, ``'gzip'``, ``'bz2'``} and other key-value pairs are
-  forwarded to ``zipfile.ZipFile``, ``gzip.GzipFile``, or ``bz2.BZ2File``.
+  set to one of {``'zip'``, ``'gzip'``, ``'bz2'``, ``'zstd'``} and other key-value pairs are
+  forwarded to ``zipfile.ZipFile``, ``gzip.GzipFile``, ``bz2.BZ2File``, or ``zstandard.ZstdDecompressor``.
   As an example, the following could be passed for faster compression and to
   create a reproducible gzip archive:
   ``compression={'method': 'gzip', 'compresslevel': 1, 'mtime': 1}``.
@@ -4022,18 +4022,18 @@ Compressed pickle files
 '''''''''''''''''''''''
 
 :func:`read_pickle`, :meth:`DataFrame.to_pickle` and :meth:`Series.to_pickle` can read
-and write compressed pickle files. The compression types of ``gzip``, ``bz2``, ``xz`` are supported for reading and writing.
+and write compressed pickle files. The compression types of ``gzip``, ``bz2``, ``xz``, ``zstd`` are supported for reading and writing.
 The ``zip`` file format only supports reading and must contain only one data file
 to be read.
 
 The compression type can be an explicit parameter or be inferred from the file extension.
-If 'infer', then use ``gzip``, ``bz2``, ``zip``, or ``xz`` if filename ends in ``'.gz'``, ``'.bz2'``, ``'.zip'``, or
-``'.xz'``, respectively.
+If 'infer', then use ``gzip``, ``bz2``, ``zip``, ``xz``, ``zstd`` if filename ends in ``'.gz'``, ``'.bz2'``, ``'.zip'``,
+``'.xz'``, or ``'.zst'``, respectively.
 
 The compression parameter can also be a ``dict`` in order to pass options to the
 compression protocol. It must have a ``'method'`` key set to the name
 of the compression protocol, which must be one of
-{``'zip'``, ``'gzip'``, ``'bz2'``}. All other key-value pairs are passed to
+{``'zip'``, ``'gzip'``, ``'bz2'``, ``'zstd'``}. All other key-value pairs are passed to
 the underlying compression library.
 
 .. ipython:: python

environment.yml

@@ -85,6 +85,7 @@ dependencies:
   - numexpr>=2.7.1
   - scipy>=1.4.1
   - numba>=0.50.1
+  - zstandard>=0.15.2
 
   # optional for io
   # ---------------

pandas/_testing/_io.py

@@ -364,7 +364,7 @@ def write_to_compressed(compression, path, data, dest="test"):
 
     Parameters
     ----------
-    compression : {'gzip', 'bz2', 'zip', 'xz'}
+    compression : {'gzip', 'bz2', 'zip', 'xz', 'zstd'}
         The compression type to use.
     path : str
         The file path to write the data.
@@ -391,6 +391,11 @@ def write_to_compressed(compression, path, data, dest="test"):
         compress_method = gzip.GzipFile
     elif compression == "bz2":
         compress_method = bz2.BZ2File
+    elif compression == "zstd":
+        import pytest
+
+        zstd = pytest.importorskip("zstandard")
+        compress_method = zstd.open
     elif compression == "xz":
         compress_method = get_lzma_file()
     else:

pandas/_testing/contexts.py

@@ -29,7 +29,7 @@ def decompress_file(path, compression):
     path : str
         The path where the file is read from.
 
-    compression : {'gzip', 'bz2', 'zip', 'xz', None}
+    compression : {'gzip', 'bz2', 'zip', 'xz', 'zstd', None}
         Name of the decompression to use
 
     Returns

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/conftest.py

@@ -44,6 +44,13 @@
     utc,
 )
 
+try:
+    import zstandard as zstd  # noqa: F401
+
+    have_zstd = True
+except ImportError:
+    have_zstd = False
+
 import pandas.util._test_decorators as td
 
 from pandas.core.dtypes.dtypes import (
@@ -267,15 +274,17 @@ def other_closed(request):
     return request.param
 
 
-@pytest.fixture(params=[None, "gzip", "bz2", "zip", "xz"])
+@pytest.fixture(
+    params=[None, "gzip", "bz2", "zip", "xz"] + (["zstd"] if have_zstd else [])
+)
 def compression(request):
     """
     Fixture for trying common compression types in compression tests.
     """
     return request.param
 
 
-@pytest.fixture(params=["gzip", "bz2", "zip", "xz"])
+@pytest.fixture(params=["gzip", "bz2", "zip", "xz"] + (["zstd"] if have_zstd else []))
 def compression_only(request):
     """
     Fixture for trying common compression types in compression tests excluding

pandas/core/frame.py

@@ -2569,11 +2569,11 @@ def to_stata(
             For on-the-fly compression of the output dta. 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
+            'bz2', 'zip', 'xz', 'zstd', None}}. If compression mode is 'infer' and
             `fname` is path-like, then detect compression from the following
-            extensions: '.gz', '.bz2', '.zip', or '.xz' (otherwise no
+            extensions: '.gz', '.bz2', '.zip', '.xz', '.zst' (otherwise no
             compression). If dict and compression mode is one of {{'zip',
-            'gzip', 'bz2'}}, or inferred as one of the above, other entries
+            'gzip', 'bz2', 'zstd'}}, or inferred as one of the above, other entries
             passed as additional compression options.
 
             .. versionadded:: 1.1.0
@@ -3020,10 +3020,11 @@ def to_xml(
             layout of elements and attributes from original output. This
             argument requires ``lxml`` to be installed. Only XSLT 1.0
             scripts and not later versions is currently supported.
-        compression : {{'infer', 'gzip', 'bz2', 'zip', 'xz', None}}, default 'infer'
+        compression : {{'infer', 'gzip', 'bz2', 'zip', 'xz', 'zstd', None}},
+            default 'infer'.
             For on-the-fly decompression of on-disk data. If 'infer', then use
-            gzip, bz2, zip or xz if path_or_buffer is a string ending in
-            '.gz', '.bz2', '.zip', or 'xz', respectively, and no decompression
+            gzip, bz2, zip, xz, or zstandard if path_or_buffer is a string ending in
+            '.gz', '.bz2', '.zip', '.xz', '.zst' respectively, and no decompression
             otherwise. If using 'zip', the ZIP file must contain only one data
             file to be read in. Set to None for no decompression.
         {storage_options}

pandas/core/generic.py

@@ -2406,7 +2406,7 @@ def to_json(
             throw ValueError if incorrect 'orient' since others are not
             list-like.
 
-        compression : {{'infer', 'gzip', 'bz2', 'zip', 'xz', None}}
+        compression : {{'infer', 'gzip', 'bz2', 'zip', 'xz', 'zstd', None}}
 
             A string representing the compression to use in the output file,
             only used when the first argument is a filename. By default, the
@@ -2933,16 +2933,16 @@ def to_pickle(
         ----------
         path : str
             File path where the pickled object will be stored.
-        compression : {{'infer', 'gzip', 'bz2', 'zip', 'xz', None}}, \
+        compression : {{'infer', 'gzip', 'bz2', 'zip', 'xz', 'zstd', None}}, \
         default 'infer'
             A string representing the compression to use in the output file. By
             default, infers from the file extension in specified path.
             Compression mode may be any of the following possible
-            values: {{‘infer’, ‘gzip’, ‘bz2’, ‘zip’, ‘xz’, None}}. If compression
-            mode is ‘infer’ and path_or_buf is path-like, then detect
+            values: {{'infer', 'gzip', 'bz2', 'zip', 'xz', 'zstd', None}}.
+            If compression mode is 'infer' and path_or_buf is path-like, then detect
             compression mode from the following extensions:
-            ‘.gz’, ‘.bz2’, ‘.zip’ or ‘.xz’. (otherwise no compression).
-            If dict given and mode is ‘zip’ or inferred as ‘zip’, other entries
+            '.gz', '.bz2', '.zip', '.xz', '.zst'. (otherwise no compression).
+            If dict given and mode is 'zip' or inferred as 'zip', other entries
             passed as additional compression options.
         protocol : int
             Int which indicates which protocol should be used by the pickler,
@@ -3406,11 +3406,11 @@ def to_csv(
         compression : str or dict, default 'infer'
             If str, represents compression mode. If dict, value at 'method' is
             the compression mode. Compression mode may be any of the following
-            possible values: {{'infer', 'gzip', 'bz2', 'zip', 'xz', None}}. If
-            compression mode is 'infer' and `path_or_buf` is path-like, then
+            possible values: {{'infer', 'gzip', 'bz2', 'zip', 'xz', 'zstd', None}}.
+            If compression mode is 'infer' and `path_or_buf` is path-like, then
             detect compression mode from the following extensions: '.gz',
-            '.bz2', '.zip' or '.xz'. (otherwise no compression). If dict given
-            and mode is one of {{'zip', 'gzip', 'bz2'}}, or inferred as
+            '.bz2', '.zip', '.xz', '.zst'. (otherwise no compression). If dict given
+            and mode is one of {{'zip', 'gzip', 'bz2', 'zstd'}}, or inferred as
             one of the above, other entries passed as
             additional compression options.
             If `path_or_buf` is omitted or `None` or is a file opened in text
@@ -3426,8 +3426,7 @@ def to_csv(
             .. versionchanged:: 1.1.0
 
                Passing compression options as keys in dict is
-               supported for compression modes 'gzip' and 'bz2'
-               as well as 'zip'.
+               supported for compression modes 'gzip', 'bz2', 'zstd', and 'zip'.
 
             .. versionchanged:: 1.2.0