diff --git a/doc/source/whatsnew/v3.0.0.rst b/doc/source/whatsnew/v3.0.0.rst index 483cf659080ea..e74967695a2ff 100644 --- a/doc/source/whatsnew/v3.0.0.rst +++ b/doc/source/whatsnew/v3.0.0.rst @@ -166,6 +166,7 @@ Performance improvements - Performance improvement in :meth:`Index.take` when ``indices`` is a full range indexer from zero to length of index (:issue:`56806`) - Performance improvement in :meth:`MultiIndex.equals` for equal length indexes (:issue:`56990`) - Performance improvement in :meth:`RangeIndex.append` when appending the same index (:issue:`57252`) +- Performance improvement in :meth:`RangeIndex.take` returning a :class:`RangeIndex` instead of a :class:`Index` when possible. (:issue:`57445`) - Performance improvement in indexing operations for string dtypes (:issue:`56997`) - diff --git a/pandas/core/indexes/range.py b/pandas/core/indexes/range.py index a8e8ae140041a..20286cad58df9 100644 --- a/pandas/core/indexes/range.py +++ b/pandas/core/indexes/range.py @@ -1006,11 +1006,13 @@ def _concat(self, indexes: list[Index], name: Hashable) -> Index: # Get the stop value from "next" or alternatively # from the last non-empty index stop = non_empty_indexes[-1].stop if next_ is None else next_ - return RangeIndex(start, stop, step).rename(name) + if len(non_empty_indexes) == 1: + step = non_empty_indexes[0].step + return RangeIndex(start, stop, step, name=name) # Here all "indexes" had 0 length, i.e. were empty. # In this case return an empty range index. - return RangeIndex(0, 0).rename(name) + return RangeIndex(_empty_range, name=name) def __len__(self) -> int: """ @@ -1168,7 +1170,7 @@ def take( # type: ignore[override] allow_fill: bool = True, fill_value=None, **kwargs, - ) -> Index: + ) -> Self | Index: if kwargs: nv.validate_take((), kwargs) if is_scalar(indices): @@ -1179,7 +1181,7 @@ def take( # type: ignore[override] self._maybe_disallow_fill(allow_fill, fill_value, indices) if len(indices) == 0: - taken = np.array([], dtype=self.dtype) + return type(self)(_empty_range, name=self.name) else: ind_max = indices.max() if ind_max >= len(self): @@ -1199,5 +1201,4 @@ def take( # type: ignore[override] if self.start != 0: taken += self.start - # _constructor so RangeIndex-> Index with an int64 dtype - return self._constructor._simple_new(taken, name=self.name) + return self._shallow_copy(taken, name=self.name) diff --git a/pandas/tests/indexes/ranges/test_range.py b/pandas/tests/indexes/ranges/test_range.py index 29d6916cc7f08..d500687763a1e 100644 --- a/pandas/tests/indexes/ranges/test_range.py +++ b/pandas/tests/indexes/ranges/test_range.py @@ -606,3 +606,20 @@ def test_range_index_rsub_by_const(self): result = 3 - RangeIndex(0, 4, 1) expected = RangeIndex(3, -1, -1) tm.assert_index_equal(result, expected) + + +def test_take_return_rangeindex(): + ri = RangeIndex(5, name="foo") + result = ri.take([]) + expected = RangeIndex(0, name="foo") + tm.assert_index_equal(result, expected, exact=True) + + result = ri.take([3, 4]) + expected = RangeIndex(3, 5, name="foo") + tm.assert_index_equal(result, expected, exact=True) + + +def test_append_one_nonempty_preserve_step(): + expected = RangeIndex(0, -1, -1) + result = RangeIndex(0).append([expected]) + tm.assert_index_equal(result, expected, exact=True)