WIP: ENH: CoW support for replace

Author: lithomas1

pandas/core/frame.py

@@ -5522,7 +5522,7 @@ def _replace_columnwise(
         DataFrame or None
         """
         # Operate column-wise
-        res = self if inplace else self.copy()
+        res = self if inplace else self.copy(deep=None)
         ax = self.columns
 
         for i, ax_value in enumerate(ax):

pandas/core/generic.py

@@ -7045,6 +7045,7 @@ def replace(
                 to_replace = [to_replace]
 
             if isinstance(to_replace, (tuple, list)):
+                # TODO: Consider copy-on-write for non-replaced columns's here
                 if isinstance(self, ABCDataFrame):
                     from pandas import Series
 
@@ -7104,7 +7105,8 @@ def replace(
             if not self.size:
                 if inplace:
                     return None
-                return self.copy()
+                # TODO: Is it even worth doing the CoW here?
+                return self.copy(deep=None)
 
             if is_dict_like(to_replace):
                 if is_dict_like(value):  # {'A' : NA} -> {'A' : 0}

pandas/core/internals/managers.py

@@ -1211,8 +1211,12 @@ def value_getitem(placement):
                 return value[placement.indexer]
 
         # Accessing public blknos ensures the public versions are initialized
+        # print(f"Loc is {loc}")
         blknos = self.blknos[loc]
+        # print(f"Blknos is {blknos}")
         blklocs = self.blklocs[loc].copy()
+        # print(f"Blklocs is {blklocs}")
+        # print(f"Blocks is {self.blocks}")
 
         unfit_mgr_locs = []
         unfit_val_locs = []
@@ -1223,8 +1227,32 @@ 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):
-                    blk.set_inplace(blk_locs, value_getitem(val_locs), copy=True)
-                    self._clear_reference_block(blkno_l)
+
+                    # print("Hit here")
+                    # print(f"Blk locs is {blk_locs}")
+                    # print(f"Blkno_l is {blkno_l}")
+                    prev_refs = self.refs[blkno_l]
+                    leftover_blocks = blk.delete(blk_locs)
+                    num_extra_blocks = len(leftover_blocks)
+                    # Preserve refs for unchanged blocks
+                    extra_refs = [prev_refs] * num_extra_blocks
+                    self.refs = self.refs[:blkno_l] + extra_refs + self.refs[blkno_l:]
+
+                    # TODO: Are we sure we want 2D?
+                    nb = new_block_2d(value, placement=blk._mgr_locs)
+                    old_blocks = self.blocks
+                    new_blocks = (
+                        old_blocks[:blkno_l]
+                        + tuple(leftover_blocks[:blkno_l])
+                        + (nb,)
+                        + tuple(leftover_blocks[blkno_l:])
+                        + old_blocks[blkno_l + num_extra_blocks :]
+                    )
+                    self.blocks = new_blocks
+
+                    self._rebuild_blknos_and_blklocs()
+                    # blk.set_inplace(blk_locs, value_getitem(val_locs), copy=True)
+                    # self._clear_reference_block(blkno_l)
                 else:
                     blk.set_inplace(blk_locs, value_getitem(val_locs))
             else:

pandas/tests/copy_view/test_methods.py

@@ -700,3 +700,33 @@ def test_squeeze(using_copy_on_write):
         # Without CoW the original will be modified
         assert np.shares_memory(series.values, get_array(df, "a"))
         assert df.loc[0, "a"] == 0
+
+
+@pytest.mark.parametrize(
+    "to_replace",
+    [
+        {"a": 1, "b": 4},
+        # Test CoW splits blocks to avoid copying unchanged columns
+        {"a": 1},
+        # TODO: broken, fix by fixing Block.replace()
+        # 1
+    ],
+)
+def test_replace(using_copy_on_write, to_replace):
+    df = DataFrame({"a": [1, 2, 3], "b": [4, 5, 6], "c": ["foo", "bar", "baz"]})
+    df_orig = df.copy()
+
+    df_replaced = df.replace(to_replace=to_replace, value=-1)
+
+    # Should share memory regardless of CoW since squeeze is just an iloc
+    if using_copy_on_write:
+        if (df_replaced["b"] == df["b"]).all():
+            assert np.shares_memory(get_array(df_replaced, "b"), get_array(df, "b"))
+        assert np.shares_memory(get_array(df_replaced, "c"), get_array(df, "c"))
+
+    # mutating squeezed df triggers a copy-on-write for that column/block
+    df_replaced.loc[0, "c"] = -1
+    if using_copy_on_write:
+        assert not np.shares_memory(get_array(df_replaced, "c"), get_array(df, "a"))
+
+    tm.assert_frame_equal(df, df_orig)