Skip to content

Commit

Permalink
ENH: Add in sort keyword to DatetimeIndex.union (pandas-dev#25110)
Browse files Browse the repository at this point in the history
  • Loading branch information
reidy-p authored and Pingviinituutti committed Feb 28, 2019
1 parent 72ddc11 commit 7042d06
Show file tree
Hide file tree
Showing 4 changed files with 123 additions and 54 deletions.
Binary file added doc/source/styled.xlsx
Binary file not shown.
1 change: 1 addition & 0 deletions doc/source/whatsnew/v0.25.0.rst
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@ Other Enhancements
- Indexing of ``DataFrame`` and ``Series`` now accepts zerodim ``np.ndarray`` (:issue:`24919`)
- :meth:`Timestamp.replace` now supports the ``fold`` argument to disambiguate DST transition times (:issue:`25017`)
- :meth:`DataFrame.at_time` and :meth:`Series.at_time` now support :meth:`datetime.time` objects with timezones (:issue:`24043`)
- :meth:`DatetimeIndex.union` now supports the ``sort`` argument. The behaviour of the sort parameter matches that of :meth:`Index.union` (:issue:`24994`)
-

.. _whatsnew_0250.api_breaking:
Expand Down
38 changes: 32 additions & 6 deletions pandas/core/indexes/datetimes.py
Original file line number Diff line number Diff line change
Expand Up @@ -460,7 +460,7 @@ def _formatter_func(self):
# --------------------------------------------------------------------
# Set Operation Methods

def union(self, other):
def union(self, other, sort=None):
"""
Specialized union for DatetimeIndex objects. If combine
overlapping ranges with the same DateOffset, will be much
Expand All @@ -469,15 +469,29 @@ def union(self, other):
Parameters
----------
other : DatetimeIndex or array-like
sort : bool or None, default None
Whether to sort the resulting Index.
* None : Sort the result, except when
1. `self` and `other` are equal.
2. `self` or `other` has length 0.
3. Some values in `self` or `other` cannot be compared.
A RuntimeWarning is issued in this case.
* False : do not sort the result
.. versionadded:: 0.25.0
Returns
-------
y : Index or DatetimeIndex
"""
self._validate_sort_keyword(sort)
self._assert_can_do_setop(other)

if len(other) == 0 or self.equals(other) or len(self) == 0:
return super(DatetimeIndex, self).union(other)
return super(DatetimeIndex, self).union(other, sort=sort)

if not isinstance(other, DatetimeIndex):
try:
Expand All @@ -488,9 +502,9 @@ def union(self, other):
this, other = self._maybe_utc_convert(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)
result = Index.union(this, other, sort=sort)
if isinstance(result, DatetimeIndex):
# TODO: we shouldn't be setting attributes like this;
# in all the tests this equality already holds
Expand Down Expand Up @@ -563,16 +577,28 @@ def _can_fast_union(self, other):
# this will raise
return False

def _fast_union(self, other):
def _fast_union(self, other, sort=None):
if len(other) == 0:
return self.view(type(self))

if len(self) == 0:
return other.view(type(self))

# to make our life easier, "sort" the two ranges
# Both DTIs are monotonic. Check if they are already
# in the "correct" order
if self[0] <= other[0]:
left, right = self, other
# DTIs are not in the "correct" order and we don't want
# to sort but want to remove overlaps
elif sort is False:
left, right = self, other
left_start = left[0]
loc = right.searchsorted(left_start, side='left')
right_chunk = right.values[:loc]
dates = _concat._concat_compat((left.values, right_chunk))
return self._shallow_copy(dates)
# DTIs are not in the "correct" order and we want
# to sort
else:
left, right = other, self

Expand Down
Loading

0 comments on commit 7042d06

Please sign in to comment.