CoW: Delay copy when enlarging df with loc (pandas-dev#52062)

Author: phofl

pandas/core/indexing.py

@@ -1979,7 +1979,10 @@ def _setitem_with_indexer_frame_value(self, indexer, value: DataFrame, name: str
                 if item in value:
                     sub_indexer[1] = item
                     val = self._align_series(
-                        tuple(sub_indexer), value[item], multiindex_indexer
+                        tuple(sub_indexer),
+                        value[item],
+                        multiindex_indexer,
+                        using_cow=using_copy_on_write(),
                     )
                 else:
                     val = np.nan
@@ -2185,7 +2188,13 @@ def _ensure_iterable_column_indexer(self, column_indexer):
             ilocs = column_indexer
         return ilocs
 
-    def _align_series(self, indexer, ser: Series, multiindex_indexer: bool = False):
+    def _align_series(
+        self,
+        indexer,
+        ser: Series,
+        multiindex_indexer: bool = False,
+        using_cow: bool = False,
+    ):
         """
         Parameters
         ----------
@@ -2254,6 +2263,8 @@ def ravel(i):
                     else:
                         new_ix = Index(new_ix)
                     if ser.index.equals(new_ix):
+                        if using_cow:
+                            return ser
                         return ser._values.copy()
 
                     return ser.reindex(new_ix)._values

pandas/tests/copy_view/test_indexing.py

@@ -1073,3 +1073,19 @@ def test_series_midx_tuples_slice(using_copy_on_write):
             index=pd.MultiIndex.from_tuples([((1, 2), 3), ((1, 2), 4), ((2, 3), 4)]),
         )
         tm.assert_series_equal(ser, expected)
+
+
+def test_loc_enlarging_with_dataframe(using_copy_on_write):
+    df = DataFrame({"a": [1, 2, 3]})
+    rhs = DataFrame({"b": [1, 2, 3], "c": [4, 5, 6]})
+    rhs_orig = rhs.copy()
+    df.loc[:, ["b", "c"]] = rhs
+    if using_copy_on_write:
+        assert np.shares_memory(get_array(df, "b"), get_array(rhs, "b"))
+        assert np.shares_memory(get_array(df, "c"), get_array(rhs, "c"))
+        assert not df._mgr._has_no_reference(1)
+    else:
+        assert not np.shares_memory(get_array(df, "b"), get_array(rhs, "b"))
+
+    df.iloc[0, 1] = 100
+    tm.assert_frame_equal(rhs, rhs_orig)