BUG: Preserve nullable index dtype in _transform_index (#65316)

Author: gautamvarmadatla

doc/source/whatsnew/v3.1.0.rst

@@ -283,6 +283,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`)

pandas/core/indexes/base.py

@@ -6892,7 +6892,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

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]})

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"])