Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

REF: share join methods for DTI/TDI #30595

Merged
merged 5 commits into from
Jan 1, 2020
Merged
Show file tree
Hide file tree
Changes from 3 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
122 changes: 95 additions & 27 deletions pandas/core/indexes/datetimelike.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,8 +6,9 @@

import numpy as np

from pandas._libs import NaT, iNaT, lib
from pandas._libs import NaT, iNaT, join as libjoin, lib
from pandas._libs.algos import unique_deltas
from pandas._libs.tslibs import timezones
from pandas.compat.numpy import function as nv
from pandas.errors import AbstractMethodError
from pandas.util._decorators import Appender, cache_readonly
Expand Down Expand Up @@ -72,6 +73,30 @@ def method(self, other):
return method


def _join_i8_wrapper(joinf, with_indexers: bool = True):
"""
Create the join wrapper methods.
"""

def wrapper(left, right):
if isinstance(left, (np.ndarray, ABCIndex, ABCSeries, DatetimeLikeArrayMixin)):
left = left.view("i8")
if isinstance(right, (np.ndarray, ABCIndex, ABCSeries, DatetimeLikeArrayMixin)):
right = right.view("i8")
results = joinf(left, right)
if with_indexers:
# dtype should be timedelta64[ns] for TimedeltaIndex
# and datetime64[ns] for DatetimeIndex
dtype = left.dtype.base

join_index, left_indexer, right_indexer = results
join_index = join_index.view(dtype)
return join_index, left_indexer, right_indexer
return results

return wrapper


class DatetimeIndexOpsMixin(ExtensionOpsMixin):
"""
Common ops mixin to support a unified interface datetimelike Index.
Expand Down Expand Up @@ -208,32 +233,6 @@ def equals(self, other):

return np.array_equal(self.asi8, other.asi8)

@staticmethod
def _join_i8_wrapper(joinf, dtype, with_indexers=True):
"""
Create the join wrapper methods.
"""
from pandas.core.arrays.datetimelike import DatetimeLikeArrayMixin

@staticmethod
def wrapper(left, right):
if isinstance(
left, (np.ndarray, ABCIndex, ABCSeries, DatetimeLikeArrayMixin)
):
left = left.view("i8")
if isinstance(
right, (np.ndarray, ABCIndex, ABCSeries, DatetimeLikeArrayMixin)
):
right = right.view("i8")
results = joinf(left, right)
if with_indexers:
join_index, left_indexer, right_indexer = results
join_index = join_index.view(dtype)
return join_index, left_indexer, right_indexer
return results

return wrapper

def _ensure_localized(
self, arg, ambiguous="raise", nonexistent="raise", from_utc=False
):
Expand Down Expand Up @@ -853,6 +852,75 @@ def _can_fast_union(self, other) -> bool:
# this will raise
return False

# --------------------------------------------------------------------
# Join Methods
_join_precedence = 10

_inner_indexer = _join_i8_wrapper(libjoin.inner_join_indexer)
_outer_indexer = _join_i8_wrapper(libjoin.outer_join_indexer)
_left_indexer = _join_i8_wrapper(libjoin.left_join_indexer)
_left_indexer_unique = _join_i8_wrapper(
libjoin.left_join_indexer_unique, with_indexers=False
)

def join(
self, other, how: str = "left", level=None, return_indexers=False, sort=False
):
"""
See Index.join
"""
if self._is_convertible_to_index_for_join(other):
try:
other = type(self)(other)
except (TypeError, ValueError):
pass

this, other = self._maybe_utc_convert(other)
return Index.join(
this,
other,
how=how,
level=level,
return_indexers=return_indexers,
sort=sort,
)

def _maybe_utc_convert(self, other):
this = self
if not hasattr(self, "tz"):
return this, other

if isinstance(other, type(self)):
if self.tz is not None:
if other.tz is None:
raise TypeError("Cannot join tz-naive with tz-aware DatetimeIndex")
elif other.tz is not None:
raise TypeError("Cannot join tz-naive with tz-aware DatetimeIndex")

if not timezones.tz_compare(self.tz, other.tz):
this = self.tz_convert("UTC")
other = other.tz_convert("UTC")
return this, other

@classmethod
def _is_convertible_to_index_for_join(cls, other: Index) -> bool:
"""
return a boolean whether I can attempt conversion to a
DatetimeIndex/TimedeltaIndex
"""
if isinstance(other, cls):
return False
elif len(other) > 0 and other.inferred_type not in (
"floating",
"mixed-integer",
"integer",
"integer-na",
"mixed-integer-float",
"mixed",
):
return True
return False


def wrap_arithmetic_op(self, other, result):
if result is NotImplemented:
Expand Down
61 changes: 0 additions & 61 deletions pandas/core/indexes/datetimes.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,6 @@
import numpy as np

from pandas._libs import NaT, Timestamp, index as libindex, lib, tslib as libts
import pandas._libs.join as libjoin
from pandas._libs.tslibs import ccalendar, fields, parsing, timezones
from pandas.util._decorators import Appender, Substitution, cache_readonly

Expand All @@ -32,7 +31,6 @@
import pandas.core.common as com
from pandas.core.indexes.base import Index, maybe_extract_name
from pandas.core.indexes.datetimelike import (
DatetimeIndexOpsMixin,
DatetimelikeDelegateMixin,
DatetimeTimedeltaMixin,
)
Expand Down Expand Up @@ -195,17 +193,6 @@ class DatetimeIndex(DatetimeTimedeltaMixin, DatetimeDelegateMixin):
"""

_typ = "datetimeindex"
_join_precedence = 10

def _join_i8_wrapper(joinf, **kwargs):
return DatetimeIndexOpsMixin._join_i8_wrapper(joinf, dtype="M8[ns]", **kwargs)

_inner_indexer = _join_i8_wrapper(libjoin.inner_join_indexer)
_outer_indexer = _join_i8_wrapper(libjoin.outer_join_indexer)
_left_indexer = _join_i8_wrapper(libjoin.left_join_indexer)
_left_indexer_unique = _join_i8_wrapper(
libjoin.left_join_indexer_unique, with_indexers=False
)

_engine_type = libindex.DatetimeEngine
_supports_partial_string_indexing = True
Expand Down Expand Up @@ -645,54 +632,6 @@ def snap(self, freq="S"):
# we know it conforms; skip check
return DatetimeIndex._simple_new(snapped, name=self.name, tz=self.tz, freq=freq)

def join(
self, other, how: str = "left", level=None, return_indexers=False, sort=False
):
"""
See Index.join
"""
if (
not isinstance(other, DatetimeIndex)
and len(other) > 0
and other.inferred_type
not in (
"floating",
"integer",
"integer-na",
"mixed-integer",
"mixed-integer-float",
"mixed",
)
):
try:
other = DatetimeIndex(other)
except (TypeError, ValueError):
pass

this, other = self._maybe_utc_convert(other)
return Index.join(
this,
other,
how=how,
level=level,
return_indexers=return_indexers,
sort=sort,
)

def _maybe_utc_convert(self, other):
this = self
if isinstance(other, DatetimeIndex):
if self.tz is not None:
if other.tz is None:
raise TypeError("Cannot join tz-naive with tz-aware DatetimeIndex")
elif other.tz is not None:
raise TypeError("Cannot join tz-naive with tz-aware DatetimeIndex")

if not timezones.tz_compare(self.tz, other.tz):
this = self.tz_convert("UTC")
other = other.tz_convert("UTC")
return this, other

def _wrap_joined_index(self, joined, other):
name = get_op_result_name(self, other)
if (
Expand Down
50 changes: 1 addition & 49 deletions pandas/core/indexes/timedeltas.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@

import numpy as np

from pandas._libs import NaT, Timedelta, index as libindex, join as libjoin, lib
from pandas._libs import NaT, Timedelta, index as libindex, lib
from pandas.util._decorators import Appender, Substitution

from pandas.core.dtypes.common import (
Expand Down Expand Up @@ -121,17 +121,6 @@ class TimedeltaIndex(
"""

_typ = "timedeltaindex"
_join_precedence = 10

def _join_i8_wrapper(joinf, **kwargs):
return DatetimeIndexOpsMixin._join_i8_wrapper(joinf, dtype="m8[ns]", **kwargs)

_inner_indexer = _join_i8_wrapper(libjoin.inner_join_indexer)
_outer_indexer = _join_i8_wrapper(libjoin.outer_join_indexer)
_left_indexer = _join_i8_wrapper(libjoin.left_join_indexer)
_left_indexer_unique = _join_i8_wrapper(
libjoin.left_join_indexer_unique, with_indexers=False
)

_engine_type = libindex.TimedeltaEngine

Expand Down Expand Up @@ -294,25 +283,6 @@ def _union(self, other, sort):
result._set_freq("infer")
return result

def join(self, other, how="left", level=None, return_indexers=False, sort=False):
"""
See Index.join
"""
if _is_convertible_to_index(other):
try:
other = TimedeltaIndex(other)
except (TypeError, ValueError):
pass

return Index.join(
self,
other,
how=how,
level=level,
return_indexers=return_indexers,
sort=sort,
)

def _wrap_joined_index(self, joined, other):
name = get_op_result_name(self, other)
if (
Expand Down Expand Up @@ -569,24 +539,6 @@ def delete(self, loc):
TimedeltaIndex._add_datetimelike_methods()


def _is_convertible_to_index(other) -> bool:
"""
return a boolean whether I can attempt conversion to a TimedeltaIndex
"""
if isinstance(other, TimedeltaIndex):
return True
elif len(other) > 0 and other.inferred_type not in (
"floating",
"mixed-integer",
"integer",
"integer-na",
"mixed-integer-float",
"mixed",
):
return True
return False


def timedelta_range(
start=None, end=None, periods=None, freq=None, name=None, closed=None
) -> TimedeltaIndex:
Expand Down