Repository: pandas-dev/pandas
Pull request: ENH: Support mask in unique
- [x] closes #30037 (Replace xxxx with the Github issue number)
- [x] [Tests added and passed](https://pandas.pydata.org/pandas-docs/dev/development/contributing_codebase.html#writing-tests) if fixing a bug or adding a new feature
- [x] All [code checks passed](https://pandas.pydata.org/pandas-docs/dev/development/contributing_codebase.html#pre-commit).
- [x] Added [type annotations](https://pandas.pydata.org/pandas-docs/dev/development/contributing_codebase.html#type-hints) to new arguments/methods/functions.
- [x] Added an entry in the latest `doc/source/whatsnew/vX.X.X.rst` file if fixing a bug or adding a new feature.


cc @jorisvandenbossche 

Problem statement:
ENH: add mask-aware implementation of factorize algos
Now we start to have mask-based dtypes/arrays (integer, boolean), we should also look into making our algos work with such masked arrays. An example for which we could explore this is `factorize` / `unique`. 

Currently, BooleanArray and IntegerArray need to convert their masked array into a single numpy array using a certain "NA sentinel" that is specified so the algo can recognize this sentinel. This happens through the `ExtensionArray._values_for_factorize`, which returns a (numpy array, NA sentinel) tuple. 
In practice this means that the boolean array is converted to integer (with NA as -1), and IntegerArray is converted to float array with NA as NaN, so the algos can handle this.

We should look into:

- Can we adapt or make a specific version of the unique/factorize hashtable class that takes a mask instead of a NA sentinel
- We could then have a variant of `ExtensionArray._values_for_factorize` that then returns (array, mask) instead of (array, NA).

Hints from issue discussion:
> and IntegerArray is converted to float array with NA as NaN

On master, IntegerArray uses an object-dtype ndarray, with NA as NaN. We probably need object to avoid floating-point imprecision for large integers.

```python
In [3]: a = pd.array([1, 2, None])

In [4]: a._values_for_factorize()
Out[4]: (array([1, 2, nan], dtype=object), nan)

```

I suppose we could avoid that object-dtype when there are no missing values, but that's straying from the original issue.
So #33064 added mask support to `HashTable.factorize`, which is used in the `factorize` methods. 

Something that is still missing is eg the `unique` support, which for ExtensionArrays has the following base implementation:

https://github.com/pandas-dev/pandas/blob/be8d1ec8808822ef606c12584440f001ad76de22/pandas/core/arrays/base.py#L862-L863

So going through an object cast. And this is also used for the masked arrays.

The actual HashTable implementation actually already supports masks, though, so I think it is mostly a matter of exposing a mask argument in `HashTable.unique` (just as we did for `HashTable.factorize`), and having a way to call this from the masked array class (for factorize, that goes through `factorize_array` in algorithms.py).

The HashTable classes also have other methods (`lookup`, `map_locations`, `get_labels_groupby`), but those can covered in separate issues I think.

Commit message: Refactor unique

Patch to review:
diff --git a/asv_bench/benchmarks/hash_functions.py b/asv_bench/benchmarks/hash_functions.py
index d9a291dc27125..da752b902b4fd 100644
--- a/asv_bench/benchmarks/hash_functions.py
+++ b/asv_bench/benchmarks/hash_functions.py
@@ -39,6 +39,21 @@ def time_unique(self, exponent):
         pd.unique(self.a2)
 
 
+class Unique:
+    params = ["Int64", "Float64"]
+    param_names = ["dtype"]
+
+    def setup(self, dtype):
+        self.ser = pd.Series(([1, pd.NA, 2] + list(range(100_000))) * 3, dtype=dtype)
+        self.ser_unique = pd.Series(list(range(300_000)) + [pd.NA], dtype=dtype)
+
+    def time_unique_with_duplicates(self, exponent):
+        pd.unique(self.ser)
+
+    def time_unique(self, exponent):
+        pd.unique(self.ser_unique)
+
+
 class NumericSeriesIndexing:
 
     params = [
diff --git a/pandas/_libs/hashtable_class_helper.pxi.in b/pandas/_libs/hashtable_class_helper.pxi.in
index 8a2b9c2f77627..6730b38ca442d 100644
--- a/pandas/_libs/hashtable_class_helper.pxi.in
+++ b/pandas/_libs/hashtable_class_helper.pxi.in
@@ -521,7 +521,7 @@ cdef class {{name}}HashTable(HashTable):
     def _unique(self, const {{dtype}}_t[:] values, {{name}}Vector uniques,
                 Py_ssize_t count_prior=0, Py_ssize_t na_sentinel=-1,
                 object na_value=None, bint ignore_na=False,
-                object mask=None, bint return_inverse=False):
+                object mask=None, bint return_inverse=False, bint use_result_mask=False):
         """
         Calculate unique values and labels (no sorting!)
 
@@ -551,6 +551,9 @@ cdef class {{name}}HashTable(HashTable):
         return_inverse : bool, default False
             Whether the mapping of the original array values to their location
             in the vector of uniques should be returned.
+        use_result_mask: bool, default False
+            Whether to create a result mask for the unique values. Not supported
+            with return_inverse=True.
 
         Returns
         -------
@@ -566,7 +569,9 @@ cdef class {{name}}HashTable(HashTable):
             {{c_type}} val, na_value2
             khiter_t k
             {{name}}VectorData *ud
-            bint use_na_value, use_mask
+            UInt8Vector result_mask
+            UInt8VectorData *rmd
+            bint use_na_value, use_mask, seen_na = False
             uint8_t[:] mask_values
 
         if return_inverse:
@@ -574,6 +579,14 @@ cdef class {{name}}HashTable(HashTable):
         ud = uniques.data
         use_na_value = na_value is not None
         use_mask = mask is not None
+        if not use_mask and use_result_mask:
+            raise NotImplementedError  # pragma: no cover
+
+        if use_result_mask and return_inverse:
+            raise NotImplementedError  # pragma: no cover
+
+        result_mask = UInt8Vector()
+        rmd = result_mask.data
 
         if use_mask:
             mask_values = mask.view("uint8")
@@ -605,6 +618,27 @@ cdef class {{name}}HashTable(HashTable):
                     # and replace the corresponding label with na_sentinel
                     labels[i] = na_sentinel
                     continue
+                elif not ignore_na and use_result_mask:
+                    if mask_values[i]:
+                        if seen_na:
+                            continue
+
+                        seen_na = True
+                        if needs_resize(ud):
+                            with gil:
+                                if uniques.external_view_exists:
+                                    raise ValueError("external reference to "
+                                                     "uniques held, but "
+                                                     "Vector.resize() needed")
+                                uniques.resize()
+                                if result_mask.external_view_exists:
+                                    raise ValueError("external reference to "
+                                                     "result_mask held, but "
+                                                     "Vector.resize() needed")
+                                result_mask.resize()
+                        append_data_{{dtype}}(ud, val)
+                        append_data_uint8(rmd, 1)
+                        continue
 
                 k = kh_get_{{dtype}}(self.table, val)
 
@@ -619,7 +653,16 @@ cdef class {{name}}HashTable(HashTable):
                                                  "uniques held, but "
                                                  "Vector.resize() needed")
                             uniques.resize()
+                            if use_result_mask:
+                                if result_mask.external_view_exists:
+                                    raise ValueError("external reference to "
+                                                     "result_mask held, but "
+                                                     "Vector.resize() needed")
+                                result_mask.resize()
                     append_data_{{dtype}}(ud, val)
+                    if use_result_mask:
+                        append_data_uint8(rmd, 0)
+
                     if return_inverse:
                         self.table.vals[k] = count
                         labels[i] = count
@@ -632,9 +675,11 @@ cdef class {{name}}HashTable(HashTable):
 
         if return_inverse:
             return uniques.to_array(), labels.base  # .base -> underlying ndarray
+        if use_result_mask:
+            return uniques.to_array(), result_mask.to_array()
         return uniques.to_array()
 
-    def unique(self, const {{dtype}}_t[:] values, bint return_inverse=False):
+    def unique(self, const {{dtype}}_t[:] values, bint return_inverse=False, object mask=None):
         """
         Calculate unique values and labels (no sorting!)
 
@@ -652,10 +697,14 @@ cdef class {{name}}HashTable(HashTable):
             Unique values of input, not sorted
         labels : ndarray[intp_t] (if return_inverse)
             The labels from values to uniques
+        mask : ndarray[bool], optional
+            If not None, the mask is used as indicator for missing values
+            (True = missing, False = valid) instead of `na_value` or
         """
         uniques = {{name}}Vector()
+        use_result_mask = True if mask is not None else False
         return self._unique(values, uniques, ignore_na=False,
-                            return_inverse=return_inverse)
+                            return_inverse=return_inverse, mask=mask, use_result_mask=use_result_mask)
 
     def factorize(self, const {{dtype}}_t[:] values, Py_ssize_t na_sentinel=-1,
                   object na_value=None, object mask=None):
@@ -1013,7 +1062,7 @@ cdef class StringHashTable(HashTable):
             return uniques.to_array(), labels.base  # .base -> underlying ndarray
         return uniques.to_array()
 
-    def unique(self, ndarray[object] values, bint return_inverse=False):
+    def unique(self, ndarray[object] values, bint return_inverse=False, object mask=None):
         """
         Calculate unique values and labels (no sorting!)
 
@@ -1024,6 +1073,8 @@ cdef class StringHashTable(HashTable):
         return_inverse : bool, default False
             Whether the mapping of the original array values to their location
             in the vector of uniques should be returned.
+        mask : ndarray[bool], optional
+            Not yet implemented for StringHashTable
 
         Returns
         -------
@@ -1266,7 +1317,7 @@ cdef class PyObjectHashTable(HashTable):
             return uniques.to_array(), labels.base  # .base -> underlying ndarray
         return uniques.to_array()
 
-    def unique(self, ndarray[object] values, bint return_inverse=False):
+    def unique(self, ndarray[object] values, bint return_inverse=False, object mask=None):
         """
         Calculate unique values and labels (no sorting!)
 
@@ -1277,6 +1328,8 @@ cdef class PyObjectHashTable(HashTable):
         return_inverse : bool, default False
             Whether the mapping of the original array values to their location
             in the vector of uniques should be returned.
+        mask : ndarray[bool], optional
+            Not yet implemented for PyObjectHashTable
 
         Returns
         -------
diff --git a/pandas/core/algorithms.py b/pandas/core/algorithms.py
index a4736c2a141a5..daa5ebd7fc28e 100644
--- a/pandas/core/algorithms.py
+++ b/pandas/core/algorithms.py
@@ -322,6 +322,8 @@ def unique(values):
     Parameters
     ----------
     values : 1d array-like
+    mask : 1d bool array
+        Mask of the values.
 
     Returns
     -------
@@ -404,6 +406,11 @@ def unique(values):
     >>> pd.unique([("a", "b"), ("b", "a"), ("a", "c"), ("b", "a")])
     array([('a', 'b'), ('b', 'a'), ('a', 'c')], dtype=object)
     """
+    return _unique(values)
+
+
+def _unique(values, mask: npt.NDArray[np.bool_] | None = None):
+    """See algorithms.unique for docs"""
     values = _ensure_arraylike(values)
 
     if is_extension_array_dtype(values.dtype):
@@ -414,9 +421,15 @@ def unique(values):
     htable, values = _get_hashtable_algo(values)
 
     table = htable(len(values))
-    uniques = table.unique(values)
-    uniques = _reconstruct_data(uniques, original.dtype, original)
-    return uniques
+    if mask is None:
+        uniques = table.unique(values)
+        uniques = _reconstruct_data(uniques, original.dtype, original)
+        return uniques
+
+    else:
+        uniques, mask = table.unique(values, mask=mask)
+        uniques = _reconstruct_data(uniques, original.dtype, original)
+        return uniques, mask.astype("bool")
 
 
 unique1d = unique
diff --git a/pandas/core/arrays/masked.py b/pandas/core/arrays/masked.py
index 128c7e44f5075..c558f70caa3e8 100644
--- a/pandas/core/arrays/masked.py
+++ b/pandas/core/arrays/masked.py
@@ -851,6 +851,17 @@ def copy(self: BaseMaskedArrayT) -> BaseMaskedArrayT:
         mask = mask.copy()
         return type(self)(data, mask, copy=False)
 
+    def unique(self: BaseMaskedArrayT) -> BaseMaskedArrayT:
+        """
+        Compute the BaseMaskedArray of unique values.
+
+        Returns
+        -------
+        uniques : BaseMaskedArray
+        """
+        uniques, mask = algos._unique(self._data, self._mask)
+        return type(self)(uniques, mask, copy=False)
+
     @doc(ExtensionArray.searchsorted)
     def searchsorted(
         self,
diff --git a/pandas/tests/test_algos.py b/pandas/tests/test_algos.py
index def63c552e059..fd617dcb4565e 100644
--- a/pandas/tests/test_algos.py
+++ b/pandas/tests/test_algos.py
@@ -834,6 +834,13 @@ def test_do_not_mangle_na_values(self, unique_nulls_fixture, unique_nulls_fixtur
         assert a[0] is unique_nulls_fixture
         assert a[1] is unique_nulls_fixture2
 
+    def test_unique_masked(self, any_numeric_ea_dtype):
+        # GH#48019
+        ser = Series([1, pd.NA, 2] * 3, dtype=any_numeric_ea_dtype)
+        result = pd.unique(ser)
+        expected = pd.array([1, pd.NA, 2], dtype=any_numeric_ea_dtype)
+        tm.assert_extension_array_equal(result, expected)
+
 
 class TestIsin:
     def test_invalid(self):


Code review comments:
diff_hunk: @@ -851,6 +851,17 @@ def copy(self: BaseMaskedArrayT) -> BaseMaskedArrayT:
         mask = mask.copy()
         return type(self)(data, mask, copy=False)
 
+    def unique(self: BaseMaskedArrayT) -> BaseMaskedArrayT:
+        """
+        Compute the BaseMaskedArray of unique values.
+
+        Returns
+        -------
+        uniques : BaseMaskedArray
+        """
+        uniques, mask = algos._unique(self._data, self._mask)
path: pandas/core/arrays/masked.py
text: @user1:
nitpick, can you rename the function to something non-private, e.g. unique_with_mask

@author:
Sure, done