Setitem-based where

Author: TomAugspurger

doc/source/whatsnew/v0.24.0.rst

@@ -1310,6 +1310,7 @@ Datetimelike
 - Bug in :class:`DatetimeIndex` where calling ``np.array(dtindex, dtype=object)`` would incorrectly return an array of ``long`` objects (:issue:`23524`)
 - Bug in :class:`Index` where passing a timezone-aware :class:`DatetimeIndex` and `dtype=object` would incorrectly raise a ``ValueError`` (:issue:`23524`)
 - Bug in :class:`Index` where calling ``np.array(dtindex, dtype=object)`` on a timezone-naive :class:`DatetimeIndex` would return an array of ``datetime`` objects instead of :class:`Timestamp` objects, potentially losing nanosecond portions of the timestamps (:issue:`23524`)
+- Bug in :class:`Categorical.__setitem__` not allowing setting with another ``Categorical`` when both are undordered and have the same categories, but in a different order (:issue:`24142`)
 
 Timedelta
 ^^^^^^^^^

pandas/core/arrays/base.py

@@ -662,42 +662,6 @@ def take(self, indices, allow_fill=False, fill_value=None):
         # pandas.api.extensions.take
         raise AbstractMethodError(self)
 
-    def where(self, cond, other):
-        """
-        Replace values where the condition is False.
-
-        .. versionadded :: 0.24.0
-
-        Parameters
-        ----------
-        cond : ndarray or ExtensionArray
-            The mask indicating which values should be kept (True)
-            or replaced from `other` (False).
-
-        other : ndarray, ExtensionArray, or scalar
-            Entries where `cond` is False are replaced with
-            corresponding value from `other`.
-
-        Notes
-        -----
-        Note that `cond` and `other` *cannot* be a Series, Index, or callable.
-        When used from, e.g., :meth:`Series.where`, pandas will unbox
-        Series and Indexes, and will apply callables before they arrive here.
-
-        Returns
-        -------
-        ExtensionArray
-            Same dtype as the original.
-
-        See Also
-        --------
-        Series.where : Similar method for Series.
-        DataFrame.where : Similar method for DataFrame.
-        """
-        return type(self)._from_sequence(np.where(cond, self, other),
-                                         dtype=self.dtype,
-                                         copy=False)
-
     def copy(self, deep=False):
         # type: (bool) -> ExtensionArray
         """

pandas/core/arrays/categorical.py

@@ -1906,49 +1906,6 @@ def take_nd(self, indexer, allow_fill=None, fill_value=None):
 
     take = take_nd
 
-    def where(self, cond, other):
-        # n.b. this now preserves the type
-        codes = self._codes
-        object_msg = (
-            "Implicitly converting categorical to object-dtype ndarray. "
-            "The values `{}' are not present in this categorical's "
-            "categories. A future version of pandas will raise a ValueError "
-            "when 'other' contains different categories.\n\n"
-            "To preserve the current behavior, add the new categories to "
-            "the categorical before calling 'where', or convert the "
-            "categorical to a different dtype."
-        )
-
-        if is_scalar(other) and isna(other):
-            other = -1
-        elif is_scalar(other):
-            item = self.categories.get_indexer([other]).item()
-
-            if item == -1:
-                # note: when removing this, also remove CategoricalBlock.where
-                warn(object_msg.format(other), FutureWarning, stacklevel=2)
-                return np.where(cond, self, other)
-
-            other = item
-
-        elif is_categorical_dtype(other):
-            if not is_dtype_equal(self, other):
-                extra = list(other.categories.difference(self.categories))
-                warn(object_msg.format(compat.reprlib.repr(extra)),
-                     FutureWarning,
-                     stacklevel=2)
-                return np.where(cond, self, other)
-            other = _get_codes_for_values(other, self.categories)
-            # get the codes from other that match our categories
-            pass
-        else:
-            other = np.where(isna(other), -1, other)
-
-        new_codes = np.where(cond, codes, other)
-        return type(self).from_codes(new_codes,
-                                     categories=self.categories,
-                                     ordered=self.ordered)
-
     def _slice(self, slicer):
         """
         Return a slice of myself.
@@ -2121,11 +2078,21 @@ def __setitem__(self, key, value):
             `Categorical` does not have the same categories
         """
 
+        if isinstance(value, (ABCIndexClass, ABCSeries)):
+            value = value.array
+
         # require identical categories set
         if isinstance(value, Categorical):
-            if not value.categories.equals(self.categories):
+            if not is_dtype_equal(self, value):
                 raise ValueError("Cannot set a Categorical with another, "
                                  "without identical categories")
+            if not self.categories.equals(value.categories):
+                new_codes = _recode_for_categories(
+                    value.codes, value.categories, self.categories
+                )
+                value = Categorical.from_codes(new_codes,
+                                               categories=self.categories,
+                                               ordered=self.ordered)
 
         rvalue = value if is_list_like(value) else [value]
 

pandas/core/arrays/interval.py

@@ -777,17 +777,6 @@ def take(self, indices, allow_fill=False, fill_value=None, axis=None,
 
         return self._shallow_copy(left_take, right_take)
 
-    def where(self, cond, other):
-        if is_scalar(other) and isna(other):
-            lother = rother = other
-        else:
-            self._check_closed_matches(other, name='other')
-            lother = other.left
-            rother = other.right
-        left = np.where(cond, self.left, lother)
-        right = np.where(cond, self.right, rother)
-        return self._shallow_copy(left, right)
-
     def value_counts(self, dropna=True):
         """
         Returns a Series containing counts of each interval.

pandas/core/arrays/period.py

@@ -4,7 +4,6 @@
 
 import numpy as np
 
-from pandas._libs import lib
 from pandas._libs.tslibs import NaT, iNaT, period as libperiod
 from pandas._libs.tslibs.fields import isleapyear_arr
 from pandas._libs.tslibs.period import (
@@ -347,22 +346,6 @@ def to_timestamp(self, freq=None, how='start'):
     # --------------------------------------------------------------------
     # Array-like / EA-Interface Methods
 
-    def where(self, cond, other):
-        # TODO(DatetimeArray): move to DatetimeLikeArrayMixin
-        # n.b. _ndarray_values candidate.
-        i8 = self.asi8
-        if lib.is_scalar(other):
-            if isna(other):
-                other = iNaT
-            elif isinstance(other, Period):
-                self._check_compatible_with(other)
-                other = other.ordinal
-        elif isinstance(other, type(self)):
-            self._check_compatible_with(other)
-            other = other.asi8
-        result = np.where(cond, i8, other)
-        return type(self)._simple_new(result, dtype=self.dtype)
-
     def _formatter(self, boxed=False):
         if boxed:
             return str

pandas/core/arrays/sparse.py

@@ -704,11 +704,6 @@ def __array__(self, dtype=None, copy=True):
         out[self.sp_index.to_int_index().indices] = self.sp_values
         return out
 
-    def __setitem__(self, key, value):
-        # I suppose we could allow setting of non-fill_value elements.
-        msg = "SparseArray does not support item assignment via setitem"
-        raise TypeError(msg)
-
     @classmethod
     def _from_sequence(cls, scalars, dtype=None, copy=False):
         return cls(scalars, dtype=dtype)
@@ -1063,20 +1058,6 @@ def take(self, indices, allow_fill=False, fill_value=None):
         return type(self)(result, fill_value=self.fill_value, kind=self.kind,
                           **kwargs)
 
-    def where(self, cond, other):
-        if is_scalar(other):
-            result_dtype = np.result_type(self.dtype.subtype, other)
-        elif isinstance(other, type(self)):
-            result_dtype = np.result_type(self.dtype.subtype,
-                                          other.dtype.subtype)
-        else:
-            result_dtype = np.result_type(self.dtype.subtype, other.dtype)
-
-        dtype = self.dtype.update_dtype(result_dtype)
-        # TODO: avoid converting to dense.
-        values = np.where(cond, self, other)
-        return type(self)(values, dtype=dtype)
-
     def _take_with_fill(self, indices, fill_value=None):
         if fill_value is None:
             fill_value = self.dtype.na_value

pandas/core/internals/blocks.py

@@ -1970,6 +1970,7 @@ def shift(self, periods, axis=0):
 
     def where(self, other, cond, align=True, errors='raise',
               try_cast=False, axis=0, transpose=False):
+        # rough attempt to see if
         if isinstance(other, (ABCIndexClass, ABCSeries)):
             other = other.array
 
@@ -1989,7 +1990,33 @@ def where(self, other, cond, align=True, errors='raise',
             # we want to replace that with the correct NA value
             # for the type
             other = self.dtype.na_value
-        result = self.values.where(cond, other)
+
+        if is_sparse(self.values):
+            # ugly workaround for ensure that the dtype is OK
+            # after we insert NaNs.
+            if is_sparse(other):
+                otype = other.dtype.subtype
+            else:
+                otype = other
+            dtype = self.dtype.update_dtype(
+                np.result_type(self.values.dtype.subtype, otype)
+            )
+        else:
+            dtype = self.dtype
+
+        if self._holder.__setitem__ is ExtensionArray.__setitem__:
+            # the array doesn't implement setitem, so convert to ndarray
+            result = self._holder._from_sequence(
+                np.where(cond, self.values, other),
+                dtype=dtype,
+            )
+        else:
+            result = self.values.copy()
+            icond = ~cond
+            if lib.is_scalar(other):
+                result[icond] = other
+            else:
+                result[icond] = other[icond]
         return self.make_block_same_class(result, placement=self.mgr_locs)
 
     @property
@@ -2673,13 +2700,55 @@ def concat_same_type(self, to_concat, placement=None):
 
     def where(self, other, cond, align=True, errors='raise',
               try_cast=False, axis=0, transpose=False):
-        result = super(CategoricalBlock, self).where(
-            other, cond, align, errors, try_cast, axis, transpose
+        # This can all be deleted in favor of ExtensionBlock.where once
+        # we enforce the deprecation.
+        object_msg = (
+            "Implicitly converting categorical to object-dtype ndarray. "
+            "The values `{}' are not present in this categorical's "
+            "categories. A future version of pandas will raise a ValueError "
+            "when 'other' contains different categories.\n\n"
+            "To preserve the current behavior, add the new categories to "
+            "the categorical before calling 'where', or convert the "
+            "categorical to a different dtype."
         )
-        if result.values.dtype != self.values.dtype:
-            # For backwards compatability, we allow upcasting to object.
-            # This fallback will be removed in the future.
-            result = result.astype(object)
+
+        scalar_other = lib.is_scalar(other)
+        categorical_other = is_categorical_dtype(other)
+        if isinstance(other, ABCDataFrame):
+            # should be 1d
+            assert other.shape[1] == 1
+            other = other.iloc[:, 0]
+
+        if isinstance(other, (ABCSeries, ABCIndexClass)):
+            other = other._values
+
+        do_as_object = (
+            # Two categoricals with different dtype (ignoring order)
+            (categorical_other and not is_dtype_equal(self.values, other)) or
+            # a not-na scalar not present in our categories
+            (scalar_other and (other not in self.values.categories
+                               and notna(other))) or
+            # an array not present in our categories
+            (not scalar_other and
+             (self.values.categories.get_indexer(
+                 other[notna(other)]) < 0).any())
+        )
+
+        if do_as_object:
+            if scalar_other:
+                msg = object_msg.format(other)
+            else:
+                msg = compat.reprlib.repr(other)
+
+            warnings.warn(msg, FutureWarning, stacklevel=6)
+            result = self.astype(object).where(other, cond, align=align,
+                                               errors=errors,
+                                               try_cast=try_cast,
+                                               axis=axis, transpose=transpose)
+        else:
+            result = super(CategoricalBlock, self).where(
+                other, cond, align, errors, try_cast, axis, transpose
+            )
         return result