REF: de-duplicate factorize and duplicated code (pandas-dev#32216)

Author: jbrockmendel

asv_bench/benchmarks/algorithms.py

@@ -31,83 +31,62 @@ def time_maybe_convert_objects(self):
 
 class Factorize:
 
-    params = [[True, False], ["int", "uint", "float", "string"]]
-    param_names = ["sort", "dtype"]
-
-    def setup(self, sort, dtype):
-        N = 10 ** 5
-        data = {
-            "int": pd.Int64Index(np.arange(N).repeat(5)),
-            "uint": pd.UInt64Index(np.arange(N).repeat(5)),
-            "float": pd.Float64Index(np.random.randn(N).repeat(5)),
-            "string": tm.makeStringIndex(N).repeat(5),
-        }
-        self.idx = data[dtype]
-
-    def time_factorize(self, sort, dtype):
-        self.idx.factorize(sort=sort)
-
-
-class FactorizeUnique:
-
-    params = [[True, False], ["int", "uint", "float", "string"]]
-    param_names = ["sort", "dtype"]
+    params = [
+        [True, False],
+        [True, False],
+        ["int", "uint", "float", "string", "datetime64[ns]", "datetime64[ns, tz]"],
+    ]
+    param_names = ["unique", "sort", "dtype"]
 
-    def setup(self, sort, dtype):
+    def setup(self, unique, sort, dtype):
         N = 10 ** 5
         data = {
             "int": pd.Int64Index(np.arange(N)),
             "uint": pd.UInt64Index(np.arange(N)),
-            "float": pd.Float64Index(np.arange(N)),
+            "float": pd.Float64Index(np.random.randn(N)),
             "string": tm.makeStringIndex(N),
-        }
-        self.idx = data[dtype]
-        assert self.idx.is_unique
-
-    def time_factorize(self, sort, dtype):
+            "datetime64[ns]": pd.date_range("2011-01-01", freq="H", periods=N),
+            "datetime64[ns, tz]": pd.date_range(
+                "2011-01-01", freq="H", periods=N, tz="Asia/Tokyo"
+            ),
+        }[dtype]
+        if not unique:
+            data = data.repeat(5)
+        self.idx = data
+
+    def time_factorize(self, unique, sort, dtype):
         self.idx.factorize(sort=sort)
 
 
 class Duplicated:
 
-    params = [["first", "last", False], ["int", "uint", "float", "string"]]
-    param_names = ["keep", "dtype"]
-
-    def setup(self, keep, dtype):
-        N = 10 ** 5
-        data = {
-            "int": pd.Int64Index(np.arange(N).repeat(5)),
-            "uint": pd.UInt64Index(np.arange(N).repeat(5)),
-            "float": pd.Float64Index(np.random.randn(N).repeat(5)),
-            "string": tm.makeStringIndex(N).repeat(5),
-        }
-        self.idx = data[dtype]
-        # cache is_unique
-        self.idx.is_unique
-
-    def time_duplicated(self, keep, dtype):
-        self.idx.duplicated(keep=keep)
-
-
-class DuplicatedUniqueIndex:
-
-    params = ["int", "uint", "float", "string"]
-    param_names = ["dtype"]
+    params = [
+        [True, False],
+        ["first", "last", False],
+        ["int", "uint", "float", "string", "datetime64[ns]", "datetime64[ns, tz]"],
+    ]
+    param_names = ["unique", "keep", "dtype"]
 
-    def setup(self, dtype):
+    def setup(self, unique, keep, dtype):
         N = 10 ** 5
         data = {
             "int": pd.Int64Index(np.arange(N)),
             "uint": pd.UInt64Index(np.arange(N)),
             "float": pd.Float64Index(np.random.randn(N)),
             "string": tm.makeStringIndex(N),
-        }
-        self.idx = data[dtype]
+            "datetime64[ns]": pd.date_range("2011-01-01", freq="H", periods=N),
+            "datetime64[ns, tz]": pd.date_range(
+                "2011-01-01", freq="H", periods=N, tz="Asia/Tokyo"
+            ),
+        }[dtype]
+        if not unique:
+            data = data.repeat(5)
+        self.idx = data
         # cache is_unique
         self.idx.is_unique
 
-    def time_duplicated_unique(self, dtype):
-        self.idx.duplicated()
+    def time_duplicated(self, unique, keep, dtype):
+        self.idx.duplicated(keep=keep)
 
 
 class Hashing:

asv_bench/benchmarks/timeseries.py

@@ -91,20 +91,6 @@ def time_reest_datetimeindex(self, tz):
         self.df.reset_index()
 
 
-class Factorize:
-
-    params = [None, "Asia/Tokyo"]
-    param_names = "tz"
-
-    def setup(self, tz):
-        N = 100000
-        self.dti = date_range("2011-01-01", freq="H", periods=N, tz=tz)
-        self.dti = self.dti.repeat(5)
-
-    def time_factorize(self, tz):
-        self.dti.factorize()
-
-
 class InferFreq:
 
     params = [None, "D", "B"]