PERF: Improve dt-like index ops

Author: sinhrks

doc/source/whatsnew/v0.16.2.txt

@@ -48,6 +48,8 @@ Performance Improvements
 
 - Improved ``Series.resample`` performance with dtype=datetime64[ns] (:issue:`7754`)
 - Modest improvement in datetime writing speed in to_csv (:issue:`10271`)
+- Improved huge ``DatetimeIndex``, ``PeriodIndex`` and ``TimedeltaIndex``'s ops performance including ``NaT`` (:issue:`10277`)
+
 
 .. _whatsnew_0162.bug_fixes:
 

pandas/tseries/base.py

@@ -130,10 +130,15 @@ def get_duplicates(self):
         values = Index.get_duplicates(self)
         return self._simple_new(values)
 
+    @cache_readonly
+    def _isnan(self):
+        """ return if each value is nan"""
+        return (self.asi8 == tslib.iNaT)
+
     @cache_readonly
     def hasnans(self):
         """ return if I have any nans; enables various perf speedups """
-        return (self.asi8 == tslib.iNaT).any()
+        return self._isnan.any()
 
     @property
     def asobject(self):
@@ -157,12 +162,11 @@ def _maybe_mask_results(self, result, fill_value=None, convert=None):
         """
 
         if self.hasnans:
-            mask = self.asi8 == tslib.iNaT
             if convert:
                 result = result.astype(convert)
             if fill_value is None:
                 fill_value = np.nan
-            result[mask] = fill_value
+            result[self._isnan] = fill_value
         return result
 
     def tolist(self):
@@ -188,8 +192,7 @@ def min(self, axis=None):
                     return self._box_func(i8[0])
 
             if self.hasnans:
-                mask = i8 == tslib.iNaT
-                min_stamp = self[~mask].asi8.min()
+                min_stamp = self[~self._isnan].asi8.min()
             else:
                 min_stamp = i8.min()
             return self._box_func(min_stamp)
@@ -207,7 +210,7 @@ def argmin(self, axis=None):
 
         i8 = self.asi8
         if self.hasnans:
-            mask = i8 == tslib.iNaT
+            mask = self._isnan
             if mask.all():
                 return -1
             i8 = i8.copy()
@@ -231,8 +234,7 @@ def max(self, axis=None):
                     return self._box_func(i8[-1])
 
             if self.hasnans:
-                mask = i8 == tslib.iNaT
-                max_stamp = self[~mask].asi8.max()
+                max_stamp = self[~self._isnan].asi8.max()
             else:
                 max_stamp = i8.max()
             return self._box_func(max_stamp)
@@ -250,7 +252,7 @@ def argmax(self, axis=None):
 
         i8 = self.asi8
         if self.hasnans:
-            mask = i8 == tslib.iNaT
+            mask = self._isnan
             if mask.all():
                 return -1
             i8 = i8.copy()
@@ -381,9 +383,8 @@ def _add_delta_td(self, other):
         # return the i8 result view
 
         inc = tslib._delta_to_nanoseconds(other)
-        mask = self.asi8 == tslib.iNaT
         new_values = (self.asi8 + inc).view(self.dtype)
-        new_values[mask] = tslib.iNaT
+        new_values[self._isnan] = tslib.iNaT
         return new_values.view(self.dtype)
 
     def _add_delta_tdi(self, other):
@@ -396,7 +397,7 @@ def _add_delta_tdi(self, other):
 
         self_i8 = self.asi8
         other_i8 = other.asi8
-        mask = (self_i8 == tslib.iNaT) | (other_i8 == tslib.iNaT)
+        mask = (self._isnan) | (other._isnan)
         new_values = self_i8 + other_i8
         new_values[mask] = tslib.iNaT
         return new_values.view(self.dtype)

pandas/tseries/index.py

@@ -85,9 +85,8 @@ def wrapper(self, other):
             if o_mask.any():
                 result[o_mask] = nat_result
 
-        mask = self.asi8 == tslib.iNaT
-        if mask.any():
-            result[mask] = nat_result
+        if self.hasnans:
+            result[self._isnan] = nat_result
 
         # support of bool dtype indexers
         if com.is_bool_dtype(result):

pandas/tseries/period.py

@@ -506,9 +506,8 @@ def shift(self, n):
         -------
         shifted : PeriodIndex
         """
-        mask = self.values == tslib.iNaT
         values = self.values + n
-        values[mask] = tslib.iNaT
+        values[self._isnan] = tslib.iNaT
         return PeriodIndex(data=values, name=self.name, freq=self.freq)
 
     @property

pandas/tseries/tdi.py

@@ -61,9 +61,8 @@ def wrapper(self, other):
             if o_mask.any():
                 result[o_mask] = nat_result
 
-        mask = self.asi8 == tslib.iNaT
-        if mask.any():
-            result[mask] = nat_result
+        if self.hasnans:
+            result[self._isnan] = nat_result
 
         # support of bool dtype indexers
         if com.is_bool_dtype(result):
@@ -334,7 +333,7 @@ def _get_field(self, m):
         hasnans = self.hasnans
         if hasnans:
             result = np.empty(len(self), dtype='float64')
-            mask = values == tslib.iNaT
+            mask = self._isnan
             imask = ~mask
             result.flat[imask] = np.array([ getattr(Timedelta(val),m) for val in values[imask] ])
             result[mask] = np.nan

pandas/tseries/tests/test_base.py

@@ -97,6 +97,8 @@ def test_minmax(self):
             for idx in [idx1, idx2]:
                 self.assertEqual(idx.min(), pd.Timestamp('2011-01-01', tz=tz))
                 self.assertEqual(idx.max(), pd.Timestamp('2011-01-03', tz=tz))
+                self.assertEqual(idx.argmin(), 0)
+                self.assertEqual(idx.argmax(), 2)
 
         for op in ['min', 'max']:
             # Return NaT
@@ -348,6 +350,8 @@ def test_minmax(self):
         for idx in [idx1, idx2]:
             self.assertEqual(idx.min(), Timedelta('1 days')),
             self.assertEqual(idx.max(), Timedelta('3 days')),
+            self.assertEqual(idx.argmin(), 0)
+            self.assertEqual(idx.argmax(), 2)
 
         for op in ['min', 'max']:
             # Return NaT
@@ -801,6 +805,10 @@ def test_minmax(self):
         for idx in [idx1, idx2]:
             self.assertEqual(idx.min(), pd.Period('2011-01-01', freq='D'))
             self.assertEqual(idx.max(), pd.Period('2011-01-03', freq='D'))
+        self.assertEqual(idx1.argmin(), 1)
+        self.assertEqual(idx2.argmin(), 0)
+        self.assertEqual(idx1.argmax(), 3)
+        self.assertEqual(idx2.argmax(), 2)
 
         for op in ['min', 'max']:
             # Return NaT