From beeb7e6afd1fcf4011b0348920c3c6bb7cde4a36 Mon Sep 17 00:00:00 2001 From: Brock Date: Sat, 15 Aug 2020 15:12:35 -0700 Subject: [PATCH] REF: insert self.on column _after_ concat --- pandas/core/window/rolling.py | 50 +++++++++++++++-------------------- 1 file changed, 21 insertions(+), 29 deletions(-) diff --git a/pandas/core/window/rolling.py b/pandas/core/window/rolling.py index 966773b7c6982..ac96258cbc3c9 100644 --- a/pandas/core/window/rolling.py +++ b/pandas/core/window/rolling.py @@ -38,7 +38,7 @@ from pandas.core.base import DataError, PandasObject, SelectionMixin, ShallowMixin import pandas.core.common as com from pandas.core.construction import extract_array -from pandas.core.indexes.api import Index, MultiIndex, ensure_index +from pandas.core.indexes.api import Index, MultiIndex from pandas.core.util.numba_ import NUMBA_FUNC_CACHE, maybe_use_numba from pandas.core.window.common import ( WindowGroupByMixin, @@ -402,36 +402,27 @@ def _wrap_results(self, results, blocks, obj, exclude=None) -> FrameOrSeries: return result final.append(result) - # if we have an 'on' column - # we want to put it back into the results - # in the same location - columns = self._selected_obj.columns - if self.on is not None and not self._on.equals(obj.index): - - name = self._on.name - final.append(Series(self._on, index=obj.index, name=name)) - - if self._selection is not None: - - selection = ensure_index(self._selection) - - # need to reorder to include original location of - # the on column (if its not already there) - if name not in selection: - columns = self.obj.columns - indexer = columns.get_indexer(selection.tolist() + [name]) - columns = columns.take(sorted(indexer)) - - # exclude nuisance columns so that they are not reindexed - if exclude is not None and exclude: - columns = [c for c in columns if c not in exclude] + exclude = exclude or [] + columns = [c for c in self._selected_obj.columns if c not in exclude] + if not columns and not len(final) and exclude: + raise DataError("No numeric types to aggregate") + elif not len(final): + return obj.astype("float64") - if not columns: - raise DataError("No numeric types to aggregate") + df = concat(final, axis=1).reindex(columns=columns, copy=False) - if not len(final): - return obj.astype("float64") - return concat(final, axis=1).reindex(columns=columns, copy=False) + # if we have an 'on' column we want to put it back into + # the results in the same location + if self.on is not None and not self._on.equals(obj.index): + name = self._on.name + extra_col = Series(self._on, index=obj.index, name=name) + if name not in df.columns and name not in df.index.names: + new_loc = len(df.columns) + df.insert(new_loc, name, extra_col) + elif name in df.columns: + # TODO: sure we want to overwrite results? + df[name] = extra_col + return df def _center_window(self, result, window) -> np.ndarray: """ @@ -2277,6 +2268,7 @@ def _get_window_indexer(self, window: int) -> GroupbyRollingIndexer: if isinstance(self.window, BaseIndexer): rolling_indexer = type(self.window) indexer_kwargs = self.window.__dict__ + assert isinstance(indexer_kwargs, dict) # We'll be using the index of each group later indexer_kwargs.pop("index_array", None) elif self.is_freq_type: