|
29 | 29 | from pandas.core.dtypes.dtypes import (
|
30 | 30 | CategoricalDtype, DatetimeTZDtype, ExtensionDtype, PandasExtensionDtype)
|
31 | 31 | from pandas.core.dtypes.generic import (
|
32 |
| - ABCDatetimeIndex, ABCExtensionArray, ABCIndexClass, ABCSeries) |
| 32 | + ABCDataFrame, ABCDatetimeIndex, ABCExtensionArray, ABCIndexClass, |
| 33 | + ABCSeries) |
33 | 34 | from pandas.core.dtypes.inference import is_scalar
|
34 | 35 | from pandas.core.dtypes.missing import (
|
35 | 36 | _isna_compat, array_equivalent, is_null_datelike_scalar, isna, notna)
|
@@ -1970,6 +1971,30 @@ def shift(self, periods, axis=0):
|
1970 | 1971 | placement=self.mgr_locs,
|
1971 | 1972 | ndim=self.ndim)]
|
1972 | 1973 |
|
| 1974 | + def where(self, other, cond, align=True, errors='raise', |
| 1975 | + try_cast=False, axis=0, transpose=False): |
| 1976 | + if isinstance(other, (ABCIndexClass, ABCSeries)): |
| 1977 | + other = other.array |
| 1978 | + |
| 1979 | + if isinstance(cond, ABCDataFrame): |
| 1980 | + assert cond.shape[1] == 1 |
| 1981 | + cond = cond.iloc[:, 0].array |
| 1982 | + |
| 1983 | + if isinstance(other, ABCDataFrame): |
| 1984 | + assert other.shape[1] == 1 |
| 1985 | + other = other.iloc[:, 0].array |
| 1986 | + |
| 1987 | + if isinstance(cond, (ABCIndexClass, ABCSeries)): |
| 1988 | + cond = cond.array |
| 1989 | + |
| 1990 | + if lib.is_scalar(other) and isna(other): |
| 1991 | + # The default `other` for Series / Frame is np.nan |
| 1992 | + # we want to replace that with the correct NA value |
| 1993 | + # for the type |
| 1994 | + other = self.dtype.na_value |
| 1995 | + result = self.values.where(cond, other) |
| 1996 | + return self.make_block_same_class(result, placement=self.mgr_locs) |
| 1997 | + |
1973 | 1998 | @property
|
1974 | 1999 | def _ftype(self):
|
1975 | 2000 | return getattr(self.values, '_pandas_ftype', Block._ftype)
|
@@ -2675,6 +2700,17 @@ def concat_same_type(self, to_concat, placement=None):
|
2675 | 2700 | values, placement=placement or slice(0, len(values), 1),
|
2676 | 2701 | ndim=self.ndim)
|
2677 | 2702 |
|
| 2703 | + def where(self, other, cond, align=True, errors='raise', |
| 2704 | + try_cast=False, axis=0, transpose=False): |
| 2705 | + result = super(CategoricalBlock, self).where( |
| 2706 | + other, cond, align, errors, try_cast, axis, transpose |
| 2707 | + ) |
| 2708 | + if result.values.dtype != self.values.dtype: |
| 2709 | + # For backwards compatability, we allow upcasting to object. |
| 2710 | + # This fallback will be removed in the future. |
| 2711 | + result = result.astype(object) |
| 2712 | + return result |
| 2713 | + |
2678 | 2714 |
|
2679 | 2715 | class DatetimeBlock(DatetimeLikeBlockMixin, Block):
|
2680 | 2716 | __slots__ = ()
|
|
0 commit comments