BUG: support count function for custom BaseIndexer rolling windows (pandas-dev#33605)

Author: AlexKirko

doc/source/whatsnew/v1.1.0.rst

@@ -174,8 +174,8 @@ Other API changes
 - Added :meth:`DataFrame.value_counts` (:issue:`5377`)
 - :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 ``count``, ``skew``, ``cov``, ``corr`` will now raise a ``NotImplementedError`` (:issue:`32865`)
-- Using a :func:`pandas.api.indexers.BaseIndexer` with ``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 ``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`)
 - Added a :func:`pandas.api.indexers.FixedForwardWindowIndexer` class to support forward-looking windows during ``rolling`` operations.
 -
 

pandas/core/window/common.py

@@ -328,6 +328,7 @@ def func(arg, window, min_periods=None):
 def validate_baseindexer_support(func_name: Optional[str]) -> None:
     # GH 32865: These functions work correctly with a BaseIndexer subclass
     BASEINDEXER_WHITELIST = {
+        "count",
         "min",
         "max",
         "mean",

pandas/core/window/rolling.py

@@ -1171,8 +1171,9 @@ class _Rolling_and_Expanding(_Rolling):
     )
 
     def count(self):
-        if isinstance(self.window, BaseIndexer):
-            validate_baseindexer_support("count")
+        # GH 32865. Using count with custom BaseIndexer subclass
+        # implementations shouldn't end up here
+        assert not isinstance(self.window, BaseIndexer)
 
         blocks, obj = self._create_blocks()
         results = []
@@ -1939,7 +1940,9 @@ def aggregate(self, func, *args, **kwargs):
     def count(self):
 
         # different impl for freq counting
-        if self.is_freq_type:
+        # GH 32865. Use a custom count function implementation
+        # when using a BaseIndexer subclass as a window
+        if self.is_freq_type or isinstance(self.window, BaseIndexer):
             window_func = self._get_roll_func("roll_count")
             return self._apply(window_func, center=self.center, name="count")
 

pandas/tests/window/test_base_indexer.py

@@ -82,7 +82,7 @@ def get_window_bounds(self, num_values, min_periods, center, closed):
         df.rolling(indexer, win_type="boxcar")
 
 
-@pytest.mark.parametrize("func", ["count", "skew", "cov", "corr"])
+@pytest.mark.parametrize("func", ["skew", "cov", "corr"])
 def test_notimplemented_functions(func):
     # GH 32865
     class CustomIndexer(BaseIndexer):
@@ -99,6 +99,7 @@ def get_window_bounds(self, num_values, min_periods, center, closed):
 @pytest.mark.parametrize(
     "func,np_func,expected,np_kwargs",
     [
+        ("count", len, [3.0, 3.0, 3.0, 3.0, 3.0, 3.0, 3.0, 3.0, 2.0, np.nan], {},),
         ("min", np.min, [0.0, 1.0, 2.0, 3.0, 4.0, 6.0, 6.0, 7.0, 8.0, np.nan], {},),
         (
             "max",