Proof of concept for Copy-on-Write implementation

pandas/core/frame.py

@@ -4930,7 +4930,7 @@ def rename(
         index: Renamer | None = None,
         columns: Renamer | None = None,
         axis: Axis | None = None,
-        copy: bool = True,
+        copy: bool | None = None,
         inplace: bool = False,
         level: Level | None = None,
         errors: str = "ignore",
@@ -5764,7 +5764,7 @@ class    max    type
         if inplace:
             new_obj = self
         else:
-            new_obj = self.copy()
+            new_obj = self.copy(deep=None)
 
         new_index = ibase.default_index(len(new_obj))
         if level is not None:

pandas/core/generic.py

@@ -992,7 +992,7 @@ def rename(
         index: Renamer | None = None,
         columns: Renamer | None = None,
         axis: Axis | None = None,
-        copy: bool_t = True,
+        copy: bool_t | None = None,
         inplace: bool_t = False,
         level: Level | None = None,
         errors: str = "ignore",
@@ -3911,6 +3911,8 @@ def _check_setitem_copy(self, stacklevel=4, t="setting", force=False):
         df.iloc[0:5]['group'] = 'a'
 
         """
+        if isinstance(self._mgr, (ArrayManager, SingleArrayManager)):
+            return
         # return early if the check is not needed
         if not (force or self._is_copy):
             return
@@ -5833,7 +5835,7 @@ def astype(
         return result
 
     @final
-    def copy(self: FrameOrSeries, deep: bool_t = True) -> FrameOrSeries:
+    def copy(self: FrameOrSeries, deep: bool_t | None = True) -> FrameOrSeries:
         """
         Make a copy of this object's indices and data.
 

pandas/core/indexing.py

@@ -1869,6 +1869,12 @@ def _setitem_single_column(self, loc: int, value, plane_indexer):
         """
         pi = plane_indexer
 
+        if not getattr(self.obj._mgr, "blocks", False):
+            # ArrayManager
+            self.obj._mgr.column_setitem(loc, plane_indexer, value)
+            self.obj._clear_item_cache()
+            return
+
         ser = self.obj._ixs(loc, axis=1)
 
         # perform the equivalent of a setitem on the info axis

pandas/core/internals/array_manager.py

@@ -9,6 +9,7 @@
     Callable,
     TypeVar,
 )
+import weakref
 
 import numpy as np
 
@@ -110,6 +111,7 @@ class BaseArrayManager(DataManager):
     ----------
     arrays : Sequence of arrays
     axes : Sequence of Index
+    refs : Sequence of weakrefs or None, optional
     verify_integrity : bool, default True
 
     """
@@ -121,11 +123,13 @@ class BaseArrayManager(DataManager):
 
     arrays: list[np.ndarray | ExtensionArray]
     _axes: list[Index]
+    refs: list[weakref.ref | None] | None
 
     def __init__(
         self,
         arrays: list[np.ndarray | ExtensionArray],
         axes: list[Index],
+        refs: list[weakref.ref | None] | None,
         verify_integrity: bool = True,
     ):
         raise NotImplementedError
@@ -176,6 +180,24 @@ def is_consolidated(self) -> bool:
     def _consolidate_inplace(self) -> None:
         pass
 
+    def _has_no_reference(self, i: int):
+        """
+        Check for column `i` if has references.
+        (whether it references another array or is itself being referenced)
+
+        Returns True if the columns has no references.
+        """
+        return (self.refs is None or self.refs[i] is None) and weakref.getweakrefcount(
+            self.arrays[i]
+        ) == 0
+
+    def _clear_reference(self, i: int):
+        """
+        Clear any reference for column `i`.
+        """
+        if self.refs is not None:
+            self.refs[i] = None
+
     def get_dtypes(self):
         return np.array([arr.dtype for arr in self.arrays], dtype="object")
 
@@ -516,6 +538,10 @@ def copy(self: T, deep=True) -> T:
         -------
         BlockManager
         """
+        if deep is None:
+            # use shallow copy
+            deep = False
+
         # this preserves the notion of view copying of axes
         if deep:
             # hit in e.g. tests.io.json.test_pandas
@@ -529,9 +555,12 @@ def copy_func(ax):
 
         if deep:
             new_arrays = [arr.copy() for arr in self.arrays]
+            refs = None
         else:
-            new_arrays = self.arrays
-        return type(self)(new_arrays, new_axes)
+            new_arrays = list(self.arrays)
+            refs = [weakref.ref(arr) for arr in self.arrays]
+
+        return type(self)(new_arrays, new_axes, refs, verify_integrity=False)
 
     def reindex_indexer(
         self: T,
@@ -599,14 +628,19 @@ def _reindex_indexer(
 
         if axis == 1:
             new_arrays = []
+            refs = []
             for i in indexer:
                 if i == -1:
                     arr = self._make_na_array(
                         fill_value=fill_value, use_na_proxy=use_na_proxy
                     )
+                    ref = None
                 else:
+                    # reusing full column array -> track with reference
                     arr = self.arrays[i]
+                    ref = weakref.ref(arr)
                 new_arrays.append(arr)
+                refs.append(ref)
 
         else:
             validate_indices(indexer, len(self._axes[0]))
@@ -625,11 +659,14 @@ def _reindex_indexer(
                 )
                 for arr in self.arrays
             ]
+            # selecting rows with take always creates a copy -> no need to
+            # track references to original arrays
+            refs = None
 
         new_axes = list(self._axes)
         new_axes[axis] = new_axis
 
-        return type(self)(new_arrays, new_axes, verify_integrity=False)
+        return type(self)(new_arrays, new_axes, refs, verify_integrity=False)
 
     def take(self: T, indexer, axis: int = 1, verify: bool = True) -> T:
         """
@@ -693,12 +730,14 @@ def __init__(
         self,
         arrays: list[np.ndarray | ExtensionArray],
         axes: list[Index],
+        refs: list[weakref.ref | None] | None = None,
         verify_integrity: bool = True,
     ):
         # Note: we are storing the axes in "_axes" in the (row, columns) order
         # which contrasts the order how it is stored in BlockManager
         self._axes = axes
         self.arrays = arrays
+        self.refs = refs
 
         if verify_integrity:
             self._axes = [ensure_index(ax) for ax in axes]
@@ -768,15 +807,20 @@ def get_slice(self, slobj: slice, axis: int = 0) -> ArrayManager:
 
         new_axes = list(self._axes)
         new_axes[axis] = new_axes[axis]._getitem_slice(slobj)
+        # slicing results in views -> track references to original arrays
+        # TODO possible to optimizate this with single ref to the full ArrayManager?
+        refs = [weakref.ref(arr) for arr in self.arrays]
 
-        return type(self)(arrays, new_axes, verify_integrity=False)
+        return type(self)(arrays, new_axes, refs, verify_integrity=False)
 
     def iget(self, i: int) -> SingleArrayManager:
         """
         Return the data as a SingleArrayManager.
         """
         values = self.arrays[i]
-        return SingleArrayManager([values], [self._axes[0]])
+        # getting single column array for Series -> track reference to original
+        ref = weakref.ref(values)
+        return SingleArrayManager([values], [self._axes[0]], [ref])
 
     def iget_values(self, i: int) -> ArrayLike:
         """
@@ -804,6 +848,7 @@ def iset(self, loc: int | slice | np.ndarray, value: ArrayLike):
             Positional location (already bounds checked)
         value : np.ndarray or ExtensionArray
         """
+        # TODO clear reference for item that is being overwritten
         # single column -> single integer index
         if lib.is_integer(loc):
 
@@ -882,6 +927,10 @@ def insert(self, loc: int, item: Hashable, value: ArrayLike) -> None:
         # TODO is this copy needed?
         arrays = self.arrays.copy()
         arrays.insert(loc, value)
+        if self.refs is not None:
+            # inserted `value` is already a copy, no need to track reference
+            # TODO can we use CoW here as well?
+            self.refs.insert(loc, None)
 
         self.arrays = arrays
         self._axes[1] = new_axis
@@ -890,13 +939,25 @@ def idelete(self, indexer):
         """
         Delete selected locations in-place (new block and array, same BlockManager)
         """
+        # TODO update refs
         to_keep = np.ones(self.shape[0], dtype=np.bool_)
         to_keep[indexer] = False
 
         self.arrays = [self.arrays[i] for i in np.nonzero(to_keep)[0]]
         self._axes = [self._axes[0], self._axes[1][to_keep]]
         return self
 
+    def column_setitem(self, loc: int, idx: int | slice | np.ndarray, value):
+        if self._has_no_reference(loc):
+            # if no reference -> set array inplace
+            self.arrays[loc][idx] = value
+        else:
+            # otherwise perform Copy-on-Write and clear the reference
+            arr = self.arrays[loc].copy()
+            arr[idx] = value
+            self.arrays[loc] = arr
+            self._clear_reference(loc)
+
     # --------------------------------------------------------------------
     # Array-wise Operation
 
@@ -1179,10 +1240,12 @@ def __init__(
         self,
         arrays: list[np.ndarray | ExtensionArray],
         axes: list[Index],
+        refs: list[weakref.ref | None] | None = None,
         verify_integrity: bool = True,
     ):
         self._axes = axes
         self.arrays = arrays
+        self.refs = refs
 
         if verify_integrity:
             assert len(axes) == 1
@@ -1267,6 +1330,7 @@ def fast_xs(self, loc: int) -> ArrayLike:
         raise NotImplementedError("Use series._values[loc] instead")
 
     def get_slice(self, slobj: slice, axis: int = 0) -> SingleArrayManager:
+        # TODO track reference
         if axis >= self.ndim:
             raise IndexError("Requested axis not found in manager")
 
@@ -1286,13 +1350,21 @@ def apply(self, func, **kwargs):
             new_array = getattr(self.array, func)(**kwargs)
         return type(self)([new_array], self._axes)
 
-    def setitem(self, indexer, value):
-        return self.apply_with_block("setitem", indexer=indexer, value=value)
+    def setitem(self, indexer, value, inplace=False):
+        if not self._has_no_reference(0):
+            # if being referenced -> perform Copy-on-Write and clear the reference
+            self.arrays[0] = self.arrays[0].copy()
+            self._clear_reference(0)
+        if inplace:
+            self.arrays[0][indexer] = value
+        else:
+            return self.apply_with_block("setitem", indexer=indexer, value=value)
 
     def idelete(self, indexer) -> SingleArrayManager:
         """
         Delete selected locations in-place (new array, same ArrayManager)
         """
+        # TODO clear reference
         to_keep = np.ones(self.shape[0], dtype=np.bool_)
         to_keep[indexer] = False
 

pandas/core/internals/managers.py

@@ -351,8 +351,12 @@ def where(self: T, other, cond, align: bool, errors: str) -> T:
             errors=errors,
         )
 
-    def setitem(self: T, indexer, value) -> T:
-        return self.apply("setitem", indexer=indexer, value=value)
+    def setitem(self: T, indexer, value, inplace=False) -> T:
+        if inplace:
+            assert self.ndim == 1
+            self._block.values[indexer] = value
+        else:
+            return self.apply("setitem", indexer=indexer, value=value)
 
     def putmask(self, mask, new, align: bool = True):
 
@@ -562,6 +566,9 @@ def copy(self: T, deep=True) -> T:
         -------
         BlockManager
         """
+        if deep is None:
+            # preserve deep copy for BlockManager with copy=None
+            deep = True
         # this preserves the notion of view copying of axes
         if deep:
             # hit in e.g. tests.io.json.test_pandas

pandas/core/series.py

@@ -1064,10 +1064,9 @@ def __setitem__(self, key, value) -> None:
         try:
             self._set_with_engine(key, value)
         except (KeyError, ValueError):
-            values = self._values
             if is_integer(key) and self.index.inferred_type != "integer":
                 # positional setter
-                values[key] = value
+                self._mgr.setitem(key, value, inplace=True)
             else:
                 # GH#12862 adding a new key to the Series
                 self.loc[key] = value
@@ -1099,7 +1098,7 @@ def _set_with_engine(self, key, value) -> None:
         # error: Argument 1 to "validate_numeric_casting" has incompatible type
         # "Union[dtype, ExtensionDtype]"; expected "dtype"
         validate_numeric_casting(self.dtype, value)  # type: ignore[arg-type]
-        self._values[loc] = value
+        self._mgr.setitem(loc, value, inplace=True)
 
     def _set_with(self, key, value):
         # other: fancy integer or otherwise