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

Author: AlexKirko

doc/source/whatsnew/v1.1.0.rst

@@ -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.
 -
 

pandas/_libs/window/aggregations.pyx

@@ -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
@@ -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()

pandas/core/window/rolling.py

@@ -1431,7 +1431,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):

pandas/tests/window/test_base_indexer.py

@@ -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):
@@ -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)