REF: reuse _combine instead of reset_dropped_locs (pandas-dev#35884)

Author: jbrockmendel

pandas/core/groupby/generic.py

@@ -21,7 +21,6 @@
     Mapping,
     Optional,
     Sequence,
-    Tuple,
     Type,
     Union,
 )
@@ -1025,16 +1024,14 @@ def _iterate_slices(self) -> Iterable[Series]:
     def _cython_agg_general(
         self, how: str, alt=None, numeric_only: bool = True, min_count: int = -1
     ) -> DataFrame:
-        agg_blocks, agg_items = self._cython_agg_blocks(
+        agg_mgr = self._cython_agg_blocks(
             how, alt=alt, numeric_only=numeric_only, min_count=min_count
         )
-        return self._wrap_agged_blocks(agg_blocks, items=agg_items)
+        return self._wrap_agged_blocks(agg_mgr.blocks, items=agg_mgr.items)
 
     def _cython_agg_blocks(
         self, how: str, alt=None, numeric_only: bool = True, min_count: int = -1
-    ) -> "Tuple[List[Block], Index]":
-        # TODO: the actual managing of mgr_locs is a PITA
-        # here, it should happen via BlockManager.combine
+    ) -> BlockManager:
 
         data: BlockManager = self._get_data_to_aggregate()
 
@@ -1124,15 +1121,14 @@ def blk_func(bvalues: ArrayLike) -> ArrayLike:
             res_values = cast_agg_result(result, bvalues, how)
             return res_values
 
-        skipped: List[int] = []
         for i, block in enumerate(data.blocks):
             try:
                 nbs = block.apply(blk_func)
             except (NotImplementedError, TypeError):
                 # TypeError -> we may have an exception in trying to aggregate
                 #  continue and exclude the block
                 # NotImplementedError -> "ohlc" with wrong dtype
-                skipped.append(i)
+                pass
             else:
                 agg_blocks.extend(nbs)
 
@@ -1141,9 +1137,8 @@ def blk_func(bvalues: ArrayLike) -> ArrayLike:
 
         # reset the locs in the blocks to correspond to our
         # current ordering
-        agg_items = data.reset_dropped_locs(agg_blocks, skipped)
-
-        return agg_blocks, agg_items
+        new_mgr = data._combine(agg_blocks)
+        return new_mgr
 
     def _aggregate_frame(self, func, *args, **kwargs) -> DataFrame:
         if self.grouper.nkeys != 1:

pandas/core/internals/managers.py

@@ -1486,38 +1486,6 @@ def unstack(self, unstacker, fill_value) -> "BlockManager":
         bm = BlockManager(new_blocks, [new_columns, new_index])
         return bm
 
-    def reset_dropped_locs(self, blocks: List[Block], skipped: List[int]) -> Index:
-        """
-        Decrement the mgr_locs of the given blocks with `skipped` removed.
-
-        Notes
-        -----
-        Alters each block's mgr_locs inplace.
-        """
-        ncols = len(self)
-
-        new_locs = [blk.mgr_locs.as_array for blk in blocks]
-        indexer = np.concatenate(new_locs)
-
-        new_items = self.items.take(np.sort(indexer))
-
-        if skipped:
-            # we need to adjust the indexer to account for the
-            #  items we have removed
-            deleted_items = [self.blocks[i].mgr_locs.as_array for i in skipped]
-            deleted = np.concatenate(deleted_items)
-            ai = np.arange(ncols)
-            mask = np.zeros(ncols)
-            mask[deleted] = 1
-            indexer = (ai - mask.cumsum())[indexer]
-
-        offset = 0
-        for blk in blocks:
-            loc = len(blk.mgr_locs)
-            blk.mgr_locs = indexer[offset : (offset + loc)]
-            offset += loc
-        return new_items
-
 
 class SingleBlockManager(BlockManager):
     """ manage a single block with """

pandas/core/window/rolling.py

@@ -561,8 +561,7 @@ def hfunc(bvalues: ArrayLike) -> ArrayLike:
         elif not len(res_blocks):
             return obj.astype("float64")
 
-        new_cols = mgr.reset_dropped_locs(res_blocks, skipped)
-        new_mgr = type(mgr).from_blocks(res_blocks, [new_cols, obj.index])
+        new_mgr = mgr._combine(res_blocks)
         out = obj._constructor(new_mgr)
         self._insert_on_column(out, obj)
         return out