Skip to content

Commit 9eaa4bc

Browse files
authored
Add support for NumpyExtensionArray in pd.unique() (#59214)
* Add support for NumpyExtensionArray in unique * Add space to end of string
1 parent 374f386 commit 9eaa4bc

File tree

2 files changed

+32
-6
lines changed

2 files changed

+32
-6
lines changed

pandas/core/algorithms.py

+15-4
Original file line numberDiff line numberDiff line change
@@ -223,13 +223,17 @@ def _ensure_arraylike(values, func_name: str) -> ArrayLike:
223223
"""
224224
ensure that we are arraylike if not already
225225
"""
226-
if not isinstance(values, (ABCIndex, ABCSeries, ABCExtensionArray, np.ndarray)):
226+
if not isinstance(
227+
values,
228+
(ABCIndex, ABCSeries, ABCExtensionArray, np.ndarray, ABCNumpyExtensionArray),
229+
):
227230
# GH#52986
228231
if func_name != "isin-targets":
229232
# Make an exception for the comps argument in isin.
230233
raise TypeError(
231234
f"{func_name} requires a Series, Index, "
232-
f"ExtensionArray, or np.ndarray, got {type(values).__name__}."
235+
f"ExtensionArray, np.ndarray or NumpyExtensionArray "
236+
f"got {type(values).__name__}."
233237
)
234238

235239
inferred = lib.infer_dtype(values, skipna=False)
@@ -325,15 +329,15 @@ def unique(values):
325329
326330
Returns
327331
-------
328-
numpy.ndarray or ExtensionArray
332+
numpy.ndarray, ExtensionArray or NumpyExtensionArray
329333
330334
The return can be:
331335
332336
* Index : when the input is an Index
333337
* Categorical : when the input is a Categorical dtype
334338
* ndarray : when the input is a Series/ndarray
335339
336-
Return numpy.ndarray or ExtensionArray.
340+
Return numpy.ndarray, ExtensionArray or NumpyExtensionArray.
337341
338342
See Also
339343
--------
@@ -405,6 +409,13 @@ def unique(values):
405409
406410
>>> pd.unique(pd.Series([("a", "b"), ("b", "a"), ("a", "c"), ("b", "a")]).values)
407411
array([('a', 'b'), ('b', 'a'), ('a', 'c')], dtype=object)
412+
413+
An NumpyExtensionArray of complex
414+
415+
>>> pd.unique(pd.array([1 + 1j, 2, 3]))
416+
<NumpyExtensionArray>
417+
[(1+1j), (2+0j), (3+0j)]
418+
Length: 3, dtype: complex128
408419
"""
409420
return unique_with_mask(values)
410421

pandas/tests/test_algos.py

+17-2
Original file line numberDiff line numberDiff line change
@@ -873,6 +873,14 @@ def test_unique_masked(self, any_numeric_ea_dtype):
873873
expected = pd.array([1, pd.NA, 2], dtype=any_numeric_ea_dtype)
874874
tm.assert_extension_array_equal(result, expected)
875875

876+
def test_unique_NumpyExtensionArray(self):
877+
arr_complex = pd.array(
878+
[1 + 1j, 2, 3]
879+
) # NumpyEADtype('complex128') => NumpyExtensionArray
880+
result = pd.unique(arr_complex)
881+
expected = pd.array([1 + 1j, 2 + 0j, 3 + 0j])
882+
tm.assert_extension_array_equal(result, expected)
883+
876884

877885
def test_nunique_ints(index_or_series_or_array):
878886
# GH#36327
@@ -1638,7 +1646,10 @@ def test_unique_tuples(self, arr, uniques):
16381646
expected = np.empty(len(uniques), dtype=object)
16391647
expected[:] = uniques
16401648

1641-
msg = "unique requires a Series, Index, ExtensionArray, or np.ndarray, got list"
1649+
msg = (
1650+
r"unique requires a Series, Index, ExtensionArray, np.ndarray "
1651+
r"or NumpyExtensionArray got list"
1652+
)
16421653
with pytest.raises(TypeError, match=msg):
16431654
# GH#52986
16441655
pd.unique(arr)
@@ -1657,7 +1668,11 @@ def test_unique_tuples(self, arr, uniques):
16571668
)
16581669
def test_unique_complex_numbers(self, array, expected):
16591670
# GH 17927
1660-
msg = "unique requires a Series, Index, ExtensionArray, or np.ndarray, got list"
1671+
msg = (
1672+
r"unique requires a Series, Index, ExtensionArray, np.ndarray "
1673+
r"or NumpyExtensionArray got list"
1674+
)
1675+
16611676
with pytest.raises(TypeError, match=msg):
16621677
# GH#52986
16631678
pd.unique(array)

0 commit comments

Comments
 (0)