Fixing shift() for ExtensionArray

doc/source/whatsnew/v0.24.0.rst

@@ -914,6 +914,7 @@ update the ``ExtensionDtype._metadata`` tuple to match the signature of your
 - Added ``ExtensionDtype._is_numeric`` for controlling whether an extension dtype is considered numeric (:issue:`22290`).
 - The ``ExtensionArray`` constructor, ``_from_sequence`` now take the keyword arg ``copy=False`` (:issue:`21185`)
 - Bug in :meth:`Series.get` for ``Series`` using ``ExtensionArray`` and integer index (:issue:`21257`)
+- :meth:`pandas.api.extensions.ExtensionArray.shift` added as part of the basic ``ExtensionArray`` interface (:issue:`22387`).
 - :meth:`~Series.shift` now dispatches to :meth:`ExtensionArray.shift` (:issue:`22386`)
 - :meth:`Series.combine()` works correctly with :class:`~pandas.api.extensions.ExtensionArray` inside of :class:`Series` (:issue:`20825`)
 - :meth:`Series.combine()` with scalar argument now works for any function type (:issue:`21248`)

pandas/core/arrays/base.py

@@ -418,13 +418,13 @@ def fillna(self, value=None, method=None, limit=None):
         return new_values
 
     def dropna(self):
-        """ Return ExtensionArray without NA values
+        """
+        Return ExtensionArray without NA values
 
         Returns
         -------
         valid : ExtensionArray
         """
-
         return self[~self.isna()]
 
     def shift(self, periods=1):
@@ -446,13 +446,24 @@ def shift(self, periods=1):
         Returns
         -------
         shifted : ExtensionArray
+
+        Notes
+        -----
+        If ``self`` is empty or ``periods`` is 0, a copy of ``self`` is
+        returned.
+
+        If ``periods > len(self)``, then an array of size
+        len(self) is returned, with all values filled with
+        ``self.dtype.na_value``.
         """
         # Note: this implementation assumes that `self.dtype.na_value` can be
         # stored in an instance of your ExtensionArray with `self.dtype`.
-        if periods == 0:
+        if len(self) == 0 or periods == 0:
             return self.copy()
-        empty = self._from_sequence([self.dtype.na_value] * abs(periods),
-                                    dtype=self.dtype)
+        empty = self._from_sequence(
+            [self.dtype.na_value] * min(abs(periods), len(self)),
+            dtype=self.dtype
+        )
         if periods > 0:
             a = empty
             b = self[:-periods]
@@ -462,7 +473,8 @@ def shift(self, periods=1):
         return self._concat_same_type([a, b])
 
     def unique(self):
-        """Compute the ExtensionArray of unique values.
+        """
+        Compute the ExtensionArray of unique values.
 
         Returns
         -------

pandas/core/arrays/sparse.py

@@ -882,7 +882,7 @@ def fillna(self, value=None, method=None, limit=None):
 
     def shift(self, periods=1):
 
-        if periods == 0:
+        if len(self) == 0 or periods == 0:
             return self.copy()
 
         subtype = np.result_type(np.nan, self.dtype.subtype)
@@ -893,8 +893,11 @@ def shift(self, periods=1):
         else:
             arr = self
 
-        empty = self._from_sequence([self.dtype.na_value] * abs(periods),
-                                    dtype=arr.dtype)
+        empty = self._from_sequence(
+            [self.dtype.na_value] * min(abs(periods), len(self)),
+            dtype=arr.dtype
+        )
+
         if periods > 0:
             a = empty
             b = arr[:-periods]

pandas/tests/extension/arrow/test_bool.py

@@ -39,6 +39,14 @@ class TestInterface(BaseArrowTests, base.BaseInterfaceTests):
     def test_repr(self, data):
         raise pytest.skip("TODO")
 
+    def test_shift_non_empty_array(self, data):
+        raise pytest.skip("Skipping because the implementation of"
+                          "`ArrowBoolArray` is quite buggy.")
+
+    def test_shift_empty_array(self, data):
+        raise pytest.skip("Skipping because the implementation of"
+                          "`ArrowBoolArray` is quite buggy.")
+
 
 class TestConstructors(BaseArrowTests, base.BaseConstructorsTests):
     def test_from_dtype(self, data):

pandas/tests/extension/base/interface.py

@@ -1,4 +1,5 @@
 import numpy as np
+import pytest
 
 from pandas.compat import StringIO
 
@@ -86,3 +87,25 @@ def test_isna_extension_array(self, data_missing):
             assert not na.all()
 
             assert na.dtype._is_boolean
+
+    @pytest.mark.parametrize('periods, indices', [
+        [-4, [-1, -1]],
+        [-1, [1, -1]],
+        [0, [0, 1]],
+        [1, [-1, 0]],
+        [4, [-1, -1]]
+    ])
+    def test_shift_non_empty_array(self, data, periods, indices):
+        subset = data[:2]
+        result = subset.shift(periods)
+        expected = subset.take(indices, allow_fill=True)
+        self.assert_extension_array_equal(result, expected)
+
+    @pytest.mark.parametrize('periods', [
+        -4, -1, 0, 1, 4
+    ])
+    def test_shift_empty_array(self, data, periods):
+        empty = data[:0]
+        result = empty.shift(periods)
+        expected = empty
+        self.assert_extension_array_equal(result, expected)