CoW: Avoid unnecessary copies for columnwise replace (#54117)

Author: phofl

pandas/core/frame.py

@@ -4188,11 +4188,15 @@ def _set_item_mgr(
         if len(self):
             self._check_setitem_copy()
 
-    def _iset_item(self, loc: int, value: Series) -> None:
+    def _iset_item(self, loc: int, value: Series, inplace: bool = True) -> None:
         # We are only called from _replace_columnwise which guarantees that
         # no reindex is necessary
-        # TODO(CoW): Optimize to avoid copy here, but have ton track refs
-        self._iset_item_mgr(loc, value._values.copy(), inplace=True)
+        if using_copy_on_write():
+            self._iset_item_mgr(
+                loc, value._values, inplace=inplace, refs=value._references
+            )
+        else:
+            self._iset_item_mgr(loc, value._values.copy(), inplace=True)
 
         # check if we are modifying a copy
         # try to set first as we want an invalid
@@ -5510,7 +5514,7 @@ def _replace_columnwise(
                 target, value = mapping[ax_value]
                 newobj = ser.replace(target, value, regex=regex)
 
-                res._iset_item(i, newobj)
+                res._iset_item(i, newobj, inplace=inplace)
 
         if inplace:
             return

pandas/core/internals/managers.py

@@ -1101,7 +1101,9 @@ def value_getitem(placement):
             if inplace and blk.should_store(value):
                 # Updating inplace -> check if we need to do Copy-on-Write
                 if using_copy_on_write() and not self._has_no_reference_block(blkno_l):
-                    self._iset_split_block(blkno_l, blk_locs, value_getitem(val_locs))
+                    self._iset_split_block(
+                        blkno_l, blk_locs, value_getitem(val_locs), refs=refs
+                    )
                 else:
                     blk.set_inplace(blk_locs, value_getitem(val_locs))
                     continue

pandas/tests/copy_view/test_replace.py

@@ -364,13 +364,23 @@ def test_replace_list_none_inplace_refs(using_copy_on_write):
         assert np.shares_memory(arr, get_array(df, "a"))
 
 
+def test_replace_columnwise_no_op_inplace(using_copy_on_write):
+    df = DataFrame({"a": [1, 2, 3], "b": [1, 2, 3]})
+    view = df[:]
+    df_orig = df.copy()
+    df.replace({"a": 10}, 100, inplace=True)
+    if using_copy_on_write:
+        assert np.shares_memory(get_array(view, "a"), get_array(df, "a"))
+        df.iloc[0, 0] = 100
+        tm.assert_frame_equal(view, df_orig)
+
+
 def test_replace_columnwise_no_op(using_copy_on_write):
     df = DataFrame({"a": [1, 2, 3], "b": [1, 2, 3]})
     df_orig = df.copy()
     df2 = df.replace({"a": 10}, 100)
     if using_copy_on_write:
-        # TODO(CoW): This should share memory
-        assert not np.shares_memory(get_array(df2, "a"), get_array(df, "a"))
+        assert np.shares_memory(get_array(df2, "a"), get_array(df, "a"))
     df2.iloc[0, 0] = 100
     tm.assert_frame_equal(df, df_orig)