Implement scalar shift_month mirroring tslib.shift_months (#18218)

Author: jbrockmendel

doc/source/whatsnew/v0.22.0.txt

@@ -71,6 +71,7 @@ Performance Improvements
 
 - Indexers on ``Series`` or ``DataFrame`` no longer create a reference cycle (:issue:`17956`)
 - Added a keyword argument, ``cache``, to :func:`to_datetime` that improved the performance of converting duplicate datetime arguments (:issue:`11665`)
+- :class`DateOffset` arithmetic performance is improved (:issue:`18218`)
 -
 
 .. _whatsnew_0220.docs:

pandas/_libs/tslibs/offsets.pyx

@@ -4,7 +4,7 @@
 cimport cython
 
 import time
-from cpython.datetime cimport timedelta, time as dt_time
+from cpython.datetime cimport datetime, timedelta, time as dt_time
 
 from dateutil.relativedelta import relativedelta
 
@@ -13,9 +13,9 @@ cimport numpy as np
 np.import_array()
 
 
-from util cimport is_string_object
+from util cimport is_string_object, is_integer_object
 
-from pandas._libs.tslib import pydt_to_i8
+from pandas._libs.tslib import pydt_to_i8, monthrange
 
 from frequencies cimport get_freq_code
 from conversion cimport tz_convert_single
@@ -375,3 +375,56 @@ class BaseOffset(_BaseOffset):
             # i.e. isinstance(other, (ABCDatetimeIndex, ABCSeries))
             return other - self
         return -self + other
+
+
+# ----------------------------------------------------------------------
+# RelativeDelta Arithmetic
+
+
+cpdef datetime shift_month(datetime stamp, int months, object day_opt=None):
+    """
+    Given a datetime (or Timestamp) `stamp`, an integer `months` and an
+    option `day_opt`, return a new datetimelike that many months later,
+    with day determined by `day_opt` using relativedelta semantics.
+
+    Scalar analogue of tslib.shift_months
+
+    Parameters
+    ----------
+    stamp : datetime or Timestamp
+    months : int
+    day_opt : None, 'start', 'end', or an integer
+        None: returned datetimelike has the same day as the input, or the
+              last day of the month if the new month is too short
+        'start': returned datetimelike has day=1
+        'end': returned datetimelike has day on the last day of the month
+        int: returned datetimelike has day equal to day_opt
+
+    Returns
+    -------
+    shifted : datetime or Timestamp (same as input `stamp`)
+    """
+    cdef:
+        int year, month, day
+        int dim, dy
+
+    dy = (stamp.month + months) // 12
+    month = (stamp.month + months) % 12
+
+    if month == 0:
+        month = 12
+        dy -= 1
+    year = stamp.year + dy
+
+    dim = monthrange(year, month)[1]
+    if day_opt is None:
+        day = min(stamp.day, dim)
+    elif day_opt == 'start':
+        day = 1
+    elif day_opt == 'end':
+        day = dim
+    elif is_integer_object(day_opt):
+        day = min(day_opt, dim)
+    else:
+        raise ValueError(day_opt)
+    return stamp.replace(year=year, month=month, day=day)

pandas/tseries/offsets.py

@@ -22,6 +22,7 @@
     _int_to_weekday, _weekday_to_int,
     _determine_offset,
     apply_index_wraps,
+    shift_month,
     BeginMixin, EndMixin,
     BaseOffset)
 
@@ -252,6 +253,8 @@ def apply_index(self, i):
                                       "applied vectorized".format(kwd=kwd))
 
     def isAnchored(self):
+        # TODO: Does this make sense for the general case?  It would help
+        # if there were a canonical docstring for what isAnchored means.
         return (self.n == 1)
 
     def _params(self):
@@ -721,6 +724,7 @@ def apply(self, other):
 
             return result
         else:
+            # TODO: Figure out the end of this sente
             raise ApplyTypeError(
                 'Only know how to combine business hour with ')
 
@@ -927,10 +931,10 @@ def apply(self, other):
         n = self.n
         _, days_in_month = tslib.monthrange(other.year, other.month)
         if other.day != days_in_month:
-            other = other + relativedelta(months=-1, day=31)
+            other = shift_month(other, -1, 'end')
             if n <= 0:
                 n = n + 1
-        other = other + relativedelta(months=n, day=31)
+        other = shift_month(other, n, 'end')
         return other
 
     @apply_index_wraps
@@ -956,7 +960,7 @@ def apply(self, other):
         if other.day > 1 and n <= 0:  # then roll forward if n<=0
             n += 1
 
-        return other + relativedelta(months=n, day=1)
+        return shift_month(other, n, 'start')
 
     @apply_index_wraps
     def apply_index(self, i):
@@ -1002,12 +1006,12 @@ def apply(self, other):
         if not self.onOffset(other):
             _, days_in_month = tslib.monthrange(other.year, other.month)
             if 1 < other.day < self.day_of_month:
-                other += relativedelta(day=self.day_of_month)
+                other = other.replace(day=self.day_of_month)
                 if n > 0:
                     # rollforward so subtract 1
                     n -= 1
             elif self.day_of_month < other.day < days_in_month:
-                other += relativedelta(day=self.day_of_month)
+                other = other.replace(day=self.day_of_month)
                 if n < 0:
                     # rollforward in the negative direction so add 1
                     n += 1
@@ -1084,19 +1088,19 @@ def onOffset(self, dt):
     def _apply(self, n, other):
         # if other.day is not day_of_month move to day_of_month and update n
         if other.day < self.day_of_month:
-            other += relativedelta(day=self.day_of_month)
+            other = other.replace(day=self.day_of_month)
             if n > 0:
                 n -= 1
         elif other.day > self.day_of_month:
-            other += relativedelta(day=self.day_of_month)
+            other = other.replace(day=self.day_of_month)
             if n == 0:
                 n = 1
             else:
                 n += 1
 
         months = n // 2
         day = 31 if n % 2 else self.day_of_month
-        return other + relativedelta(months=months, day=day)
+        return shift_month(other, months, day)
 
     def _get_roll(self, i, before_day_of_month, after_day_of_month):
         n = self.n
@@ -1141,21 +1145,21 @@ def onOffset(self, dt):
     def _apply(self, n, other):
         # if other.day is not day_of_month move to day_of_month and update n
         if other.day < self.day_of_month:
-            other += relativedelta(day=self.day_of_month)
+            other = other.replace(day=self.day_of_month)
             if n == 0:
                 n = -1
             else:
                 n -= 1
         elif other.day > self.day_of_month:
-            other += relativedelta(day=self.day_of_month)
+            other = other.replace(day=self.day_of_month)
             if n == 0:
                 n = 1
             elif n < 0:
                 n += 1
 
         months = n // 2 + n % 2
         day = 1 if n % 2 else self.day_of_month
-        return other + relativedelta(months=months, day=day)
+        return shift_month(other, months, day)
 
     def _get_roll(self, i, before_day_of_month, after_day_of_month):
         n = self.n
@@ -1191,7 +1195,7 @@ def apply(self, other):
             n = n - 1
         elif n <= 0 and other.day > lastBDay:
             n = n + 1
-        other = other + relativedelta(months=n, day=31)
+        other = shift_month(other, n, 'end')
 
         if other.weekday() > 4:
             other = other - BDay()
@@ -1215,7 +1219,7 @@ def apply(self, other):
             other = other + timedelta(days=first - other.day)
             n -= 1
 
-        other = other + relativedelta(months=n)
+        other = shift_month(other, n, None)
         wkday, _ = tslib.monthrange(other.year, other.month)
         first = _get_firstbday(wkday)
         result = datetime(other.year, other.month, first,
@@ -1520,8 +1524,7 @@ def apply(self, other):
             else:
                 months = self.n + 1
 
-        other = self.getOffsetOfMonth(
-            other + relativedelta(months=months, day=1))
+        other = self.getOffsetOfMonth(shift_month(other, months, 'start'))
         other = datetime(other.year, other.month, other.day, base.hour,
                          base.minute, base.second, base.microsecond)
         return other
@@ -1612,8 +1615,7 @@ def apply(self, other):
             else:
                 months = self.n + 1
 
-        return self.getOffsetOfMonth(
-            other + relativedelta(months=months, day=1))
+        return self.getOffsetOfMonth(shift_month(other, months, 'start'))
 
     def getOffsetOfMonth(self, dt):
         m = MonthEnd()
@@ -1716,7 +1718,7 @@ def apply(self, other):
         elif n <= 0 and other.day > lastBDay and monthsToGo == 0:
             n = n + 1
 
-        other = other + relativedelta(months=monthsToGo + 3 * n, day=31)
+        other = shift_month(other, monthsToGo + 3 * n, 'end')
         other = tslib._localize_pydatetime(other, base.tzinfo)
         if other.weekday() > 4:
             other = other - BDay()
@@ -1761,7 +1763,7 @@ def apply(self, other):
             n = n - 1
 
         # get the first bday for result
-        other = other + relativedelta(months=3 * n - monthsSince)
+        other = shift_month(other, 3 * n - monthsSince, None)
         wkday, _ = tslib.monthrange(other.year, other.month)
         first = _get_firstbday(wkday)
         result = datetime(other.year, other.month, first,
@@ -1795,7 +1797,7 @@ def apply(self, other):
         if n > 0 and not (other.day >= days_in_month and monthsToGo == 0):
             n = n - 1
 
-        other = other + relativedelta(months=monthsToGo + 3 * n, day=31)
+        other = shift_month(other, monthsToGo + 3 * n, 'end')
         return other
 
     @apply_index_wraps
@@ -1830,7 +1832,7 @@ def apply(self, other):
             # after start, so come back an extra period as if rolled forward
             n = n + 1
 
-        other = other + relativedelta(months=3 * n - monthsSince, day=1)
+        other = shift_month(other, 3 * n - monthsSince, 'start')
         return other
 
     @apply_index_wraps
@@ -1889,7 +1891,7 @@ def apply(self, other):
                     (other.month == self.month and other.day > lastBDay)):
                 years += 1
 
-        other = other + relativedelta(years=years)
+        other = shift_month(other, 12 * years, None)
 
         _, days_in_month = tslib.monthrange(other.year, self.month)
         result = datetime(other.year, self.month, days_in_month,
@@ -1927,7 +1929,7 @@ def apply(self, other):
                 years += 1
 
         # set first bday for result
-        other = other + relativedelta(years=years)
+        other = shift_month(other, years * 12, None)
         wkday, days_in_month = tslib.monthrange(other.year, self.month)
         first = _get_firstbday(wkday)
         return datetime(other.year, self.month, first, other.hour,
@@ -2145,8 +2147,8 @@ def onOffset(self, dt):
 
         if self.variation == "nearest":
             # We have to check the year end of "this" cal year AND the previous
-            return year_end == dt or \
-                self.get_year_end(dt - relativedelta(months=1)) == dt
+            return (year_end == dt or
+                    self.get_year_end(shift_month(dt, -1, None)) == dt)
         else:
             return year_end == dt
 
@@ -2226,8 +2228,8 @@ def get_year_end(self, dt):
     def get_target_month_end(self, dt):
         target_month = datetime(
             dt.year, self.startingMonth, 1, tzinfo=dt.tzinfo)
-        next_month_first_of = target_month + relativedelta(months=+1)
-        return next_month_first_of + relativedelta(days=-1)
+        next_month_first_of = shift_month(target_month, 1, None)
+        return next_month_first_of + timedelta(days=-1)
 
     def _get_year_end_nearest(self, dt):
         target_date = self.get_target_month_end(dt)
@@ -2382,7 +2384,7 @@ def apply(self, other):
                     qtr_lens = self.get_weeks(other + self._offset)
 
                 for weeks in qtr_lens:
-                    start += relativedelta(weeks=weeks)
+                    start += timedelta(weeks=weeks)
                     if start > other:
                         other = start
                         n -= 1
@@ -2399,7 +2401,7 @@ def apply(self, other):
                     qtr_lens = self.get_weeks(other)
 
                 for weeks in reversed(qtr_lens):
-                    end -= relativedelta(weeks=weeks)
+                    end -= timedelta(weeks=weeks)
                     if end < other:
                         other = end
                         n -= 1
@@ -2442,7 +2444,7 @@ def onOffset(self, dt):
 
         current = next_year_end
         for qtr_len in qtr_lens[0:4]:
-            current += relativedelta(weeks=qtr_len)
+            current += timedelta(weeks=qtr_len)
             if dt == current:
                 return True
         return False