Skip to content

Commit

Permalink
REF: Simplifying creation of window indexers internally (#37177)
Browse files Browse the repository at this point in the history
Co-authored-by: Matt Roeschke <[email protected]>
  • Loading branch information
mroeschke and Matt Roeschke authored Oct 17, 2020
1 parent 009ffa8 commit a8e2f92
Show file tree
Hide file tree
Showing 2 changed files with 31 additions and 39 deletions.
19 changes: 10 additions & 9 deletions pandas/core/window/expanding.py
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
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

0 comments on commit a8e2f92

Please sign in to comment.