diff --git a/pandas/core/algorithms.py b/pandas/core/algorithms.py index 92bd55cac9c5e..948836bf6a51d 100644 --- a/pandas/core/algorithms.py +++ b/pandas/core/algorithms.py @@ -223,13 +223,17 @@ def _ensure_arraylike(values, func_name: str) -> ArrayLike: """ ensure that we are arraylike if not already """ - if not isinstance(values, (ABCIndex, ABCSeries, ABCExtensionArray, np.ndarray)): + if not isinstance( + values, + (ABCIndex, ABCSeries, ABCExtensionArray, np.ndarray, ABCNumpyExtensionArray), + ): # GH#52986 if func_name != "isin-targets": # Make an exception for the comps argument in isin. raise TypeError( f"{func_name} requires a Series, Index, " - f"ExtensionArray, or np.ndarray, got {type(values).__name__}." + f"ExtensionArray, np.ndarray or NumpyExtensionArray " + f"got {type(values).__name__}." ) inferred = lib.infer_dtype(values, skipna=False) @@ -325,7 +329,7 @@ def unique(values): Returns ------- - numpy.ndarray or ExtensionArray + numpy.ndarray, ExtensionArray or NumpyExtensionArray The return can be: @@ -333,7 +337,7 @@ def unique(values): * Categorical : when the input is a Categorical dtype * ndarray : when the input is a Series/ndarray - Return numpy.ndarray or ExtensionArray. + Return numpy.ndarray, ExtensionArray or NumpyExtensionArray. See Also -------- @@ -405,6 +409,13 @@ def unique(values): >>> pd.unique(pd.Series([("a", "b"), ("b", "a"), ("a", "c"), ("b", "a")]).values) array([('a', 'b'), ('b', 'a'), ('a', 'c')], dtype=object) + + An NumpyExtensionArray of complex + + >>> pd.unique(pd.array([1 + 1j, 2, 3])) + + [(1+1j), (2+0j), (3+0j)] + Length: 3, dtype: complex128 """ return unique_with_mask(values) diff --git a/pandas/tests/test_algos.py b/pandas/tests/test_algos.py index 134ebededd163..cdcd36846c560 100644 --- a/pandas/tests/test_algos.py +++ b/pandas/tests/test_algos.py @@ -873,6 +873,14 @@ def test_unique_masked(self, any_numeric_ea_dtype): expected = pd.array([1, pd.NA, 2], dtype=any_numeric_ea_dtype) tm.assert_extension_array_equal(result, expected) + def test_unique_NumpyExtensionArray(self): + arr_complex = pd.array( + [1 + 1j, 2, 3] + ) # NumpyEADtype('complex128') => NumpyExtensionArray + result = pd.unique(arr_complex) + expected = pd.array([1 + 1j, 2 + 0j, 3 + 0j]) + tm.assert_extension_array_equal(result, expected) + def test_nunique_ints(index_or_series_or_array): # GH#36327 @@ -1638,7 +1646,10 @@ def test_unique_tuples(self, arr, uniques): expected = np.empty(len(uniques), dtype=object) expected[:] = uniques - msg = "unique requires a Series, Index, ExtensionArray, or np.ndarray, got list" + msg = ( + r"unique requires a Series, Index, ExtensionArray, np.ndarray " + r"or NumpyExtensionArray got list" + ) with pytest.raises(TypeError, match=msg): # GH#52986 pd.unique(arr) @@ -1657,7 +1668,11 @@ def test_unique_tuples(self, arr, uniques): ) def test_unique_complex_numbers(self, array, expected): # GH 17927 - msg = "unique requires a Series, Index, ExtensionArray, or np.ndarray, got list" + msg = ( + r"unique requires a Series, Index, ExtensionArray, np.ndarray " + r"or NumpyExtensionArray got list" + ) + with pytest.raises(TypeError, match=msg): # GH#52986 pd.unique(array)