Skip to content

Commit

Permalink
Make Rolling Reducible.
Browse files Browse the repository at this point in the history
  • Loading branch information
vyasr committed Dec 16, 2021
1 parent 6bd3547 commit 3f63397
Show file tree
Hide file tree
Showing 2 changed files with 23 additions and 11 deletions.
2 changes: 2 additions & 0 deletions python/cudf/cudf/core/groupby/groupby.py
Original file line number Diff line number Diff line change
Expand Up @@ -49,6 +49,8 @@ class GroupBy(Serializable, Reducible):
"nunique",
"first",
"last",
"var",
"std",
}

_MAX_GROUPS_BEFORE_WARN = 100
Expand Down
32 changes: 21 additions & 11 deletions python/cudf/cudf/core/window/rolling.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,11 +11,12 @@
from cudf.api.types import is_integer, is_number
from cudf.core import column
from cudf.core.column.column import as_column
from cudf.core.reductions import Reducible
from cudf.utils import cudautils
from cudf.utils.utils import GetAttrGetItemMixin


class Rolling(GetAttrGetItemMixin):
class Rolling(GetAttrGetItemMixin, Reducible):
"""
Rolling window calculations.
Expand Down Expand Up @@ -163,6 +164,15 @@ class Rolling(GetAttrGetItemMixin):

_time_window = False

_VALID_REDUCTIONS = {
"sum",
"min",
"max",
"mean",
"var",
"std",
}

def __init__(
self,
obj,
Expand Down Expand Up @@ -256,17 +266,17 @@ def _apply_agg(self, agg_name):
else:
return self._apply_agg_dataframe(self.obj, agg_name)

def sum(self):
return self._apply_agg("sum")

def min(self):
return self._apply_agg("min")

def max(self):
return self._apply_agg("max")
def _reduce(
self, op: str, *args, **kwargs,
):
"""Calculate the rolling {op}.
def mean(self):
return self._apply_agg("mean")
Returns
-------
Series or DataFrame
Return type is the same as the original object.
"""
return self._apply_agg(op)

def var(self, ddof=1):
self.agg_params["ddof"] = ddof
Expand Down

0 comments on commit 3f63397

Please sign in to comment.