diff --git a/pandas/core/generic.py b/pandas/core/generic.py index 04b69a78dad42..65100972ab57f 100644 --- a/pandas/core/generic.py +++ b/pandas/core/generic.py @@ -3895,6 +3895,7 @@ class max_speed return self._take(indices, axis) + @final def _take( self: NDFrameT, indices, @@ -3915,6 +3916,11 @@ def _take( and is_range_indexer(indices, len(self)) ): return self.copy(deep=None) + else: + # We can get here with a slice via DataFrame.__geittem__ + indices = np.arange( + indices.start, indices.stop, indices.step, dtype=np.intp + ) new_data = self._mgr.take( indices, diff --git a/pandas/core/internals/array_manager.py b/pandas/core/internals/array_manager.py index 1e3d88d42ef9b..f4b8daab403df 100644 --- a/pandas/core/internals/array_manager.py +++ b/pandas/core/internals/array_manager.py @@ -630,7 +630,7 @@ def _reindex_indexer( def take( self: T, - indexer, + indexer: npt.NDArray[np.intp], axis: AxisInt = 1, verify: bool = True, convert_indices: bool = True, @@ -638,13 +638,10 @@ def take( """ Take items along any axis. """ - axis = self._normalize_axis(axis) + assert isinstance(indexer, np.ndarray), type(indexer) + assert indexer.dtype == np.intp, indexer.dtype - indexer = ( - np.arange(indexer.start, indexer.stop, indexer.step, dtype="int64") - if isinstance(indexer, slice) - else np.asanyarray(indexer, dtype="int64") - ) + axis = self._normalize_axis(axis) if not indexer.ndim == 1: raise ValueError("indexer should be 1-dimensional") diff --git a/pandas/core/internals/managers.py b/pandas/core/internals/managers.py index 664a122015ba5..416a565bd3522 100644 --- a/pandas/core/internals/managers.py +++ b/pandas/core/internals/managers.py @@ -897,7 +897,7 @@ def _make_na_block( def take( self: T, - indexer, + indexer: npt.NDArray[np.intp], axis: AxisInt = 1, verify: bool = True, convert_indices: bool = True, @@ -905,7 +905,7 @@ def take( """ Take items along any axis. - indexer : np.ndarray or slice + indexer : np.ndarray[np.intp] axis : int, default 1 verify : bool, default True Check that all entries are between 0 and len(self) - 1, inclusive. @@ -917,12 +917,8 @@ def take( ------- BlockManager """ - # We have 6 tests that get here with a slice - indexer = ( - np.arange(indexer.start, indexer.stop, indexer.step, dtype=np.intp) - if isinstance(indexer, slice) - else np.asanyarray(indexer, dtype=np.intp) - ) + assert isinstance(indexer, np.ndarray), type(indexer) + assert indexer.dtype == np.intp, indexer.dtype n = self.shape[axis] if convert_indices: diff --git a/pandas/tests/internals/test_internals.py b/pandas/tests/internals/test_internals.py index f60c38c041fcf..964a4c4982481 100644 --- a/pandas/tests/internals/test_internals.py +++ b/pandas/tests/internals/test_internals.py @@ -1003,13 +1003,15 @@ def assert_take_ok(mgr, axis, indexer): for ax in range(mgr.ndim): # take/fancy indexer - assert_take_ok(mgr, ax, indexer=[]) - assert_take_ok(mgr, ax, indexer=[0, 0, 0]) - assert_take_ok(mgr, ax, indexer=list(range(mgr.shape[ax]))) + assert_take_ok(mgr, ax, indexer=np.array([], dtype=np.intp)) + assert_take_ok(mgr, ax, indexer=np.array([0, 0, 0], dtype=np.intp)) + assert_take_ok( + mgr, ax, indexer=np.array(list(range(mgr.shape[ax])), dtype=np.intp) + ) if mgr.shape[ax] >= 3: - assert_take_ok(mgr, ax, indexer=[0, 1, 2]) - assert_take_ok(mgr, ax, indexer=[-1, -2, -3]) + assert_take_ok(mgr, ax, indexer=np.array([0, 1, 2], dtype=np.intp)) + assert_take_ok(mgr, ax, indexer=np.array([-1, -2, -3], dtype=np.intp)) @pytest.mark.parametrize("mgr", MANAGERS) @pytest.mark.parametrize("fill_value", [None, np.nan, 100.0])