Skip to content

Commit 05d1667

Browse files
reidy-preidy-p
committed
ENH: Add sort parameter to RangeIndex.union (pandas-dev#24471)
1 parent e8d951d commit 05d1667

File tree

5 files changed

+99
-32
lines changed

5 files changed

+99
-32
lines changed

‎pandas/core/indexes/base.py

+1-1
Original file line numberDiff line numberDiff line change
@@ -2319,7 +2319,7 @@ def union(self, other, sort=None):
23192319
else:
23202320
rvals = other._values
23212321

2322-
if self.is_monotonic and other.is_monotonic:
2322+
if self.is_monotonic and other.is_monotonic and sort is None:
23232323
try:
23242324
result = self._outer_indexer(lvals, rvals)[0]
23252325
except TypeError:

‎pandas/core/indexes/range.py

+4-4
Original file line numberDiff line numberDiff line change
@@ -463,7 +463,7 @@ def _extended_gcd(self, a, b):
463463
old_t, t = t, old_t - quotient * t
464464
return old_r, old_s, old_t
465465

466-
def union(self, other):
466+
def union(self, other, sort=None):
467467
"""
468468
Form the union of two Index objects and sorts if possible
469469
@@ -477,9 +477,9 @@ def union(self, other):
477477
"""
478478
self._assert_can_do_setop(other)
479479
if len(other) == 0 or self.equals(other) or len(self) == 0:
480-
return super(RangeIndex, self).union(other)
480+
return super(RangeIndex, self).union(other, sort=sort)
481481

482-
if isinstance(other, RangeIndex):
482+
if isinstance(other, RangeIndex) and sort is None:
483483
start_s, step_s = self._start, self._step
484484
end_s = self._start + self._step * (len(self) - 1)
485485
start_o, step_o = other._start, other._step
@@ -516,7 +516,7 @@ def union(self, other):
516516
(end_s - step_o <= end_o)):
517517
return RangeIndex(start_r, end_r + step_o, step_o)
518518

519-
return self._int64index.union(other)
519+
return self._int64index.union(other, sort=sort)
520520

521521
@Appender(_index_shared_docs['join'])
522522
def join(self, other, how='left', level=None, return_indexers=False,

‎pandas/tests/indexes/datetimes/test_setops.py

+10-2
Original file line numberDiff line numberDiff line change
@@ -86,7 +86,11 @@ def test_union_bug_1730(self, sort):
8686
rng_b = date_range('1/1/2012', periods=4, freq='4H')
8787

8888
result = rng_a.union(rng_b, sort=sort)
89-
exp = DatetimeIndex(sorted(set(list(rng_a)) | set(list(rng_b))))
89+
exp = list(rng_a) + list(rng_b[1:])
90+
if sort is None:
91+
exp = DatetimeIndex(sorted(exp))
92+
else:
93+
exp = DatetimeIndex(exp)
9094
tm.assert_index_equal(result, exp)
9195

9296
@pytest.mark.parametrize("sort", [None, False])
@@ -112,7 +116,11 @@ def test_union_bug_4564(self, sort):
112116
right = left + DateOffset(minutes=15)
113117

114118
result = left.union(right, sort=sort)
115-
exp = DatetimeIndex(sorted(set(list(left)) | set(list(right))))
119+
exp = list(left) + list(right)
120+
if sort is None:
121+
exp = DatetimeIndex(sorted(exp))
122+
else:
123+
exp = DatetimeIndex(exp)
116124
tm.assert_index_equal(result, exp)
117125

118126
@pytest.mark.parametrize("sort", [None, False])

‎pandas/tests/indexes/period/test_setops.py

+9-2
Original file line numberDiff line numberDiff line change
@@ -43,7 +43,12 @@ def test_union(self, sort):
4343
# union
4444
other1 = pd.period_range('1/1/2000', freq='D', periods=5)
4545
rng1 = pd.period_range('1/6/2000', freq='D', periods=5)
46-
expected1 = pd.period_range('1/1/2000', freq='D', periods=10)
46+
expected1 = pd.PeriodIndex(['2000-01-06', '2000-01-07',
47+
'2000-01-08', '2000-01-09',
48+
'2000-01-10', '2000-01-01',
49+
'2000-01-02', '2000-01-03',
50+
'2000-01-04', '2000-01-05'],
51+
freq='D')
4752

4853
rng2 = pd.period_range('1/1/2000', freq='D', periods=5)
4954
other2 = pd.period_range('1/4/2000', freq='D', periods=5)
@@ -77,7 +82,9 @@ def test_union(self, sort):
7782

7883
rng7 = pd.period_range('2003-01-01', freq='A', periods=5)
7984
other7 = pd.period_range('1998-01-01', freq='A', periods=8)
80-
expected7 = pd.period_range('1998-01-01', freq='A', periods=10)
85+
expected7 = pd.PeriodIndex(['2003', '2004', '2005', '2006', '2007',
86+
'1998', '1999', '2000', '2001', '2002'],
87+
freq='A')
8188

8289
rng8 = pd.PeriodIndex(['1/3/2000', '1/2/2000', '1/1/2000',
8390
'1/5/2000', '1/4/2000'], freq='D')

‎pandas/tests/indexes/test_range.py

+75-23
Original file line numberDiff line numberDiff line change
@@ -586,36 +586,88 @@ def test_union_noncomparable(self):
586586
def test_union(self):
587587
RI = RangeIndex
588588
I64 = Int64Index
589-
cases = [(RI(0, 10, 1), RI(0, 10, 1), RI(0, 10, 1)),
590-
(RI(0, 10, 1), RI(5, 20, 1), RI(0, 20, 1)),
591-
(RI(0, 10, 1), RI(10, 20, 1), RI(0, 20, 1)),
592-
(RI(0, -10, -1), RI(0, -10, -1), RI(0, -10, -1)),
593-
(RI(0, -10, -1), RI(-10, -20, -1), RI(-19, 1, 1)),
594-
(RI(0, 10, 2), RI(1, 10, 2), RI(0, 10, 1)),
595-
(RI(0, 11, 2), RI(1, 12, 2), RI(0, 12, 1)),
596-
(RI(0, 21, 4), RI(-2, 24, 4), RI(-2, 24, 2)),
597-
(RI(0, -20, -2), RI(-1, -21, -2), RI(-19, 1, 1)),
598-
(RI(0, 100, 5), RI(0, 100, 20), RI(0, 100, 5)),
599-
(RI(0, -100, -5), RI(5, -100, -20), RI(-95, 10, 5)),
600-
(RI(0, -11, -1), RI(1, -12, -4), RI(-11, 2, 1)),
601-
(RI(0), RI(0), RI(0)),
602-
(RI(0, -10, -2), RI(0), RI(0, -10, -2)),
603-
(RI(0, 100, 2), RI(100, 150, 200), RI(0, 102, 2)),
604-
(RI(0, -100, -2), RI(-100, 50, 102), RI(-100, 4, 2)),
605-
(RI(0, -100, -1), RI(0, -50, -3), RI(-99, 1, 1)),
606-
(RI(0, 1, 1), RI(5, 6, 10), RI(0, 6, 5)),
607-
(RI(0, 10, 5), RI(-5, -6, -20), RI(-5, 10, 5)),
608-
(RI(0, 3, 1), RI(4, 5, 1), I64([0, 1, 2, 4])),
609-
(RI(0, 10, 1), I64([]), RI(0, 10, 1)),
610-
(RI(0), I64([1, 5, 6]), I64([1, 5, 6]))]
611-
for idx1, idx2, expected in cases:
589+
590+
inputs = [(RI(0, 10, 1), RI(0, 10, 1)),
591+
(RI(0, 10, 1), RI(5, 20, 1)),
592+
(RI(0, 10, 1), RI(10, 20, 1)),
593+
(RI(0, -10, -1), RI(0, -10, -1)),
594+
(RI(0, -10, -1), RI(-10, -20, -1)),
595+
(RI(0, 10, 2), RI(1, 10, 2)),
596+
(RI(0, 11, 2), RI(1, 12, 2)),
597+
(RI(0, 21, 4), RI(-2, 24, 4)),
598+
(RI(0, -20, -2), RI(-1, -21, -2)),
599+
(RI(0, 100, 5), RI(0, 100, 20)),
600+
(RI(0, -100, -5), RI(5, -100, -20)),
601+
(RI(0, -11, -1), RI(1, -12, -4)),
602+
(RI(0), RI(0)),
603+
(RI(0, -10, -2), RI(0)),
604+
(RI(0, 100, 2), RI(100, 150, 200)),
605+
(RI(0, -100, -2), RI(-100, 50, 102)),
606+
(RI(0, -100, -1), RI(0, -50, -3)),
607+
(RI(0, 1, 1), RI(5, 6, 10)),
608+
(RI(0, 10, 5), RI(-5, -6, -20)),
609+
(RI(0, 3, 1), RI(4, 5, 1)),
610+
(RI(0, 10, 1), I64([])),
611+
(RI(0), I64([1, 5, 6]))]
612+
613+
expected_sorted = [RI(0, 10, 1),
614+
RI(0, 20, 1),
615+
RI(0, 20, 1),
616+
RI(0, -10, -1),
617+
RI(-19, 1, 1),
618+
RI(0, 10, 1),
619+
RI(0, 12, 1),
620+
RI(-2, 24, 2),
621+
RI(-19, 1, 1),
622+
RI(0, 100, 5),
623+
RI(-95, 10, 5),
624+
RI(-11, 2, 1),
625+
RI(0),
626+
RI(0, -10, -2),
627+
RI(0, 102, 2),
628+
RI(-100, 4, 2),
629+
RI(-99, 1, 1),
630+
RI(0, 6, 5),
631+
RI(-5, 10, 5),
632+
I64([0, 1, 2, 4]),
633+
RI(0, 10, 1),
634+
I64([1, 5, 6])]
635+
636+
for ((idx1, idx2), expected) in zip(inputs, expected_sorted):
612637
res1 = idx1.union(idx2)
613638
res2 = idx2.union(idx1)
614639
res3 = idx1._int64index.union(idx2)
615640
tm.assert_index_equal(res1, expected, exact=True)
616641
tm.assert_index_equal(res2, expected, exact=True)
617642
tm.assert_index_equal(res3, expected)
618643

644+
expected_notsorted = [RI(0, 10, 1),
645+
I64(range(20)),
646+
I64(range(20)),
647+
RI(0, -10, -1),
648+
I64(range(0, -20, -1)),
649+
I64(list(range(0, 10, 2)) + list(range(1, 10, 2))),
650+
I64(list(range(0, 11, 2)) + list(range(1, 12, 2))),
651+
I64(list(range(0, 21, 4)) + list(range(-2, 24, 4))),
652+
I64(list(range(0, -20, -2)) + list(range(-1, -21, -2))),
653+
I64(range(0, 100, 5)),
654+
I64(list(range(0, -100, -5)) + [5]),
655+
I64(list(range(0, -11, -1)) + [1, -11]),
656+
RI(0),
657+
RI(0, -10, -2),
658+
I64(range(0, 102, 2)),
659+
I64(list(range(0, -100, -2)) + [-100, 2]),
660+
I64(list(range(0, -100, -1))),
661+
I64([0, 5]),
662+
I64([0, 5, -5]),
663+
I64([0, 1, 2, 4]),
664+
RI(0, 10, 1),
665+
I64([1, 5, 6])]
666+
667+
for ((idx1, idx2), expected) in zip(inputs, expected_notsorted):
668+
res1 = idx1.union(idx2, sort=False)
669+
tm.assert_index_equal(res1, expected, exact=True)
670+
619671
def test_nbytes(self):
620672

621673
# memory savings vs int index

0 commit comments

Comments
 (0)