Add Zstandard compression support

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

pandas/_testing/_io.py

@@ -15,6 +15,7 @@
     ReadPickleBuffer,
 )
 from pandas.compat import get_lzma_file
+from pandas.compat._optional import import_optional_dependency
 
 import pandas as pd
 from pandas._testing._random import rands
@@ -364,7 +365,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 +392,8 @@ def write_to_compressed(compression, path, data, dest="test"):
         compress_method = gzip.GzipFile
     elif compression == "bz2":
         compress_method = bz2.BZ2File
+    elif compression == "zstd":
+        compress_method = import_optional_dependency("zstandard").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

@@ -267,15 +267,32 @@ def other_closed(request):
     return request.param
 
 
-@pytest.fixture(params=[None, "gzip", "bz2", "zip", "xz"])
+@pytest.fixture(
+    params=[
+        None,
+        "gzip",
+        "bz2",
+        "zip",
+        "xz",
+        pytest.param("zstd", marks=td.skip_if_no("zstandard")),
+    ]
+)
 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",
+        pytest.param("zstd", marks=td.skip_if_no("zstandard")),
+    ]
+)
 def compression_only(request):
     """
     Fixture for trying common compression types in compression tests excluding

pandas/core/describe.py

@@ -35,8 +35,6 @@
 
 from pandas.core.reshape.concat import concat
 
-from pandas.io.formats.format import format_percentiles
-
 if TYPE_CHECKING:
     from pandas import (
         DataFrame,
@@ -230,6 +228,8 @@ def describe_numeric_1d(series: Series, percentiles: Sequence[float]) -> Series:
     """
     from pandas import Series
 
+    from pandas.io.formats.format import format_percentiles
+
     # error: Argument 1 to "format_percentiles" has incompatible type "Sequence[float]";
     # expected "Union[ndarray, List[Union[int, float]], List[float], List[Union[str,
     # float]]]"

pandas/core/frame.py

@@ -2486,7 +2486,10 @@ def _from_arrays(
         )
         return cls(mgr)
 
-    @doc(storage_options=generic._shared_docs["storage_options"])
+    @doc(
+        storage_options=generic._shared_docs["storage_options"],
+        compression_options=generic._shared_docs["compression_options"] % "path",
+    )
     @deprecate_kwarg(old_arg_name="fname", new_arg_name="path")
     def to_stata(
         self,
@@ -2565,16 +2568,7 @@ def to_stata(
             format. Only available if version is 117.  Storing strings in the
             StrL format can produce smaller dta files if strings have more than
             8 characters and values are repeated.
-        compression : str or dict, default 'infer'
-            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
-            `fname` 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, other entries
-            passed as additional compression options.
+        {compression_options}
 
             .. versionadded:: 1.1.0
 
@@ -2943,7 +2937,11 @@ def to_html(
             render_links=render_links,
         )
 
-    @doc(storage_options=generic._shared_docs["storage_options"])
+    @doc(
+        storage_options=generic._shared_docs["storage_options"],
+        compression_options=generic._shared_docs["compression_options"]
+        % "path_or_buffer",
+    )
     def to_xml(
         self,
         path_or_buffer: FilePath | WriteBuffer[bytes] | WriteBuffer[str] | None = None,
@@ -3020,12 +3018,7 @@ 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'
-            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
-            otherwise. If using 'zip', the ZIP file must contain only one data
-            file to be read in. Set to None for no decompression.
+        {compression_options}
         {storage_options}
 
         Returns

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
 

pandas/core/shared_docs.py

@@ -402,6 +402,35 @@
     starting with "s3://", and "gcs://") the key-value pairs are forwarded to
     ``fsspec``. Please see ``fsspec`` and ``urllib`` for more details."""
 
+_shared_docs[
+    "compression_options"
+] = """compression : str or dict, default 'infer'
+    For on-the-fly compression of the output data. If 'infer' and '%s'
+    path-like, then detect compression from the following extensions: '.gz',
+    '.bz2', '.zip', '.xz', or '.zst' (otherwise no compression). Set to
+    ``None`` for no compression. Can also be a dict with key ``'method'`` 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``, respectively. 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}``.
+"""
+
+_shared_docs[
+    "decompression_options"
+] = """compression : str or dict, default 'infer'
+    For on-the-fly decompression of on-disk data. If 'infer' and '%s' is
+    path-like, then detect compression from the following extensions: '.gz',
+    '.bz2', '.zip', '.xz', or '.zst' (otherwise no compression). 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'``, ``'zstd'``} and other
+    key-value pairs are forwarded to ``zipfile.ZipFile``, ``gzip.GzipFile``,
+    ``bz2.BZ2File``, or ``zstandard.ZstdDecompressor``, respectively. As an
+    example, the following could be passed for Zstandard decompression using a
+    custom compression dictionary: ``compression={'method': 'zstd', 'dict_data': my_compression_dict}``.
+"""
+
 _shared_docs[
     "replace"
 ] = """