Skip to content

Commit 72327f3

Browse files
authored
ENH/PERF: RangeIndex.argsort accept kind; pd.isna(rangeindex) fastpath (pandas-dev#44263)
1 parent a35ae69 commit 72327f3

File tree

3 files changed

+13
-10
lines changed

3 files changed

+13
-10
lines changed

pandas/core/dtypes/missing.py

+9-5
Original file line numberDiff line numberDiff line change
@@ -169,13 +169,17 @@ def _isna(obj, inf_as_na: bool = False):
169169
return False
170170
elif isinstance(obj, (np.ndarray, ABCExtensionArray)):
171171
return _isna_array(obj, inf_as_na=inf_as_na)
172-
elif isinstance(obj, (ABCSeries, ABCIndex)):
172+
elif isinstance(obj, ABCIndex):
173+
# Try to use cached isna, which also short-circuits for integer dtypes
174+
# and avoids materializing RangeIndex._values
175+
if not obj._can_hold_na:
176+
return obj.isna()
177+
return _isna_array(obj._values, inf_as_na=inf_as_na)
178+
179+
elif isinstance(obj, ABCSeries):
173180
result = _isna_array(obj._values, inf_as_na=inf_as_na)
174181
# box
175-
if isinstance(obj, ABCSeries):
176-
result = obj._constructor(
177-
result, index=obj.index, name=obj.name, copy=False
178-
)
182+
result = obj._constructor(result, index=obj.index, name=obj.name, copy=False)
179183
return result
180184
elif isinstance(obj, ABCDataFrame):
181185
return obj.isna()

pandas/core/indexes/range.py

+1
Original file line numberDiff line numberDiff line change
@@ -496,6 +496,7 @@ def argsort(self, *args, **kwargs) -> npt.NDArray[np.intp]:
496496
numpy.ndarray.argsort
497497
"""
498498
ascending = kwargs.pop("ascending", True) # EA compat
499+
kwargs.pop("kind", None) # e.g. "mergesort" is irrelevant
499500
nv.validate_argsort(args, kwargs)
500501

501502
if self._range.step > 0:

pandas/tests/indexes/common.py

+3-5
Original file line numberDiff line numberDiff line change
@@ -333,11 +333,9 @@ def test_numpy_argsort(self, index):
333333
expected = index.argsort()
334334
tm.assert_numpy_array_equal(result, expected)
335335

336-
if not isinstance(index, RangeIndex):
337-
# TODO: add compatibility to RangeIndex?
338-
result = np.argsort(index, kind="mergesort")
339-
expected = index.argsort(kind="mergesort")
340-
tm.assert_numpy_array_equal(result, expected)
336+
result = np.argsort(index, kind="mergesort")
337+
expected = index.argsort(kind="mergesort")
338+
tm.assert_numpy_array_equal(result, expected)
341339

342340
# these are the only two types that perform
343341
# pandas compatibility input validation - the

0 commit comments

Comments
 (0)