ENH: implement EA._where (#44187)

Author: jbrockmendel

pandas/core/arrays/_mixins.py

@@ -320,7 +320,7 @@ def putmask(self, mask: np.ndarray, value) -> None:
 
         np.putmask(self._ndarray, mask, value)
 
-    def where(
+    def _where(
         self: NDArrayBackedExtensionArrayT, mask: np.ndarray, value
     ) -> NDArrayBackedExtensionArrayT:
         """

pandas/core/arrays/base.py

@@ -1411,6 +1411,31 @@ def insert(self: ExtensionArrayT, loc: int, item) -> ExtensionArrayT:
 
         return type(self)._concat_same_type([self[:loc], item_arr, self[loc:]])
 
+    def _where(
+        self: ExtensionArrayT, mask: npt.NDArray[np.bool_], value
+    ) -> ExtensionArrayT:
+        """
+        Analogue to np.where(mask, self, value)
+
+        Parameters
+        ----------
+        mask : np.ndarray[bool]
+        value : scalar or listlike
+
+        Returns
+        -------
+        same type as self
+        """
+        result = self.copy()
+
+        if is_list_like(value):
+            val = value[~mask]
+        else:
+            val = value
+
+        result[~mask] = val
+        return result
+
     @classmethod
     def _empty(cls, shape: Shape, dtype: ExtensionDtype):
         """

pandas/core/arrays/sparse/array.py

@@ -1307,6 +1307,13 @@ def to_dense(self) -> np.ndarray:
 
     _internal_get_values = to_dense
 
+    def _where(self, mask, value):
+        # NB: may not preserve dtype, e.g. result may be Sparse[float64]
+        #  while self is Sparse[int64]
+        naive_implementation = np.where(mask, self, value)
+        result = type(self)._from_sequence(naive_implementation)
+        return result
+
     # ------------------------------------------------------------------------
     # IO
     # ------------------------------------------------------------------------

pandas/core/internals/blocks.py

@@ -50,7 +50,6 @@
     is_extension_array_dtype,
     is_interval_dtype,
     is_list_like,
-    is_sparse,
     is_string_dtype,
 )
 from pandas.core.dtypes.dtypes import (
@@ -1596,30 +1595,9 @@ def where(self, other, cond, errors="raise") -> list[Block]:
             # for the type
             other = self.dtype.na_value
 
-        if is_sparse(self.values):
-            # TODO(SparseArray.__setitem__): remove this if condition
-            # We need to re-infer the type of the data after doing the
-            # where, for cases where the subtypes don't match
-            dtype = None
-        else:
-            dtype = self.dtype
-
-        result = self.values.copy()
-        icond = ~cond
-        if lib.is_scalar(other):
-            set_other = other
-        else:
-            set_other = other[icond]
         try:
-            result[icond] = set_other
-        except (NotImplementedError, TypeError):
-            # NotImplementedError for class not implementing `__setitem__`
-            # TypeError for SparseArray, which implements just to raise
-            # a TypeError
-            if isinstance(result, Categorical):
-                # TODO: don't special-case
-                raise
-
+            result = self.values._where(cond, other)
+        except TypeError:
             if is_interval_dtype(self.dtype):
                 # TestSetitemFloatIntervalWithIntIntervalValues
                 blk = self.coerce_to_target_dtype(other)
@@ -1628,10 +1606,7 @@ def where(self, other, cond, errors="raise") -> list[Block]:
                     #  Interval[int64]->Interval[float64]
                     raise
                 return blk.where(other, cond, errors)
-
-            result = type(self.values)._from_sequence(
-                np.where(cond, self.values, other), dtype=dtype
-            )
+            raise
 
         return [self.make_block_same_class(result)]
 
@@ -1721,7 +1696,7 @@ def where(self, other, cond, errors="raise") -> list[Block]:
         cond = extract_bool_array(cond)
 
         try:
-            res_values = arr.T.where(cond, other).T
+            res_values = arr.T._where(cond, other).T
         except (ValueError, TypeError):
             return Block.where(self, other, cond, errors=errors)
 

pandas/tests/indexes/categorical/test_indexing.py

@@ -325,7 +325,7 @@ def test_where_non_categories(self):
         msg = "Cannot setitem on a Categorical with a new category"
         with pytest.raises(TypeError, match=msg):
             # Test the Categorical method directly
-            ci._data.where(mask, 2)
+            ci._data._where(mask, 2)
 
 
 class TestContains:

pandas/tests/series/indexing/test_where.py

@@ -459,7 +459,7 @@ def test_where_datetimelike_categorical(tz_naive_fixture):
     tm.assert_index_equal(res, dr)
 
     # DatetimeArray.where
-    res = lvals._data.where(mask, rvals)
+    res = lvals._data._where(mask, rvals)
     tm.assert_datetime_array_equal(res, dr._data)
 
     # Series.where