From 9241c074e0b21563e9f4e4e61525415c8a645855 Mon Sep 17 00:00:00 2001 From: Brock Date: Sun, 25 Jun 2023 12:06:01 -0700 Subject: [PATCH 1/3] REF: remove axis keyword from Manager/Block.shift --- pandas/core/arrays/_mixins.py | 4 ++- pandas/core/arrays/base.py | 2 ++ pandas/core/generic.py | 5 ++-- pandas/core/internals/array_manager.py | 10 ++----- pandas/core/internals/blocks.py | 38 ++++++++++---------------- pandas/core/internals/managers.py | 5 ++-- 6 files changed, 26 insertions(+), 38 deletions(-) diff --git a/pandas/core/arrays/_mixins.py b/pandas/core/arrays/_mixins.py index 0100c17805d76..f586de3d2bdee 100644 --- a/pandas/core/arrays/_mixins.py +++ b/pandas/core/arrays/_mixins.py @@ -238,7 +238,9 @@ def searchsorted( return self._ndarray.searchsorted(npvalue, side=side, sorter=sorter) @doc(ExtensionArray.shift) - def shift(self, periods: int = 1, fill_value=None, axis: AxisInt = 0): + def shift(self, periods: int = 1, fill_value=None): + # NB: shift is always along axis=0 + axis = 0 fill_value = self._validate_scalar(fill_value) new_values = shift(self._ndarray, periods, axis, fill_value) diff --git a/pandas/core/arrays/base.py b/pandas/core/arrays/base.py index 4ebc312985aed..78ecd93d5cc75 100644 --- a/pandas/core/arrays/base.py +++ b/pandas/core/arrays/base.py @@ -851,6 +851,8 @@ def shift(self, periods: int = 1, fill_value: object = None) -> ExtensionArray: If ``periods > len(self)``, then an array of size len(self) is returned, with all values filled with ``self.dtype.na_value``. + + For 2-dimensional ExtensionArrays, we are always shifting along axis=0. """ # Note: this implementation assumes that `self.dtype.na_value` can be # stored in an instance of your ExtensionArray with `self.dtype`. diff --git a/pandas/core/generic.py b/pandas/core/generic.py index 0be840f9a4ef1..2805581beb740 100644 --- a/pandas/core/generic.py +++ b/pandas/core/generic.py @@ -10538,9 +10538,8 @@ def shift( if freq is None: # when freq is None, data is shifted, index is not axis = self._get_axis_number(axis) - new_data = self._mgr.shift( - periods=periods, axis=axis, fill_value=fill_value - ) + assert axis == 0 # axis == 1 cases handled in DataFrame.shift + new_data = self._mgr.shift(periods=periods, fill_value=fill_value) return self._constructor(new_data).__finalize__(self, method="shift") # when freq is given, index is shifted, data is not diff --git a/pandas/core/internals/array_manager.py b/pandas/core/internals/array_manager.py index ccb181c0537dc..25a8f2f73a793 100644 --- a/pandas/core/internals/array_manager.py +++ b/pandas/core/internals/array_manager.py @@ -360,17 +360,11 @@ def diff(self, n: int) -> Self: def interpolate(self, **kwargs) -> Self: return self.apply_with_block("interpolate", swap_axis=False, **kwargs) - def shift(self, periods: int, axis: AxisInt, fill_value) -> Self: + def shift(self, periods: int, fill_value) -> Self: if fill_value is lib.no_default: fill_value = None - if axis == 1 and self.ndim == 2: - # TODO column-wise shift - raise NotImplementedError - - return self.apply_with_block( - "shift", periods=periods, axis=axis, fill_value=fill_value - ) + return self.apply_with_block("shift", periods=periods, fill_value=fill_value) def fillna(self, value, limit: int | None, inplace: bool, downcast) -> Self: if limit is not None: diff --git a/pandas/core/internals/blocks.py b/pandas/core/internals/blocks.py index e61d233a0ae84..b96a7d1525c7e 100644 --- a/pandas/core/internals/blocks.py +++ b/pandas/core/internals/blocks.py @@ -1434,12 +1434,11 @@ def diff(self, n: int) -> list[Block]: new_values = algos.diff(self.values.T, n, axis=0).T return [self.make_block(values=new_values)] - def shift( - self, periods: int, axis: AxisInt = 0, fill_value: Any = None - ) -> list[Block]: + def shift(self, periods: int, fill_value: Any = None) -> list[Block]: """shift the block by periods, possibly upcast""" # convert integer to float if necessary. need to do a lot more than # that, handle boolean etc also + axis = self.ndim - 1 # Note: periods is never 0 here, as that is handled at the top of # NDFrame.shift. If that ever changes, we can do a check for periods=0 @@ -1461,7 +1460,7 @@ def shift( ) except LossySetitemError: nb = self.coerce_to_target_dtype(fill_value) - return nb.shift(periods, axis=axis, fill_value=fill_value) + return nb.shift(periods, fill_value=fill_value) else: values = cast(np.ndarray, self.values) @@ -1617,6 +1616,18 @@ class EABackedBlock(Block): values: ExtensionArray + def shift(self, periods: int, fill_value: Any = None) -> list[Self]: + """ + Shift the block by `periods`. + + Dispatches to underlying ExtensionArray and re-boxes in an + ExtensionBlock. + """ + # Transpose since EA.shift is always along axis=0, while we want to shift + # along rows. + new_values = self.values.T.shift(periods=periods, fill_value=fill_value).T + return [self.make_block_same_class(new_values)] + def setitem(self, indexer, value, using_cow: bool = False): """ Attempt self.values[indexer] = value, possibly creating a new array. @@ -2069,18 +2080,6 @@ def slice_block_rows(self, slicer: slice) -> Self: new_values = self.values[slicer] return type(self)(new_values, self._mgr_locs, ndim=self.ndim, refs=self.refs) - def shift( - self, periods: int, axis: AxisInt = 0, fill_value: Any = None - ) -> list[Block]: - """ - Shift the block by `periods`. - - Dispatches to underlying ExtensionArray and re-boxes in an - ExtensionBlock. - """ - new_values = self.values.shift(periods=periods, fill_value=fill_value) - return [self.make_block_same_class(new_values)] - def _unstack( self, unstacker, @@ -2187,13 +2186,6 @@ def is_view(self) -> bool: # check the ndarray values of the DatetimeIndex values return self.values._ndarray.base is not None - def shift( - self, periods: int, axis: AxisInt = 0, fill_value: Any = None - ) -> list[Block]: - values = self.values - new_values = values.shift(periods, fill_value=fill_value, axis=axis) - return [self.make_block_same_class(new_values)] - def _catch_deprecated_value_error(err: Exception) -> None: """ diff --git a/pandas/core/internals/managers.py b/pandas/core/internals/managers.py index f11c909a5df41..a9b6002c4e0b7 100644 --- a/pandas/core/internals/managers.py +++ b/pandas/core/internals/managers.py @@ -420,12 +420,11 @@ def interpolate(self, inplace: bool, **kwargs) -> Self: "interpolate", inplace=inplace, **kwargs, using_cow=using_copy_on_write() ) - def shift(self, periods: int, axis: AxisInt, fill_value) -> Self: - axis = self._normalize_axis(axis) + def shift(self, periods: int, fill_value) -> Self: if fill_value is lib.no_default: fill_value = None - return self.apply("shift", periods=periods, axis=axis, fill_value=fill_value) + return self.apply("shift", periods=periods, fill_value=fill_value) def fillna(self, value, limit: int | None, inplace: bool, downcast) -> Self: if limit is not None: From 0766ffbc0baa692d861e14e25e8fdb45b77c0313 Mon Sep 17 00:00:00 2001 From: Brock Date: Mon, 26 Jun 2023 09:21:50 -0700 Subject: [PATCH 2/3] mypy fixup --- pandas/core/internals/blocks.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pandas/core/internals/blocks.py b/pandas/core/internals/blocks.py index 94026b7282b65..653a5a7ffc84a 100644 --- a/pandas/core/internals/blocks.py +++ b/pandas/core/internals/blocks.py @@ -1463,7 +1463,7 @@ def diff(self, n: int) -> list[Block]: new_values = algos.diff(self.values.T, n, axis=0).T return [self.make_block(values=new_values)] - def shift(self, periods: int, fill_value: Any = None) -> list[Block]: + def shift(self, periods: int, fill_value: Any = None) -> list[Self]: """shift the block by periods, possibly upcast""" # convert integer to float if necessary. need to do a lot more than # that, handle boolean etc also From a1ed7cb68b173f40fa1f6c2fe494632c637a53cc Mon Sep 17 00:00:00 2001 From: Brock Date: Mon, 26 Jun 2023 12:33:22 -0700 Subject: [PATCH 3/3] mypy fixup --- pandas/core/internals/blocks.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/pandas/core/internals/blocks.py b/pandas/core/internals/blocks.py index 653a5a7ffc84a..a3195d9e82d58 100644 --- a/pandas/core/internals/blocks.py +++ b/pandas/core/internals/blocks.py @@ -1463,7 +1463,7 @@ def diff(self, n: int) -> list[Block]: new_values = algos.diff(self.values.T, n, axis=0).T return [self.make_block(values=new_values)] - def shift(self, periods: int, fill_value: Any = None) -> list[Self]: + def shift(self, periods: int, fill_value: Any = None) -> list[Block]: """shift the block by periods, possibly upcast""" # convert integer to float if necessary. need to do a lot more than # that, handle boolean etc also @@ -1494,7 +1494,7 @@ def shift(self, periods: int, fill_value: Any = None) -> list[Self]: else: values = cast(np.ndarray, self.values) new_values = shift(values, periods, axis, casted) - return [self.make_block(new_values)] + return [self.make_block_same_class(new_values)] @final def quantile( @@ -1645,7 +1645,7 @@ class EABackedBlock(Block): values: ExtensionArray - def shift(self, periods: int, fill_value: Any = None) -> list[Self]: + def shift(self, periods: int, fill_value: Any = None) -> list[Block]: """ Shift the block by `periods`.