diff --git a/pandas/core/window/ewm.py b/pandas/core/window/ewm.py index 7a2d8e84bec76..c57c434dd3040 100644 --- a/pandas/core/window/ewm.py +++ b/pandas/core/window/ewm.py @@ -12,9 +12,7 @@ from pandas.util._decorators import Appender, Substitution, doc from pandas.core.dtypes.common import is_datetime64_ns_dtype -from pandas.core.dtypes.generic import ABCDataFrame -from pandas.core.base import DataError import pandas.core.common as common from pandas.core.window.common import _doc_template, _shared_docs, zsqrt from pandas.core.window.rolling import _flex_binary_moment, _Rolling @@ -302,30 +300,13 @@ def _apply(self, func): ------- y : same type as input argument """ - blocks, obj = self._create_blocks(self._selected_obj) - block_list = list(blocks) - - results = [] - exclude = [] - for i, b in enumerate(blocks): - try: - values = self._prep_values(b.values) - - except (TypeError, NotImplementedError) as err: - if isinstance(obj, ABCDataFrame): - exclude.extend(b.columns) - del block_list[i] - continue - else: - raise DataError("No numeric types to aggregate") from err + def homogeneous_func(values: np.ndarray): if values.size == 0: - results.append(values.copy()) - continue + return values.copy() + return np.apply_along_axis(func, self.axis, values) - results.append(np.apply_along_axis(func, self.axis, values)) - - return self._wrap_results(results, block_list, obj, exclude) + return self._apply_blockwise(homogeneous_func) @Substitution(name="ewm", func_name="mean") @Appender(_doc_template) diff --git a/pandas/core/window/rolling.py b/pandas/core/window/rolling.py index 445f179248226..e1be7c72f274f 100644 --- a/pandas/core/window/rolling.py +++ b/pandas/core/window/rolling.py @@ -12,7 +12,7 @@ from pandas._libs.tslibs import BaseOffset, to_offset import pandas._libs.window.aggregations as window_aggregations -from pandas._typing import Axis, FrameOrSeries, Scalar +from pandas._typing import ArrayLike, Axis, FrameOrSeries, Scalar from pandas.compat._optional import import_optional_dependency from pandas.compat.numpy import function as nv from pandas.util._decorators import Appender, Substitution, cache_readonly, doc @@ -487,6 +487,38 @@ def _get_window_indexer(self, window: int) -> BaseIndexer: return VariableWindowIndexer(index_array=self._on.asi8, window_size=window) return FixedWindowIndexer(window_size=window) + def _apply_blockwise( + self, homogeneous_func: Callable[..., ArrayLike] + ) -> FrameOrSeries: + """ + Apply the given function to the DataFrame broken down into homogeneous + sub-frames. + """ + # This isn't quite blockwise, since `blocks` is actually a collection + # of homogenenous DataFrames. + blocks, obj = self._create_blocks(self._selected_obj) + + skipped: List[int] = [] + results: List[ArrayLike] = [] + exclude: List[Scalar] = [] + for i, b in enumerate(blocks): + try: + values = self._prep_values(b.values) + + except (TypeError, NotImplementedError) as err: + if isinstance(obj, ABCDataFrame): + skipped.append(i) + exclude.extend(b.columns) + continue + else: + raise DataError("No numeric types to aggregate") from err + + result = homogeneous_func(values) + results.append(result) + + block_list = [blk for i, blk in enumerate(blocks) if i not in skipped] + return self._wrap_results(results, block_list, obj, exclude) + def _apply( self, func: Callable, @@ -524,30 +556,14 @@ def _apply( """ win_type = self._get_win_type(kwargs) window = self._get_window(win_type=win_type) - - blocks, obj = self._create_blocks(self._selected_obj) - block_list = list(blocks) window_indexer = self._get_window_indexer(window) - results = [] - exclude: List[Scalar] = [] - for i, b in enumerate(blocks): - try: - values = self._prep_values(b.values) - - except (TypeError, NotImplementedError) as err: - if isinstance(obj, ABCDataFrame): - exclude.extend(b.columns) - del block_list[i] - continue - else: - raise DataError("No numeric types to aggregate") from err + def homogeneous_func(values: np.ndarray): + # calculation function if values.size == 0: - results.append(values.copy()) - continue + return values.copy() - # calculation function offset = calculate_center_offset(window) if center else 0 additional_nans = np.array([np.nan] * offset) @@ -594,9 +610,9 @@ def calc(x): if center: result = self._center_window(result, window) - results.append(result) + return result - return self._wrap_results(results, block_list, obj, exclude) + return self._apply_blockwise(homogeneous_func) def aggregate(self, func, *args, **kwargs): result, how = self._aggregate(func, *args, **kwargs)