Proof of concept for Copy-on-Write implementation

pandas/core/frame.py

@@ -4913,7 +4913,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",
@@ -5747,7 +5747,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",
@@ -3887,6 +3887,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
@@ -5817,7 +5819,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

@@ -1894,6 +1894,16 @@ def _setitem_single_column(self, loc: int, value, plane_indexer):
         """
         pi = plane_indexer
 
+        if not hasattr(self.obj._mgr, "blocks"):
+            # ArrayManager
+            if com.is_null_slice(pi) or com.is_full_slice(pi, len(self.obj)):
+                arr = self.obj._sanitize_column(value)
+                self.obj._mgr.iset(loc, arr)
+            else:
+                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,11 +180,33 @@ 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")
 
     # TODO setstate getstate
 
+    def __reduce__(self):
+        # to avoid pickling the refs
+        return type(self), (self.arrays, self._axes)
+
     def __repr__(self) -> str:
         output = type(self).__name__
         output += f"\nIndex: {self._axes[0]}"
@@ -352,6 +378,7 @@ def where(self: T, other, cond, align: bool, errors: str) -> T:
     #     return self.apply_with_block("setitem", indexer=indexer, value=value)
 
     def putmask(self, mask, new, align: bool = True):
+        # TODO(CoW) check references to copy if needed
         if align:
             align_keys = ["new", "mask"]
         else:
@@ -516,6 +543,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 +560,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 +633,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 +664,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 +735,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]
@@ -728,6 +772,12 @@ def _verify_integrity(self) -> None:
                     "Passed arrays should be 1-dimensional, got array with "
                     f"{arr.ndim} dimensions instead."
                 )
+        if self.refs is not None:
+            if not len(self.refs) == n_columns:
+                raise ValueError(
+                    "Number of passed refs must equal the size of the column Index: "
+                    f"{len(self.refs)} refs vs {n_columns} columns."
+                )
 
     # --------------------------------------------------------------------
     # Indexing
@@ -763,20 +813,27 @@ def get_slice(self, slobj: slice, axis: int = 0) -> ArrayManager:
 
         if axis == 0:
             arrays = [arr[slobj] for arr in self.arrays]
+            # 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]
         elif axis == 1:
             arrays = self.arrays[slobj]
+            # track reference to subset of column arrays
+            refs = [weakref.ref(arr) for arr in arrays]
 
         new_axes = list(self._axes)
         new_axes[axis] = new_axes[axis]._getitem_slice(slobj)
 
-        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 +861,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):
 
@@ -880,6 +938,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
@@ -893,8 +955,28 @@ def idelete(self, indexer):
 
         self.arrays = [self.arrays[i] for i in np.nonzero(to_keep)[0]]
         self._axes = [self._axes[0], self._axes[1][to_keep]]
+        if self.refs is not None:
+            self.refs = [ref for i, ref in enumerate(self.refs) if to_keep[i]]
         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 (potentially) inplace
+            arr = self.arrays[loc]
+            # TODO we should try to avoid this (indexing.py::_setitem_single_column
+            # does a copy for the BM path as well)
+            arr = arr.copy()
+        else:
+            # otherwise perform Copy-on-Write and clear the reference
+            arr = self.arrays[loc].copy()
+            self._clear_reference(loc)
+
+        # create temporary SingleArrayManager without ref to use setitem implementation
+        mgr = SingleArrayManager([arr], [self._axes[0]])
+        new_mgr = mgr.setitem((idx,), value)
+        # update existing ArrayManager in-place
+        self.arrays[loc] = new_mgr.arrays[0]
+
     # --------------------------------------------------------------------
     # Array-wise Operation
 
@@ -1177,10 +1259,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
@@ -1265,6 +1349,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")
 
@@ -1275,7 +1360,9 @@ def get_slice(self, slobj: slice, axis: int = 0) -> SingleArrayManager:
     def getitem_mgr(self, indexer) -> SingleArrayManager:
         new_array = self.array[indexer]
         new_index = self.index[indexer]
-        return type(self)([new_array], [new_index])
+        # TODO in theory only need to track reference if new_array is a view
+        ref = weakref.ref(self.array)
+        return type(self)([new_array], [new_index], [ref])
 
     def apply(self, func, **kwargs):
         if callable(func):
@@ -1284,8 +1371,15 @@ 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:
         """
@@ -1296,6 +1390,8 @@ def idelete(self, indexer) -> SingleArrayManager:
 
         self.arrays = [self.arrays[0][to_keep]]
         self._axes = [self._axes[0][to_keep]]
+        # clear reference since we are backed by new array
+        self.refs = None
         return self
 
     def _get_data_subset(self, predicate: Callable) -> SingleArrayManager:

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):
 
@@ -566,6 +570,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