Skip to content

Commit cd67a82

Browse files
authored
[Data] Support pandas 2.2.0 for pandas tensor extensions (#42892)
`ExtensionArrayFormatter` is replaced with `_ ExtensionArrayFormatter` in pandas 2.2.0. If Ray Data users have this version of pandas installed, it will result in the error in #42842. This PR adds the case to support the change to `ExtensionArrayFormatter`. Tested locally with pandas 2.2.0 and Ray Data usage works without the error. Signed-off-by: Scott Lee <sjl@anyscale.com>
1 parent 9fd7a7f commit cd67a82

File tree

1 file changed

+12
-7
lines changed
  • python/ray/air/util/tensor_extensions

1 file changed

+12
-7
lines changed

python/ray/air/util/tensor_extensions/pandas.py

Lines changed: 12 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -42,7 +42,6 @@
4242
from pandas.compat import set_function_name
4343
from pandas.core.dtypes.generic import ABCDataFrame, ABCSeries
4444
from pandas.core.indexers import check_array_indexer, validate_indices
45-
from pandas.io.formats.format import ExtensionArrayFormatter
4645

4746
from ray.air.util.tensor_extensions.utils import (
4847
_create_possibly_ragged_ndarray,
@@ -167,14 +166,20 @@ def _format(x):
167166
_FORMATTER_ENABLED_ENV_VAR = "TENSOR_COLUMN_EXTENSION_FORMATTER_ENABLED"
168167

169168
if os.getenv(_FORMATTER_ENABLED_ENV_VAR, "1") == "1":
170-
ExtensionArrayFormatter._format_strings_orig = (
171-
ExtensionArrayFormatter._format_strings
172-
)
169+
if Version(pd.__version__) < Version("2.2.0"):
170+
from pandas.io.formats.format import ExtensionArrayFormatter
171+
172+
formatter_cls = ExtensionArrayFormatter
173+
else:
174+
from pandas.io.formats.format import _ExtensionArrayFormatter
175+
176+
formatter_cls = _ExtensionArrayFormatter
177+
formatter_cls._format_strings_orig = formatter_cls._format_strings
173178
if Version("1.1.0") <= Version(pd.__version__) < Version("1.3.0"):
174-
ExtensionArrayFormatter._format_strings = _format_strings_patched
179+
formatter_cls._format_strings = _format_strings_patched
175180
else:
176-
ExtensionArrayFormatter._format_strings = _format_strings_patched_v1_0_0
177-
ExtensionArrayFormatter._patched_by_ray_datasets = True
181+
formatter_cls._format_strings = _format_strings_patched_v1_0_0
182+
formatter_cls._patched_by_ray_datasets = True
178183

179184
###########################################
180185
# End patching of ExtensionArrayFormatter #

0 commit comments

Comments
 (0)