diff --git a/pandas/core/internals/blocks.py b/pandas/core/internals/blocks.py index 06bf2e5d7b18e..b3dce41514ed9 100644 --- a/pandas/core/internals/blocks.py +++ b/pandas/core/internals/blocks.py @@ -1935,64 +1935,59 @@ def to_native_types( return self.make_block(res) -class DatetimeLikeBlockMixin(HybridMixin, Block): - """Mixin class for DatetimeBlock, DatetimeTZBlock, and TimedeltaBlock.""" - - _dtype: np.dtype - _holder: Type[Union[DatetimeArray, TimedeltaArray]] - - @classmethod - def _maybe_coerce_values(cls, values): - """ - Input validation for values passed to __init__. Ensure that - we have nanosecond datetime64/timedelta64, coercing if necessary. - - Parameters - ---------- - values : array-like - Must be convertible to datetime64/timedelta64 - - Returns - ------- - values : ndarray[datetime64ns/timedelta64ns] - - Overridden by DatetimeTZBlock. - """ - if values.dtype != cls._dtype: - # non-nano we will convert to nano - if values.dtype.kind != cls._dtype.kind: - # caller is responsible for ensuring td64/dt64 dtype - raise TypeError(values.dtype) # pragma: no cover - - values = cls._holder._from_sequence(values)._data - - if isinstance(values, cls._holder): - values = values._data +class NDArrayBackedExtensionBlock(HybridMixin, Block): + """ + Block backed by an NDArrayBackedExtensionArray + """ - assert isinstance(values, np.ndarray), type(values) - return values + def internal_values(self): + # Override to return DatetimeArray and TimedeltaArray + return self.array_values() def get_values(self, dtype: Optional[DtypeObj] = None) -> np.ndarray: """ return object dtype as boxed values, such as Timestamps/Timedelta """ + values = self.array_values() if is_object_dtype(dtype): # DTA/TDA constructor and astype can handle 2D - return self._holder(self.values).astype(object) - return self.values - - def internal_values(self): - # Override to return DatetimeArray and TimedeltaArray - return self.array_values() - - def array_values(self): - return self._holder._simple_new(self.values) + values = values.astype(object) + # TODO(EA2D): reshape not needed with 2D EAs + return np.asarray(values).reshape(self.shape) def iget(self, key): # GH#31649 we need to wrap scalars in Timestamp/Timedelta # TODO(EA2D): this can be removed if we ever have 2D EA return self.array_values().reshape(self.shape)[key] + def putmask(self, mask, new) -> List[Block]: + mask = extract_bool_array(mask) + + if not self._can_hold_element(new): + return self.astype(object).putmask(mask, new) + + # TODO(EA2D): reshape unnecessary with 2D EAs + arr = self.array_values().reshape(self.shape) + arr = cast("NDArrayBackedExtensionArray", arr) + arr.T.putmask(mask, new) + return [self] + + def where(self, other, cond, errors="raise", axis: int = 0) -> List[Block]: + # TODO(EA2D): reshape unnecessary with 2D EAs + arr = self.array_values().reshape(self.shape) + + cond = extract_bool_array(cond) + + try: + res_values = arr.T.where(cond, other).T + except (ValueError, TypeError): + return super().where(other, cond, errors=errors, axis=axis) + + # TODO(EA2D): reshape not needed with 2D EAs + res_values = res_values.reshape(self.values.shape) + nb = self.make_block_same_class(res_values) + return [nb] + def diff(self, n: int, axis: int = 0) -> List[Block]: """ 1st discrete difference. @@ -2017,50 +2012,61 @@ def diff(self, n: int, axis: int = 0) -> List[Block]: values = self.array_values().reshape(self.shape) new_values = values - values.shift(n, axis=axis) - return [ - TimeDeltaBlock(new_values, placement=self.mgr_locs.indexer, ndim=self.ndim) - ] + return [self.make_block(new_values)] def shift(self, periods: int, axis: int = 0, fill_value: Any = None) -> List[Block]: # TODO(EA2D) this is unnecessary if these blocks are backed by 2D EAs - values = self.array_values() + values = self.array_values().reshape(self.shape) new_values = values.shift(periods, fill_value=fill_value, axis=axis) return [self.make_block_same_class(new_values)] - def to_native_types(self, na_rep="NaT", **kwargs): - """ convert to our native types format """ - arr = self.array_values() - result = arr._format_native_types(na_rep=na_rep, **kwargs) - return self.make_block(result) +class DatetimeLikeBlockMixin(NDArrayBackedExtensionBlock): + """Mixin class for DatetimeBlock, DatetimeTZBlock, and TimedeltaBlock.""" - def putmask(self, mask, new) -> List[Block]: - mask = extract_bool_array(mask) + _dtype: np.dtype + _holder: Type[Union[DatetimeArray, TimedeltaArray]] - if not self._can_hold_element(new): - return self.astype(object).putmask(mask, new) + @classmethod + def _maybe_coerce_values(cls, values): + """ + Input validation for values passed to __init__. Ensure that + we have nanosecond datetime64/timedelta64, coercing if necessary. - # TODO(EA2D): reshape unnecessary with 2D EAs - arr = self.array_values().reshape(self.shape) - arr = cast("NDArrayBackedExtensionArray", arr) - arr.T.putmask(mask, new) - return [self] + Parameters + ---------- + values : array-like + Must be convertible to datetime64/timedelta64 - def where(self, other, cond, errors="raise", axis: int = 0) -> List[Block]: - # TODO(EA2D): reshape unnecessary with 2D EAs - arr = self.array_values().reshape(self.shape) + Returns + ------- + values : ndarray[datetime64ns/timedelta64ns] - cond = extract_bool_array(cond) + Overridden by DatetimeTZBlock. + """ + if values.dtype != cls._dtype: + # non-nano we will convert to nano + if values.dtype.kind != cls._dtype.kind: + # caller is responsible for ensuring td64/dt64 dtype + raise TypeError(values.dtype) # pragma: no cover - try: - res_values = arr.T.where(cond, other).T - except (ValueError, TypeError): - return super().where(other, cond, errors=errors, axis=axis) + values = cls._holder._from_sequence(values)._data - # TODO(EA2D): reshape not needed with 2D EAs - res_values = res_values.reshape(self.values.shape) - nb = self.make_block_same_class(res_values) - return [nb] + if isinstance(values, cls._holder): + values = values._data + + assert isinstance(values, np.ndarray), type(values) + return values + + def array_values(self): + return self._holder._simple_new(self.values) + + def to_native_types(self, na_rep="NaT", **kwargs): + """ convert to our native types format """ + arr = self.array_values() + + result = arr._format_native_types(na_rep=na_rep, **kwargs) + return self.make_block(result) class DatetimeBlock(DatetimeLikeBlockMixin): @@ -2133,37 +2139,6 @@ def is_view(self) -> bool: # check the ndarray values of the DatetimeIndex values return self.values._data.base is not None - def get_values(self, dtype: Optional[DtypeObj] = None) -> np.ndarray: - """ - Returns an ndarray of values. - - Parameters - ---------- - dtype : np.dtype - Only `object`-like dtypes are respected here (not sure - why). - - Returns - ------- - values : ndarray - When ``dtype=object``, then and object-dtype ndarray of - boxed values is returned. Otherwise, an M8[ns] ndarray - is returned. - - DatetimeArray is always 1-d. ``get_values`` will reshape - the return value to be the same dimensionality as the - block. - """ - values = self.values - if is_object_dtype(dtype): - values = values.astype(object) - - # TODO(EA2D): reshape unnecessary with 2D EAs - # Ensure that our shape is correct for DataFrame. - # ExtensionArrays are always 1-D, even in a DataFrame when - # the analogous NumPy-backed column would be a 2-D ndarray. - return np.asarray(values).reshape(self.shape) - def external_values(self): # NB: this is different from np.asarray(self.values), since that # return an object-dtype ndarray of Timestamps.