Skip to content

PERF: RangeIndex.round returns RangeIndex when possible #57824

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 7 commits into from
Mar 13, 2024
Merged
Show file tree
Hide file tree
Changes from 5 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions doc/source/whatsnew/v3.0.0.rst
Original file line number Diff line number Diff line change
Expand Up @@ -274,6 +274,7 @@ Performance improvements
- Performance improvement in :meth:`MultiIndex.equals` for equal length indexes (:issue:`56990`)
- 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`)
- Performance improvement in :meth:`RangeIndex.append` when appending the same index (:issue:`57252`)
- Performance improvement in :meth:`RangeIndex.round` returning a :class:`RangeIndex` instead of a :class:`Index` when possible. (:issue:`57824`)
- Performance improvement in :meth:`RangeIndex.join` returning a :class:`RangeIndex` instead of a :class:`Index` when possible. (:issue:`57651`, :issue:`57752`)
- Performance improvement in :meth:`RangeIndex.reindex` returning a :class:`RangeIndex` instead of a :class:`Index` when possible. (:issue:`57647`, :issue:`57752`)
- Performance improvement in :meth:`RangeIndex.take` returning a :class:`RangeIndex` instead of a :class:`Index` when possible. (:issue:`57445`, :issue:`57752`)
Expand Down
1 change: 0 additions & 1 deletion pandas/core/indexes/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -6737,7 +6737,6 @@ def diff(self, periods: int = 1) -> Index:
"""
return Index(self.to_series().diff(periods))

@final
def round(self, decimals: int = 0) -> Self:
"""
Round each value in the Index to the given number of decimals.
Expand Down
37 changes: 37 additions & 0 deletions pandas/core/indexes/range.py
Original file line number Diff line number Diff line change
Expand Up @@ -1165,6 +1165,43 @@ def any(self, *args, **kwargs) -> bool:

# --------------------------------------------------------------------

# error: Return type "RangeIndex | Index" of "round" incompatible with
# return type "RangeIndex" in supertype "Index"
def round(self, decimals: int = 0) -> Self | Index: # type: ignore[override]
"""
Round each value in the Index to the given number of decimals.

Parameters
----------
decimals : int, optional
Number of decimal places to round to. If decimals is negative,
it specifies the number of positions to the left of the decimal point.
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I know this is from the original docstring, but for me it's not obvious what this means by reading it, maybe adding an example or two here would be helpful.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sure thing. I added an example for the negative case.


Returns
-------
Index or RangeIndex
A new Index with the rounded values.

Examples
--------
>>> import pandas as pd
>>> idx = pd.RangeIndex(10, 30, 10)
>>> idx.round(decimals=-1)
RangeIndex(start=10, stop=30, step=10)
>>> idx = pd.RangeIndex(10, 15, 1)
>>> idx.round(decimals=-1)
Index([10, 10, 10, 10, 10], dtype='int64')
"""
if decimals >= 0:
return self.copy()
elif all(
getattr(self, attr) % 10**-decimals == 0 for attr in ("start", "step")
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Personal preference, but I'd find more readable to simply use self.start % 10 ... and self.stop % 10 ....

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fair point. I'll change this to your suggestion

):
# e.g. Range(10, 30, 10).round(-1) doesn't need rounding
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
# e.g. Range(10, 30, 10).round(-1) doesn't need rounding
# e.g. RangeIndex(10, 30, 10).round(-1) doesn't need rounding

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks!

return self.copy()
else:
return super().round(decimals=decimals)

def _cmp_method(self, other, op):
if isinstance(other, RangeIndex) and self._range == other._range:
# Both are immutable so if ._range attr. are equal, shortcut is possible
Expand Down
31 changes: 31 additions & 0 deletions pandas/tests/indexes/ranges/test_range.py
Original file line number Diff line number Diff line change
Expand Up @@ -608,6 +608,37 @@ def test_range_index_rsub_by_const(self):
tm.assert_index_equal(result, expected)


@pytest.mark.parametrize(
"rng, decimals",
[
[range(5), 0],
[range(5), 2],
[range(10, 30, 10), -1],
[range(30, 10, -10), -1],
],
)
def test_range_round_returns_rangeindex(rng, decimals):
ri = RangeIndex(rng)
expected = ri.copy()
result = ri.round(decimals=decimals)
tm.assert_index_equal(result, expected, exact=True)


@pytest.mark.parametrize(
"rng, decimals",
[
[range(10, 30, 1), -1],
[range(30, 10, -1), -1],
[range(11, 14), -10],
],
)
def test_range_round_returns_index(rng, decimals):
ri = RangeIndex(rng)
expected = Index(list(rng)).round(decimals=decimals)
result = ri.round(decimals=decimals)
tm.assert_index_equal(result, expected, exact=True)


def test_reindex_1_value_returns_rangeindex():
ri = RangeIndex(0, 10, 2, name="foo")
result, result_indexer = ri.reindex([2])
Expand Down