Skip to content

Commit 0641183

Browse files
authored
ENH/CoW: use lazy copy in set_axis method (#49600)
1 parent 428c2a6 commit 0641183

File tree

5 files changed

+58
-12
lines changed

5 files changed

+58
-12
lines changed

pandas/core/frame.py

+1-1
Original file line numberDiff line numberDiff line change
@@ -4942,7 +4942,7 @@ def set_axis(
49424942
labels,
49434943
*,
49444944
axis: Axis = 0,
4945-
copy: bool = True,
4945+
copy: bool | None = None,
49464946
) -> DataFrame:
49474947
return super().set_axis(labels, axis=axis, copy=copy)
49484948

pandas/core/generic.py

+4-2
Original file line numberDiff line numberDiff line change
@@ -700,7 +700,7 @@ def set_axis(
700700
labels,
701701
*,
702702
axis: Axis = 0,
703-
copy: bool_t = True,
703+
copy: bool_t | None = None,
704704
) -> NDFrameT:
705705
"""
706706
Assign desired index to given axis.
@@ -734,7 +734,9 @@ def set_axis(
734734
return self._set_axis_nocheck(labels, axis, inplace=False, copy=copy)
735735

736736
@final
737-
def _set_axis_nocheck(self, labels, axis: Axis, inplace: bool_t, copy: bool_t):
737+
def _set_axis_nocheck(
738+
self, labels, axis: Axis, inplace: bool_t, copy: bool_t | None
739+
):
738740
if inplace:
739741
setattr(self, self._get_axis_name(axis), labels)
740742
else:

pandas/core/series.py

+1-1
Original file line numberDiff line numberDiff line change
@@ -4809,7 +4809,7 @@ def set_axis(
48094809
labels,
48104810
*,
48114811
axis: Axis = 0,
4812-
copy: bool = True,
4812+
copy: bool | None = None,
48134813
) -> Series:
48144814
return super().set_axis(labels, axis=axis, copy=copy)
48154815

pandas/tests/copy_view/test_methods.py

+34
Original file line numberDiff line numberDiff line change
@@ -356,3 +356,37 @@ def test_reorder_levels(using_copy_on_write):
356356
if using_copy_on_write:
357357
assert not np.shares_memory(get_array(df2, "a"), get_array(df, "a"))
358358
tm.assert_frame_equal(df, df_orig)
359+
360+
361+
def test_frame_set_axis(using_copy_on_write):
362+
# GH 49473
363+
df = DataFrame({"a": [1, 2, 3], "b": [4, 5, 6], "c": [0.1, 0.2, 0.3]})
364+
df_orig = df.copy()
365+
df2 = df.set_axis(["a", "b", "c"], axis="index")
366+
367+
if using_copy_on_write:
368+
assert np.shares_memory(get_array(df2, "a"), get_array(df, "a"))
369+
else:
370+
assert not np.shares_memory(get_array(df2, "a"), get_array(df, "a"))
371+
372+
# mutating df2 triggers a copy-on-write for that column / block
373+
df2.iloc[0, 0] = 0
374+
assert not np.shares_memory(get_array(df2, "a"), get_array(df, "a"))
375+
tm.assert_frame_equal(df, df_orig)
376+
377+
378+
def test_series_set_axis(using_copy_on_write):
379+
# GH 49473
380+
ser = Series([1, 2, 3])
381+
ser_orig = ser.copy()
382+
ser2 = ser.set_axis(["a", "b", "c"], axis="index")
383+
384+
if using_copy_on_write:
385+
assert np.shares_memory(ser, ser2)
386+
else:
387+
assert not np.shares_memory(ser, ser2)
388+
389+
# mutating ser triggers a copy-on-write for the column / block
390+
ser2.iloc[0] = 0
391+
assert not np.shares_memory(ser2, ser)
392+
tm.assert_series_equal(ser, ser_orig)

pandas/tests/frame/methods/test_set_axis.py

+18-8
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@ def test_set_axis(self, obj):
2121
result = obj.set_axis(new_index, axis=0)
2222
tm.assert_equal(expected, result)
2323

24-
def test_set_axis_copy(self, obj):
24+
def test_set_axis_copy(self, obj, using_copy_on_write):
2525
# Test copy keyword GH#47932
2626
new_index = list("abcd")[: len(obj)]
2727

@@ -57,14 +57,24 @@ def test_set_axis_copy(self, obj):
5757
result = obj.set_axis(new_index, axis=0)
5858
tm.assert_equal(expected, result)
5959
assert result is not obj
60-
# check we DID make a copy
61-
if obj.ndim == 1:
62-
assert not tm.shares_memory(result, obj)
60+
if using_copy_on_write:
61+
# check we DID NOT make a copy
62+
if obj.ndim == 1:
63+
assert tm.shares_memory(result, obj)
64+
else:
65+
assert any(
66+
tm.shares_memory(result.iloc[:, i], obj.iloc[:, i])
67+
for i in range(obj.shape[1])
68+
)
6369
else:
64-
assert not any(
65-
tm.shares_memory(result.iloc[:, i], obj.iloc[:, i])
66-
for i in range(obj.shape[1])
67-
)
70+
# check we DID make a copy
71+
if obj.ndim == 1:
72+
assert not tm.shares_memory(result, obj)
73+
else:
74+
assert not any(
75+
tm.shares_memory(result.iloc[:, i], obj.iloc[:, i])
76+
for i in range(obj.shape[1])
77+
)
6878

6979
res = obj.set_axis(new_index, copy=False)
7080
tm.assert_equal(expected, res)

0 commit comments

Comments
 (0)