Skip to content

Commit 9d2eca8

Browse files
authored
CLN: freq inference in DTI/TDI set ops (#33839)
1 parent 3ee9000 commit 9d2eca8

File tree

1 file changed

+37
-19
lines changed

1 file changed

+37
-19
lines changed

pandas/core/indexes/datetimelike.py

+37-19
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22
Base and utility classes for tseries type pandas objects.
33
"""
44
from datetime import datetime
5-
from typing import Any, List, Optional, Union, cast
5+
from typing import Any, List, Optional, TypeVar, Union, cast
66

77
import numpy as np
88

@@ -45,6 +45,8 @@
4545

4646
_index_doc_kwargs = dict(ibase._index_doc_kwargs)
4747

48+
_T = TypeVar("_T", bound="DatetimeIndexOpsMixin")
49+
4850

4951
def _join_i8_wrapper(joinf, with_indexers: bool = True):
5052
"""
@@ -655,13 +657,7 @@ def intersection(self, other, sort=False):
655657
result = result._with_freq("infer")
656658
return result
657659

658-
elif (
659-
other.freq is None
660-
or self.freq is None
661-
or other.freq != self.freq
662-
or not other.freq.is_anchored()
663-
or (not self.is_monotonic or not other.is_monotonic)
664-
):
660+
elif not self._can_fast_intersect(other):
665661
result = Index.intersection(self, other, sort=sort)
666662
result = result._with_freq("infer")
667663
return result
@@ -684,7 +680,28 @@ def intersection(self, other, sort=False):
684680
left_chunk = left._values[lslice]
685681
return self._shallow_copy(left_chunk)
686682

687-
def _can_fast_union(self, other) -> bool:
683+
def _can_fast_intersect(self: _T, other: _T) -> bool:
684+
if self.freq is None:
685+
return False
686+
687+
if other.freq != self.freq:
688+
return False
689+
690+
if not self.is_monotonic_increasing:
691+
# Because freq is not None, we must then be monotonic decreasing
692+
return False
693+
694+
if not self.freq.is_anchored():
695+
# If freq is not anchored, then despite having matching freqs,
696+
# we might not "line up"
697+
return False
698+
699+
return True
700+
701+
def _can_fast_union(self: _T, other: _T) -> bool:
702+
# Assumes that type(self) == type(other), as per the annotation
703+
# The ability to fast_union also implies that `freq` should be
704+
# retained on union.
688705
if not isinstance(other, type(self)):
689706
return False
690707

@@ -693,7 +710,9 @@ def _can_fast_union(self, other) -> bool:
693710
if freq is None or freq != other.freq:
694711
return False
695712

696-
if not self.is_monotonic or not other.is_monotonic:
713+
if not self.is_monotonic_increasing:
714+
# Because freq is not None, we must then be monotonic decreasing
715+
# TODO: do union on the reversed indexes?
697716
return False
698717

699718
if len(self) == 0 or len(other) == 0:
@@ -709,12 +728,7 @@ def _can_fast_union(self, other) -> bool:
709728
left_end = left[-1]
710729

711730
# Only need to "adjoin", not overlap
712-
try:
713-
return (right_start == left_end + freq) or right_start in left
714-
except ValueError:
715-
# if we are comparing a freq that does not propagate timezones
716-
# this will raise
717-
return False
731+
return (right_start == left_end + freq) or right_start in left
718732

719733
def _fast_union(self, other, sort=None):
720734
if len(other) == 0:
@@ -734,7 +748,7 @@ def _fast_union(self, other, sort=None):
734748
loc = right.searchsorted(left_start, side="left")
735749
right_chunk = right._values[:loc]
736750
dates = concat_compat((left._values, right_chunk))
737-
# TODO: can we infer that it has self.freq?
751+
# With sort being False, we can't infer that result.freq == self.freq
738752
result = self._shallow_copy(dates)._with_freq("infer")
739753
return result
740754
else:
@@ -748,8 +762,10 @@ def _fast_union(self, other, sort=None):
748762
loc = right.searchsorted(left_end, side="right")
749763
right_chunk = right._values[loc:]
750764
dates = concat_compat([left._values, right_chunk])
751-
# TODO: can we infer that it has self.freq?
752-
result = self._shallow_copy(dates)._with_freq("infer")
765+
# The can_fast_union check ensures that the result.freq
766+
# should match self.freq
767+
dates = type(self._data)(dates, freq=self.freq)
768+
result = type(self)._simple_new(dates, name=self.name)
753769
return result
754770
else:
755771
return left
@@ -766,6 +782,8 @@ def _union(self, other, sort):
766782
if this._can_fast_union(other):
767783
result = this._fast_union(other, sort=sort)
768784
if result.freq is None:
785+
# In the case where sort is None, _can_fast_union
786+
# implies that result.freq should match self.freq
769787
result = result._with_freq("infer")
770788
return result
771789
else:

0 commit comments

Comments
 (0)