diff --git a/pandas/core/array_algos/__init__.py b/pandas/core/array_algos/__init__.py new file mode 100644 index 0000000000000..a7655a013c6cf --- /dev/null +++ b/pandas/core/array_algos/__init__.py @@ -0,0 +1,9 @@ +""" +core.array_algos is for algorithms that operate on ndarray and ExtensionArray. +These should: + +- Assume that any Index, Series, or DataFrame objects have already been unwrapped. +- Assume that any list arguments have already been cast to ndarray/EA. +- Not depend on Index, Series, or DataFrame, nor import any of these. +- May dispatch to ExtensionArray methods, but should not import from core.arrays. +""" diff --git a/pandas/core/array_algos/transforms.py b/pandas/core/array_algos/transforms.py new file mode 100644 index 0000000000000..f775b6d733d9c --- /dev/null +++ b/pandas/core/array_algos/transforms.py @@ -0,0 +1,33 @@ +""" +transforms.py is for shape-preserving functions. +""" + +import numpy as np + +from pandas.core.dtypes.common import ensure_platform_int + + +def shift(values: np.ndarray, periods: int, axis: int, fill_value) -> np.ndarray: + new_values = values + + # make sure array sent to np.roll is c_contiguous + f_ordered = values.flags.f_contiguous + if f_ordered: + new_values = new_values.T + axis = new_values.ndim - axis - 1 + + if np.prod(new_values.shape): + new_values = np.roll(new_values, ensure_platform_int(periods), axis=axis) + + axis_indexer = [slice(None)] * values.ndim + if periods > 0: + axis_indexer[axis] = slice(None, periods) + else: + axis_indexer[axis] = slice(periods, None) + new_values[tuple(axis_indexer)] = fill_value + + # restore original order + if f_ordered: + new_values = new_values.T + + return new_values diff --git a/pandas/core/arrays/datetimelike.py b/pandas/core/arrays/datetimelike.py index 105d9581b1a25..7510bfd1f67ad 100644 --- a/pandas/core/arrays/datetimelike.py +++ b/pandas/core/arrays/datetimelike.py @@ -40,6 +40,7 @@ from pandas.core import missing, nanops, ops from pandas.core.algorithms import checked_add_with_arr, take, unique1d, value_counts +from pandas.core.array_algos.transforms import shift from pandas.core.arrays.base import ExtensionArray, ExtensionOpsMixin import pandas.core.common as com from pandas.core.construction import array, extract_array @@ -773,26 +774,7 @@ def shift(self, periods=1, fill_value=None, axis=0): fill_value = self._unbox_scalar(fill_value) - new_values = self._data - - # make sure array sent to np.roll is c_contiguous - f_ordered = new_values.flags.f_contiguous - if f_ordered: - new_values = new_values.T - axis = new_values.ndim - axis - 1 - - new_values = np.roll(new_values, periods, axis=axis) - - axis_indexer = [slice(None)] * self.ndim - if periods > 0: - axis_indexer[axis] = slice(None, periods) - else: - axis_indexer[axis] = slice(periods, None) - new_values[tuple(axis_indexer)] = fill_value - - # restore original order - if f_ordered: - new_values = new_values.T + new_values = shift(self._data, periods, axis, fill_value) return type(self)._simple_new(new_values, dtype=self.dtype) diff --git a/pandas/core/internals/blocks.py b/pandas/core/internals/blocks.py index cab2bd5146745..adeb1ae04a58d 100644 --- a/pandas/core/internals/blocks.py +++ b/pandas/core/internals/blocks.py @@ -29,7 +29,6 @@ from pandas.core.dtypes.common import ( _NS_DTYPE, _TD_DTYPE, - ensure_platform_int, is_bool_dtype, is_categorical, is_categorical_dtype, @@ -66,6 +65,7 @@ ) import pandas.core.algorithms as algos +from pandas.core.array_algos.transforms import shift from pandas.core.arrays import ( Categorical, DatetimeArray, @@ -1316,25 +1316,7 @@ def shift(self, periods, axis: int = 0, fill_value=None): # that, handle boolean etc also new_values, fill_value = maybe_upcast(self.values, fill_value) - # make sure array sent to np.roll is c_contiguous - f_ordered = new_values.flags.f_contiguous - if f_ordered: - new_values = new_values.T - axis = new_values.ndim - axis - 1 - - if np.prod(new_values.shape): - new_values = np.roll(new_values, ensure_platform_int(periods), axis=axis) - - axis_indexer = [slice(None)] * self.ndim - if periods > 0: - axis_indexer[axis] = slice(None, periods) - else: - axis_indexer[axis] = slice(periods, None) - new_values[tuple(axis_indexer)] = fill_value - - # restore original order - if f_ordered: - new_values = new_values.T + new_values = shift(new_values, periods, axis, fill_value) return [self.make_block(new_values)]