ENH/CoW: use lazy copy in set_axis method

pandas/core/frame.py

@@ -4942,7 +4942,7 @@ def set_axis(
         labels,
         *,
         axis: Axis = 0,
-        copy: bool = True,
+        copy: bool | None = None,
     ) -> DataFrame:
         return super().set_axis(labels, axis=axis, copy=copy)
 

pandas/core/generic.py

@@ -700,7 +700,7 @@ def set_axis(
         labels,
         *,
         axis: Axis = 0,
-        copy: bool_t = True,
+        copy: bool_t | None = None,
     ) -> NDFrameT:
         """
         Assign desired index to given axis.
@@ -734,7 +734,9 @@ def set_axis(
         return self._set_axis_nocheck(labels, axis, inplace=False, copy=copy)
 
     @final
-    def _set_axis_nocheck(self, labels, axis: Axis, inplace: bool_t, copy: bool_t):
+    def _set_axis_nocheck(
+        self, labels, axis: Axis, inplace: bool_t, copy: bool_t | None
+    ):
         if inplace:
             setattr(self, self._get_axis_name(axis), labels)
         else:

pandas/core/series.py

@@ -4809,7 +4809,7 @@ def set_axis(
         labels,
         *,
         axis: Axis = 0,
-        copy: bool = True,
+        copy: bool | None = None,
     ) -> Series:
         return super().set_axis(labels, axis=axis, copy=copy)
 

pandas/tests/copy_view/test_methods.py

@@ -356,3 +356,37 @@ def test_reorder_levels(using_copy_on_write):
     if using_copy_on_write:
         assert not np.shares_memory(get_array(df2, "a"), get_array(df, "a"))
     tm.assert_frame_equal(df, df_orig)
+
+
+def test_frame_set_axis(using_copy_on_write):
+    # GH 49473
+    df = DataFrame({"a": [1, 2, 3], "b": [4, 5, 6], "c": [0.1, 0.2, 0.3]})
+    df_orig = df.copy()
+    df2 = df.set_axis(["a", "b", "c"], axis="index")
+
+    if using_copy_on_write:
+        assert np.shares_memory(get_array(df2, "a"), get_array(df, "a"))
+    else:
+        assert not np.shares_memory(get_array(df2, "a"), get_array(df, "a"))
+
+    # mutating df2 triggers a copy-on-write for that column / block
+    df2.iloc[0, 0] = 0
+    assert not np.shares_memory(get_array(df2, "a"), get_array(df, "a"))
+    tm.assert_frame_equal(df, df_orig)
+
+
+def test_series_set_axis(using_copy_on_write):
+    # GH 49473
+    ser = Series([1, 2, 3])
+    ser_orig = ser.copy()
+    ser2 = ser.set_axis(["a", "b", "c"], axis="index")
+
+    if using_copy_on_write:
+        assert np.shares_memory(ser, ser2)
+    else:
+        assert not np.shares_memory(ser, ser2)
+
+    # mutating ser triggers a copy-on-write for the column / block
+    ser2.iloc[0] = 0
+    assert not np.shares_memory(ser2, ser)
+    tm.assert_series_equal(ser, ser_orig)

pandas/tests/frame/methods/test_set_axis.py

@@ -21,7 +21,7 @@ def test_set_axis(self, obj):
         result = obj.set_axis(new_index, axis=0)
         tm.assert_equal(expected, result)
 
-    def test_set_axis_copy(self, obj):
+    def test_set_axis_copy(self, obj, using_copy_on_write):
         # Test copy keyword GH#47932
         new_index = list("abcd")[: len(obj)]
 
@@ -57,14 +57,24 @@ def test_set_axis_copy(self, obj):
         result = obj.set_axis(new_index, axis=0)
         tm.assert_equal(expected, result)
         assert result is not obj
-        # check we DID make a copy
-        if obj.ndim == 1:
-            assert not tm.shares_memory(result, obj)
+        if using_copy_on_write:
+            # check we DID NOT make a copy
+            if obj.ndim == 1:
+                assert tm.shares_memory(result, obj)
+            else:
+                assert any(
+                    tm.shares_memory(result.iloc[:, i], obj.iloc[:, i])
+                    for i in range(obj.shape[1])
+                )
         else:
-            assert not any(
-                tm.shares_memory(result.iloc[:, i], obj.iloc[:, i])
-                for i in range(obj.shape[1])
-            )
+            # check we DID make a copy
+            if obj.ndim == 1:
+                assert not tm.shares_memory(result, obj)
+            else:
+                assert not any(
+                    tm.shares_memory(result.iloc[:, i], obj.iloc[:, i])
+                    for i in range(obj.shape[1])
+                )
 
         res = obj.set_axis(new_index, copy=False)
         tm.assert_equal(expected, res)