Skip to content

Commit d5267d7

Browse files
mroeschkepmhatre1
authored andcommitted
PERF: RangeIndex.round returns RangeIndex when possible (pandas-dev#57824)
* PERF: RangeIndex.round returns RangeIndex when possible * Add whatsnew * Add typing * Address feedback
1 parent e8d237a commit d5267d7

File tree

4 files changed

+68
-1
lines changed

4 files changed

+68
-1
lines changed

doc/source/whatsnew/v3.0.0.rst

+1
Original file line numberDiff line numberDiff line change
@@ -274,6 +274,7 @@ Performance improvements
274274
- Performance improvement in :meth:`MultiIndex.equals` for equal length indexes (:issue:`56990`)
275275
- Performance improvement in :meth:`RangeIndex.__getitem__` with a boolean mask or integers returning a :class:`RangeIndex` instead of a :class:`Index` when possible. (:issue:`57588`)
276276
- Performance improvement in :meth:`RangeIndex.append` when appending the same index (:issue:`57252`)
277+
- Performance improvement in :meth:`RangeIndex.round` returning a :class:`RangeIndex` instead of a :class:`Index` when possible. (:issue:`57824`)
277278
- Performance improvement in :meth:`RangeIndex.join` returning a :class:`RangeIndex` instead of a :class:`Index` when possible. (:issue:`57651`, :issue:`57752`)
278279
- Performance improvement in :meth:`RangeIndex.reindex` returning a :class:`RangeIndex` instead of a :class:`Index` when possible. (:issue:`57647`, :issue:`57752`)
279280
- Performance improvement in :meth:`RangeIndex.take` returning a :class:`RangeIndex` instead of a :class:`Index` when possible. (:issue:`57445`, :issue:`57752`)

pandas/core/indexes/base.py

-1
Original file line numberDiff line numberDiff line change
@@ -6718,7 +6718,6 @@ def diff(self, periods: int = 1) -> Index:
67186718
"""
67196719
return Index(self.to_series().diff(periods))
67206720

6721-
@final
67226721
def round(self, decimals: int = 0) -> Self:
67236722
"""
67246723
Round each value in the Index to the given number of decimals.

pandas/core/indexes/range.py

+36
Original file line numberDiff line numberDiff line change
@@ -1165,6 +1165,42 @@ def any(self, *args, **kwargs) -> bool:
11651165

11661166
# --------------------------------------------------------------------
11671167

1168+
# error: Return type "RangeIndex | Index" of "round" incompatible with
1169+
# return type "RangeIndex" in supertype "Index"
1170+
def round(self, decimals: int = 0) -> Self | Index: # type: ignore[override]
1171+
"""
1172+
Round each value in the Index to the given number of decimals.
1173+
1174+
Parameters
1175+
----------
1176+
decimals : int, optional
1177+
Number of decimal places to round to. If decimals is negative,
1178+
it specifies the number of positions to the left of the decimal point
1179+
e.g. ``round(11.0, -1) == 10.0``.
1180+
1181+
Returns
1182+
-------
1183+
Index or RangeIndex
1184+
A new Index with the rounded values.
1185+
1186+
Examples
1187+
--------
1188+
>>> import pandas as pd
1189+
>>> idx = pd.RangeIndex(10, 30, 10)
1190+
>>> idx.round(decimals=-1)
1191+
RangeIndex(start=10, stop=30, step=10)
1192+
>>> idx = pd.RangeIndex(10, 15, 1)
1193+
>>> idx.round(decimals=-1)
1194+
Index([10, 10, 10, 10, 10], dtype='int64')
1195+
"""
1196+
if decimals >= 0:
1197+
return self.copy()
1198+
elif self.start % 10**-decimals == 0 and self.step % 10**-decimals == 0:
1199+
# e.g. RangeIndex(10, 30, 10).round(-1) doesn't need rounding
1200+
return self.copy()
1201+
else:
1202+
return super().round(decimals=decimals)
1203+
11681204
def _cmp_method(self, other, op):
11691205
if isinstance(other, RangeIndex) and self._range == other._range:
11701206
# Both are immutable so if ._range attr. are equal, shortcut is possible

pandas/tests/indexes/ranges/test_range.py

+31
Original file line numberDiff line numberDiff line change
@@ -608,6 +608,37 @@ def test_range_index_rsub_by_const(self):
608608
tm.assert_index_equal(result, expected)
609609

610610

611+
@pytest.mark.parametrize(
612+
"rng, decimals",
613+
[
614+
[range(5), 0],
615+
[range(5), 2],
616+
[range(10, 30, 10), -1],
617+
[range(30, 10, -10), -1],
618+
],
619+
)
620+
def test_range_round_returns_rangeindex(rng, decimals):
621+
ri = RangeIndex(rng)
622+
expected = ri.copy()
623+
result = ri.round(decimals=decimals)
624+
tm.assert_index_equal(result, expected, exact=True)
625+
626+
627+
@pytest.mark.parametrize(
628+
"rng, decimals",
629+
[
630+
[range(10, 30, 1), -1],
631+
[range(30, 10, -1), -1],
632+
[range(11, 14), -10],
633+
],
634+
)
635+
def test_range_round_returns_index(rng, decimals):
636+
ri = RangeIndex(rng)
637+
expected = Index(list(rng)).round(decimals=decimals)
638+
result = ri.round(decimals=decimals)
639+
tm.assert_index_equal(result, expected, exact=True)
640+
641+
611642
def test_reindex_1_value_returns_rangeindex():
612643
ri = RangeIndex(0, 10, 2, name="foo")
613644
result, result_indexer = ri.reindex([2])

0 commit comments

Comments
 (0)