From cb90c588a9e74fd768c21c291cf20c1306e4b398 Mon Sep 17 00:00:00 2001 From: Brock Date: Thu, 19 Nov 2020 19:23:11 -0800 Subject: [PATCH 1/2] PERF: IntervalArray.argsort --- pandas/core/arrays/interval.py | 15 +++++++++++++++ pandas/core/indexes/interval.py | 5 ----- 2 files changed, 15 insertions(+), 5 deletions(-) diff --git a/pandas/core/arrays/interval.py b/pandas/core/arrays/interval.py index d007bb112c86c..90b872032b89a 100644 --- a/pandas/core/arrays/interval.py +++ b/pandas/core/arrays/interval.py @@ -619,6 +619,21 @@ def __lt__(self, other): def __le__(self, other): return self._cmp_method(other, operator.le) + def argsort( + self, + ascending: bool = True, + kind: str = "quicksort", + na_position: str = "last", + *args, + **kwargs, + ) -> np.ndarray: + if ascending and kind == "quicksort" and na_position == "last": + return np.lexsort((self.right, self.left)) + # TODO: other cases we can use lexsort for? much more performant. + return super().argsort( + *args, ascending=ascending, kind=kind, na_position=na_position, **kwargs + ) + def fillna(self, value=None, method=None, limit=None): """ Fill NA/NaN values using the specified method. diff --git a/pandas/core/indexes/interval.py b/pandas/core/indexes/interval.py index b0f8be986fe5d..de0ad03ee7105 100644 --- a/pandas/core/indexes/interval.py +++ b/pandas/core/indexes/interval.py @@ -957,11 +957,6 @@ def _format_space(self) -> str: space = " " * (len(type(self).__name__) + 1) return f"\n{space}" - # -------------------------------------------------------------------- - - def argsort(self, *args, **kwargs) -> np.ndarray: - return np.lexsort((self.right, self.left)) - # -------------------------------------------------------------------- # Set Operations From 27b7d64a7cf12aa88022a4bc8dbf48d9942bf2b6 Mon Sep 17 00:00:00 2001 From: Brock Date: Thu, 19 Nov 2020 19:54:01 -0800 Subject: [PATCH 2/2] mypy fixup --- pandas/core/arrays/interval.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/pandas/core/arrays/interval.py b/pandas/core/arrays/interval.py index 90b872032b89a..500e96a0c2784 100644 --- a/pandas/core/arrays/interval.py +++ b/pandas/core/arrays/interval.py @@ -627,11 +627,14 @@ def argsort( *args, **kwargs, ) -> np.ndarray: + ascending = nv.validate_argsort_with_ascending(ascending, args, kwargs) + if ascending and kind == "quicksort" and na_position == "last": return np.lexsort((self.right, self.left)) + # TODO: other cases we can use lexsort for? much more performant. return super().argsort( - *args, ascending=ascending, kind=kind, na_position=na_position, **kwargs + ascending=ascending, kind=kind, na_position=na_position, **kwargs ) def fillna(self, value=None, method=None, limit=None):