ENH: Support mask in Int64Factorizer (#49549)

* ENH: Avoid object path for merge and masked arrays

* ENH: Support mask in Int64Factorizer

* Fix mypy

* Add gh ref

* Fix tests

* Fix tests

* Fix tests

Author: phofl

pandas/_libs/hashtable.pyi

@@ -36,6 +36,7 @@ class Int64Factorizer(Factorizer):
         sort: bool = ...,
         na_sentinel=...,
         na_value=...,
+        mask=...,
     ) -> npt.NDArray[np.intp]: ...
 
 class Int64Vector:
@@ -137,6 +138,7 @@ class HashTable:
         count_prior: int = ...,
         na_sentinel: int = ...,
         na_value: object = ...,
+        mask=...,
     ) -> npt.NDArray[np.intp]: ...
     def unique(
         self,

pandas/_libs/hashtable.pyx

@@ -129,7 +129,7 @@ cdef class Int64Factorizer(Factorizer):
         self.uniques = Int64Vector()
 
     def factorize(self, const int64_t[:] values,
-                  na_sentinel=-1, na_value=None) -> np.ndarray:
+                  na_sentinel=-1, na_value=None, object mask=None) -> np.ndarray:
         """
         Returns
         -------
@@ -152,6 +152,6 @@ cdef class Int64Factorizer(Factorizer):
             self.uniques = uniques
         labels = self.table.get_labels(values, self.uniques,
                                        self.count, na_sentinel,
-                                       na_value=na_value)
+                                       na_value=na_value, mask=mask)
         self.count = len(self.uniques)
         return labels

pandas/_libs/hashtable_class_helper.pxi.in

@@ -821,11 +821,11 @@ cdef class {{name}}HashTable(HashTable):
 
     def get_labels(self, const {{dtype}}_t[:] values, {{name}}Vector uniques,
                    Py_ssize_t count_prior=0, Py_ssize_t na_sentinel=-1,
-                   object na_value=None):
+                   object na_value=None, object mask=None):
         # -> np.ndarray[np.intp]
         _, labels = self._unique(values, uniques, count_prior=count_prior,
                                  na_sentinel=na_sentinel, na_value=na_value,
-                                 ignore_na=True, return_inverse=True)
+                                 ignore_na=True, return_inverse=True, mask=mask)
         return labels
 
     {{if dtype == 'int64'}}

pandas/tests/test_algos.py

@@ -222,6 +222,27 @@ def test_factorize_nan(self):
             tm.assert_numpy_array_equal(pd.isna(key), expected == na_sentinel)
             tm.assert_numpy_array_equal(ids, expected)
 
+    def test_factorizer_with_mask(self):
+        # GH#49549
+        data = np.array([1, 2, 3, 1, 1, 0], dtype="int64")
+        mask = np.array([False, False, False, False, False, True])
+        rizer = ht.Int64Factorizer(len(data))
+        result = rizer.factorize(data, mask=mask)
+        expected = np.array([0, 1, 2, 0, 0, -1], dtype=np.intp)
+        tm.assert_numpy_array_equal(result, expected)
+        expected_uniques = np.array([1, 2, 3], dtype="int64")
+        tm.assert_numpy_array_equal(rizer.uniques.to_array(), expected_uniques)
+
+    def test_factorizer_object_with_nan(self):
+        # GH#49549
+        data = np.array([1, 2, 3, 1, np.nan])
+        rizer = ht.ObjectFactorizer(len(data))
+        result = rizer.factorize(data.astype(object))
+        expected = np.array([0, 1, 2, 0, -1], dtype=np.intp)
+        tm.assert_numpy_array_equal(result, expected)
+        expected_uniques = np.array([1, 2, 3], dtype=object)
+        tm.assert_numpy_array_equal(rizer.uniques.to_array(), expected_uniques)
+
     @pytest.mark.parametrize(
         "data, expected_codes, expected_uniques",
         [