CoW: Delay copy when setting Series or Frame with isetitem (pandas-dev#52051)

Author: phofl

pandas/core/frame.py

@@ -3913,6 +3913,7 @@ def isetitem(self, loc, value) -> None:
         In cases where ``frame.columns`` is unique, this is equivalent to
         ``frame[frame.columns[i]] = value``.
         """
+        using_cow = using_copy_on_write()
         if isinstance(value, DataFrame):
             if is_integer(loc):
                 loc = [loc]
@@ -3924,12 +3925,14 @@ def isetitem(self, loc, value) -> None:
                 )
 
             for i, idx in enumerate(loc):
-                arraylike, _ = self._sanitize_column(value.iloc[:, i])
-                self._iset_item_mgr(idx, arraylike, inplace=False)
+                arraylike, refs = self._sanitize_column(
+                    value.iloc[:, i], using_cow=using_cow
+                )
+                self._iset_item_mgr(idx, arraylike, inplace=False, refs=refs)
             return
 
-        arraylike, _ = self._sanitize_column(value)
-        self._iset_item_mgr(loc, arraylike, inplace=False)
+        arraylike, refs = self._sanitize_column(value, using_cow=using_cow)
+        self._iset_item_mgr(loc, arraylike, inplace=False, refs=refs)
 
     def __setitem__(self, key, value):
         if not PYPY and using_copy_on_write():

pandas/tests/copy_view/test_methods.py

@@ -1521,8 +1521,8 @@ def test_isetitem_series(using_copy_on_write, dtype):
     df.isetitem(0, ser)
 
     if using_copy_on_write:
-        # TODO(CoW) this can share memory
-        assert not np.shares_memory(get_array(df, "a"), get_array(ser))
+        assert np.shares_memory(get_array(df, "a"), get_array(ser))
+        assert not df._mgr._has_no_reference(0)
 
     # mutating dataframe doesn't update series
     df.loc[0, "a"] = 0
@@ -1538,6 +1538,23 @@ def test_isetitem_series(using_copy_on_write, dtype):
     tm.assert_frame_equal(df, expected)
 
 
+def test_isetitem_frame(using_copy_on_write):
+    df = DataFrame({"a": [1, 2, 3], "b": 1, "c": 2})
+    rhs = DataFrame({"a": [4, 5, 6], "b": 2})
+    df.isetitem([0, 1], rhs)
+    if using_copy_on_write:
+        assert np.shares_memory(get_array(df, "a"), get_array(rhs, "a"))
+        assert np.shares_memory(get_array(df, "b"), get_array(rhs, "b"))
+        assert not df._mgr._has_no_reference(0)
+    else:
+        assert not np.shares_memory(get_array(df, "a"), get_array(rhs, "a"))
+        assert not np.shares_memory(get_array(df, "b"), get_array(rhs, "b"))
+    expected = df.copy()
+    rhs.iloc[0, 0] = 100
+    rhs.iloc[0, 1] = 100
+    tm.assert_frame_equal(df, expected)
+
+
 @pytest.mark.parametrize("key", ["a", ["a"]])
 def test_get(using_copy_on_write, key):
     df = DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]})