Skip to content

Commit bfaf917

Browse files
authored
PERF: Avoid np.divmod in maybe_sequence_to_range (pandas-dev#57812)
* PERF: Avoid np.divmod in RangeIndex._shallow_copy * Make is_range * pyi error * Use step * Switch back to int6432 * try int64_t * Revert "try int64_t" This reverts commit b8ea98c. * Adjust maybe_sequence_to_range * Access first element once
1 parent 97c3a45 commit bfaf917

File tree

3 files changed

+28
-8
lines changed

3 files changed

+28
-8
lines changed

pandas/_libs/lib.pyi

+4
Original file line numberDiff line numberDiff line change
@@ -231,3 +231,7 @@ def is_range_indexer(
231231
left: np.ndarray,
232232
n: int, # np.ndarray[np.int64, ndim=1]
233233
) -> bool: ...
234+
def is_sequence_range(
235+
sequence: np.ndarray,
236+
step: int, # np.ndarray[np.int64, ndim=1]
237+
) -> bool: ...

pandas/_libs/lib.pyx

+22
Original file line numberDiff line numberDiff line change
@@ -678,6 +678,28 @@ def is_range_indexer(ndarray[int6432_t, ndim=1] left, Py_ssize_t n) -> bool:
678678
return True
679679

680680

681+
@cython.wraparound(False)
682+
@cython.boundscheck(False)
683+
def is_sequence_range(ndarray[int6432_t, ndim=1] sequence, int64_t step) -> bool:
684+
"""
685+
Check if sequence is equivalent to a range with the specified step.
686+
"""
687+
cdef:
688+
Py_ssize_t i, n = len(sequence)
689+
int6432_t first_element
690+
691+
if step == 0:
692+
return False
693+
if n == 0:
694+
return True
695+
696+
first_element = sequence[0]
697+
for i in range(1, n):
698+
if sequence[i] != first_element + i * step:
699+
return False
700+
return True
701+
702+
681703
ctypedef fused ndarr_object:
682704
ndarray[object, ndim=1]
683705
ndarray[object, ndim=2]

pandas/core/indexes/base.py

+2-8
Original file line numberDiff line numberDiff line change
@@ -7169,7 +7169,7 @@ def maybe_sequence_to_range(sequence) -> Any | range:
71697169
-------
71707170
Any : input or range
71717171
"""
7172-
if isinstance(sequence, (ABCSeries, Index)):
7172+
if isinstance(sequence, (ABCSeries, Index, range)):
71737173
return sequence
71747174
np_sequence = np.asarray(sequence)
71757175
if np_sequence.dtype.kind != "i" or len(np_sequence) == 1:
@@ -7179,13 +7179,7 @@ def maybe_sequence_to_range(sequence) -> Any | range:
71797179
diff = np_sequence[1] - np_sequence[0]
71807180
if diff == 0:
71817181
return sequence
7182-
elif len(np_sequence) == 2:
7183-
return range(np_sequence[0], np_sequence[1] + diff, diff)
7184-
maybe_range_indexer, remainder = np.divmod(np_sequence - np_sequence[0], diff)
7185-
if (
7186-
lib.is_range_indexer(maybe_range_indexer, len(maybe_range_indexer))
7187-
and not remainder.any()
7188-
):
7182+
elif len(np_sequence) == 2 or lib.is_sequence_range(np_sequence, diff):
71897183
return range(np_sequence[0], np_sequence[-1] + diff, diff)
71907184
else:
71917185
return sequence

0 commit comments

Comments
 (0)