PERF/BUG: improve factorize for datetimetz (#13750)

Author: sinhrks

asv_bench/benchmarks/algorithms.py

@@ -0,0 +1,17 @@
+import numpy as np
+import pandas as pd
+
+
+class algorithm(object):
+    goal_time = 0.2
+
+    def setup(self):
+        N = 100000
+        self.int = pd.Int64Index(np.arange(N).repeat(5))
+        self.float = pd.Float64Index(np.random.randn(N).repeat(5))
+
+    def time_int_factorize(self):
+        self.int.factorize()
+
+    def time_float_factorize(self):
+        self.int.factorize()

asv_bench/benchmarks/period.py

@@ -9,6 +9,22 @@ def time_period_index(self):
         PeriodIndex(date_range('1985', periods=1000).to_pydatetime(), freq='D')
 
 
+class period_setitem(object):
+    goal_time = 0.2
+
+    def setup(self):
+        self.N = 100000
+        self.rng = date_range(start='1/1/2000', periods=self.N, freq='T')
+        if hasattr(Series, 'convert'):
+            Series.resample = Series.convert
+        self.ts = Series(np.random.randn(self.N), index=self.rng)
+        self.rng = period_range(start='1/1/1990', freq='S', periods=20000)
+        self.df = DataFrame(index=range(len(self.rng)))
+
+    def time_period_setitem(self):
+        self.df['col'] = self.rng
+
+
 class period_algorithm(object):
     goal_time = 0.2
 

asv_bench/benchmarks/timeseries.py

@@ -218,20 +218,20 @@ def time_dti_reset_index_tz(self):
         self.df.reset_index()
 
 
-class period_setitem(object):
+class datetime_algorithm(object):
     goal_time = 0.2
 
     def setup(self):
-        self.N = 100000
-        self.rng = date_range(start='1/1/2000', periods=self.N, freq='T')
-        if hasattr(Series, 'convert'):
-            Series.resample = Series.convert
-        self.ts = Series(np.random.randn(self.N), index=self.rng)
-        self.rng = period_range(start='1/1/1990', freq='S', periods=20000)
-        self.df = DataFrame(index=range(len(self.rng)))
+        N = 100000
+        self.dti = pd.date_range('2011-01-01', freq='H', periods=N).repeat(5)
+        self.dti_tz = pd.date_range('2011-01-01', freq='H', periods=N,
+                                    tz='Asia/Tokyo').repeat(5)
+
+    def time_dti_factorize(self):
+        self.dti.factorize()
 
-    def time_period_setitem(self):
-        self.df['col'] = self.rng
+    def time_dti_tz_factorize(self):
+        self.dti_tz.factorize()
 
 
 class timeseries_1min_5min_mean(object):

doc/source/whatsnew/v0.19.0.txt

@@ -652,6 +652,8 @@ Performance Improvements
 - Improved performance of ``Index.difference`` (:issue:`12044`)
 - Improved performance of datetime string parsing in ``DatetimeIndex`` (:issue:`13692`)
 - Improved performance of hashing ``Period`` (:issue:`12817`)
+- Improved performance of ``factorize`` of datetime with timezone (:issue:`13750`)
+
 
 
 .. _whatsnew_0190.bug_fixes:
@@ -738,6 +740,7 @@ Bug Fixes
 - Bug in ``pd.set_eng_float_format()`` that would prevent NaN's from formatting (:issue:`11981`)
 - Bug in ``.unstack`` with ``Categorical`` dtype resets ``.ordered`` to ``True`` (:issue:`13249`)
 - Clean some compile time warnings in datetime parsing (:issue:`13607`)
+- Bug in ``factorize`` raises ``AmbiguousTimeError`` if data contains datetime near DST boundary (:issue:`13750`)
 
 - Bug in ``Series`` comparison operators when dealing with zero dim NumPy arrays (:issue:`13006`)
 - Bug in ``groupby`` where ``apply`` returns different result depending on whether first result is ``None`` or not (:issue:`12824`)

pandas/core/algorithms.py

@@ -293,7 +293,7 @@ def factorize(values, sort=False, order=None, na_sentinel=-1, size_hint=None):
     is_datetimetz_type = is_datetimetz(values)
     if is_datetimetz_type:
         values = DatetimeIndex(values)
-        vals = values.tz_localize(None)
+        vals = values.asi8
 
     is_datetime = is_datetime64_dtype(vals)
     is_timedelta = is_timedelta64_dtype(vals)
@@ -313,8 +313,7 @@ def factorize(values, sort=False, order=None, na_sentinel=-1, size_hint=None):
 
     if is_datetimetz_type:
         # reset tz
-        uniques = DatetimeIndex(uniques.astype('M8[ns]')).tz_localize(
-            values.tz)
+        uniques = values._shallow_copy(uniques)
     elif is_datetime:
         uniques = uniques.astype('M8[ns]')
     elif is_timedelta:

pandas/tseries/tests/test_timeseries.py

@@ -3826,6 +3826,37 @@ def test_factorize(self):
         self.assert_numpy_array_equal(arr, exp_arr)
         tm.assert_index_equal(idx, idx3)
 
+    def test_factorize_tz(self):
+        # GH 13750
+        for tz in [None, 'UTC', 'US/Eastern', 'Asia/Tokyo']:
+            base = pd.date_range('2016-11-05', freq='H', periods=100, tz=tz)
+            idx = base.repeat(5)
+
+            exp_arr = np.arange(100).repeat(5)
+
+            for obj in [idx, pd.Series(idx)]:
+                arr, res = obj.factorize()
+                self.assert_numpy_array_equal(arr, exp_arr)
+                tm.assert_index_equal(res, base)
+
+    def test_factorize_dst(self):
+        # GH 13750
+        idx = pd.date_range('2016-11-06', freq='H', periods=12,
+                            tz='US/Eastern')
+
+        for obj in [idx, pd.Series(idx)]:
+            arr, res = obj.factorize()
+            self.assert_numpy_array_equal(arr, np.arange(12))
+            tm.assert_index_equal(res, idx)
+
+        idx = pd.date_range('2016-06-13', freq='H', periods=12,
+                            tz='US/Eastern')
+
+        for obj in [idx, pd.Series(idx)]:
+            arr, res = obj.factorize()
+            self.assert_numpy_array_equal(arr, np.arange(12))
+            tm.assert_index_equal(res, idx)
+
     def test_slice_with_negative_step(self):
         ts = Series(np.arange(20),
                     date_range('2014-01-01', periods=20, freq='MS'))