ENH: Add CoW optimization for fillna

Author: phofl

doc/source/whatsnew/v2.0.0.rst

@@ -223,6 +223,7 @@ Copy-on-Write improvements
   - :meth:`DataFrame.to_period` / :meth:`Series.to_period`
   - :meth:`DataFrame.truncate`
   - :meth:`DataFrame.tz_convert` / :meth:`Series.tz_localize`
+  - :meth:`DataFrame.fillna` / :meth:`Series.fillna`
   - :meth:`DataFrame.infer_objects` / :meth:`Series.infer_objects`
   - :meth:`DataFrame.astype` / :meth:`Series.astype`
   - :func:`concat`

pandas/core/internals/blocks.py

@@ -421,7 +421,9 @@ def coerce_to_target_dtype(self, other) -> Block:
         return self.astype(new_dtype, copy=False)
 
     @final
-    def _maybe_downcast(self, blocks: list[Block], downcast=None) -> list[Block]:
+    def _maybe_downcast(
+        self, blocks: list[Block], downcast=None, using_cow: bool = False
+    ) -> list[Block]:
         if downcast is False:
             return blocks
 
@@ -431,23 +433,30 @@ def _maybe_downcast(self, blocks: list[Block], downcast=None) -> list[Block]:
             #  but ATM it breaks too much existing code.
             # split and convert the blocks
 
-            return extend_blocks([blk.convert() for blk in blocks])
+            return extend_blocks(
+                [blk.convert(copy=not using_cow, using_cow=using_cow) for blk in blocks]
+            )
 
         if downcast is None:
             return blocks
 
-        return extend_blocks([b._downcast_2d(downcast) for b in blocks])
+        return extend_blocks(
+            [b._downcast_2d(downcast, using_cow=using_cow) for b in blocks]
+        )
 
     @final
     @maybe_split
-    def _downcast_2d(self, dtype) -> list[Block]:
+    def _downcast_2d(self, dtype, using_cow: bool = False) -> list[Block]:
         """
         downcast specialized to 2D case post-validation.
 
         Refactored to allow use of maybe_split.
         """
         new_values = maybe_downcast_to_dtype(self.values, dtype=dtype)
-        return [self.make_block(new_values)]
+        refs = None
+        if using_cow and new_values is self.values:
+            refs = self.refs
+        return [self.make_block(new_values, refs=refs)]
 
     def convert(
         self,
@@ -1152,7 +1161,12 @@ def where(self, other, cond, _downcast: str | bool = "infer") -> list[Block]:
         return [self.make_block(result)]
 
     def fillna(
-        self, value, limit: int | None = None, inplace: bool = False, downcast=None
+        self,
+        value,
+        limit: int | None = None,
+        inplace: bool = False,
+        downcast=None,
+        using_cow: bool = False,
     ) -> list[Block]:
         """
         fillna on the block with the value. If we fail, then convert to
@@ -1171,19 +1185,25 @@ def fillna(
         if noop:
             # we can't process the value, but nothing to do
             if inplace:
+                if using_cow:
+                    return [self.copy(deep=False)]
                 # Arbitrarily imposing the convention that we ignore downcast
                 #  on no-op when inplace=True
                 return [self]
             else:
                 # GH#45423 consistent downcasting on no-ops.
-                nb = self.copy()
-                nbs = nb._maybe_downcast([nb], downcast=downcast)
+                nb = self.copy(deep=not using_cow)
+                nbs = nb._maybe_downcast([nb], downcast=downcast, using_cow=using_cow)
                 return nbs
 
         if limit is not None:
             mask[mask.cumsum(self.ndim - 1) > limit] = False
 
         if inplace:
+            if using_cow and self.refs.has_reference():
+                # TODO(CoW): If using_cow is implemented for putmask we can defer
+                # the copy
+                self = self.copy()
             nbs = self.putmask(mask.T, value)
         else:
             # without _downcast, we would break
@@ -1194,7 +1214,10 @@ def fillna(
         #  makes a difference bc blk may have object dtype, which has
         #  different behavior in _maybe_downcast.
         return extend_blocks(
-            [blk._maybe_downcast([blk], downcast=downcast) for blk in nbs]
+            [
+                blk._maybe_downcast([blk], downcast=downcast, using_cow=using_cow)
+                for blk in nbs
+            ]
         )
 
     def interpolate(
@@ -1662,12 +1685,21 @@ class ExtensionBlock(libinternals.Block, EABackedBlock):
     values: ExtensionArray
 
     def fillna(
-        self, value, limit: int | None = None, inplace: bool = False, downcast=None
+        self,
+        value,
+        limit: int | None = None,
+        inplace: bool = False,
+        downcast=None,
+        using_cow: bool = False,
     ) -> list[Block]:
         if is_interval_dtype(self.dtype):
             # Block.fillna handles coercion (test_fillna_interval)
             return super().fillna(
-                value=value, limit=limit, inplace=inplace, downcast=downcast
+                value=value,
+                limit=limit,
+                inplace=inplace,
+                downcast=downcast,
+                using_cow=using_cow,
             )
         new_values = self.values.fillna(value=value, method=None, limit=limit)
         nb = self.make_block_same_class(new_values)

pandas/core/internals/managers.py

@@ -410,15 +410,14 @@ def fillna(self: T, value, limit, inplace: bool, downcast) -> T:
         if limit is not None:
             # Do this validation even if we go through one of the no-op paths
             limit = libalgos.validate_limit(None, limit=limit)
-        if inplace:
-            # TODO(CoW) can be optimized to only copy those blocks that have refs
-            if using_copy_on_write() and any(
-                not self._has_no_reference_block(i) for i in range(len(self.blocks))
-            ):
-                self = self.copy()
 
         return self.apply(
-            "fillna", value=value, limit=limit, inplace=inplace, downcast=downcast
+            "fillna",
+            value=value,
+            limit=limit,
+            inplace=inplace,
+            downcast=downcast,
+            using_cow=using_copy_on_write(),
         )
 
     def astype(self: T, dtype, copy: bool | None = False, errors: str = "raise") -> T:

pandas/tests/copy_view/test_interp_fillna.py

@@ -0,0 +1,79 @@
+import numpy as np
+import pytest
+
+from pandas import (
+    DataFrame,
+    Interval,
+    Series,
+    interval_range,
+)
+import pandas._testing as tm
+from pandas.tests.copy_view.util import get_array
+
+
+def test_fillna(using_copy_on_write):
+    df = DataFrame({"a": [1.5, np.nan], "b": 1})
+    df_orig = df.copy()
+
+    df2 = df.fillna(5.5)
+    if using_copy_on_write:
+        assert np.shares_memory(get_array(df, "b"), get_array(df2, "b"))
+    else:
+        assert not np.shares_memory(get_array(df, "b"), get_array(df2, "b"))
+
+    df2.iloc[0, 1] = 100
+    tm.assert_frame_equal(df_orig, df)
+
+
+@pytest.mark.parametrize("downcast", [None, False])
+def test_fillna_inplace(using_copy_on_write, downcast):
+    df = DataFrame({"a": [1.5, np.nan], "b": 1})
+    arr_a = get_array(df, "a")
+    arr_b = get_array(df, "b")
+
+    df.fillna(5.5, inplace=True, downcast=downcast)
+    assert np.shares_memory(get_array(df, "a"), arr_a)
+    assert np.shares_memory(get_array(df, "b"), arr_b)
+    if using_copy_on_write:
+        assert df._mgr._has_no_reference(0)
+        assert df._mgr._has_no_reference(1)
+
+
+def test_fillna_inplace_reference(using_copy_on_write):
+    df = DataFrame({"a": [1.5, np.nan], "b": 1})
+    df_orig = df.copy()
+    arr_a = get_array(df, "a")
+    arr_b = get_array(df, "b")
+    view = df[:]
+
+    df.fillna(5.5, inplace=True)
+    if using_copy_on_write:
+        assert not np.shares_memory(get_array(df, "a"), arr_a)
+        assert np.shares_memory(get_array(df, "b"), arr_b)
+        assert view._mgr._has_no_reference(0)
+        assert df._mgr._has_no_reference(0)
+        tm.assert_frame_equal(view, df_orig)
+    else:
+        assert np.shares_memory(get_array(df, "a"), arr_a)
+        assert np.shares_memory(get_array(df, "b"), arr_b)
+    expected = DataFrame({"a": [1.5, 5.5], "b": 1})
+    tm.assert_frame_equal(df, expected)
+
+
+def test_fillna_interval_inplace_reference(using_copy_on_write):
+    ser = Series(interval_range(start=0, end=5), name="a")
+    ser.iloc[1] = np.nan
+
+    ser_orig = ser.copy()
+    view = ser[:]
+    ser.fillna(value=Interval(left=0, right=5), inplace=True)
+
+    if using_copy_on_write:
+        assert not np.shares_memory(
+            get_array(ser, "a").left.values, get_array(view, "a").left.values
+        )
+        tm.assert_series_equal(view, ser_orig)
+    else:
+        assert np.shares_memory(
+            get_array(ser, "a").left.values, get_array(view, "a").left.values
+        )