[ArrayManager] Array version of putmask logic

pandas/core/internals/array_manager.py

@@ -89,6 +89,7 @@
     new_block,
     to_native_types,
 )
+from pandas.core.internals.methods import putmask_flexible
 
 if TYPE_CHECKING:
     from pandas import Float64Index
@@ -187,13 +188,7 @@ def __repr__(self) -> str:
             output += f"\n{arr.dtype}"
         return output
 
-    def apply(
-        self: T,
-        f,
-        align_keys: list[str] | None = None,
-        ignore_failures: bool = False,
-        **kwargs,
-    ) -> T:
+    def apply(self: T, f, ignore_failures: bool = False, **kwargs) -> T:
         """
         Iterate over the arrays, collect and create a new ArrayManager.
 
@@ -212,32 +207,14 @@ def apply(
         """
         assert "filter" not in kwargs
 
-        align_keys = align_keys or []
         result_arrays: list[np.ndarray] = []
         result_indices: list[int] = []
         # fillna: Series/DataFrame is responsible for making sure value is aligned
 
-        aligned_args = {k: kwargs[k] for k in align_keys}
-
         if f == "apply":
             f = kwargs.pop("func")
 
         for i, arr in enumerate(self.arrays):
-
-            if aligned_args:
-
-                for k, obj in aligned_args.items():
-                    if isinstance(obj, (ABCSeries, ABCDataFrame)):
-                        # The caller is responsible for ensuring that
-                        #  obj.axes[-1].equals(self.items)
-                        if obj.ndim == 1:
-                            kwargs[k] = obj.iloc[i]
-                        else:
-                            kwargs[k] = obj.iloc[:, i]._values
-                    else:
-                        # otherwise we have an array-like
-                        kwargs[k] = obj[i]
-
             try:
                 if callable(f):
                     applied = f(arr, **kwargs)
@@ -352,12 +329,28 @@ def putmask(self, mask, new, align: bool = True):
             align_keys = ["mask"]
             new = extract_array(new, extract_numpy=True)
 
-        return self.apply_with_block(
-            "putmask",
-            align_keys=align_keys,
-            mask=mask,
-            new=new,
-        )
+        kwargs = {"mask": mask, "new": new}
+        aligned_kwargs = {k: kwargs[k] for k in align_keys}
+
+        for i, arr in enumerate(self.arrays):
+
+            for k, obj in aligned_kwargs.items():
+                if isinstance(obj, (ABCSeries, ABCDataFrame)):
+                    # The caller is responsible for ensuring that
+                    #  obj.axes[-1].equals(self.items)
+                    if obj.ndim == 1:
+                        kwargs[k] = obj._values
+                    else:
+                        kwargs[k] = obj.iloc[:, i]._values
+                else:
+                    # otherwise we have an ndarray
+                    if self.ndim == 2:
+                        kwargs[k] = obj[i]
+
+            new = putmask_flexible(arr, **kwargs)
+            self.arrays[i] = new
+
+        return self
 
     def diff(self: T, n: int, axis: int) -> T:
         if axis == 1:

pandas/core/internals/methods.py

@@ -0,0 +1,114 @@
+"""
+Wrappers around array_algos with internals-specific logic
+"""
+from __future__ import annotations
+
+import numpy as np
+
+from pandas.core.dtypes.cast import (
+    can_hold_element,
+    find_common_type,
+    infer_dtype_from,
+)
+from pandas.core.dtypes.common import is_interval_dtype
+from pandas.core.dtypes.generic import (
+    ABCDataFrame,
+    ABCIndex,
+    ABCSeries,
+)
+from pandas.core.dtypes.missing import (
+    is_valid_na_for_dtype,
+    na_value_for_dtype,
+)
+
+from pandas.core.array_algos.putmask import (
+    extract_bool_array,
+    putmask_smart,
+    putmask_without_repeat,
+    setitem_datetimelike_compat,
+    validate_putmask,
+)
+from pandas.core.arrays import ExtensionArray
+from pandas.core.arrays._mixins import NDArrayBackedExtensionArray
+
+
+def putmask_flexible(array: np.ndarray | ExtensionArray, mask, new):
+    """
+    Putmask implementation for ArrayManager.putmask.
+
+    Flexible version that will upcast if needed.
+    """
+    if isinstance(array, np.ndarray):
+        return putmask_flexible_ndarray(array, mask=mask, new=new)
+    else:
+        return putmask_flexible_ea(array, mask=mask, new=new)
+
+
+def putmask_flexible_ndarray(array: np.ndarray, mask, new):
+    """
+    Putmask implementation for ArrayManager putmask for ndarray.
+
+    Flexible version that will upcast if needed.
+    """
+    mask, noop = validate_putmask(array, mask)
+    assert not isinstance(new, (ABCIndex, ABCSeries, ABCDataFrame))
+
+    # if we are passed a scalar None, convert it here
+    if not array.dtype == "object" and is_valid_na_for_dtype(new, array.dtype):
+        new = na_value_for_dtype(array.dtype, compat=False)
+
+    if can_hold_element(array, new):
+        putmask_without_repeat(array, mask, new)
+        return array
+
+    elif noop:
+        return array
+
+    dtype, _ = infer_dtype_from(new)
+    if dtype.kind in ["m", "M"]:
+        array = array.astype(object)
+        # convert to list to avoid numpy coercing datetimelikes to integers
+        new = setitem_datetimelike_compat(array, mask.sum(), new)
+        # putmask_smart below converts it back to array
+        np.putmask(array, mask, new)
+        return array
+
+    new_values = putmask_smart(array, mask, new)
+    return new_values
+
+
+def _coerce_to_target_dtype(array, new):
+    dtype, _ = infer_dtype_from(new, pandas_dtype=True)
+    new_dtype = find_common_type([array.dtype, dtype])
+    return array.astype(new_dtype, copy=False)
+
+
+def putmask_flexible_ea(array: ExtensionArray, mask, new):
+    """
+    Putmask implementation for ArrayManager putmask for EA.
+
+    Flexible version that will upcast if needed.
+    """
+    mask = extract_bool_array(mask)
+
+    if isinstance(array, NDArrayBackedExtensionArray):
+        if not can_hold_element(array, new):
+            array = _coerce_to_target_dtype(array, new)
+            return putmask_flexible(array, mask, new)
+
+    try:
+        array._putmask(mask, new)
+    except TypeError:
+        if not is_interval_dtype(array.dtype):
+            # Discussion about what we want to support in the general
+            #  case GH#39584
+            raise
+
+        array = _coerce_to_target_dtype(array, new)
+        if array.dtype == np.dtype("object"):
+            # For now at least, only support casting e.g.
+            #  Interval[int64]->Interval[float64],
+            raise
+        return putmask_flexible(array, mask, new)
+
+    return array

pandas/tests/internals/test_api.py

@@ -19,6 +19,7 @@ def test_namespace():
         "blocks",
         "concat",
         "managers",
+        "methods",
         "construction",
         "array_manager",
         "base",