diff --git a/pandas/_libs/internals.pyi b/pandas/_libs/internals.pyi index a46a1747d1d8d..f3436e9c7afba 100644 --- a/pandas/_libs/internals.pyi +++ b/pandas/_libs/internals.pyi @@ -79,3 +79,5 @@ class BlockManager: _blklocs: np.ndarray def __init__(self, blocks: tuple[B, ...], axes: list[Index], verify_integrity=True): ... + + def get_slice(self: T, slobj: slice, axis: int=...) -> T: ... diff --git a/pandas/_libs/internals.pyx b/pandas/_libs/internals.pyx index 3fd580684a6a2..f3bc70ad8a26b 100644 --- a/pandas/_libs/internals.pyx +++ b/pandas/_libs/internals.pyx @@ -515,7 +515,7 @@ cdef class NumpyBlock(SharedBlock): self.values = values # @final # not useful in cython, but we _would_ annotate with @final - def getitem_block_index(self, slicer: slice) -> NumpyBlock: + cpdef NumpyBlock getitem_block_index(self, slice slicer): """ Perform __getitem__-like specialized to slicing along index. @@ -610,3 +610,30 @@ cdef class BlockManager: self._rebuild_blknos_and_blklocs() # ------------------------------------------------------------------- + # Indexing + + cdef BlockManager _get_index_slice(self, slobj): + cdef: + SharedBlock blk, nb + + nbs = [] + for blk in self.blocks: + nb = blk.getitem_block_index(slobj) + nbs.append(nb) + + new_axes = [self.axes[0], self.axes[1]._getitem_slice(slobj)] + return type(self)(tuple(nbs), new_axes, verify_integrity=False) + + def get_slice(self, slobj: slice, axis: int = 0) -> BlockManager: + + if axis == 0: + new_blocks = self._slice_take_blocks_ax0(slobj) + elif axis == 1: + return self._get_index_slice(slobj) + else: + raise IndexError("Requested axis not found in manager") + + new_axes = list(self.axes) + new_axes[axis] = new_axes[axis]._getitem_slice(slobj) + + return type(self)(tuple(new_blocks), new_axes, verify_integrity=False) diff --git a/pandas/core/internals/managers.py b/pandas/core/internals/managers.py index 373d3566e1e8a..97d605e2fa2d1 100644 --- a/pandas/core/internals/managers.py +++ b/pandas/core/internals/managers.py @@ -1116,21 +1116,6 @@ def fast_xs(self, loc: int) -> ArrayLike: return result - def get_slice(self, slobj: slice, axis: int = 0) -> BlockManager: - assert isinstance(slobj, slice), type(slobj) - - if axis == 0: - new_blocks = self._slice_take_blocks_ax0(slobj) - elif axis == 1: - new_blocks = [blk.getitem_block_index(slobj) for blk in self.blocks] - else: - raise IndexError("Requested axis not found in manager") - - new_axes = list(self.axes) - new_axes[axis] = new_axes[axis]._getitem_slice(slobj) - - return type(self)(tuple(new_blocks), new_axes, verify_integrity=False) - def iget(self, i: int) -> SingleBlockManager: """ Return the data as a SingleBlockManager.