ENH: Add copy-on-write to DataFrame.drop (pandas-dev#49689)

Author: seljaks

pandas/core/generic.py

@@ -4483,6 +4483,7 @@ def _drop_axis(
             indexer,
             axis=bm_axis,
             allow_dups=True,
+            copy=None,
             only_slice=only_slice,
         )
         result = self._constructor(new_mgr)

pandas/core/internals/array_manager.py

@@ -546,7 +546,7 @@ def reindex_indexer(
         axis: AxisInt,
         fill_value=None,
         allow_dups: bool = False,
-        copy: bool = True,
+        copy: bool | None = True,
         # ignored keywords
         only_slice: bool = False,
         # ArrayManager specific keywords
@@ -570,7 +570,7 @@ def _reindex_indexer(
         axis: AxisInt,
         fill_value=None,
         allow_dups: bool = False,
-        copy: bool = True,
+        copy: bool | None = True,
         use_na_proxy: bool = False,
     ) -> T:
         """

pandas/tests/copy_view/test_methods.py

@@ -132,6 +132,25 @@ def test_reindex_columns(using_copy_on_write):
     tm.assert_frame_equal(df, df_orig)
 
 
+def test_drop_on_column(using_copy_on_write):
+    df = DataFrame({"a": [1, 2, 3], "b": [4, 5, 6], "c": [0.1, 0.2, 0.3]})
+    df_orig = df.copy()
+    df2 = df.drop(columns="a")
+    df2._mgr._verify_integrity()
+
+    if using_copy_on_write:
+        assert np.shares_memory(get_array(df2, "b"), get_array(df, "b"))
+        assert np.shares_memory(get_array(df2, "c"), get_array(df, "c"))
+    else:
+        assert not np.shares_memory(get_array(df2, "b"), get_array(df, "b"))
+        assert not np.shares_memory(get_array(df2, "c"), get_array(df, "c"))
+    df2.iloc[0, 0] = 0
+    assert not np.shares_memory(get_array(df2, "b"), get_array(df, "b"))
+    if using_copy_on_write:
+        assert np.shares_memory(get_array(df2, "c"), get_array(df, "c"))
+    tm.assert_frame_equal(df, df_orig)
+
+
 def test_select_dtypes(using_copy_on_write):
     # Case: selecting columns using `select_dtypes()` returns a new dataframe
     # + afterwards modifying the result