Example with _from_factorize

Author: TomAugspurger

pandas/core/arrays/base.py

@@ -77,7 +77,21 @@ def _constructor_from_sequence(cls, scalars):
         raise AbstractMethodError(cls)
 
     @classmethod
-    def _constructor_from_simple_ndarray(cls, values, instance):
+    def _from_factorized(cls, values, original):
+        """Reconstruct an ExtensionArray after factorization.
+
+        Parameters
+        ----------
+        values : ndarray
+            An integer ndarray with the factorized values.
+        original : ExtensionArray
+            The original ndarray that was factorized.
+
+        See Also
+        --------
+        pandas.factorize
+        ExtensionArray.factorize
+        """
         raise AbstractMethodError(cls)
 
     # ------------------------------------------------------------------------
@@ -305,7 +319,16 @@ def unique(self):
         uniques = unique(self.astype(object))
         return self._constructor_from_sequence(uniques)
 
-    def _simple_ndarray(self):
+    def _values_for_factorize(self):
+        """Return an array suitable for factorization.
+
+        Returns
+        -------
+        ndarray
+            An array suitable for factoraization. This should maintain order
+            and be a supported dtype.
+
+        """
         return self.astype(object)
 
     def factorize(self, na_sentinel=-1):
@@ -337,17 +360,17 @@ def factorize(self, na_sentinel=-1):
         -----
         :meth:`pandas.factorize` offers a `sort` keyword as well.
         """
-        # Implementor note: make sure to exclude missing values from your
-        # `uniques`. It should only contain non-NA values.
+        # Implementor notes: There are two options for overriding the
+        # behavior of `factorize`: here and `_values_for_factorize`.
         from pandas.core.algorithms import _factorize_array
 
         mask = self.isna()
-        arr = self._simple_ndarray()
+        arr = self._values_for_factorize()
         arr[mask] = np.nan
 
         labels, uniques = _factorize_array(arr, check_nulls=True,
                                            na_sentinel=na_sentinel)
-        uniques = self._constructor_from_simple_ndarray(uniques, instance=arr)
+        uniques = self._from_factorized(uniques, arr)
         return labels, uniques
 
     # ------------------------------------------------------------------------

pandas/tests/extension/decimal/array.py

@@ -36,6 +36,10 @@ def __init__(self, values):
     def _constructor_from_sequence(cls, scalars):
         return cls(scalars)
 
+    @classmethod
+    def _from_factorized(cls, values, original):
+        return cls(values)
+
     def __getitem__(self, item):
         if isinstance(item, numbers.Integral):
             return self.values[item]

pandas/tests/extension/json/array.py

@@ -7,7 +7,6 @@
 
 import numpy as np
 
-import pandas as pd
 from pandas.core.dtypes.base import ExtensionDtype
 from pandas.core.arrays import ExtensionArray
 
@@ -38,6 +37,10 @@ def __init__(self, values):
     def _constructor_from_sequence(cls, scalars):
         return cls(scalars)
 
+    @classmethod
+    def _from_factorized(cls, values, original):
+        return cls([collections.UserDict(x) for x in values if x != ()])
+
     def __getitem__(self, item):
         if isinstance(item, numbers.Integral):
             return self.data[item]
@@ -105,20 +108,9 @@ def _concat_same_type(cls, to_concat):
         data = list(itertools.chain.from_iterable([x.data for x in to_concat]))
         return cls(data)
 
-    def factorize(self, na_sentinel=-1):
+    def _values_for_factorize(self):
         frozen = tuple(tuple(x.items()) for x in self)
-        labels, uniques = pd.factorize(frozen)
-
-        # fixup NA
-        if self.isna().any():
-            na_code = labels[self.isna()][0]
-
-            labels[labels == na_code] = na_sentinel
-            labels[labels > na_code] -= 1
-
-        uniques = JSONArray([collections.UserDict(x)
-                             for x in uniques if x != ()])
-        return labels, uniques
+        return np.array(frozen, dtype=object)
 
 
 def make_data():