Skip to content

BUG: support median function for custom BaseIndexer rolling windows #33626

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 10 commits into from
Apr 21, 2020
2 changes: 1 addition & 1 deletion doc/source/whatsnew/v1.1.0.rst
Original file line number Diff line number Diff line change
Expand Up @@ -175,7 +175,7 @@ Other API changes
- :meth:`Groupby.groups` now returns an abbreviated representation when called on large dataframes (:issue:`1135`)
- ``loc`` lookups with an object-dtype :class:`Index` and an integer key will now raise ``KeyError`` instead of ``TypeError`` when key is missing (:issue:`31905`)
- Using a :func:`pandas.api.indexers.BaseIndexer` with ``skew``, ``cov``, ``corr`` will now raise a ``NotImplementedError`` (:issue:`32865`)
- Using a :func:`pandas.api.indexers.BaseIndexer` with ``count``, ``min``, ``max`` will now return correct results for any monotonic :func:`pandas.api.indexers.BaseIndexer` descendant (:issue:`32865`)
- Using a :func:`pandas.api.indexers.BaseIndexer` with ``count``, ``min``, ``max``, ``median`` will now return correct results for any monotonic :func:`pandas.api.indexers.BaseIndexer` descendant (:issue:`32865`)
- Added a :func:`pandas.api.indexers.FixedForwardWindowIndexer` class to support forward-looking windows during ``rolling`` operations.
-

Expand Down
5 changes: 3 additions & 2 deletions pandas/_libs/window/aggregations.pyx
Original file line number Diff line number Diff line change
Expand Up @@ -843,7 +843,8 @@ def roll_kurt_variable(ndarray[float64_t] values, ndarray[int64_t] start,


def roll_median_c(ndarray[float64_t] values, ndarray[int64_t] start,
ndarray[int64_t] end, int64_t minp, int64_t win):
ndarray[int64_t] end, int64_t minp, int64_t win=0):
# GH 32865. win argument kept for compatibility
cdef:
float64_t val, res, prev
bint err = False
Expand All @@ -858,7 +859,7 @@ def roll_median_c(ndarray[float64_t] values, ndarray[int64_t] start,
# actual skiplist ops outweigh any window computation costs
output = np.empty(N, dtype=float)

if win == 0 or (end - start).max() == 0:
if (end - start).max() == 0:
output[:] = NaN
return output
win = (end - start).max()
Expand Down
3 changes: 2 additions & 1 deletion pandas/core/window/rolling.py
Original file line number Diff line number Diff line change
Expand Up @@ -1429,7 +1429,8 @@ def mean(self, *args, **kwargs):

def median(self, **kwargs):
window_func = self._get_roll_func("roll_median_c")
window_func = partial(window_func, win=self._get_window())
# GH 32865. Move max window size calculation to
# the median function implementation
return self._apply(window_func, center=self.center, name="median", **kwargs)

def std(self, ddof=1, *args, **kwargs):
Expand Down
18 changes: 18 additions & 0 deletions pandas/tests/window/test_base_indexer.py
Original file line number Diff line number Diff line change
Expand Up @@ -141,6 +141,12 @@ def get_window_bounds(self, num_values, min_periods, center, closed):
],
{"ddof": 1},
),
(
"median",
np.median,
[1.0, 2.0, 3.0, 4.0, 6.0, 7.0, 7.0, 8.0, 8.5, np.nan],
{},
),
],
)
def test_rolling_forward_window(constructor, func, np_func, expected, np_kwargs):
Expand All @@ -162,7 +168,19 @@ def test_rolling_forward_window(constructor, func, np_func, expected, np_kwargs)

rolling = constructor(values).rolling(window=indexer, min_periods=2)
result = getattr(rolling, func)()

# Check that the function output matches the explicitly provided array
expected = constructor(expected)
tm.assert_equal(result, expected)

# Check that the rolling function output matches applying an alternative
# function to the rolling window object
expected2 = constructor(rolling.apply(lambda x: np_func(x, **np_kwargs)))
tm.assert_equal(result, expected2)

# Check that the function output matches applying an alternative function
# if min_periods isn't specified
rolling3 = constructor(values).rolling(window=indexer)
result3 = getattr(rolling3, func)()
expected3 = constructor(rolling3.apply(lambda x: np_func(x, **np_kwargs)))
tm.assert_equal(result3, expected3)