ENH: at_time/between_time work with tz-localized time series. refactoring and cleanup close #1481

Author: wesm

pandas/core/series.py

@@ -2773,26 +2773,23 @@ def _repr_footer(self):
         namestr = "Name: %s, " % str(self.name) if self.name else ""
         return '%s%sLength: %d' % (freqstr, namestr, len(self))
 
-    def at_time(self, time, tz=None, asof=False):
+    def at_time(self, time, asof=False):
         """
         Select values at particular time of day (e.g. 9:30AM)
 
         Parameters
         ----------
         time : datetime.time or string
-        tz : string or pytz.timezone
-            Time zone for time. Corresponding timestamps would be converted to
-            time zone of the TimeSeries
 
         Returns
         -------
         values_at_time : TimeSeries
         """
-        from pandas.tseries.resample import values_at_time
-        return values_at_time(self, time, tz=tz, asof=asof)
+        indexer = self.index.indexer_at_time(time, asof=asof)
+        return self.take(indexer)
 
     def between_time(self, start_time, end_time, include_start=True,
-                     include_end=True, tz=None):
+                     include_end=True):
         """
         Select values between particular times of the day (e.g., 9:00-9:30 AM)
 
@@ -2802,16 +2799,15 @@ def between_time(self, start_time, end_time, include_start=True,
         end_time : datetime.time or string
         include_start : boolean, default True
         include_end : boolean, default True
-        tz : string or pytz.timezone, default None
 
         Returns
         -------
         values_between_time : TimeSeries
         """
-        from pandas.tseries.resample import values_between_time
-        return values_between_time(self, start_time, end_time, tz=tz,
-                                   include_start=include_start,
-                                   include_end=include_end)
+        indexer = self.index.indexer_between_time(
+            start_time, end_time, include_start=include_start,
+            include_end=include_end)
+        return self.take(indexer)
 
     def to_timestamp(self, freq=None, how='start', copy=True):
         """

pandas/src/datetime.pyx

@@ -409,8 +409,7 @@ cdef class _Timestamp(datetime):
             return datetime.__sub__(self, other)
 
     cpdef _get_field(self, field):
-        out = fast_field_accessor(np.array([self.value], dtype=np.int64),
-                                  field)
+        out = get_date_field(np.array([self.value], dtype=np.int64), field)
         return out[0]
 
 
@@ -1049,8 +1048,26 @@ def build_field_sarray(ndarray[int64_t] dtindex):
 
     return out
 
+def get_time_micros(ndarray[int64_t] dtindex):
+    '''
+    Datetime as int64 representation to a structured array of fields
+    '''
+    cdef:
+        Py_ssize_t i, n = len(dtindex)
+        pandas_datetimestruct dts
+        ndarray[int64_t] micros
+
+    micros = np.empty(n, dtype=np.int64)
+
+    for i in range(n):
+        pandas_datetime_to_datetimestruct(dtindex[i], PANDAS_FR_ns, &dts)
+        micros[i] = 1000000LL * (dts.hour * 60 * 60 +
+                                 60 * dts.min + dts.sec) + dts.us
+
+    return micros
+
 @cython.wraparound(False)
-def fast_field_accessor(ndarray[int64_t] dtindex, object field):
+def get_date_field(ndarray[int64_t] dtindex, object field):
     '''
     Given a int64-based datetime index, extract the year, month, etc.,
     field and return an array of these values.
@@ -1178,107 +1195,6 @@ cdef inline int m8_weekday(int64_t val):
 
 cdef int64_t DAY_NS = 86400000000000LL
 
-def values_at_time(ndarray[int64_t] stamps, int64_t time):
-    cdef:
-        Py_ssize_t i, j, count, n = len(stamps)
-        ndarray[int64_t] indexer, times
-        int64_t last, cur
-
-    # Assumes stamps is sorted
-
-    if len(stamps) == 0:
-        return np.empty(0, dtype=np.int64)
-
-    # is this OK?
-    # days = stamps // DAY_NS
-    times = stamps % DAY_NS
-
-    # Nanosecond resolution
-    count = 0
-    for i in range(n):
-        if times[i] == time:
-            count += 1
-
-    indexer = np.empty(count, dtype=np.int64)
-
-    j = 0
-    # last = days[0]
-    for i in range(n):
-        if times[i] == time:
-            indexer[j] = i
-            j += 1
-
-    return indexer
-
-def values_between_time(ndarray[int64_t] stamps, int64_t stime, int64_t etime,
-                        bint include_start, bint include_end):
-    cdef:
-        Py_ssize_t i, j, count, n = len(stamps)
-        ndarray[int64_t] indexer, times
-        int64_t last, cur
-
-    # Assumes stamps is sorted
-
-    if len(stamps) == 0:
-        return np.empty(0, dtype=np.int64)
-
-    # is this OK?
-    # days = stamps // DAY_NS
-    times = stamps % DAY_NS
-
-    # Nanosecond resolution
-    count = 0
-    if include_start and include_end:
-        for i in range(n):
-            cur = times[i]
-            if cur >= stime and cur <= etime:
-                count += 1
-    elif include_start:
-        for i in range(n):
-            cur = times[i]
-            if cur >= stime and cur < etime:
-                count += 1
-    elif include_end:
-        for i in range(n):
-            cur = times[i]
-            if cur > stime and cur <= etime:
-                count += 1
-    else:
-        for i in range(n):
-            cur = times[i]
-            if cur > stime and cur < etime:
-                count += 1
-
-    indexer = np.empty(count, dtype=np.int64)
-
-    j = 0
-    # last = days[0]
-    if include_start and include_end:
-        for i in range(n):
-            cur = times[i]
-            if cur >= stime and cur <= etime:
-                indexer[j] = i
-                j += 1
-    elif include_start:
-        for i in range(n):
-            cur = times[i]
-            if cur >= stime and cur < etime:
-                indexer[j] = i
-                j += 1
-    elif include_end:
-        for i in range(n):
-            cur = times[i]
-            if cur > stime and cur <= etime:
-                indexer[j] = i
-                j += 1
-    else:
-        for i in range(n):
-            cur = times[i]
-            if cur > stime and cur < etime:
-                indexer[j] = i
-                j += 1
-
-    return indexer
 
 def date_normalize(ndarray[int64_t] stamps):
     cdef:

pandas/tseries/index.py

@@ -39,7 +39,7 @@ def f(self):
             utc = _utc()
             if self.tz is not utc:
                 values = lib.tz_convert(values, utc, self.tz)
-        return lib.fast_field_accessor(values, field)
+        return lib.get_date_field(values, field)
     f.__name__ = name
     return property(f)
 
@@ -578,9 +578,23 @@ def asi8(self):
     @property
     def asstruct(self):
         if self._sarr_cache is None:
-            self._sarr_cache = lib.build_field_sarray(self.asi8)
+            self._sarr_cache = self._get_field_sarr()
         return self._sarr_cache
 
+    def _get_field_sarr(self):
+        utc = _utc()
+        values = self.asi8
+        if self.tz is not None and self.tz is not utc:
+            values = lib.tz_convert(values, utc, self.tz)
+        return lib.build_field_sarray(values)
+
+    def _get_time_micros(self):
+        utc = _utc()
+        values = self.asi8
+        if self.tz is not None and self.tz is not utc:
+            values = lib.tz_convert(values, utc, self.tz)
+        return lib.get_time_micros(values)
+
     @property
     def asobject(self):
         """
@@ -712,7 +726,7 @@ def take(self, indices, axis=0):
             return self[maybe_slice]
         indices = com._ensure_platform_int(indices)
         taken = self.values.take(indices, axis=axis)
-        return DatetimeIndex(taken, tz=self.tz, name=self.name)
+        return self._simple_new(taken, self.name, None, self.tz)
 
     def union(self, other):
         """
@@ -943,7 +957,7 @@ def get_value(self, series, key):
                 pass
 
             if isinstance(key, time):
-                locs = self._indices_at_time(key)
+                locs = self.indexer_at_time(key)
                 return series.take(locs)
 
             stamp = Timestamp(key)
@@ -969,22 +983,13 @@ def get_loc(self, key):
                 pass
 
             if isinstance(key, time):
-                return self._indices_at_time(key)
+                return self.indexer_at_time(key)
 
             try:
                 return self._engine.get_loc(Timestamp(key))
             except (KeyError, ValueError):
                 raise KeyError(key)
 
-    def _indices_at_time(self, key):
-        from dateutil.parser import parse
-
-        # TODO: time object with tzinfo?
-
-        nanos = _time_to_nanosecond(key)
-        indexer = lib.values_at_time(self.asi8, nanos)
-        return com._ensure_platform_int(indexer)
-
     def _get_string_slice(self, key):
         freq = getattr(self, 'freqstr',
                        getattr(self, 'inferred_freq', None))
@@ -1246,6 +1251,84 @@ def tz_validate(self):
 
         return True
 
+    def indexer_at_time(self, time, asof=False):
+        """
+        Select values at particular time of day (e.g. 9:30AM)
+
+        Parameters
+        ----------
+        time : datetime.time or string
+        tz : string or pytz.timezone
+            Time zone for time. Corresponding timestamps would be converted to
+            time zone of the TimeSeries
+
+        Returns
+        -------
+        values_at_time : TimeSeries
+        """
+        from dateutil.parser import parse
+
+        if asof:
+            raise NotImplementedError
+
+        if isinstance(time, basestring):
+            time = parse(time).time()
+
+        if time.tzinfo:
+            # TODO
+            raise NotImplementedError
+
+        time_micros = self._get_time_micros()
+        micros = _time_to_micros(time)
+        return (micros == time_micros).nonzero()[0]
+
+    def indexer_between_time(self, start_time, end_time, include_start=True,
+                             include_end=True):
+        """
+        Select values between particular times of day (e.g., 9:00-9:30AM)
+
+        Parameters
+        ----------
+        start_time : datetime.time or string
+        end_time : datetime.time or string
+        include_start : boolean, default True
+        include_end : boolean, default True
+        tz : string or pytz.timezone, default None
+
+        Returns
+        -------
+        values_between_time : TimeSeries
+        """
+        from dateutil.parser import parse
+
+        if isinstance(start_time, basestring):
+            start_time = parse(start_time).time()
+
+        if isinstance(end_time, basestring):
+            end_time = parse(end_time).time()
+
+        if start_time.tzinfo or end_time.tzinfo:
+            raise NotImplementedError
+
+        time_micros = self._get_time_micros()
+        start_micros = _time_to_micros(start_time)
+        end_micros = _time_to_micros(end_time)
+
+        if include_start and include_end:
+            mask = ((start_micros <= time_micros) &
+                    (time_micros <= end_micros))
+        elif include_start:
+            mask = ((start_micros <= time_micros) &
+                    (time_micros < end_micros))
+        elif include_end:
+            mask = ((start_micros < time_micros) &
+                    (time_micros <= end_micros))
+        else:
+            mask = ((start_micros < time_micros) &
+                    (time_micros < end_micros))
+
+        return mask.nonzero()[0]
+
 def _generate_regular_range(start, end, periods, offset):
     if com._count_not_none(start, end, periods) < 2:
         raise ValueError('Must specify two of start, end, or periods')
@@ -1399,8 +1482,11 @@ def _in_range(start, end, rng_start, rng_end):
     return start > rng_start and end < rng_end
 
 def _time_to_nanosecond(time):
+    return _time_to_micros(time) * 1000
+
+def _time_to_micros(time):
     seconds = time.hour * 60 * 60 + 60 * time.minute + time.second
-    return (1000000 * seconds + time.microsecond) * 1000
+    return 1000000 * seconds + time.microsecond
 
 
 def _concat(to_concat):
@@ -1410,3 +1496,4 @@ def _concat(to_concat):
         return new_values.view(_NS_DTYPE)
     else:
         return np.concatenate(to_concat)
+