ENH: Support mask in duplicated algorithm (pandas-dev#48150)

* ENH: Support mask in duplicated algorithm

* Fix mypy

* Add tests

* Improve test

* Fix bug

Author: phofl

asv_bench/benchmarks/algorithms.py

@@ -95,6 +95,29 @@ def time_duplicated(self, unique, keep, dtype):
         self.idx.duplicated(keep=keep)
 
 
+class DuplicatedMaskedArray:
+
+    params = [
+        [True, False],
+        ["first", "last", False],
+        ["Int64", "Float64"],
+    ]
+    param_names = ["unique", "keep", "dtype"]
+
+    def setup(self, unique, keep, dtype):
+        N = 10**5
+        data = pd.Series(np.arange(N), dtype=dtype)
+        data[list(range(1, N, 100))] = pd.NA
+        if not unique:
+            data = data.repeat(5)
+        self.ser = data
+        # cache is_unique
+        self.ser.is_unique
+
+    def time_duplicated(self, unique, keep, dtype):
+        self.ser.duplicated(keep=keep)
+
+
 class Hashing:
     def setup_cache(self):
         N = 10**5

pandas/_libs/hashtable.pyi

@@ -183,6 +183,7 @@ class IntpHashTable(HashTable): ...
 def duplicated(
     values: np.ndarray,
     keep: Literal["last", "first", False] = ...,
+    mask: npt.NDArray[np.bool_] | None = ...,
 ) -> npt.NDArray[np.bool_]: ...
 def mode(
     values: np.ndarray, dropna: bool, mask: npt.NDArray[np.bool_] | None = ...

pandas/_libs/hashtable_func_helper.pxi.in

@@ -118,9 +118,9 @@ cdef value_count_{{dtype}}(const {{dtype}}_t[:] values, bint dropna, const uint8
 @cython.wraparound(False)
 @cython.boundscheck(False)
 {{if dtype == 'object'}}
-cdef duplicated_{{dtype}}(ndarray[{{dtype}}] values, object keep='first'):
+cdef duplicated_{{dtype}}(ndarray[{{dtype}}] values, object keep='first', const uint8_t[:] mask=None):
 {{else}}
-cdef duplicated_{{dtype}}(const {{dtype}}_t[:] values, object keep='first'):
+cdef duplicated_{{dtype}}(const {{dtype}}_t[:] values, object keep='first', const uint8_t[:] mask=None):
 {{endif}}
     cdef:
         int ret = 0
@@ -129,10 +129,12 @@ cdef duplicated_{{dtype}}(const {{dtype}}_t[:] values, object keep='first'):
         {{else}}
         PyObject* value
         {{endif}}
-        Py_ssize_t i, n = len(values)
+        Py_ssize_t i, n = len(values), first_na = -1
         khiter_t k
         kh_{{ttype}}_t *table = kh_init_{{ttype}}()
         ndarray[uint8_t, ndim=1, cast=True] out = np.empty(n, dtype='bool')
+        bint seen_na = False, uses_mask = mask is not None
+        bint seen_multiple_na = False
 
     kh_resize_{{ttype}}(table, min(kh_needed_n_buckets(n), SIZE_HINT_LIMIT))
 
@@ -147,9 +149,16 @@ cdef duplicated_{{dtype}}(const {{dtype}}_t[:] values, object keep='first'):
         {{endif}}
             for i in range(n - 1, -1, -1):
                 # equivalent: range(n)[::-1], which cython doesn't like in nogil
-                value = {{to_c_type}}(values[i])
-                kh_put_{{ttype}}(table, value, &ret)
-                out[i] = ret == 0
+                if uses_mask and mask[i]:
+                    if seen_na:
+                        out[i] = True
+                    else:
+                        out[i] = False
+                        seen_na = True
+                else:
+                    value = {{to_c_type}}(values[i])
+                    kh_put_{{ttype}}(table, value, &ret)
+                    out[i] = ret == 0
 
     elif keep == 'first':
         {{if dtype == 'object'}}
@@ -158,9 +167,16 @@ cdef duplicated_{{dtype}}(const {{dtype}}_t[:] values, object keep='first'):
         with nogil:
         {{endif}}
             for i in range(n):
-                value = {{to_c_type}}(values[i])
-                kh_put_{{ttype}}(table, value, &ret)
-                out[i] = ret == 0
+                if uses_mask and mask[i]:
+                    if seen_na:
+                        out[i] = True
+                    else:
+                        out[i] = False
+                        seen_na = True
+                else:
+                    value = {{to_c_type}}(values[i])
+                    kh_put_{{ttype}}(table, value, &ret)
+                    out[i] = ret == 0
 
     else:
         {{if dtype == 'object'}}
@@ -169,15 +185,28 @@ cdef duplicated_{{dtype}}(const {{dtype}}_t[:] values, object keep='first'):
         with nogil:
         {{endif}}
             for i in range(n):
-                value = {{to_c_type}}(values[i])
-                k = kh_get_{{ttype}}(table, value)
-                if k != table.n_buckets:
-                    out[table.vals[k]] = 1
-                    out[i] = 1
+                if uses_mask and mask[i]:
+                    if not seen_na:
+                        first_na = i
+                        seen_na = True
+                        out[i] = 0
+                    elif not seen_multiple_na:
+                        out[i] = 1
+                        out[first_na] = 1
+                        seen_multiple_na = True
+                    else:
+                        out[i] = 1
+
                 else:
-                    k = kh_put_{{ttype}}(table, value, &ret)
-                    table.vals[k] = i
-                    out[i] = 0
+                    value = {{to_c_type}}(values[i])
+                    k = kh_get_{{ttype}}(table, value)
+                    if k != table.n_buckets:
+                        out[table.vals[k]] = 1
+                        out[i] = 1
+                    else:
+                        k = kh_put_{{ttype}}(table, value, &ret)
+                        table.vals[k] = i
+                        out[i] = 0
 
     kh_destroy_{{ttype}}(table)
     return out
@@ -301,37 +330,37 @@ cpdef value_count(ndarray[htfunc_t] values, bint dropna, const uint8_t[:] mask=N
         raise TypeError(values.dtype)
 
 
-cpdef duplicated(ndarray[htfunc_t] values, object keep="first"):
+cpdef duplicated(ndarray[htfunc_t] values, object keep="first", const uint8_t[:] mask=None):
     if htfunc_t is object:
-        return duplicated_object(values, keep)
+        return duplicated_object(values, keep, mask=mask)
 
     elif htfunc_t is int8_t:
-        return duplicated_int8(values, keep)
+        return duplicated_int8(values, keep, mask=mask)
     elif htfunc_t is int16_t:
-        return duplicated_int16(values, keep)
+        return duplicated_int16(values, keep, mask=mask)
     elif htfunc_t is int32_t:
-        return duplicated_int32(values, keep)
+        return duplicated_int32(values, keep, mask=mask)
     elif htfunc_t is int64_t:
-        return duplicated_int64(values, keep)
+        return duplicated_int64(values, keep, mask=mask)
 
     elif htfunc_t is uint8_t:
-        return duplicated_uint8(values, keep)
+        return duplicated_uint8(values, keep, mask=mask)
     elif htfunc_t is uint16_t:
-        return duplicated_uint16(values, keep)
+        return duplicated_uint16(values, keep, mask=mask)
     elif htfunc_t is uint32_t:
-        return duplicated_uint32(values, keep)
+        return duplicated_uint32(values, keep, mask=mask)
     elif htfunc_t is uint64_t:
-        return duplicated_uint64(values, keep)
+        return duplicated_uint64(values, keep, mask=mask)
 
     elif htfunc_t is float64_t:
-        return duplicated_float64(values, keep)
+        return duplicated_float64(values, keep, mask=mask)
     elif htfunc_t is float32_t:
-        return duplicated_float32(values, keep)
+        return duplicated_float32(values, keep, mask=mask)
 
     elif htfunc_t is complex128_t:
-        return duplicated_complex128(values, keep)
+        return duplicated_complex128(values, keep, mask=mask)
     elif htfunc_t is complex64_t:
-        return duplicated_complex64(values, keep)
+        return duplicated_complex64(values, keep, mask=mask)
 
     else:
         raise TypeError(values.dtype)

pandas/core/algorithms.py

@@ -1043,6 +1043,10 @@ def duplicated(
     -------
     duplicated : ndarray[bool]
     """
+    if hasattr(values, "dtype") and isinstance(values.dtype, BaseMaskedDtype):
+        values = cast("BaseMaskedArray", values)
+        return htable.duplicated(values._data, keep=keep, mask=values._mask)
+
     values = _ensure_data(values)
     return htable.duplicated(values, keep=keep)
 

pandas/tests/series/methods/test_duplicated.py

@@ -2,6 +2,7 @@
 import pytest
 
 from pandas import (
+    NA,
     Categorical,
     Series,
 )
@@ -50,3 +51,28 @@ def test_duplicated_categorical_bool_na(nulls_fixture):
     result = ser.duplicated()
     expected = Series([False, False, True, True, False])
     tm.assert_series_equal(result, expected)
+
+
+@pytest.mark.parametrize(
+    "keep, vals",
+    [
+        ("last", [True, True, False]),
+        ("first", [False, True, True]),
+        (False, [True, True, True]),
+    ],
+)
+def test_duplicated_mask(keep, vals):
+    # GH#48150
+    ser = Series([1, 2, NA, NA, NA], dtype="Int64")
+    result = ser.duplicated(keep=keep)
+    expected = Series([False, False] + vals)
+    tm.assert_series_equal(result, expected)
+
+
+@pytest.mark.parametrize("keep", ["last", "first", False])
+def test_duplicated_mask_no_duplicated_na(keep):
+    # GH#48150
+    ser = Series([1, 2, NA], dtype="Int64")
+    result = ser.duplicated(keep=keep)
+    expected = Series([False, False, False])
+    tm.assert_series_equal(result, expected)