Skip to content

Commit 9aa319b

Browse files
mroeschkepmhatre1
authored andcommitted
PERF: Unary methods on RangeIndex returns RangeIndex (pandas-dev#57825)
* PERF: Unary methods on RangeIndex returns RangeIndex * Whatsnew number
1 parent 514ba85 commit 9aa319b

File tree

3 files changed

+83
-0
lines changed

3 files changed

+83
-0
lines changed

doc/source/whatsnew/v3.0.0.rst

+1
Original file line numberDiff line numberDiff line change
@@ -282,6 +282,7 @@ Performance improvements
282282
- Performance improvement in :meth:`RangeIndex.take` returning a :class:`RangeIndex` instead of a :class:`Index` when possible. (:issue:`57445`, :issue:`57752`)
283283
- Performance improvement in ``DataFrameGroupBy.__len__`` and ``SeriesGroupBy.__len__`` (:issue:`57595`)
284284
- Performance improvement in indexing operations for string dtypes (:issue:`56997`)
285+
- Performance improvement in unary methods on a :class:`RangeIndex` returning a :class:`RangeIndex` instead of a :class:`Index` when possible. (:issue:`57825`)
285286

286287
.. ---------------------------------------------------------------------------
287288
.. _whatsnew_300.bug_fixes:

pandas/core/indexes/range.py

+21
Original file line numberDiff line numberDiff line change
@@ -1314,6 +1314,27 @@ def _arith_method(self, other, op):
13141314
# test_arithmetic_explicit_conversions
13151315
return super()._arith_method(other, op)
13161316

1317+
def __abs__(self) -> Self | Index:
1318+
if len(self) == 0 or self.min() >= 0:
1319+
return self.copy()
1320+
elif self.max() <= 0:
1321+
return -self
1322+
else:
1323+
return super().__abs__()
1324+
1325+
def __neg__(self) -> Self:
1326+
rng = range(-self.start, -self.stop, -self.step)
1327+
return self._simple_new(rng, name=self.name)
1328+
1329+
def __pos__(self) -> Self:
1330+
return self.copy()
1331+
1332+
def __invert__(self) -> Self:
1333+
if len(self) == 0:
1334+
return self.copy()
1335+
rng = range(~self.start, ~self.stop, -self.step)
1336+
return self._simple_new(rng, name=self.name)
1337+
13171338
# error: Return type "Index" of "take" incompatible with return type
13181339
# "RangeIndex" in supertype "Index"
13191340
def take( # type: ignore[override]

pandas/tests/indexes/ranges/test_range.py

+61
Original file line numberDiff line numberDiff line change
@@ -763,6 +763,67 @@ def test_getitem_boolmask_wrong_length():
763763
ri[[True]]
764764

765765

766+
def test_pos_returns_rangeindex():
767+
ri = RangeIndex(2, name="foo")
768+
expected = ri.copy()
769+
result = +ri
770+
tm.assert_index_equal(result, expected, exact=True)
771+
772+
773+
def test_neg_returns_rangeindex():
774+
ri = RangeIndex(2, name="foo")
775+
result = -ri
776+
expected = RangeIndex(0, -2, -1, name="foo")
777+
tm.assert_index_equal(result, expected, exact=True)
778+
779+
ri = RangeIndex(-2, 2, name="foo")
780+
result = -ri
781+
expected = RangeIndex(2, -2, -1, name="foo")
782+
tm.assert_index_equal(result, expected, exact=True)
783+
784+
785+
@pytest.mark.parametrize(
786+
"rng, exp_rng",
787+
[
788+
[range(0), range(0)],
789+
[range(10), range(10)],
790+
[range(-2, 1, 1), range(2, -1, -1)],
791+
[range(0, -10, -1), range(0, 10, 1)],
792+
],
793+
)
794+
def test_abs_returns_rangeindex(rng, exp_rng):
795+
ri = RangeIndex(rng, name="foo")
796+
expected = RangeIndex(exp_rng, name="foo")
797+
result = abs(ri)
798+
tm.assert_index_equal(result, expected, exact=True)
799+
800+
801+
def test_abs_returns_index():
802+
ri = RangeIndex(-2, 2, name="foo")
803+
result = abs(ri)
804+
expected = Index([2, 1, 0, 1], name="foo")
805+
tm.assert_index_equal(result, expected, exact=True)
806+
807+
808+
@pytest.mark.parametrize(
809+
"rng",
810+
[
811+
range(0),
812+
range(5),
813+
range(0, -5, -1),
814+
range(-2, 2, 1),
815+
range(2, -2, -2),
816+
range(0, 5, 2),
817+
],
818+
)
819+
def test_invert_returns_rangeindex(rng):
820+
ri = RangeIndex(rng, name="foo")
821+
result = ~ri
822+
assert isinstance(result, RangeIndex)
823+
expected = ~Index(list(rng), name="foo")
824+
tm.assert_index_equal(result, expected, exact=False)
825+
826+
766827
@pytest.mark.parametrize(
767828
"rng",
768829
[

0 commit comments

Comments
 (0)