Review (jorisvandenbossche)

pandas/_libs/hashtable_class_helper.pxi.in

@@ -357,9 +357,9 @@ cdef class {{name}}HashTable(HashTable):
 
     @cython.boundscheck(False)
     @cython.wraparound(False)
-    def _unique_with_inverse(self, const {{dtype}}_t[:] values,
-                             {{name}}Vector uniques, Py_ssize_t count_prior=0,
-                             Py_ssize_t na_sentinel=-1, object na_value=None):
+    def _factorize(self, const {{dtype}}_t[:] values, {{name}}Vector uniques,
+                   Py_ssize_t count_prior=0, Py_ssize_t na_sentinel=-1,
+                   object na_value=None):
         """
         Calculate unique values and labels (no sorting); ignores all NA-values
 
@@ -437,20 +437,20 @@ cdef class {{name}}HashTable(HashTable):
                     labels[i] = count
                     count += 1
 
-        return uniques.to_array(), np.asarray(labels)
+        return np.asarray(labels)
 
-    def factorize(self, {{dtype}}_t[:] values):
+    def factorize(self, const {{dtype}}_t[:] values, Py_ssize_t na_sentinel=-1,
+                  object na_value=None):
         uniques = {{name}}Vector()
-        return self._unique_with_inverse(values, uniques=uniques)
+        labels = self._factorize(values, uniques=uniques,
+                                 na_sentinel=na_sentinel, na_value=na_value)
+        return labels, uniques.to_array()
 
     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):
-        _, labels = self._unique_with_inverse(values, uniques,
-                                              count_prior=count_prior,
-                                              na_sentinel=na_sentinel,
-                                              na_value=na_value)
-        return labels
+        return self._factorize(values, uniques, count_prior=count_prior,
+                               na_sentinel=na_sentinel, na_value=na_value)
 
     @cython.boundscheck(False)
     def get_labels_groupby(self, const {{dtype}}_t[:] values):
@@ -727,9 +727,9 @@ cdef class StringHashTable(HashTable):
 
     @cython.boundscheck(False)
     @cython.wraparound(False)
-    def _unique_with_inverse(self, ndarray[object] values,
-                             ObjectVector uniques, Py_ssize_t count_prior=0,
-                             Py_ssize_t na_sentinel=-1, object na_value=None):
+    def _factorize(self, ndarray[object] values, ObjectVector uniques,
+                   Py_ssize_t count_prior=0, Py_ssize_t na_sentinel=-1,
+                   object na_value=None):
         """
         Calculate unique values and labels (no sorting); ignores all NA-values
 
@@ -806,20 +806,20 @@ cdef class StringHashTable(HashTable):
         for i in range(count):
             uniques.append(values[uindexer[i]])
 
-        return uniques.to_array(), np.asarray(labels)
+        return np.asarray(labels)
 
-    def factorize(self, ndarray[object] values):
+    def factorize(self, ndarray[object] values, Py_ssize_t na_sentinel=-1,
+                  object na_value=None):
         uniques = ObjectVector()
-        return self._unique_with_inverse(values, uniques=uniques)
+        labels = self._factorize(values, uniques=uniques,
+                                 na_sentinel=na_sentinel, na_value=na_value)
+        return labels, uniques.to_array()
 
     def get_labels(self, ndarray[object] values, ObjectVector uniques,
                    Py_ssize_t count_prior=0, Py_ssize_t na_sentinel=-1,
                    object na_value=None):
-        _, labels = self._unique_with_inverse(values, uniques,
-                                              count_prior=count_prior,
-                                              na_sentinel=na_sentinel,
-                                              na_value=na_value)
-        return labels
+        return self._factorize(values, uniques, count_prior=count_prior,
+                               na_sentinel=na_sentinel, na_value=na_value)
 
 
 cdef class PyObjectHashTable(HashTable):
@@ -942,9 +942,9 @@ cdef class PyObjectHashTable(HashTable):
 
     @cython.boundscheck(False)
     @cython.wraparound(False)
-    def _unique_with_inverse(self, ndarray[object] values,
-                             ObjectVector uniques, Py_ssize_t count_prior=0,
-                             Py_ssize_t na_sentinel=-1, object na_value=None):
+    def _factorize(self, ndarray[object] values, ObjectVector uniques,
+                   Py_ssize_t count_prior=0, Py_ssize_t na_sentinel=-1,
+                   object na_value=None):
         """
         Calculate unique values and labels (no sorting); ignores all NA-values
 
@@ -1002,17 +1002,17 @@ cdef class PyObjectHashTable(HashTable):
                 labels[i] = count
                 count += 1
 
-        return uniques.to_array(), np.asarray(labels)
+        return np.asarray(labels)
 
-    def factorize(self, ndarray[object] values):
+    def factorize(self, ndarray[object] values, Py_ssize_t na_sentinel=-1,
+                  object na_value=None):
         uniques = ObjectVector()
-        return self._unique_with_inverse(values, uniques=uniques)
+        labels = self._factorize(values, uniques=uniques,
+                                 na_sentinel=na_sentinel, na_value=na_value)
+        return labels, uniques.to_array()
 
     def get_labels(self, ndarray[object] values, ObjectVector uniques,
                    Py_ssize_t count_prior=0, Py_ssize_t na_sentinel=-1,
                    object na_value=None):
-        _, labels = self._unique_with_inverse(values, uniques,
-                                              count_prior=count_prior,
-                                              na_sentinel=na_sentinel,
-                                              na_value=na_value)
-        return labels
+        return self._factorize(values, uniques, count_prior=count_prior,
+                               na_sentinel=na_sentinel, na_value=na_value)

pandas/core/algorithms.py

@@ -468,15 +468,13 @@ def _factorize_array(values, na_sentinel=-1, size_hint=None,
     -------
     labels, uniques : ndarray
     """
-    (hash_klass, vec_klass), values = _get_data_algo(values, _hashtables)
+    (hash_klass, _), values = _get_data_algo(values, _hashtables)
 
     table = hash_klass(size_hint or len(values))
-    uniques = vec_klass()
-    labels = table.get_labels(values, uniques, 0, na_sentinel,
-                              na_value=na_value)
+    labels, uniques = table.factorize(values, na_sentinel=na_sentinel,
+                                      na_value=na_value)
 
     labels = ensure_platform_int(labels)
-    uniques = uniques.to_array()
     return labels, uniques
 
 

pandas/tests/test_algos.py

@@ -1322,14 +1322,64 @@ def test_hashtable_unique(self, htable, tm_dtype):
             s.loc[500:502] = [np.nan, None, pd.NaT]
 
         # create duplicated selection
-        s_duplicated = s.sample(frac=3, replace=True)
+        s_duplicated = s.sample(frac=3, replace=True).reset_index(drop=True)
 
         # drop_duplicates has own cython code (hash_table_func_helper.pxi)
-        #  and is tested separately; keeps first occurrence like ht.unique()
+        # and is tested separately; keeps first occurrence like ht.unique()
         expected_unique = s_duplicated.drop_duplicates(keep='first').values
         result_unique = htable().unique(s_duplicated.values)
         tm.assert_numpy_array_equal(result_unique, expected_unique)
 
+    @pytest.mark.parametrize('na_sentinel', [-1])
+    @pytest.mark.parametrize('htable, tm_dtype', [
+        (ht.PyObjectHashTable, 'String'),
+        (ht.StringHashTable, 'String'),
+        (ht.Float64HashTable, 'Float'),
+        (ht.Int64HashTable, 'Int'),
+        (ht.UInt64HashTable, 'UInt')])
+    def test_hashtable_factorize(self, htable, tm_dtype, na_sentinel):
+        # output of maker has guaranteed unique elements
+        maker = getattr(tm, 'make' + tm_dtype + 'Index')
+        s = Series(maker(1000))
+        if htable == ht.Float64HashTable:
+            # add NaN for float column
+            s.loc[500] = np.nan
+        elif htable == ht.PyObjectHashTable:
+            # use different NaN types for object column
+            s.loc[500:502] = [np.nan, None, pd.NaT]
+
+        # create duplicated selection
+        idx_duplicated = pd.Series(s.index).sample(frac=3, replace=True)
+        s_duplicated = s[idx_duplicated.values].reset_index(drop=True)
+        na_mask = s_duplicated.isna().values
+
+        result_inverse, result_unique = htable().factorize(s_duplicated.values)
+
+        # drop_duplicates has own cython code (hash_table_func_helper.pxi)
+        # and is tested separately; keeps first occurrence like ht.unique()
+        expected_unique = s_duplicated.dropna().drop_duplicates(keep='first')
+        expected_unique = expected_unique.values
+        tm.assert_numpy_array_equal(result_unique, expected_unique)
+
+        # ignore NaNs for calculating inverse
+        _, values2unique, unique2values = np.unique(idx_duplicated[~na_mask],
+                                                    return_inverse=True,
+                                                    return_index=True)
+        expected_inverse = np.ones(s_duplicated.shape,
+                                   dtype=np.intp) * na_sentinel
+
+        # np.unique yields a __SORTED__ list of uniques, and values2unique
+        # resp. unique2values are relative to this order. To restore the
+        # original order, we argsort values2unique, because values2unique would
+        # be ordered if np.unique had not sorted implicitly. The first argsort
+        # gives the permutation from values2unique to its sorted form, but we
+        # need the inverse permutation (the map from the unsorted uniques to
+        # values2unique, from which we can continue with unique2values).
+        # This inversion (as a permutation) is achieved by the second argsort.
+        inverse_no_na = np.argsort(np.argsort(values2unique))[unique2values]
+        expected_inverse[~na_mask] = inverse_no_na
+        tm.assert_numpy_array_equal(result_inverse, expected_inverse)
+
 
 def test_quantile():
     s = Series(np.random.randn(100))