REF: simplify BlockManager.idelete (#40548)

Author: jbrockmendel

pandas/core/generic.py

@@ -4063,7 +4063,7 @@ def __delitem__(self, key) -> None:
             # there was no match, this call should raise the appropriate
             # exception:
             loc = self.axes[-1].get_loc(key)
-            self._mgr.idelete(loc)
+            self._mgr = self._mgr.idelete(loc)
 
         # delete from the caches
         try:

pandas/core/internals/array_manager.py

@@ -848,6 +848,7 @@ def idelete(self, indexer):
 
         self.arrays = [self.arrays[i] for i in np.nonzero(to_keep)[0]]
         self._axes = [self._axes[0], self._axes[1][to_keep]]
+        return self
 
     def iset(self, loc: Union[int, slice, np.ndarray], value):
         """
@@ -1247,7 +1248,7 @@ def apply(self, func, **kwargs):
     def setitem(self, indexer, value):
         return self.apply_with_block("setitem", indexer=indexer, value=value)
 
-    def idelete(self, indexer):
+    def idelete(self, indexer) -> SingleArrayManager:
         """
         Delete selected locations in-place (new array, same ArrayManager)
         """
@@ -1256,6 +1257,7 @@ def idelete(self, indexer):
 
         self.arrays = [self.arrays[0][to_keep]]
         self._axes = [self._axes[0][to_keep]]
+        return self
 
     def _get_data_subset(self, predicate: Callable) -> ArrayManager:
         # used in get_numeric_data / get_bool_data

pandas/core/internals/managers.py

@@ -1046,41 +1046,18 @@ def iget_values(self, i: int) -> ArrayLike:
         values = block.iget(self.blklocs[i])
         return values
 
-    def idelete(self, indexer):
+    def idelete(self, indexer) -> BlockManager:
         """
-        Delete selected locations in-place (new block and array, same BlockManager)
+        Delete selected locations, returning a new BlockManager.
         """
         is_deleted = np.zeros(self.shape[0], dtype=np.bool_)
         is_deleted[indexer] = True
-        ref_loc_offset = -is_deleted.cumsum()
+        taker = (~is_deleted).nonzero()[0]
 
-        is_blk_deleted = [False] * len(self.blocks)
-
-        if isinstance(indexer, int):
-            affected_start = indexer
-        else:
-            affected_start = is_deleted.nonzero()[0][0]
-
-        for blkno, _ in _fast_count_smallints(self.blknos[affected_start:]):
-            blk = self.blocks[blkno]
-            bml = blk.mgr_locs
-            blk_del = is_deleted[bml.indexer].nonzero()[0]
-
-            if len(blk_del) == len(bml):
-                is_blk_deleted[blkno] = True
-                continue
-            elif len(blk_del) != 0:
-                blk.delete(blk_del)
-                bml = blk.mgr_locs
-
-            blk.mgr_locs = bml.add(ref_loc_offset[bml.indexer])
-
-        # FIXME: use Index.delete as soon as it uses fastpath=True
-        self.axes[0] = self.items[~is_deleted]
-        self.blocks = tuple(
-            b for blkno, b in enumerate(self.blocks) if not is_blk_deleted[blkno]
-        )
-        self._rebuild_blknos_and_blklocs()
+        nbs = self._slice_take_blocks_ax0(taker, only_slice=True)
+        new_columns = self.items[~is_deleted]
+        axes = [new_columns, self.axes[1]]
+        return type(self)._simple_new(tuple(nbs), axes)
 
     def iset(self, loc: Union[int, slice, np.ndarray], value):
         """
@@ -1683,14 +1660,15 @@ def _consolidate_check(self):
     def _consolidate_inplace(self):
         pass
 
-    def idelete(self, indexer):
+    def idelete(self, indexer) -> SingleBlockManager:
         """
         Delete single location from SingleBlockManager.
 
         Ensures that self.blocks doesn't become empty.
         """
         self._block.delete(indexer)
         self.axes[0] = self.axes[0].delete(indexer)
+        return self
 
     def fast_xs(self, loc):
         """