API: rolling.apply will pass Series to function

closes #5071

Author: jreback

doc/source/whatsnew/v0.23.0.txt

@@ -479,6 +479,33 @@ If you wish to retain the old behavior while using Python >= 3.6, you can use
                'Taxes': -200,
                'Net result': 300}).sort_index()
 
+.. _whatsnew_0230.api_breaking.window_raw:
+
+Rolling/Expanding.apply() will by default send a ``Series`` to the function
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+The :func`Series.rolling`, :func:`DataFrame.rolling`, :func`Series.expanding`, :func:`DataFrame.expanding` methods when used with ``.apply()`` have gained a ``raw=False`` parameter.
+This is similar to :func:`DataFame.apply`. This parameter, ``False`` by default allows one to send a ``np.ndarray`` to the applied function, rather than the default of a ``Series``.
+This is a change from prior versions, when the applied function would *always* received an ndarray. This allow one to use pandas operations generically. (:issue:`5071`)
+
+.. ipython:: python
+
+   s = pd.Series(np.arange(5), np.arange(5) + 1)
+   s
+
+Pass a ``Series``:
+
+.. ipython:: python
+
+   s.rolling(2, min_periods=1).apply(lambda x: x.iloc[-1])
+
+Mimic the original behavior of passing a ndarray:
+
+.. ipython:: python
+
+   s.rolling(2, min_periods=1).apply(lambda x: x[-1], raw=True)
+
+
 .. _whatsnew_0230.api_breaking.deprecate_panel:
 
 Deprecate Panel

pandas/_libs/window.pyx

@@ -1432,39 +1432,44 @@ def roll_quantile(ndarray[float64_t, cast=True] input, int64_t win,
     return output
 
 
-def roll_generic(ndarray[float64_t, cast=True] input,
+def roll_generic(object obj,
                  int64_t win, int64_t minp, object index, object closed,
-                 int offset, object func,
+                 int offset, object func, bint raw,
                  object args, object kwargs):
     cdef:
         ndarray[double_t] output, counts, bufarr
+        ndarray[float64_t, cast=True] arr
         float64_t *buf
         float64_t *oldbuf
         int64_t nobs = 0, i, j, s, e, N
         bint is_variable
         ndarray[int64_t] start, end
 
-    if not input.flags.c_contiguous:
-        input = input.copy('C')
-
-    n = len(input)
+    n = len(obj)
     if n == 0:
-        return input
+        return obj
+
+    arr = np.asarray(obj)
+
+    # ndarray input
+    if not raw:
+        if not arr.flags.c_contiguous:
+            arr = arr.copy('C')
 
-    counts = roll_sum(np.concatenate([np.isfinite(input).astype(float),
+    counts = roll_sum(np.concatenate([np.isfinite(arr).astype(float),
                                       np.array([0.] * offset)]),
                       win, minp, index, closed)[offset:]
 
-    start, end, N, win, minp, is_variable = get_window_indexer(input, win,
+    start, end, N, win, minp, is_variable = get_window_indexer(arr, win,
                                                                minp, index,
                                                                closed,
                                                                floor=0)
 
     output = np.empty(N, dtype=float)
 
     if is_variable:
+        # variable window arr or series
 
-        # variable window
         if offset != 0:
             raise ValueError("unable to roll_generic with a non-zero offset")
 
@@ -1473,7 +1478,20 @@ def roll_generic(ndarray[float64_t, cast=True] input,
             e = end[i]
 
             if counts[i] >= minp:
-                output[i] = func(input[s:e], *args, **kwargs)
+                if raw:
+                    output[i] = func(arr[s:e], *args, **kwargs)
+                else:
+                    output[i] = func(obj.iloc[s:e], *args, **kwargs)
+            else:
+                output[i] = NaN
+
+    elif not raw:
+        # series
+        for i from 0 <= i < N:
+            if counts[i] >= minp:
+                sl = slice(int_max(i + offset - win + 1, 0),
+                           int_min(i + offset + 1, N))
+                output[i] = func(obj.iloc[sl], *args, **kwargs)
             else:
                 output[i] = NaN
 
@@ -1482,12 +1500,12 @@ def roll_generic(ndarray[float64_t, cast=True] input,
         # truncated windows at the beginning, through first full-length window
         for i from 0 <= i < (int_min(win, N) - offset):
             if counts[i] >= minp:
-                output[i] = func(input[0: (i + offset + 1)], *args, **kwargs)
+                output[i] = func(arr[0: (i + offset + 1)], *args, **kwargs)
             else:
                 output[i] = NaN
 
         # remaining full-length windows
-        buf = <float64_t *> input.data
+        buf = <float64_t *> arr.data
         bufarr = np.empty(win, dtype=float)
         oldbuf = <float64_t *> bufarr.data
         for i from (win - offset) <= i < (N - offset):
@@ -1502,7 +1520,7 @@ def roll_generic(ndarray[float64_t, cast=True] input,
         # truncated windows at the end
         for i from int_max(N - offset, 0) <= i < N:
             if counts[i] >= minp:
-                output[i] = func(input[int_max(i + offset - win + 1, 0): N],
+                output[i] = func(arr[int_max(i + offset - win + 1, 0): N],
                                  *args,
                                  **kwargs)
             else:

pandas/core/window.py

@@ -955,9 +955,20 @@ def count(self):
     ----------
     func : function
         Must produce a single value from an ndarray input
-        \*args and \*\*kwargs are passed to the function""")
+    raw : bool, default False
+        * ``False`` : passes each row or column as a Series to the
+          function.
+        * ``True`` : the passed function will receive ndarray objects
+          instead.
+          If you are just applying a NumPy reduction function this will
+          achieve much better performance.
+        .. versionadded:: 0.23.0
+
+    \*args and \*\*kwargs are passed to the function""")
+
+    def apply(self, func, raw=False, args=(), kwargs={}):
+        from pandas import Series
 
-    def apply(self, func, args=(), kwargs={}):
         # TODO: _level is unused?
         _level = kwargs.pop('_level', None)  # noqa
         window = self._get_window()
@@ -966,8 +977,11 @@ def apply(self, func, args=(), kwargs={}):
 
         def f(arg, window, min_periods, closed):
             minp = _use_window(min_periods, window)
-            return _window.roll_generic(arg, window, minp, indexi, closed,
-                                        offset, func, args, kwargs)
+            if not raw:
+                arg = Series(arg, index=self.obj.index)
+            return _window.roll_generic(
+                arg, window, minp, indexi,
+                closed, offset, func, raw, args, kwargs)
 
         return self._apply(f, func, args=args, kwargs=kwargs,
                            center=False)
@@ -1498,8 +1512,9 @@ def count(self):
     @Substitution(name='rolling')
     @Appender(_doc_template)
     @Appender(_shared_docs['apply'])
-    def apply(self, func, args=(), kwargs={}):
-        return super(Rolling, self).apply(func, args=args, kwargs=kwargs)
+    def apply(self, func, raw=False, args=(), kwargs={}):
+        return super(Rolling, self).apply(
+            func, raw=raw, args=args, kwargs=kwargs)
 
     @Substitution(name='rolling')
     @Appender(_shared_docs['sum'])
@@ -1756,8 +1771,9 @@ def count(self, **kwargs):
     @Substitution(name='expanding')
     @Appender(_doc_template)
     @Appender(_shared_docs['apply'])
-    def apply(self, func, args=(), kwargs={}):
-        return super(Expanding, self).apply(func, args=args, kwargs=kwargs)
+    def apply(self, func, raw=False, args=(), kwargs={}):
+        return super(Expanding, self).apply(
+            func, raw=raw, args=args, kwargs=kwargs)
 
     @Substitution(name='expanding')
     @Appender(_shared_docs['sum'])

pandas/tests/test_window.py

@@ -29,6 +29,11 @@ def assert_equal(left, right):
         tm.assert_frame_equal(left, right)
 
 
+@pytest.fixture(params=[True, False])
+def raw(request):
+    return request.param
+
+
 class Base(object):
 
     _nan_locs = np.arange(20, 40)
@@ -1150,17 +1155,17 @@ def test_rolling_quantile_param(self):
         with pytest.raises(TypeError):
             ser.rolling(3).quantile('foo')
 
-    def test_rolling_apply(self):
+    def test_rolling_apply(self, raw):
         # suppress warnings about empty slices, as we are deliberately testing
         # with a 0-length Series
         with warnings.catch_warnings():
             warnings.filterwarnings("ignore",
                                     message=".*(empty slice|0 for slice).*",
                                     category=RuntimeWarning)
 
-            ser = Series([])
-            tm.assert_series_equal(ser,
-                                   ser.rolling(10).apply(lambda x: x.mean()))
+            expected = Series([])
+            result = expected.rolling(10).apply(lambda x: x.mean(), raw=raw)
+            tm.assert_series_equal(result, expected)
 
             def f(x):
                 return x[np.isfinite(x)].mean()
@@ -1169,24 +1174,37 @@ def f(x):
 
         # GH 8080
         s = Series([None, None, None])
-        result = s.rolling(2, min_periods=0).apply(lambda x: len(x))
+        result = s.rolling(2, min_periods=0).apply(lambda x: len(x), raw=raw)
         expected = Series([1., 2., 2.])
         tm.assert_series_equal(result, expected)
 
-        result = s.rolling(2, min_periods=0).apply(len)
+        result = s.rolling(2, min_periods=0).apply(len, raw=raw)
         tm.assert_series_equal(result, expected)
 
-    def test_rolling_apply_out_of_bounds(self):
+    def test_rolling_apply_out_of_bounds(self, raw):
         # #1850
         vals = pd.Series([1, 2, 3, 4])
 
-        result = vals.rolling(10).apply(np.sum)
+        result = vals.rolling(10).apply(np.sum, raw=raw)
         assert result.isna().all()
 
-        result = vals.rolling(10, min_periods=1).apply(np.sum)
+        result = vals.rolling(10, min_periods=1).apply(np.sum, raw=raw)
         expected = pd.Series([1, 3, 6, 10], dtype=float)
         tm.assert_almost_equal(result, expected)
 
+    def test_rolling_apply_with_pandas_objects(self):
+        # 5071
+        df = pd.DataFrame({'A': np.random.randn(5),
+                           'B': np.random.randint(0, 10, size=5)})
+
+        def f(x):
+            return x.iloc[-1]
+
+        df.rolling(2).apply(f)
+        result = df.rolling(2).apply(f)
+        expected = df[1:].reindex_like(df)
+        tm.assert_frame_equal(result, expected)
+
     def test_rolling_std(self):
         self._check_moment_func(lambda x: np.std(x, ddof=1),
                                 name='std')
@@ -2805,20 +2823,25 @@ def expanding_func(x, min_periods=1, center=False, axis=0):
             return getattr(exp, func)()
         self._check_expanding(expanding_func, static_comp, preserve_nan=False)
 
-    def test_expanding_apply(self):
+    def test_expanding_apply(self, raw):
 
         def expanding_mean(x, min_periods=1):
+
             exp = x.expanding(min_periods=min_periods)
-            return exp.apply(lambda x: x.mean())
+            result = exp.apply(lambda x: x.mean(), raw=raw)
+            return result
 
-        self._check_expanding(expanding_mean, np.mean)
+        # TODO(jreback), needed to add preserve_nan=False
+        # here to make this pass
+        self._check_expanding(expanding_mean, np.mean, preserve_nan=False)
 
         ser = Series([])
-        tm.assert_series_equal(ser, ser.expanding().apply(lambda x: x.mean()))
+        tm.assert_series_equal(ser, ser.expanding().apply(
+            lambda x: x.mean(), raw=raw))
 
         # GH 8080
         s = Series([None, None, None])
-        result = s.expanding(min_periods=0).apply(lambda x: len(x))
+        result = s.expanding(min_periods=0).apply(lambda x: len(x), raw=raw)
         expected = Series([1., 2., 3.])
         tm.assert_series_equal(result, expected)
 
@@ -3057,7 +3080,7 @@ def func(x):
             expected = g.apply(func)
             tm.assert_series_equal(result, expected)
 
-    def test_rolling_apply(self):
+    def test_rolling_apply(self, raw):
         g = self.frame.groupby('A')
         r = g.rolling(window=4)