Skip to content

REF: Simplify creation of rolling window indexers internally #37177

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.

Sign up for GitHub

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 1 commit into from
Oct 17, 2020
Merged

REF: Simplify creation of rolling window indexers internally #37177

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 10 additions & 9 deletions pandas/core/window/expanding.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@
from pandas.util._decorators import Appender, Substitution, doc

from pandas.core.window.common import _doc_template, _shared_docs
from pandas.core.window.indexers import ExpandingIndexer, GroupbyIndexer
from pandas.core.window.indexers import BaseIndexer, ExpandingIndexer, GroupbyIndexer
from pandas.core.window.rolling import BaseWindowGroupby, RollingAndExpandingMixin


Expand Down Expand Up @@ -68,11 +68,17 @@ def __init__(self, obj, min_periods=1, center=None, axis=0, **kwargs):
def _constructor(self):
return Expanding

def _get_window(
def _get_window_indexer(self) -> BaseIndexer:
"""
Return an indexer class that will compute the window start and end bounds
"""
return ExpandingIndexer()

def _get_cov_corr_window(
self, other: Optional[Union[np.ndarray, FrameOrSeries]] = None, **kwargs
) -> int:
"""
Get the window length over which to perform some operation.
Get the window length over which to perform cov and corr operations.

Parameters
----------
Expand Down Expand Up @@ -275,15 +281,10 @@ class ExpandingGroupby(BaseWindowGroupby, Expanding):
Provide a expanding groupby implementation.
"""

def _get_window_indexer(self, window: int) -> GroupbyIndexer:
def _get_window_indexer(self) -> GroupbyIndexer:
"""
Return an indexer class that will compute the window start and end bounds

Parameters
----------
window : int
window size for FixedWindowIndexer (unused)

Returns
-------
GroupbyIndexer
Expand Down
51 changes: 21 additions & 30 deletions pandas/core/window/rolling.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -213,9 +213,9 @@ def __getattr__(self, attr: str):
def _dir_additions(self):
return self.obj._dir_additions()

def _get_window(
def _get_cov_corr_window(
self, other: Optional[Union[np.ndarray, FrameOrSeries]] = None
) -> int:
) -> Optional[Union[int, timedelta, BaseOffset, BaseIndexer]]:
"""
Return window length.

Expand All @@ -228,8 +228,6 @@ def _get_window(
-------
window : int
"""
if isinstance(self.window, BaseIndexer):
return self.min_periods or 0
return self.window

@property
Expand All @@ -249,11 +247,10 @@ def __repr__(self) -> str:
return f"{self._window_type} [{attrs}]"

def __iter__(self):
window = self._get_window()
obj = self._create_data(self._selected_obj)
index = self._get_window_indexer(window=window)
indexer = self._get_window_indexer()

start, end = index.get_window_bounds(
start, end = indexer.get_window_bounds(
num_values=len(obj),
min_periods=self.min_periods,
center=self.center,
Expand Down Expand Up @@ -340,15 +337,17 @@ def _get_roll_func(self, func_name: str) -> Callable[..., Any]:
)
return window_func

def _get_window_indexer(self, window: int) -> BaseIndexer:
def _get_window_indexer(self) -> BaseIndexer:
"""
Return an indexer class that will compute the window start and end bounds
"""
if isinstance(self.window, BaseIndexer):
return self.window
if self.is_freq_type:
return VariableWindowIndexer(index_array=self._on.asi8, window_size=window)
return FixedWindowIndexer(window_size=window)
return VariableWindowIndexer(
index_array=self._on.asi8, window_size=self.window
)
return FixedWindowIndexer(window_size=self.window)

def _apply_series(
self, homogeneous_func: Callable[..., ArrayLike], name: Optional[str] = None
Expand Down Expand Up @@ -428,8 +427,7 @@ def _apply(
-------
y : type of input
"""
window = self._get_window()
window_indexer = self._get_window_indexer(window)
window_indexer = self._get_window_indexer()
min_periods = (
self.min_periods
if self.min_periods is not None
Expand Down Expand Up @@ -1750,14 +1748,10 @@ def cov(self, other=None, pairwise=None, ddof=1, **kwargs):
# GH 32865. We leverage rolling.mean, so we pass
# to the rolling constructors the data used when constructing self:
# window width, frequency data, or a BaseIndexer subclass
if isinstance(self.window, BaseIndexer):
window = self.window
else:
# GH 16058: offset window
if self.is_freq_type:
window = self.win_freq
else:
window = self._get_window(other)
# GH 16058: offset window
window = (
self._get_cov_corr_window(other) if not self.is_freq_type else self.win_freq
)

def _get_cov(X, Y):
# GH #12373 : rolling functions error on float32 data
Expand Down Expand Up @@ -1899,10 +1893,10 @@ def corr(self, other=None, pairwise=None, **kwargs):
# GH 32865. We leverage rolling.cov and rolling.std here, so we pass
# to the rolling constructors the data used when constructing self:
# window width, frequency data, or a BaseIndexer subclass
if isinstance(self.window, BaseIndexer):
window = self.window
else:
window = self._get_window(other) if not self.is_freq_type else self.win_freq
# GH 16058: offset window
window = (
self._get_cov_corr_window(other) if not self.is_freq_type else self.win_freq
)

def _get_corr(a, b):
a = a.rolling(
Expand Down Expand Up @@ -2208,28 +2202,25 @@ class RollingGroupby(BaseWindowGroupby, Rolling):
Provide a rolling groupby implementation.
"""

def _get_window_indexer(self, window: int) -> GroupbyIndexer:
def _get_window_indexer(self) -> GroupbyIndexer:
"""
Return an indexer class that will compute the window start and end bounds

Parameters
----------
window : int
window size for FixedWindowIndexer

Returns
-------
GroupbyIndexer
"""
rolling_indexer: Type[BaseIndexer]
indexer_kwargs: Optional[Dict[str, Any]] = None
index_array = self._on.asi8
window = self.window
if isinstance(self.window, BaseIndexer):
rolling_indexer = type(self.window)
indexer_kwargs = self.window.__dict__
assert isinstance(indexer_kwargs, dict) # for mypy
# We'll be using the index of each group later
indexer_kwargs.pop("index_array", None)
window = 0
elif self.is_freq_type:
rolling_indexer = VariableWindowIndexer
else:
Expand Down