CLN: Unify Window._apply_window and Rolling._apply functions

pandas/_libs/window.pyx

@@ -1675,9 +1675,13 @@ def roll_generic(object obj,
     return output
 
 
-def roll_window(ndarray[float64_t, ndim=1, cast=True] values,
-                ndarray[float64_t, ndim=1, cast=True] weights,
-                int minp, bint avg=True):
+# ----------------------------------------------------------------------
+# Rolling mean for weighted window
+
+
+def roll_window_mean(ndarray[float64_t, ndim=1, cast=True] values,
+                     ndarray[float64_t, ndim=1, cast=True] weights,
+                     int minp):
     """
     Assume len(weights) << len(values)
     """
@@ -1688,57 +1692,82 @@ def roll_window(ndarray[float64_t, ndim=1, cast=True] values,
 
     in_n = len(values)
     win_n = len(weights)
+
     output = np.zeros(in_n, dtype=float)
     counts = np.zeros(in_n, dtype=float)
-    if avg:
-        tot_wgt = np.zeros(in_n, dtype=float)
+    tot_wgt = np.zeros(in_n, dtype=float)
 
     minp = _check_minp(len(weights), minp, in_n)
 
-    if avg:
-        for win_i in range(win_n):
-            val_win = weights[win_i]
-            if val_win != val_win:
-                continue
-
-            for in_i from 0 <= in_i < in_n - (win_n - win_i) + 1:
-                val_in = values[in_i]
-                if val_in == val_in:
-                    output[in_i + (win_n - win_i) - 1] += val_in * val_win
-                    counts[in_i + (win_n - win_i) - 1] += 1
-                    tot_wgt[in_i + (win_n - win_i) - 1] += val_win
-
-        for in_i in range(in_n):
-            c = counts[in_i]
-            if c < minp:
+    for win_i in range(win_n):
+        val_win = weights[win_i]
+        if val_win != val_win:
+            continue
+
+        for in_i from 0 <= in_i < in_n - (win_n - win_i) + 1:
+            val_in = values[in_i]
+            if val_in == val_in:
+                output[in_i + (win_n - win_i) - 1] += val_in * val_win
+                counts[in_i + (win_n - win_i) - 1] += 1
+                tot_wgt[in_i + (win_n - win_i) - 1] += val_win
+
+    for in_i in range(in_n):
+        c = counts[in_i]
+        if c < minp:
+            output[in_i] = NaN
+        else:
+            w = tot_wgt[in_i]
+            if w == 0:
                 output[in_i] = NaN
             else:
-                w = tot_wgt[in_i]
-                if w == 0:
-                    output[in_i] = NaN
-                else:
-                    output[in_i] /= tot_wgt[in_i]
+                output[in_i] /= tot_wgt[in_i]
 
-    else:
-        for win_i in range(win_n):
-            val_win = weights[win_i]
-            if val_win != val_win:
-                continue
+    return output
 
-            for in_i from 0 <= in_i < in_n - (win_n - win_i) + 1:
-                val_in = values[in_i]
 
-                if val_in == val_in:
-                    output[in_i + (win_n - win_i) - 1] += val_in * val_win
-                    counts[in_i + (win_n - win_i) - 1] += 1
+# ----------------------------------------------------------------------
+# Rolling sum for weighted window
 
-        for in_i in range(in_n):
-            c = counts[in_i]
-            if c < minp:
-                output[in_i] = NaN
+
+def roll_window_sum(ndarray[float64_t, ndim=1, cast=True] values,
+                    ndarray[float64_t, ndim=1, cast=True] weights,
+                    int minp):
+    """
+    Assume len(weights) << len(values)
+    """
+    cdef:
+        ndarray[float64_t] output, counts
+        Py_ssize_t in_i, win_i, win_n, in_n
+        float64_t val_in, val_win
+
+    in_n = len(values)
+    win_n = len(weights)
+
+    output = np.zeros(in_n, dtype=float)
+    counts = np.zeros(in_n, dtype=float)
+
+    minp = _check_minp(len(weights), minp, in_n)
+
+    for win_i in range(win_n):
+        val_win = weights[win_i]
+        if val_win != val_win:
+            continue
+
+        for in_i from 0 <= in_i < in_n - (win_n - win_i) + 1:
+            val_in = values[in_i]
+
+            if val_in == val_in:
+                output[in_i + (win_n - win_i) - 1] += val_in * val_win
+                counts[in_i + (win_n - win_i) - 1] += 1
+
+    for in_i in range(in_n):
+        c = counts[in_i]
+        if c < minp:
+            output[in_i] = NaN
 
     return output
 
+
 # ----------------------------------------------------------------------
 # Exponentially weighted moving average
 

pandas/core/window.py

@@ -707,14 +707,14 @@ def _pop_args(win_type, arg_names, kwargs):
             # GH #15662. `False` makes symmetric window, rather than periodic.
             return sig.get_window(win_type, window, False).astype(float)
 
-    def _apply_window(self, mean=True, **kwargs):
+    def _apply_window(self, func, **kwargs):
         """
         Applies a moving window of type ``window_type`` on the data.
 
         Parameters
         ----------
-        mean : bool, default True
-            If True computes weighted mean, else weighted sum
+        func : str
+            Name of function to apply
 
         Returns
         -------
@@ -749,12 +749,13 @@ def _apply_window(self, mean=True, **kwargs):
             additional_nans = np.array([np.NaN] * offset)
 
             def f(arg, *args, **kwargs):
+                cfunc = getattr(libwindow, func)
                 minp = _use_window(self.min_periods, len(window))
-                return libwindow.roll_window(
+
+                return cfunc(
                     np.concatenate((arg, additional_nans)) if center else arg,
                     window,
                     minp,
-                    avg=mean,
                 )
 
             result = np.apply_along_axis(f, self.axis, values)
@@ -831,13 +832,13 @@ def aggregate(self, arg, *args, **kwargs):
     @Appender(_shared_docs["sum"])
     def sum(self, *args, **kwargs):
         nv.validate_window_func("sum", args, kwargs)
-        return self._apply_window(mean=False, **kwargs)
+        return self._apply_window("roll_window_sum", **kwargs)
 
     @Substitution(name="window")
     @Appender(_shared_docs["mean"])
     def mean(self, *args, **kwargs):
         nv.validate_window_func("mean", args, kwargs)
-        return self._apply_window(mean=True, **kwargs)
+        return self._apply_window("roll_window_mean", **kwargs)
 
 
 class _GroupByMixin(GroupByMixin):