WIP

Author: rhshadrach

pandas/_libs/hashtable_class_helper.pxi.in

@@ -560,14 +560,13 @@ cdef class {{name}}HashTable(HashTable):
             The labels from values to uniques
         """
         cdef:
-            Py_ssize_t i, idx, count = count_prior, n = len(values), na_index = -1, na_code = -1
+            Py_ssize_t i, idx, count = count_prior, n = len(values)
             intp_t[::1] labels
             int ret = 0
             {{c_type}} val, na_value2
             khiter_t k
             {{name}}VectorData *ud
-            bint use_na_value, use_mask, seen_na = False
-            uint8_t is_na
+            bint use_na_value, use_mask
             uint8_t[:] mask_values
 
         if return_inverse:
@@ -593,43 +592,22 @@ cdef class {{name}}HashTable(HashTable):
             for i in range(n):
                 val = {{to_c_type}}(values[i])
 
-                if use_mask:
-                    is_na = mask_values[i]
-                else:
-                    is_na = (
-                       is_nan_{{c_type}}(val) or
-                       (use_na_value and are_equivalent_{{c_type}}(val, na_value2))
-                    )
-
-                if is_na:
-                    if ignore_na:
-                        # if missing values do not count as unique values (i.e. if
-                        # ignore_na is True), skip the hashtable entry for them,
-                        # and replace the corresponding label with na_sentinel
+                if ignore_na and use_mask:
+                    if mask_values[i]:
                         labels[i] = na_sentinel
                         continue
-
-                    if not seen_na:
-                        if needs_resize(ud):
-                            with gil:
-                                if uniques.external_view_exists:
-                                    raise ValueError("external reference to "
-                                                     "uniques held, but "
-                                                     "Vector.resize() needed")
-                                uniques.resize()
-                        append_data_{{dtype}}(ud, val)
-                        na_index = i
-                        if return_inverse:
-                            labels[i] = count
-                            na_code = count
-                            count += 1
-                        seen_na = True
-                    if return_inverse:
-                        idx = na_code
-                        labels[i] = idx
+                elif ignore_na and (
+                   is_nan_{{c_type}}(val) or
+                   (use_na_value and are_equivalent_{{c_type}}(val, na_value2))
+                ):
+                    # if missing values do not count as unique values (i.e. if
+                    # ignore_na is True), skip the hashtable entry for them,
+                    # and replace the corresponding label with na_sentinel
+                    labels[i] = na_sentinel
                     continue
 
                 k = kh_get_{{dtype}}(self.table, val)
+
                 if k == self.table.n_buckets:
                     # k hasn't been seen yet
                     k = kh_put_{{dtype}}(self.table, val, &ret)
@@ -653,8 +631,8 @@ cdef class {{name}}HashTable(HashTable):
                     labels[i] = idx
 
         if return_inverse:
-            return uniques.to_array(), na_index, labels.base  # .base -> underlying ndarray
-        return uniques.to_array(), na_index
+            return uniques.to_array(), labels.base  # .base -> underlying ndarray
+        return uniques.to_array()
 
     def unique(self, const {{dtype}}_t[:] values, bint return_inverse=False):
         """
@@ -719,7 +697,7 @@ cdef class {{name}}HashTable(HashTable):
                    Py_ssize_t count_prior=0, Py_ssize_t na_sentinel=-1,
                    object na_value=None):
         # -> np.ndarray[np.intp]
-        _, _, labels = self._unique(values, uniques, count_prior=count_prior,
+        _, labels = self._unique(values, uniques, count_prior=count_prior,
                                  na_sentinel=na_sentinel, na_value=na_value,
                                  ignore_na=True, return_inverse=True)
         return labels
@@ -1032,8 +1010,8 @@ cdef class StringHashTable(HashTable):
             uniques.append(values[uindexer[i]])
 
         if return_inverse:
-            return uniques.to_array(), -1, labels.base  # .base -> underlying ndarray
-        return uniques.to_array(), -1
+            return uniques.to_array(), labels.base  # .base -> underlying ndarray
+        return uniques.to_array()
 
     def unique(self, ndarray[object] values, bint return_inverse=False):
         """
@@ -1096,7 +1074,7 @@ cdef class StringHashTable(HashTable):
                    Py_ssize_t count_prior=0, Py_ssize_t na_sentinel=-1,
                    object na_value=None):
         # -> np.ndarray[np.intp]
-        _, _, labels = self._unique(values, uniques, count_prior=count_prior,
+        _, labels = self._unique(values, uniques, count_prior=count_prior,
                                  na_sentinel=na_sentinel, na_value=na_value,
                                  ignore_na=True, return_inverse=True)
         return labels
@@ -1244,13 +1222,12 @@ cdef class PyObjectHashTable(HashTable):
             The labels from values to uniques
         """
         cdef:
-            Py_ssize_t i, idx, count = count_prior, n = len(values), na_index = -1
+            Py_ssize_t i, idx, count = count_prior, n = len(values)
             intp_t[::1] labels
             int ret = 0
             object val
             khiter_t k
-            bint use_na_value, seen_na = False, is_na
-            intp_t na_code
+            bint use_na_value
 
         if return_inverse:
             labels = np.empty(n, dtype=np.intp)
@@ -1260,26 +1237,14 @@ cdef class PyObjectHashTable(HashTable):
             val = values[i]
             hash(val)
 
-            is_na = checknull(val) or (use_na_value and val == na_value)
-
-            if is_na:
-                if ignore_na:
-                    # if missing values do not count as unique values (i.e. if
-                    # ignore_na is True), skip the hashtable entry for them, and
-                    # replace the corresponding label with na_sentinel
-                    labels[i] = na_sentinel
-                    continue
-
-                if not seen_na:
-                    seen_na = True
-                    uniques.append(val)
-                    na_index = i
-                    if return_inverse:
-                        labels[i] = count
-                        na_code = count
-                        count += 1
-                elif return_inverse:
-                    labels[i] = na_code
+            if ignore_na and (
+                checknull(val)
+                or (use_na_value and val == na_value)
+            ):
+                # if missing values do not count as unique values (i.e. if
+                # ignore_na is True), skip the hashtable entry for them, and
+                # replace the corresponding label with na_sentinel
+                labels[i] = na_sentinel
                 continue
 
             k = kh_get_pymap(self.table, <PyObject*>val)
@@ -1298,8 +1263,8 @@ cdef class PyObjectHashTable(HashTable):
                 labels[i] = idx
 
         if return_inverse:
-            return uniques.to_array(), na_index, labels.base  # .base -> underlying ndarray
-        return uniques.to_array(), na_index
+            return uniques.to_array(), labels.base  # .base -> underlying ndarray
+        return uniques.to_array()
 
     def unique(self, ndarray[object] values, bint return_inverse=False):
         """
@@ -1362,7 +1327,7 @@ cdef class PyObjectHashTable(HashTable):
                    Py_ssize_t count_prior=0, Py_ssize_t na_sentinel=-1,
                    object na_value=None):
         # -> np.ndarray[np.intp]
-        _, _, labels = self._unique(values, uniques, count_prior=count_prior,
+        _, labels = self._unique(values, uniques, count_prior=count_prior,
                                  na_sentinel=na_sentinel, na_value=na_value,
                                  ignore_na=True, return_inverse=True)
         return labels

pandas/core/algorithms.py

@@ -412,7 +412,7 @@ def unique(values):
     htable, values = _get_hashtable_algo(values)
 
     table = htable(len(values))
-    uniques, _ = table.unique(values)
+    uniques = table.unique(values)
     uniques = _reconstruct_data(uniques, original.dtype, original)
     return uniques
 
@@ -497,13 +497,13 @@ def f(c, v):
 
 
 def factorize_array(
-    arr: ArrayLike,
+    values: np.ndarray,
     na_sentinel: int = -1,
     size_hint: int | None = None,
     na_value=None,
     mask: npt.NDArray[np.bool_] | None = None,
     ignore_na: bool = True,
-) -> tuple[npt.NDArray[np.intp], ArrayLike]:
+) -> tuple[npt.NDArray[np.intp], np.ndarray]:
     """
     Factorize a numpy array to codes and uniques.
 
@@ -530,11 +530,6 @@ def factorize_array(
     codes : ndarray[np.intp]
     uniques : ndarray
     """
-    if is_extension_array_dtype(arr):
-        values, na_value = arr._values_for_factorize()
-    else:
-        values = arr
-
     original = values
     if values.dtype.kind in ["m", "M"]:
         # _get_hashtable_algo will cast dt64/td64 to i8 via _ensure_data, so we
@@ -546,7 +541,7 @@ def factorize_array(
     hash_klass, values = _get_hashtable_algo(values)
 
     table = hash_klass(size_hint or len(values))
-    uniques, na_index, codes = table.factorize(
+    uniques, codes = table.factorize(
         values,
         na_sentinel=na_sentinel,
         na_value=na_value,
@@ -558,7 +553,7 @@ def factorize_array(
     uniques = _reconstruct_data(uniques, original.dtype, original)
 
     codes = ensure_platform_int(codes)
-    return codes, uniques, na_index
+    return codes, uniques
 
 
 @doc(
@@ -744,7 +739,7 @@ def factorize(
         )
     else:
         values = np.asarray(values)  # convert DTA/TDA/MultiIndex
-        codes, uniques, _ = factorize_array(
+        codes, uniques = factorize_array(
             values,
             na_sentinel=na_sentinel,
             size_hint=size_hint,

pandas/core/arrays/arrow/array.py

@@ -108,7 +108,8 @@ def copy(self: ArrowExtensionArrayT) -> ArrowExtensionArrayT:
     def factorize(
         self, na_sentinel: int = -1, ignore_na=True
     ) -> tuple[np.ndarray, ExtensionArray]:
-        encoded = self._data.dictionary_encode()
+        null_encoding = "mask" if ignore_na else "encode"
+        encoded = self._data.dictionary_encode(null_encoding=null_encoding)
         indices = pa.chunked_array(
             [c.indices for c in encoded.chunks], type=encoded.type.index_type
         ).to_pandas()

pandas/core/arrays/base.py

@@ -1045,12 +1045,11 @@ def factorize(
         #    Complete control over factorization.
         arr, na_value = self._values_for_factorize()
 
-        codes, uniques, _ = factorize_array(
+        codes, uniques = factorize_array(
             arr, na_sentinel=na_sentinel, na_value=na_value, ignore_na=ignore_na
         )
 
         uniques_ea = self._from_factorized(uniques, self)
-        # TODO: Use na_index here?
         return codes, uniques_ea
 
     _extension_array_shared_docs[

pandas/core/arrays/masked.py

@@ -874,17 +874,28 @@ def factorize(
         arr = self._data
         mask = self._mask
 
-        codes, uniques, na_index = factorize_array(
-            arr, na_sentinel=na_sentinel, mask=mask, ignore_na=ignore_na
+        codes, uniques = factorize_array(
+            arr, na_sentinel=na_sentinel, mask=mask, ignore_na=True
         )
 
         # check that factorize_array correctly preserves dtype.
         assert uniques.dtype == self.dtype.numpy_dtype, (uniques.dtype, self.dtype)
 
-        mask = np.zeros(len(uniques), dtype=bool)
-        if na_index >= 0:
-            mask[na_index] = True
-        uniques_ea = type(self)(uniques, mask)
+        size = len(uniques) if ignore_na else len(uniques) + 1
+        uniques_mask = np.zeros(size, dtype=bool)
+        if not ignore_na:
+            na_index = mask.argmax()
+            if mask[na_index]:
+                # Insert na with the proper code
+                # TODO: This only works with na_sentinel being -1
+                na_code = codes[:na_index].argmax() + 1
+                codes[codes >= na_code] += 1
+                codes[codes == na_sentinel] = na_code
+                # dummy value for uniques
+                uniques = np.insert(uniques, na_code, 0)
+                uniques_mask[na_code] = True
+        uniques_ea = type(self)(uniques, uniques_mask)
+
         return codes, uniques_ea
 
     @doc(ExtensionArray._values_for_argsort)

pandas/core/arrays/string_.py

@@ -382,6 +382,12 @@ def _values_for_factorize(self):
         arr[mask] = -1
         return arr, -1
 
+    @classmethod
+    def _from_factorized(cls, values, original):
+        assert values.dtype == original._ndarray.dtype
+        values[values == -1] = None
+        return original._from_backing_data(values)
+
     def __setitem__(self, key, value):
         value = extract_array(value, extract_numpy=True)
         if isinstance(value, type(self)):