REF: de-duplicate NDFrame.take, remove Manager.take keyword (#51482)

Author: jbrockmendel

pandas/core/generic.py

@@ -3816,6 +3816,7 @@ def _clear_item_cache(self) -> None:
     # ----------------------------------------------------------------------
     # Indexing Methods
 
+    @final
     def take(self: NDFrameT, indices, axis: Axis = 0, **kwargs) -> NDFrameT:
         """
         Return the elements in the given *positional* indices along an axis.
@@ -3893,20 +3894,6 @@ class  max_speed
 
         nv.validate_take((), kwargs)
 
-        return self._take(indices, axis)
-
-    @final
-    def _take(
-        self: NDFrameT,
-        indices,
-        axis: Axis = 0,
-        convert_indices: bool_t = True,
-    ) -> NDFrameT:
-        """
-        Internal version of the `take` allowing specification of additional args.
-
-        See the docstring of `take` for full explanation of the parameters.
-        """
         if not isinstance(indices, slice):
             indices = np.asarray(indices, dtype=np.intp)
             if (
@@ -3916,8 +3903,14 @@ def _take(
                 and is_range_indexer(indices, len(self))
             ):
                 return self.copy(deep=None)
+        elif self.ndim == 1:
+            # TODO: be consistent here for DataFrame vs Series
+            raise TypeError(
+                f"{type(self).__name__}.take requires a sequence of integers, "
+                "not slice."
+            )
         else:
-            # We can get here with a slice via DataFrame.__geittem__
+            # We can get here with a slice via DataFrame.__getitem__
             indices = np.arange(
                 indices.start, indices.stop, indices.step, dtype=np.intp
             )
@@ -3926,21 +3919,23 @@ def _take(
             indices,
             axis=self._get_block_manager_axis(axis),
             verify=True,
-            convert_indices=convert_indices,
         )
         return self._constructor(new_data).__finalize__(self, method="take")
 
+    @final
     def _take_with_is_copy(self: NDFrameT, indices, axis: Axis = 0) -> NDFrameT:
         """
         Internal version of the `take` method that sets the `_is_copy`
         attribute to keep track of the parent dataframe (using in indexing
         for the SettingWithCopyWarning).
 
+        For Series this does the same as the public take (it never sets `_is_copy`).
+
         See the docstring of `take` for full explanation of the parameters.
         """
-        result = self._take(indices=indices, axis=axis)
+        result = self.take(indices=indices, axis=axis)
         # Maybe set copy if we didn't actually change the index.
-        if not result._get_axis(axis).equals(self._get_axis(axis)):
+        if self.ndim == 2 and not result._get_axis(axis).equals(self._get_axis(axis)):
             result._set_is_copy(self)
         return result
 

pandas/core/groupby/groupby.py

@@ -1572,7 +1572,10 @@ def _wrap_transform_fast_result(self, result: NDFrameT) -> NDFrameT:
             # GH#46209
             # Don't convert indices: negative indices need to give rise
             # to null values in the result
-            output = result._take(ids, axis=axis, convert_indices=False)
+            new_ax = result.axes[axis].take(ids)
+            output = result._reindex_with_indexers(
+                {axis: (new_ax, ids)}, allow_dups=True, copy=False
+            )
             output = output.set_axis(obj._get_axis(self.axis), axis=axis)
         return output
 

pandas/core/internals/array_manager.py

@@ -633,7 +633,6 @@ def take(
         indexer: npt.NDArray[np.intp],
         axis: AxisInt = 1,
         verify: bool = True,
-        convert_indices: bool = True,
     ) -> T:
         """
         Take items along any axis.
@@ -647,8 +646,7 @@ def take(
             raise ValueError("indexer should be 1-dimensional")
 
         n = self.shape_proper[axis]
-        if convert_indices:
-            indexer = maybe_convert_indices(indexer, n, verify=verify)
+        indexer = maybe_convert_indices(indexer, n, verify=verify)
 
         new_labels = self._axes[axis].take(indexer)
         return self._reindex_indexer(

pandas/core/internals/managers.py

@@ -910,7 +910,6 @@ def take(
         indexer: npt.NDArray[np.intp],
         axis: AxisInt = 1,
         verify: bool = True,
-        convert_indices: bool = True,
     ) -> T:
         """
         Take items along any axis.
@@ -920,8 +919,6 @@ def take(
         verify : bool, default True
             Check that all entries are between 0 and len(self) - 1, inclusive.
             Pass verify=False if this check has been done by the caller.
-        convert_indices : bool, default True
-            Whether to attempt to convert indices to positive values.
 
         Returns
         -------
@@ -931,8 +928,7 @@ def take(
         assert indexer.dtype == np.intp, indexer.dtype
 
         n = self.shape[axis]
-        if convert_indices:
-            indexer = maybe_convert_indices(indexer, n, verify=verify)
+        indexer = maybe_convert_indices(indexer, n, verify=verify)
 
         new_labels = self.axes[axis].take(indexer)
         return self.reindex_indexer(

pandas/core/series.py

@@ -96,7 +96,6 @@
     maybe_cast_pointwise_result,
 )
 from pandas.core.dtypes.common import (
-    ensure_platform_int,
     is_dict_like,
     is_extension_array_dtype,
     is_integer,
@@ -908,36 +907,6 @@ def axes(self) -> list[Index]:
     # ----------------------------------------------------------------------
     # Indexing Methods
 
-    @Appender(NDFrame.take.__doc__)
-    def take(self, indices, axis: Axis = 0, **kwargs) -> Series:
-        nv.validate_take((), kwargs)
-
-        indices = ensure_platform_int(indices)
-
-        if (
-            indices.ndim == 1
-            and using_copy_on_write()
-            and is_range_indexer(indices, len(self))
-        ):
-            return self.copy(deep=None)
-
-        new_index = self.index.take(indices)
-        new_values = self._values.take(indices)
-
-        result = self._constructor(new_values, index=new_index, fastpath=True)
-        return result.__finalize__(self, method="take")
-
-    def _take_with_is_copy(self, indices, axis: Axis = 0) -> Series:
-        """
-        Internal version of the `take` method that sets the `_is_copy`
-        attribute to keep track of the parent dataframe (using in indexing
-        for the SettingWithCopyWarning). For Series this does the same
-        as the public take (it never sets `_is_copy`).
-
-        See the docstring of `take` for full explanation of the parameters.
-        """
-        return self.take(indices=indices, axis=axis)
-
     def _ixs(self, i: int, axis: AxisInt = 0) -> Any:
         """
         Return the i-th value or values in the Series by location.

pandas/tests/series/indexing/test_take.py

@@ -16,10 +16,10 @@ def test_take():
     expected = Series([4, 2, 4], index=[4, 3, 4])
     tm.assert_series_equal(actual, expected)
 
-    msg = lambda x: f"index {x} is out of bounds for( axis 0 with)? size 5"
-    with pytest.raises(IndexError, match=msg(10)):
+    msg = "indices are out-of-bounds"
+    with pytest.raises(IndexError, match=msg):
         ser.take([1, 10])
-    with pytest.raises(IndexError, match=msg(5)):
+    with pytest.raises(IndexError, match=msg):
         ser.take([2, 5])
 
 
@@ -31,3 +31,11 @@ def test_take_categorical():
         pd.Categorical(["b", "b", "a"], categories=["a", "b", "c"]), index=[1, 1, 0]
     )
     tm.assert_series_equal(result, expected)
+
+
+def test_take_slice_raises():
+    ser = Series([-1, 5, 6, 2, 4])
+
+    msg = "Series.take requires a sequence of integers, not slice"
+    with pytest.raises(TypeError, match=msg):
+        ser.take(slice(0, 3, 1))