ENH: speed up Series.__getitem__ with timestamps, much more efficient implementation of Series.asof, close #1168

Author: wesm

pandas/core/algorithms.py

@@ -140,6 +140,7 @@ def value_counts(values, sort=True, ascending=False):
 
     return result
 
+
 def rank(values, axis=0, method='average', na_option='keep',
          ascending=True):
     """

pandas/core/daterange.py

@@ -43,8 +43,3 @@ def __setstate__(self, aug_state):
         self.offset = offset
         self.tzinfo = tzinfo
         Index.__setstate__(self, *index_state)
-
-def interval_range():
-    """
-    Return a fixed frequency interval index
-    """

pandas/core/index.py

@@ -11,7 +11,6 @@
 from pandas.util import py3compat
 import pandas.core.common as com
 import pandas._tseries as lib
-import pandas._engines as _gin
 
 
 __all__ = ['Index']
@@ -65,7 +64,7 @@ class Index(np.ndarray):
     name = None
     asi8 = None
 
-    _engine_type = _gin.ObjectEngine
+    _engine_type = lib.ObjectEngine
 
     def __new__(cls, data, dtype=None, copy=False, name=None):
         if isinstance(data, np.ndarray):
@@ -385,6 +384,20 @@ def asof(self, label):
 
         return label
 
+    def asof_locs(self, where, mask):
+        """
+
+        """
+        locs = self.values[mask].searchsorted(where.values, side='right')
+        locs = np.where(locs > 0, locs - 1, 0)
+
+        result = np.arange(len(self))[mask].take(locs)
+
+        first = mask.argmax()
+        result[(locs == 0) & (where < self.values[first])] = -1
+
+        return result
+
     def order(self, return_indexer=False, ascending=True):
         """
         Return sorted copy of Index
@@ -611,7 +624,7 @@ def get_value(self, series, key):
                 raise
 
             try:
-                return _gin.get_value_at(series, key)
+                return lib.get_value_at(series, key)
             except IndexError:
                 raise
             except TypeError:
@@ -1053,7 +1066,7 @@ class Int64Index(Index):
     _inner_indexer = lib.inner_join_indexer_int64
     _outer_indexer = lib.outer_join_indexer_int64
 
-    _engine_type = _gin.Int64Engine
+    _engine_type = lib.Int64Engine
 
     def __new__(cls, data, dtype=None, copy=False, name=None):
         if not isinstance(data, np.ndarray):
@@ -1283,7 +1296,7 @@ def get_value(self, series, key):
                 pass
 
             try:
-                return _gin.get_value_at(series, key)
+                return lib.get_value_at(series, key)
             except IndexError:
                 raise
             except TypeError:
@@ -2340,3 +2353,4 @@ def _maybe_box_dtindex(idx):
     if isinstance(idx, DatetimeIndex):
         return Index(_dt_box_array(idx.asi8), dtype='object')
     return idx
+

pandas/core/series.py

@@ -219,7 +219,7 @@ def f(self, *args, **kwargs):
     return f
 
 _stat_doc = """
-Return %(name)s  of values
+Return %(name)s of values
 %(na_action)s
 
 Parameters
@@ -238,7 +238,21 @@ def f(self, *args, **kwargs):
 _doc_ndarray_interface = ("Extra parameters are to preserve ndarray"
                           "interface.\n")
 
-#-------------------------------------------------------------------------------
+
+def _make_stat_func(nanop, name, shortname, na_action=_doc_exclude_na,
+                    extras=_doc_ndarray_interface):
+
+    @Substitution(name=name, shortname=shortname,
+                  na_action=na_action, extras=extras)
+    @Appender(_stat_doc)
+    def f(self, axis=0, dtype=None, out=None, skipna=True, level=None):
+        if level is not None:
+            return self._agg_by_level(shortname, level=level, skipna=skipna)
+        return nanop(self.values, skipna=skipna)
+    f.__name__ = shortname
+    return f
+
+#----------------------------------------------------------------------
 # Series class
 
 class Series(np.ndarray, generic.PandasObject):
@@ -985,21 +999,10 @@ def nunique(self):
         """
         return len(self.value_counts())
 
-    @Substitution(name='sum', shortname='sum', na_action=_doc_exclude_na,
-                  extras=_doc_ndarray_interface)
-    @Appender(_stat_doc)
-    def sum(self, axis=0, dtype=None, out=None, skipna=True, level=None):
-        if level is not None:
-            return self._agg_by_level('sum', level=level, skipna=skipna)
-        return nanops.nansum(self.values, skipna=skipna)
-
-    @Substitution(name='mean', shortname='mean', na_action=_doc_exclude_na,
-                  extras=_doc_ndarray_interface)
-    @Appender(_stat_doc)
-    def mean(self, axis=0, dtype=None, out=None, skipna=True, level=None):
-        if level is not None:
-            return self._agg_by_level('mean', level=level, skipna=skipna)
-        return nanops.nanmean(self.values, skipna=skipna)
+    sum = _make_stat_func(nanops.nansum, 'sum', 'sum')
+    mean = _make_stat_func(nanops.nanmean, 'mean', 'mean')
+    median = _make_stat_func(nanops.nanmedian, 'median', 'median', extras='')
+    prod = _make_stat_func(nanops.nanprod, 'product', 'prod', extras='')
 
     @Substitution(name='mean absolute deviation', shortname='mad',
                   na_action=_doc_exclude_na, extras='')
@@ -1011,22 +1014,6 @@ def mad(self, skipna=True, level=None):
         demeaned = self - self.mean(skipna=skipna)
         return np.abs(demeaned).mean(skipna=skipna)
 
-    @Substitution(name='median', shortname='median',
-                  na_action=_doc_exclude_na, extras='')
-    @Appender(_stat_doc)
-    def median(self, skipna=True, level=None):
-        if level is not None:
-            return self._agg_by_level('median', level=level, skipna=skipna)
-        return nanops.nanmedian(self.values, skipna=skipna)
-
-    @Substitution(name='product', shortname='product',
-                  na_action=_doc_exclude_na, extras='')
-    @Appender(_stat_doc)
-    def prod(self, axis=None, dtype=None, out=None, skipna=True, level=None):
-        if level is not None:
-            return self._agg_by_level('prod', level=level, skipna=skipna)
-        return nanops.nanprod(self.values, skipna=skipna)
-
     @Substitution(name='minimum', shortname='min',
                   na_action=_doc_exclude_na, extras='')
     @Appender(_stat_doc)
@@ -2278,7 +2265,7 @@ def shift(self, periods=1, freq=None, **kwds):
             return Series(self, index=self.index.shift(periods, offset),
                           name=self.name)
 
-    def asof(self, date):
+    def asof(self, where):
         """
         Return last good (non-NaN) value in TimeSeries if value is NaN for
         requested date.
@@ -2287,7 +2274,7 @@ def asof(self, date):
 
         Parameters
         ----------
-        date : datetime or similar value
+        wehre : date or array of dates
 
         Notes
         -----
@@ -2297,46 +2284,27 @@ def asof(self, date):
         -------
         value or NaN
         """
-        if isinstance(date, basestring):
-            date = datetools.to_datetime(date)
-
-        if not isinstance(date, (list, tuple, np.ndarray, Index)):
-            # treat scalar values differently
-            v = self.get(date)
-            if isnull(v):
-                candidates = self.index[notnull(self)]
-                index = candidates.searchsorted(lib.Timestamp(date))
-                if index > 0:
-                    asOfDate = candidates[index - 1]
-                    return self.get(asOfDate)
-                return nan
-            return v
+        if isinstance(where, basestring):
+            where = datetools.to_datetime(where)
 
-        if not isinstance(date, Index):
-            date = Index(date)
-
-        candidates = self.index[notnull(self)]
-
-        mask = date.isin(candidates)
-
-        there = self[date[mask]]
-        todo = date[-mask]
-
-        if len(there) == len(date):
-            if len(there) == 1:
-                return there[0]
-            return there
-
-        index = candidates.searchsorted(todo)
-        index = index - 1
-        asof_mask = index >= 0
-        asof = self.ix[candidates[index[asof_mask]]]
-        asof.index = todo[asof_mask]
-
-        if len(date) == 1 and len(asof) > 0:
-            return asof[0]
+        values = self.values
 
-        return there.combine_first(asof).reindex(date)
+        if not hasattr(where, '__iter__'):
+            if where < self.index[0]:
+                return np.nan
+            loc = self.index.searchsorted(where, side='right')
+            if loc > 0:
+                loc -= 1
+            while isnull(values[loc]) and loc > 0:
+                loc -= 1
+            return values[loc]
+
+        if not isinstance(where, Index):
+            where = Index(where)
+
+        locs = self.index.asof_locs(where, notnull(values))
+        new_values = com.take_1d(values, locs)
+        return Series(new_values, index=where, name=self.name)
 
     def interpolate(self, method='linear'):
         """

pandas/sparse/array.py

@@ -14,7 +14,7 @@
 
 from pandas._sparse import BlockIndex, IntIndex
 import pandas._sparse as splib
-import pandas._engines as _gin
+import pandas._tseries as lib
 
 
 def _sparse_op_wrap(op, name):
@@ -261,7 +261,7 @@ def _get_val_at(self, loc):
         if sp_loc == -1:
             return self.fill_value
         else:
-            return _gin.get_value_at(self, sp_loc)
+            return lib.get_value_at(self, sp_loc)
 
     def take(self, indices, axis=0):
         """