diff --git a/doc/source/whatsnew/v0.24.0.rst b/doc/source/whatsnew/v0.24.0.rst index 364c6f2a2daf2a..78f864f0dcb735 100644 --- a/doc/source/whatsnew/v0.24.0.rst +++ b/doc/source/whatsnew/v0.24.0.rst @@ -987,7 +987,8 @@ update the ``ExtensionDtype._metadata`` tuple to match the signature of your **Other changes** -- ``ExtensionArray`` has gained the abstract methods ``.dropna()`` (:issue:`21185`) +- :meth:`~pandas.api.types.ExtensionArray.dropna` has been added (:issue:`21185`) +- :meth:`~pandas.api.types.ExtensionArray.repeat` has been added (:issue:`24349`) - ``ExtensionDtype`` has gained the ability to instantiate from string dtypes, e.g. ``decimal`` would instantiate a registered ``DecimalDtype``; furthermore the ``ExtensionDtype`` has gained the method ``construct_array_type`` (:issue:`21185`) - An ``ExtensionArray`` with a boolean dtype now works correctly as a boolean indexer. :meth:`pandas.api.types.is_bool_dtype` now properly considers them boolean (:issue:`22326`) diff --git a/pandas/core/arrays/base.py b/pandas/core/arrays/base.py index cf145064fd7b17..a848dafbb06ef1 100644 --- a/pandas/core/arrays/base.py +++ b/pandas/core/arrays/base.py @@ -580,6 +580,35 @@ def factorize(self, na_sentinel=-1): uniques = self._from_factorized(uniques, self) return labels, uniques + def repeat(self, repeats, axis=None): + """ + Repeat elements of an array. + + .. versionadded:: 0.24.0 + + Parameters + ---------- + repeats : int + This should be a non-negative integer. Repeating 0 times + will return an empty array. + + Returns + ------- + repeated_array : ExtensionArray + Same type as the input, with elements repeated `repeats` times. + + See Also + -------- + numpy.repeat : Similar method for :class:`numpy.ndarray`. + ExtensionArray.take : Take arbitrary positions. + """ + if axis is not None: + raise ValueError("'axis' must be None.") + if repeats < 0: + raise ValueError("negative repeats are not allowed.") + ind = np.arange(len(self)).repeat(repeats) + return self.take(ind) + # ------------------------------------------------------------------------ # Indexing methods # ------------------------------------------------------------------------ diff --git a/pandas/tests/extension/base/methods.py b/pandas/tests/extension/base/methods.py index 4a409a84f3db4a..3403d0e9e02f1b 100644 --- a/pandas/tests/extension/base/methods.py +++ b/pandas/tests/extension/base/methods.py @@ -264,3 +264,34 @@ def test_where_series(self, data, na_value, as_frame): if as_frame: expected = expected.to_frame(name='a') self.assert_equal(result, expected) + + @pytest.mark.parametrize("as_series", [True, False]) + @pytest.mark.parametrize("repeats", [0, 1, 2]) + def test_repeat(self, data, repeats, as_series): + a, b, c = data[:3] + arr = type(data)._from_sequence([a, b, c], dtype=data.dtype) + + if as_series: + arr = pd.Series(arr) + + result = arr.repeat(repeats) + + if repeats == 0: + expected = [] + elif repeats == 1: + expected = [a, b, c] + else: + expected = [a, a, b, b, c, c] + expected = type(data)._from_sequence(expected, dtype=data.dtype) + if as_series: + index = pd.Series(np.arange(len(arr))).repeat(repeats).index + expected = pd.Series(expected, index=index) + self.assert_equal(result, expected) + + def test_repeat_raises(self, data): + with pytest.raises(ValueError, match="'axis'"): + data.repeat(2, axis=1) + + with pytest.raises(ValueError, + match="negative"): + data.repeat(-1)