Backport PR #53495 on branch 2.0.x (BUG: RangeIndex.union(sort=True) …

…with another RangeIndex) (#53650) Backport PR #53495: BUG: RangeIndex.union(sort=True) with another RangeIndex Co-authored-by: Matthew Roeschke <10647082+mroeschke@users.noreply.github.com>
pandas-dev · Jun 13, 2023 · c504ed6 · c504ed6
1 parent 4b2b518
commit c504ed6
Show file tree

Hide file tree

Showing 4 changed files with 45 additions and 6 deletions.
diff --git a/doc/source/whatsnew/v2.0.3.rst b/doc/source/whatsnew/v2.0.3.rst
@@ -21,6 +21,7 @@ Fixed regressions
 
 Bug fixes
 ~~~~~~~~~
+- Bug in :func:`RangeIndex.union` when using ``sort=True`` with another :class:`RangeIndex` (:issue:`53490`)
 - Bug in :func:`read_csv` when defining ``dtype`` with ``bool[pyarrow]`` for the ``"c"`` and ``"python"`` engines (:issue:`53390`)
 - Bug in :meth:`Series.str.split` and :meth:`Series.str.rsplit` with ``expand=True`` for :class:`ArrowDtype` with ``pyarrow.string`` (:issue:`53532`)
 -

diff --git a/pandas/core/indexes/base.py b/pandas/core/indexes/base.py
@@ -3201,7 +3201,7 @@ def union(self, other, sort=None):
 
         return self._wrap_setop_result(other, result)
 
-    def _union(self, other: Index, sort):
+    def _union(self, other: Index, sort: bool | None):
         """
         Specific union logic should go here. In subclasses, union behavior
         should be overwritten here rather than in `self.union`.
@@ -3212,6 +3212,7 @@ def _union(self, other: Index, sort):
         sort : False or None, default False
             Whether to sort the resulting index.
 
+            * True : sort the result
             * False : do not sort the result.
             * None : sort the result, except when `self` and `other` are equal
               or when the values cannot be compared.
@@ -3224,7 +3225,7 @@ def _union(self, other: Index, sort):
         rvals = other._values
 
         if (
-            sort is None
+            sort in (None, True)
             and self.is_monotonic_increasing
             and other.is_monotonic_increasing
             and not (self.has_duplicates and other.has_duplicates)

diff --git a/pandas/core/indexes/range.py b/pandas/core/indexes/range.py
@@ -607,17 +607,17 @@ def _range_in_self(self, other: range) -> bool:
             return False
         return other.start in self._range and other[-1] in self._range
 
-    def _union(self, other: Index, sort):
+    def _union(self, other: Index, sort: bool | None):
         """
         Form the union of two Index objects and sorts if possible
 
         Parameters
         ----------
         other : Index or array-like
 
-        sort : False or None, default None
+        sort : bool or None, default None
             Whether to sort (monotonically increasing) the resulting index.
-            ``sort=None`` returns a ``RangeIndex`` if possible or a sorted
+            ``sort=None|True`` returns a ``RangeIndex`` if possible or a sorted
             ``Index`` with a int64 dtype if not.
             ``sort=False`` can return a ``RangeIndex`` if self is monotonically
             increasing and other is fully contained in self. Otherwise, returns
@@ -628,7 +628,7 @@ def _union(self, other: Index, sort):
         union : Index
         """
         if isinstance(other, RangeIndex):
-            if sort is None or (
+            if sort in (None, True) or (
                 sort is False and self.step > 0 and self._range_in_self(other._range)
             ):
                 # GH 47557: Can still return a RangeIndex

diff --git a/pandas/tests/indexes/test_setops.py b/pandas/tests/indexes/test_setops.py
@@ -578,6 +578,43 @@ def test_union_nan_in_both(dup):
     tm.assert_index_equal(result, expected)
 
 
+def test_union_rangeindex_sort_true():
+    # GH 53490
+    idx1 = RangeIndex(1, 100, 6)
+    idx2 = RangeIndex(1, 50, 3)
+    result = idx1.union(idx2, sort=True)
+    expected = Index(
+        [
+            1,
+            4,
+            7,
+            10,
+            13,
+            16,
+            19,
+            22,
+            25,
+            28,
+            31,
+            34,
+            37,
+            40,
+            43,
+            46,
+            49,
+            55,
+            61,
+            67,
+            73,
+            79,
+            85,
+            91,
+            97,
+        ]
+    )
+    tm.assert_index_equal(result, expected)
+
+
 def test_union_with_duplicate_index_not_subset_and_non_monotonic(
     any_dtype_for_small_pos_integer_indexes,
 ):