Skip to content

Commit eaacf83

Browse files
authored
BUG: CoW not tracking references when indexing midx with slice (#51944)
1 parent ffc55a0 commit eaacf83

File tree

3 files changed

+22
-1
lines changed

3 files changed

+22
-1
lines changed

doc/source/whatsnew/v2.0.0.rst

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -206,6 +206,9 @@ Copy-on-Write improvements
206206
- Arithmetic operations that can be inplace, e.g. ``ser *= 2`` will now respect the
207207
Copy-on-Write mechanism.
208208

209+
- :meth:`DataFrame.__getitem__` will now respect the Copy-on-Write mechanism when the
210+
:class:`DataFrame` has :class:`MultiIndex` columns.
211+
209212
- :meth:`Series.view` will now respect the Copy-on-Write mechanism.
210213

211214
Copy-on-Write can be enabled through one of

pandas/core/frame.py

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3832,10 +3832,13 @@ def _getitem_multilevel(self, key):
38323832
result = self.reindex(columns=new_columns)
38333833
result.columns = result_columns
38343834
else:
3835-
new_values = self.values[:, loc]
3835+
new_values = self._values[:, loc]
38363836
result = self._constructor(
38373837
new_values, index=self.index, columns=result_columns
38383838
)
3839+
if using_copy_on_write() and isinstance(loc, slice):
3840+
result._mgr.add_references(self._mgr) # type: ignore[arg-type]
3841+
38393842
result = result.__finalize__(self)
38403843

38413844
# If there is only one column being returned, and its name is

pandas/tests/copy_view/test_indexing.py

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1034,3 +1034,18 @@ def test_set_value_copy_only_necessary_column(
10341034
assert not np.shares_memory(get_array(df, "a"), get_array(view, "a"))
10351035
else:
10361036
assert np.shares_memory(get_array(df, "a"), get_array(view, "a"))
1037+
1038+
1039+
def test_getitem_midx_slice(using_copy_on_write, using_array_manager):
1040+
df = DataFrame({("a", "x"): [1, 2], ("a", "y"): 1, ("b", "x"): 2})
1041+
df_orig = df.copy()
1042+
new_df = df[("a",)]
1043+
1044+
if using_copy_on_write:
1045+
assert not new_df._mgr._has_no_reference(0)
1046+
1047+
if not using_array_manager:
1048+
assert np.shares_memory(get_array(df, ("a", "x")), get_array(new_df, "x"))
1049+
if using_copy_on_write:
1050+
new_df.iloc[0, 0] = 100
1051+
tm.assert_frame_equal(df_orig, df)

0 commit comments

Comments
 (0)