diff --git a/doc/source/whatsnew/v3.1.0.rst b/doc/source/whatsnew/v3.1.0.rst index f3055427be5e0..251c6f288c548 100644 --- a/doc/source/whatsnew/v3.1.0.rst +++ b/doc/source/whatsnew/v3.1.0.rst @@ -276,6 +276,7 @@ Indexing - Bug in :meth:`DataFrame.loc` raising ``ValueError`` when setting a row with a list-like value on a single-column :class:`DataFrame` with :class:`ExtensionArray` dtype (:issue:`44103`) - Bug in :meth:`DataFrame.loc` with a :class:`MultiIndex` returning wrong results instead of raising ``KeyError`` when passing string keys for numeric index levels (:issue:`60104`) - Bug in :meth:`DataFrame.mask` with ``inplace=True`` where incorrect values were produced when ``other`` was a :class:`Series` with :class:`ExtensionArray` values (:issue:`64635`) +- Bug in :meth:`DataFrame.rename` and :meth:`Series.rename` not preserving nullable extension dtype (e.g. ``Int64``, ``Float64``) when relabeling index or column labels (:issue:`65315`) - Bug in :meth:`DataFrame.where` and :meth:`DataFrame.mask` raising ``TypeError`` when ``cond`` is a :class:`Series` and ``axis=1`` (:issue:`58190`) - Bug in :meth:`DataFrame.xs` where ``drop_level=False`` was ignored for fully specified :class:`MultiIndex` keys when ``level`` was not explicitly provided (:issue:`6507`) - Bug in :meth:`Index.get_level_values` mishandling boolean, NA-like (``np.nan``, ``pd.NA``, ``pd.NaT``) and integer index names (:issue:`62169`) diff --git a/pandas/core/indexes/base.py b/pandas/core/indexes/base.py index 504ea3b6f96bf..2d38e1f9d1648 100644 --- a/pandas/core/indexes/base.py +++ b/pandas/core/indexes/base.py @@ -6889,7 +6889,18 @@ def _transform_index(self, func: Callable, *, level: int | None = None) -> Index return type(self).from_arrays(values) else: items = [func(x) for x in self] - return Index(items, name=self.name, tupleize_cols=False) + if not items: + return Index( + items, dtype=self.dtype, name=self.name, tupleize_cols=False + ) + new_values = self.array._cast_pointwise_result(items) + return Index( + new_values, + dtype=new_values.dtype, + copy=False, + name=self.name, + tupleize_cols=False, + ) def isin( self, values: Axes | set, level: str_t | int | None = None diff --git a/pandas/tests/frame/methods/test_rename.py b/pandas/tests/frame/methods/test_rename.py index 8c02e28bc138c..5271f304e9763 100644 --- a/pandas/tests/frame/methods/test_rename.py +++ b/pandas/tests/frame/methods/test_rename.py @@ -5,10 +5,12 @@ import pytest from pandas import ( + NA, DataFrame, Index, MultiIndex, Series, + array, merge, ) import pandas._testing as tm @@ -411,6 +413,64 @@ def test_rename_boolean_index(self): ) tm.assert_frame_equal(res, exp) + def test_rename_preserves_nullable_int_index(self): + # GH#65315 + df = DataFrame( + {"val": [1, 2, 3]}, + index=Index(array([1, 2, 3], dtype="Int64"), name="id"), + ) + result = df.rename({1: 9}) + expected = Index(array([9, 2, 3], dtype="Int64"), name="id") + tm.assert_index_equal(result.index, expected) + + def test_rename_preserves_nullable_float_columns(self): + # GH#65315 + df = DataFrame( + [[1, 2, 3]], + columns=Index(array([1.0, 2.0, 3.0], dtype="Float64")), + ) + result = df.rename(columns={1.0: 9.0}) + expected = Index(array([9.0, 2.0, 3.0], dtype="Float64")) + tm.assert_index_equal(result.columns, expected) + + def test_rename_nullable_index_to_na(self): + # GH#65315 + df = DataFrame( + {"val": [1, 2, 3]}, + index=Index(array([1, 2, 3], dtype="Int64")), + ) + result = df.rename({1: NA}) + expected = Index(array([NA, 2, 3], dtype="Int64")) + tm.assert_index_equal(result.index, expected) + + def test_rename_empty_nullable_index(self): + # GH#65315 + df = DataFrame({"val": []}, index=Index(array([], dtype="Int64"))) + result = df.rename(index={}) + tm.assert_index_equal(result.index, df.index) + + def test_rename_nullable_index_type_change_widens(self): + # GH#65315 + df = DataFrame( + {"val": [1, 2, 3]}, + index=Index(array([1, 2, 3], dtype="Int64")), + ) + result = df.rename(lambda x: f"label_{x}") + assert list(result.index) == ["label_1", "label_2", "label_3"] + assert result.index.dtype != "Int64" + + def test_rename_tuple_columns_not_multiindex(self): + # GH#65315 + df = DataFrame( + [[1, 2]], + columns=Index(array([1, 2], dtype="Int64")), + ) + result = df.rename(columns=lambda x: (x, x)) + assert not isinstance(result.columns, MultiIndex) + tm.assert_index_equal( + result.columns, Index([(1, 1), (2, 2)], tupleize_cols=False) + ) + def test_rename_non_unique_index_series(self): # GH#58621 df = DataFrame({"A": [1, 2, 3], "B": [4, 5, 6], "C": [7, 8, 9]}) diff --git a/pandas/tests/series/methods/test_rename.py b/pandas/tests/series/methods/test_rename.py index c912b9abd22ca..8b8792479810e 100644 --- a/pandas/tests/series/methods/test_rename.py +++ b/pandas/tests/series/methods/test_rename.py @@ -162,6 +162,16 @@ def test_rename_series_with_multiindex_keeps_ea_dtypes(self): tm.assert_series_equal(result, expected) + def test_rename_preserves_nullable_index_dtype(self): + # GH#65315 + ser = Series( + [1, 2, 3], + index=Index(array([1, 2, 3], dtype="Int64"), name="id"), + ) + result = ser.rename({1: 9}) + expected = Index(array([9, 2, 3], dtype="Int64"), name="id") + tm.assert_index_equal(result.index, expected) + def test_rename_error_arg(self): # GH 46889 ser = Series(["foo", "bar"])