Fix bugs, should work for 1-D

Author: lithomas1

pandas/core/internals/managers.py

@@ -1211,12 +1211,8 @@ 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 = []
@@ -1227,30 +1223,38 @@ 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):
-
-                    # 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)
+                    leftover_blocks = tuple(blk.delete(blk_locs))
                     # Preserve refs for unchanged blocks
-                    extra_refs = [prev_refs] * num_extra_blocks
-                    self.refs = self.refs[:blkno_l] + extra_refs + self.refs[blkno_l:]
+                    # Add new block where old block was and remaining blocks at
+                    # the end to avoid updating all block numbers
+                    extra_refs = [prev_refs] * len(leftover_blocks)
+                    self.refs = (
+                        self.refs[:blkno_l]
+                        + [None]
+                        + self.refs[blkno_l + 1 :]
+                        + extra_refs
+                    )
 
                     # TODO: Are we sure we want 2D?
-                    nb = new_block_2d(value, placement=blk._mgr_locs)
+                    # Also generalize this to bigger than 1-D locs
+                    nb = new_block_2d(
+                        value,
+                        placement=BlockPlacement(slice(blk_locs[0], blk_locs[0] + 1)),
+                    )
                     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 :]
+                        + old_blocks[blkno_l + 1 :]
+                        + leftover_blocks
                     )
+                    self._blklocs[nb.mgr_locs.indexer] = np.arange(len(nb))
+                    for i, nb in enumerate(leftover_blocks):
+                        self._blklocs[nb.mgr_locs.indexer] = np.arange(len(nb))
+                        self._blknos[nb.mgr_locs.indexer] = i + len(old_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:

pandas/tests/copy_view/test_methods.py

@@ -703,22 +703,24 @@ def test_squeeze(using_copy_on_write):
 
 
 @pytest.mark.parametrize(
-    "to_replace",
+    "replace_kwargs",
     [
-        {"a": 1, "b": 4},
+        {"to_replace": {"a": 1, "b": 4}, "value": -1},
         # Test CoW splits blocks to avoid copying unchanged columns
-        {"a": 1},
+        {"to_replace": {"a": 1}, "value": -1},
+        {"to_replace": {"b": 4}, "value": -1},
+        # TODO: broken, fix by fixing Block.replace_list()
+        # {"to_replace": {"b": 1}},
         # TODO: broken, fix by fixing Block.replace()
         # 1
     ],
 )
-def test_replace(using_copy_on_write, to_replace):
+def test_replace(using_copy_on_write, replace_kwargs):
     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)
+    df_replaced = df.replace(**replace_kwargs)
 
-    # 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"))