PERF: Improve pickle support with BZ2 & LZMA

pandas/compat/__init__.py

@@ -10,9 +10,16 @@
 from __future__ import annotations
 
 import os
+from pickle import PickleBuffer
 import platform
 import sys
-from typing import TYPE_CHECKING
+
+try:
+    import pandas.compat.lzma
+
+    has_lzma = True
+except ImportError:
+    has_lzma = False
 
 from pandas._typing import F
 from pandas.compat.numpy import (
@@ -26,9 +33,6 @@
     pa_version_under9p0,
 )
 
-if TYPE_CHECKING:
-    import lzma
-
 PY39 = sys.version_info >= (3, 9)
 PY310 = sys.version_info >= (3, 10)
 PY311 = sys.version_info >= (3, 11)
@@ -121,7 +125,7 @@ def is_ci_environment() -> bool:
     return os.environ.get("PANDAS_CI", "0") == "1"
 
 
-def get_lzma_file() -> type[lzma.LZMAFile]:
+def get_lzma_file() -> type[pandas.compat.lzma.LZMAFile]:
     """
     Importing the `LZMAFile` class from the `lzma` module.
 
@@ -135,15 +139,13 @@ def get_lzma_file() -> type[lzma.LZMAFile]:
     RuntimeError
         If the `lzma` module was not imported correctly, or didn't exist.
     """
-    try:
-        import lzma
-    except ImportError:
+    if not has_lzma:
         raise RuntimeError(
             "lzma module not available. "
             "A Python re-install with the proper dependencies, "
             "might be required to solve this issue."
         )
-    return lzma.LZMAFile
+    return pandas.compat.lzma.LZMAFile
 
 
 __all__ = [

pandas/compat/_utils.py

@@ -0,0 +1,29 @@
+"""
+Utilities used in various `compat` components.
+"""
+
+from __future__ import annotations
+
+import pickle
+
+
+def flatten_buffer(b: bytes | bytearray | memoryview | pickle.PickleBuffer):
+    """
+    Return some 1-D `uint8` typed buffer.
+
+    Coerces anything that does not match that description to one that does
+    without copying if possible (otherwise will copy).
+    """
+
+    if isinstance(b, (bytes, bytearray)):
+        return b
+
+    if not isinstance(b, pickle.PickleBuffer):
+        b = pickle.PickleBuffer(b)
+
+    try:
+        # coerce to 1-D `uint8` C-contiguous `memoryview` zero-copy
+        return b.raw()
+    except BufferError:
+        # perform in-memory copy if buffer is not contiguous
+        return memoryview(b).tobytes()

pandas/compat/bz2.py

@@ -0,0 +1,25 @@
+"""
+Patched ``BZ2File`` to handle pickle protocol 5.
+"""
+
+from __future__ import annotations
+
+import bz2
+import sys
+
+from pandas.compat._utils import flatten_buffer
+
+
+if sys.version_info < (3, 10):
+    class BZ2File(bz2.BZ2File):
+        def write(self, b) -> int:
+            # Workaround issue where `bz2.BZ2File` expects `len`
+            # to return the number of bytes in `b` by converting
+            # `b` into something that meets that constraint with
+            # minimal copying.
+            #
+            # Note: This is fixed in Python 3.10.
+            return super().write(flatten_buffer(b))
+else:
+    class BZ2File(bz2.BZ2File):
+        pass

pandas/compat/lzma.py

@@ -0,0 +1,25 @@
+"""
+Patched ``LZMAFile`` to handle pickle protocol 5.
+"""
+
+from __future__ import annotations
+
+import lzma
+import sys
+
+from pandas.compat._utils import flatten_buffer
+
+
+if sys.version_info < (3, 10):
+    class LZMAFile(lzma.LZMAFile):
+        def write(self, b) -> int:
+            # Workaround issue where `lzma.LZMAFile` expects `len`
+            # to return the number of bytes in `b` by converting
+            # `b` into something that meets that constraint with
+            # minimal copying.
+            #
+            # Note: This is fixed in Python 3.10.
+            return super().write(flatten_buffer(b))
+else:
+    class LZMAFile(lzma.LZMAFile):
+        pass

pandas/io/common.py

@@ -5,7 +5,6 @@
     ABC,
     abstractmethod,
 )
-import bz2
 import codecs
 import dataclasses
 import functools
@@ -21,6 +20,7 @@
 import mmap
 import os
 from pathlib import Path
+from pickle import PickleBuffer
 import re
 import tarfile
 from typing import (
@@ -56,6 +56,7 @@
 )
 from pandas.compat import get_lzma_file
 from pandas.compat._optional import import_optional_dependency
+from pandas.compat.bz2 import BZ2File as _BZ2File
 from pandas.util._decorators import doc
 from pandas.util._exceptions import find_stack_level
 
@@ -761,9 +762,9 @@ def get_handle(
 
         # BZ Compression
         elif compression == "bz2":
-            # No overload variant of "BZ2File" matches argument types
+            # Overload of "BZ2File" to handle pickle protocol 5
             # "Union[str, BaseBuffer]", "str", "Dict[str, Any]"
-            handle = bz2.BZ2File(  # type: ignore[call-overload]
+            handle = _BZ2File(  # type: ignore[call-overload]
                 handle,
                 mode=ioargs.mode,
                 **compression_args,

pandas/io/pickle.py

@@ -101,15 +101,8 @@ def to_pickle(
         is_text=False,
         storage_options=storage_options,
     ) as handles:
-        if handles.compression["method"] in ("bz2", "xz") and protocol >= 5:
-            # some weird TypeError GH#39002 with pickle 5: fallback to letting
-            # pickle create the entire object and then write it to the buffer.
-            # "zip" would also be here if pandas.io.common._BytesZipFile
-            # wouldn't buffer write calls
-            handles.handle.write(pickle.dumps(obj, protocol=protocol))
-        else:
-            # letting pickle write directly to the buffer is more memory-efficient
-            pickle.dump(obj, handles.handle, protocol=protocol)
+        # letting pickle write directly to the buffer is more memory-efficient
+        pickle.dump(obj, handles.handle, protocol=protocol)
 
 
 @doc(

pandas/tests/io/test_pickle.py

@@ -10,6 +10,7 @@
 
 3. Move the created pickle to "data/legacy_pickle/<version>" directory.
 """
+from array import array
 import bz2
 import datetime
 import functools
@@ -38,6 +39,7 @@
     is_platform_little_endian,
 )
 from pandas.compat._optional import import_optional_dependency
+from pandas.compat._utils import flatten_buffer
 import pandas.util._test_decorators as td
 
 import pandas as pd
@@ -105,6 +107,34 @@ def legacy_pickle(request, datapath):
 # ---------------------
 # tests
 # ---------------------
+
+
+@pytest.mark.parametrize(
+    "data",
+    [
+        b"123",
+        b"123456",
+        bytearray(b"123"),
+        memoryview(b"123"),
+        pickle.PickleBuffer(b"123"),
+        array("I", [1, 2, 3]),
+        memoryview(b"123456").cast("B", (3, 2)),
+        memoryview(b"123456").cast("B", (3, 2))[::2],
+        np.arange(12).reshape((3, 4), order="C"),
+        np.arange(12).reshape((3, 4), order="F"),
+        np.arange(12).reshape((3, 4), order="C")[:, ::2],
+    ],
+)
+def test_flatten_buffer(data):
+    result = flatten_buffer(data)
+    assert result == bytes(data)
+    if isinstance(result, memoryview):
+        assert result.ndim == 1
+        assert result.format == "B"
+        assert result.contiguous
+        assert result.shape == (result.nbytes,)
+
+
 def test_pickles(legacy_pickle):
     if not is_platform_little_endian():
         pytest.skip("known failure on non-little endian")