From 213e1fd7761b3206e4d53c81a398d8c899c3e072 Mon Sep 17 00:00:00 2001 From: jbrockmendel Date: Sat, 4 Jan 2020 14:32:39 -0800 Subject: [PATCH] BUG: TimedeltaIndex.union with sort=False --- pandas/core/indexes/timedeltas.py | 13 +++++++++++-- pandas/tests/indexes/timedeltas/test_setops.py | 16 ++++++++++++++++ 2 files changed, 27 insertions(+), 2 deletions(-) diff --git a/pandas/core/indexes/timedeltas.py b/pandas/core/indexes/timedeltas.py index 6fa7a7d3d4fb6..66b4101a92197 100644 --- a/pandas/core/indexes/timedeltas.py +++ b/pandas/core/indexes/timedeltas.py @@ -260,7 +260,7 @@ def _union(self, other, sort): this, other = self, other if this._can_fast_union(other): - return this._fast_union(other) + return this._fast_union(other, sort=sort) else: result = Index._union(this, other, sort=sort) if isinstance(result, TimedeltaIndex): @@ -268,7 +268,7 @@ def _union(self, other, sort): result._set_freq("infer") return result - def _fast_union(self, other): + def _fast_union(self, other, sort=None): if len(other) == 0: return self.view(type(self)) @@ -278,6 +278,15 @@ def _fast_union(self, other): # to make our life easier, "sort" the two ranges if self[0] <= other[0]: left, right = self, other + elif sort is False: + # TDIs are not in the "correct" order and we don't want + # to sort but want to remove overlaps + left, right = self, other + left_start = left[0] + loc = right.searchsorted(left_start, side="left") + right_chunk = right.values[:loc] + dates = concat_compat((left.values, right_chunk)) + return self._shallow_copy(dates) else: left, right = other, self diff --git a/pandas/tests/indexes/timedeltas/test_setops.py b/pandas/tests/indexes/timedeltas/test_setops.py index b2024d04efc66..e6d5ca97a0bdf 100644 --- a/pandas/tests/indexes/timedeltas/test_setops.py +++ b/pandas/tests/indexes/timedeltas/test_setops.py @@ -22,6 +22,22 @@ def test_union(self): i1.union(i2) # Works i2.union(i1) # Fails with "AttributeError: can't set attribute" + def test_union_sort_false(self): + tdi = timedelta_range("1day", periods=5) + + left = tdi[3:] + right = tdi[:3] + + # Check that we are testing the desired code path + assert left._can_fast_union(right) + + result = left.union(right) + tm.assert_index_equal(result, tdi) + + result = left.union(right, sort=False) + expected = pd.TimedeltaIndex(["4 Days", "5 Days", "1 Days", "2 Day", "3 Days"]) + tm.assert_index_equal(result, expected) + def test_union_coverage(self): idx = TimedeltaIndex(["3d", "1d", "2d"])