Skip to content

REF: de-duplicate ndarray[datetimelike] wrapping #38129

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

Merged
merged 4 commits into from
Nov 29, 2020
Merged
Show file tree
Hide file tree
Changes from all 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
12 changes: 7 additions & 5 deletions pandas/core/arrays/interval.py
Original file line number Diff line number Diff line change
Expand Up @@ -44,7 +44,11 @@
from pandas.core.arrays.base import ExtensionArray, _extension_array_shared_docs
from pandas.core.arrays.categorical import Categorical
import pandas.core.common as com
from pandas.core.construction import array, extract_array
from pandas.core.construction import (
array,
ensure_wrapped_if_datetimelike,
extract_array,
)
from pandas.core.indexers import check_array_indexer
from pandas.core.indexes.base import ensure_index
from pandas.core.ops import invalid_comparison, unpack_zerodim_and_defer
Expand Down Expand Up @@ -251,11 +255,9 @@ def _simple_new(
raise ValueError(msg)

# For dt64/td64 we want DatetimeArray/TimedeltaArray instead of ndarray
from pandas.core.ops.array_ops import maybe_upcast_datetimelike_array

left = maybe_upcast_datetimelike_array(left)
left = ensure_wrapped_if_datetimelike(left)
left = extract_array(left, extract_numpy=True)
right = maybe_upcast_datetimelike_array(right)
right = ensure_wrapped_if_datetimelike(right)
right = extract_array(right, extract_numpy=True)

lbase = getattr(left, "_ndarray", left).base
Expand Down
18 changes: 18 additions & 0 deletions pandas/core/construction.py
Original file line number Diff line number Diff line change
Expand Up @@ -402,6 +402,24 @@ def extract_array(obj: object, extract_numpy: bool = False) -> Union[Any, ArrayL
return obj


def ensure_wrapped_if_datetimelike(arr):
"""
Wrap datetime64 and timedelta64 ndarrays in DatetimeArray/TimedeltaArray.
"""
if isinstance(arr, np.ndarray):
if arr.dtype.kind == "M":
from pandas.core.arrays import DatetimeArray

return DatetimeArray._from_sequence(arr)

elif arr.dtype.kind == "m":
from pandas.core.arrays import TimedeltaArray

return TimedeltaArray._from_sequence(arr)

return arr


def sanitize_array(
data,
index: Optional[Index],
Expand Down
22 changes: 5 additions & 17 deletions pandas/core/dtypes/concat.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@

from pandas.core.arrays import ExtensionArray
from pandas.core.arrays.sparse import SparseArray
from pandas.core.construction import array
from pandas.core.construction import array, ensure_wrapped_if_datetimelike


def _get_dtype_kinds(arrays) -> Set[str]:
Expand Down Expand Up @@ -360,12 +360,14 @@ def _concat_datetime(to_concat, axis=0):
-------
a single array, preserving the combined dtypes
"""
to_concat = [_wrap_datetimelike(x) for x in to_concat]
to_concat = [ensure_wrapped_if_datetimelike(x) for x in to_concat]

single_dtype = len({x.dtype for x in to_concat}) == 1

# multiple types, need to coerce to object
if not single_dtype:
# wrap_datetimelike ensures that astype(object) wraps in Timestamp/Timedelta
# ensure_wrapped_if_datetimelike ensures that astype(object) wraps
# in Timestamp/Timedelta
return _concatenate_2d([x.astype(object) for x in to_concat], axis=axis)

if axis == 1:
Expand All @@ -379,17 +381,3 @@ def _concat_datetime(to_concat, axis=0):
assert result.shape[0] == 1
result = result[0]
return result


def _wrap_datetimelike(arr):
"""
Wrap datetime64 and timedelta64 ndarrays in DatetimeArray/TimedeltaArray.

DTA/TDA handle .astype(object) correctly.
"""
from pandas.core.construction import array as pd_array, extract_array

arr = extract_array(arr, extract_numpy=True)
if isinstance(arr, np.ndarray) and arr.dtype.kind in ["m", "M"]:
arr = pd_array(arr)
return arr
34 changes: 5 additions & 29 deletions pandas/core/ops/array_ops.py
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,7 @@
from pandas.core.dtypes.generic import ABCExtensionArray, ABCIndexClass, ABCSeries
from pandas.core.dtypes.missing import isna, notna

from pandas.core.construction import ensure_wrapped_if_datetimelike
from pandas.core.ops import missing
from pandas.core.ops.dispatch import should_extension_dispatch
from pandas.core.ops.invalid import invalid_comparison
Expand Down Expand Up @@ -175,8 +176,8 @@ def arithmetic_op(left: ArrayLike, right: Any, op):

# NB: We assume that extract_array has already been called
# on `left` and `right`.
lvalues = maybe_upcast_datetimelike_array(left)
rvalues = maybe_upcast_datetimelike_array(right)
lvalues = ensure_wrapped_if_datetimelike(left)
rvalues = ensure_wrapped_if_datetimelike(right)
rvalues = _maybe_upcast_for_op(rvalues, lvalues.shape)

if should_extension_dispatch(lvalues, rvalues) or isinstance(rvalues, Timedelta):
Expand Down Expand Up @@ -206,7 +207,7 @@ def comparison_op(left: ArrayLike, right: Any, op) -> ArrayLike:
ndarray or ExtensionArray
"""
# NB: We assume extract_array has already been called on left and right
lvalues = maybe_upcast_datetimelike_array(left)
lvalues = ensure_wrapped_if_datetimelike(left)
rvalues = right

rvalues = lib.item_from_zerodim(rvalues)
Expand Down Expand Up @@ -331,7 +332,7 @@ def fill_bool(x, left=None):
right = construct_1d_object_array_from_listlike(right)

# NB: We assume extract_array has already been called on left and right
lvalues = maybe_upcast_datetimelike_array(left)
lvalues = ensure_wrapped_if_datetimelike(left)
rvalues = right

if should_extension_dispatch(lvalues, rvalues):
Expand Down Expand Up @@ -400,31 +401,6 @@ def get_array_op(op):
raise NotImplementedError(op_name)


def maybe_upcast_datetimelike_array(obj: ArrayLike) -> ArrayLike:
"""
If we have an ndarray that is either datetime64 or timedelta64, wrap in EA.

Parameters
----------
obj : ndarray or ExtensionArray

Returns
-------
ndarray or ExtensionArray
"""
if isinstance(obj, np.ndarray):
if obj.dtype.kind == "m":
from pandas.core.arrays import TimedeltaArray

return TimedeltaArray._from_sequence(obj)
if obj.dtype.kind == "M":
from pandas.core.arrays import DatetimeArray

return DatetimeArray._from_sequence(obj)

return obj


def _maybe_upcast_for_op(obj, shape: Shape):
"""
Cast non-pandas objects to pandas types to unify behavior of arithmetic
Expand Down