"""
|
Functions for arithmetic and comparison operations on NumPy arrays and
|
ExtensionArrays.
|
"""
|
from __future__ import annotations
|
|
import datetime
|
from functools import partial
|
import operator
|
from typing import Any
|
|
import numpy as np
|
|
from pandas._libs import (
|
NaT,
|
Timedelta,
|
Timestamp,
|
lib,
|
ops as libops,
|
)
|
from pandas._libs.tslibs import (
|
BaseOffset,
|
get_supported_reso,
|
get_unit_from_dtype,
|
is_supported_unit,
|
is_unitless,
|
npy_unit_to_abbrev,
|
)
|
from pandas._typing import (
|
ArrayLike,
|
Shape,
|
)
|
|
from pandas.core.dtypes.cast import (
|
construct_1d_object_array_from_listlike,
|
find_common_type,
|
)
|
from pandas.core.dtypes.common import (
|
ensure_object,
|
is_bool_dtype,
|
is_integer_dtype,
|
is_list_like,
|
is_numeric_v_string_like,
|
is_object_dtype,
|
is_scalar,
|
)
|
from pandas.core.dtypes.generic import (
|
ABCExtensionArray,
|
ABCIndex,
|
ABCSeries,
|
)
|
from pandas.core.dtypes.missing import (
|
isna,
|
notna,
|
)
|
|
from pandas.core.computation import expressions
|
from pandas.core.construction import ensure_wrapped_if_datetimelike
|
from pandas.core.ops import (
|
missing,
|
roperator,
|
)
|
from pandas.core.ops.dispatch import should_extension_dispatch
|
from pandas.core.ops.invalid import invalid_comparison
|
|
|
def comp_method_OBJECT_ARRAY(op, x, y):
|
if isinstance(y, list):
|
y = construct_1d_object_array_from_listlike(y)
|
|
if isinstance(y, (np.ndarray, ABCSeries, ABCIndex)):
|
if not is_object_dtype(y.dtype):
|
y = y.astype(np.object_)
|
|
if isinstance(y, (ABCSeries, ABCIndex)):
|
y = y._values
|
|
if x.shape != y.shape:
|
raise ValueError("Shapes must match", x.shape, y.shape)
|
result = libops.vec_compare(x.ravel(), y.ravel(), op)
|
else:
|
result = libops.scalar_compare(x.ravel(), y, op)
|
return result.reshape(x.shape)
|
|
|
def _masked_arith_op(x: np.ndarray, y, op):
|
"""
|
If the given arithmetic operation fails, attempt it again on
|
only the non-null elements of the input array(s).
|
|
Parameters
|
----------
|
x : np.ndarray
|
y : np.ndarray, Series, Index
|
op : binary operator
|
"""
|
# For Series `x` is 1D so ravel() is a no-op; calling it anyway makes
|
# the logic valid for both Series and DataFrame ops.
|
xrav = x.ravel()
|
assert isinstance(x, np.ndarray), type(x)
|
if isinstance(y, np.ndarray):
|
dtype = find_common_type([x.dtype, y.dtype])
|
result = np.empty(x.size, dtype=dtype)
|
|
if len(x) != len(y):
|
raise ValueError(x.shape, y.shape)
|
ymask = notna(y)
|
|
# NB: ravel() is only safe since y is ndarray; for e.g. PeriodIndex
|
# we would get int64 dtype, see GH#19956
|
yrav = y.ravel()
|
mask = notna(xrav) & ymask.ravel()
|
|
# See GH#5284, GH#5035, GH#19448 for historical reference
|
if mask.any():
|
result[mask] = op(xrav[mask], yrav[mask])
|
|
else:
|
if not is_scalar(y):
|
raise TypeError(
|
f"Cannot broadcast np.ndarray with operand of type { type(y) }"
|
)
|
|
# mask is only meaningful for x
|
result = np.empty(x.size, dtype=x.dtype)
|
mask = notna(xrav)
|
|
# 1 ** np.nan is 1. So we have to unmask those.
|
if op is pow:
|
mask = np.where(x == 1, False, mask)
|
elif op is roperator.rpow:
|
mask = np.where(y == 1, False, mask)
|
|
if mask.any():
|
result[mask] = op(xrav[mask], y)
|
|
np.putmask(result, ~mask, np.nan)
|
result = result.reshape(x.shape) # 2D compat
|
return result
|
|
|
def _na_arithmetic_op(left: np.ndarray, right, op, is_cmp: bool = False):
|
"""
|
Return the result of evaluating op on the passed in values.
|
|
If native types are not compatible, try coercion to object dtype.
|
|
Parameters
|
----------
|
left : np.ndarray
|
right : np.ndarray or scalar
|
Excludes DataFrame, Series, Index, ExtensionArray.
|
is_cmp : bool, default False
|
If this a comparison operation.
|
|
Returns
|
-------
|
array-like
|
|
Raises
|
------
|
TypeError : invalid operation
|
"""
|
if isinstance(right, str):
|
# can never use numexpr
|
func = op
|
else:
|
func = partial(expressions.evaluate, op)
|
|
try:
|
result = func(left, right)
|
except TypeError:
|
if not is_cmp and (is_object_dtype(left.dtype) or is_object_dtype(right)):
|
# For object dtype, fallback to a masked operation (only operating
|
# on the non-missing values)
|
# Don't do this for comparisons, as that will handle complex numbers
|
# incorrectly, see GH#32047
|
result = _masked_arith_op(left, right, op)
|
else:
|
raise
|
|
if is_cmp and (is_scalar(result) or result is NotImplemented):
|
# numpy returned a scalar instead of operating element-wise
|
# e.g. numeric array vs str
|
# TODO: can remove this after dropping some future numpy version?
|
return invalid_comparison(left, right, op)
|
|
return missing.dispatch_fill_zeros(op, left, right, result)
|
|
|
def arithmetic_op(left: ArrayLike, right: Any, op):
|
"""
|
Evaluate an arithmetic operation `+`, `-`, `*`, `/`, `//`, `%`, `**`, ...
|
|
Note: the caller is responsible for ensuring that numpy warnings are
|
suppressed (with np.errstate(all="ignore")) if needed.
|
|
Parameters
|
----------
|
left : np.ndarray or ExtensionArray
|
right : object
|
Cannot be a DataFrame or Index. Series is *not* excluded.
|
op : {operator.add, operator.sub, ...}
|
Or one of the reversed variants from roperator.
|
|
Returns
|
-------
|
ndarray or ExtensionArray
|
Or a 2-tuple of these in the case of divmod or rdivmod.
|
"""
|
# NB: We assume that extract_array and ensure_wrapped_if_datetimelike
|
# have already been called on `left` and `right`,
|
# and `maybe_prepare_scalar_for_op` has already been called on `right`
|
# We need to special-case datetime64/timedelta64 dtypes (e.g. because numpy
|
# casts integer dtypes to timedelta64 when operating with timedelta64 - GH#22390)
|
|
if (
|
should_extension_dispatch(left, right)
|
or isinstance(right, (Timedelta, BaseOffset, Timestamp))
|
or right is NaT
|
):
|
# Timedelta/Timestamp and other custom scalars are included in the check
|
# because numexpr will fail on it, see GH#31457
|
res_values = op(left, right)
|
else:
|
# TODO we should handle EAs consistently and move this check before the if/else
|
# (https://github.com/pandas-dev/pandas/issues/41165)
|
_bool_arith_check(op, left, right)
|
|
# error: Argument 1 to "_na_arithmetic_op" has incompatible type
|
# "Union[ExtensionArray, ndarray[Any, Any]]"; expected "ndarray[Any, Any]"
|
res_values = _na_arithmetic_op(left, right, op) # type: ignore[arg-type]
|
|
return res_values
|
|
|
def comparison_op(left: ArrayLike, right: Any, op) -> ArrayLike:
|
"""
|
Evaluate a comparison operation `=`, `!=`, `>=`, `>`, `<=`, or `<`.
|
|
Note: the caller is responsible for ensuring that numpy warnings are
|
suppressed (with np.errstate(all="ignore")) if needed.
|
|
Parameters
|
----------
|
left : np.ndarray or ExtensionArray
|
right : object
|
Cannot be a DataFrame, Series, or Index.
|
op : {operator.eq, operator.ne, operator.gt, operator.ge, operator.lt, operator.le}
|
|
Returns
|
-------
|
ndarray or ExtensionArray
|
"""
|
# NB: We assume extract_array has already been called on left and right
|
lvalues = ensure_wrapped_if_datetimelike(left)
|
rvalues = ensure_wrapped_if_datetimelike(right)
|
|
rvalues = lib.item_from_zerodim(rvalues)
|
if isinstance(rvalues, list):
|
# We don't catch tuple here bc we may be comparing e.g. MultiIndex
|
# to a tuple that represents a single entry, see test_compare_tuple_strs
|
rvalues = np.asarray(rvalues)
|
|
if isinstance(rvalues, (np.ndarray, ABCExtensionArray)):
|
# TODO: make this treatment consistent across ops and classes.
|
# We are not catching all listlikes here (e.g. frozenset, tuple)
|
# The ambiguous case is object-dtype. See GH#27803
|
if len(lvalues) != len(rvalues):
|
raise ValueError(
|
"Lengths must match to compare", lvalues.shape, rvalues.shape
|
)
|
|
if should_extension_dispatch(lvalues, rvalues) or (
|
(isinstance(rvalues, (Timedelta, BaseOffset, Timestamp)) or right is NaT)
|
and not is_object_dtype(lvalues.dtype)
|
):
|
# Call the method on lvalues
|
res_values = op(lvalues, rvalues)
|
|
elif is_scalar(rvalues) and isna(rvalues): # TODO: but not pd.NA?
|
# numpy does not like comparisons vs None
|
if op is operator.ne:
|
res_values = np.ones(lvalues.shape, dtype=bool)
|
else:
|
res_values = np.zeros(lvalues.shape, dtype=bool)
|
|
elif is_numeric_v_string_like(lvalues, rvalues):
|
# GH#36377 going through the numexpr path would incorrectly raise
|
return invalid_comparison(lvalues, rvalues, op)
|
|
elif is_object_dtype(lvalues.dtype) or isinstance(rvalues, str):
|
res_values = comp_method_OBJECT_ARRAY(op, lvalues, rvalues)
|
|
else:
|
res_values = _na_arithmetic_op(lvalues, rvalues, op, is_cmp=True)
|
|
return res_values
|
|
|
def na_logical_op(x: np.ndarray, y, op):
|
try:
|
# For exposition, write:
|
# yarr = isinstance(y, np.ndarray)
|
# yint = is_integer(y) or (yarr and y.dtype.kind == "i")
|
# ybool = is_bool(y) or (yarr and y.dtype.kind == "b")
|
# xint = x.dtype.kind == "i"
|
# xbool = x.dtype.kind == "b"
|
# Then Cases where this goes through without raising include:
|
# (xint or xbool) and (yint or bool)
|
result = op(x, y)
|
except TypeError:
|
if isinstance(y, np.ndarray):
|
# bool-bool dtype operations should be OK, should not get here
|
assert not (is_bool_dtype(x.dtype) and is_bool_dtype(y.dtype))
|
x = ensure_object(x)
|
y = ensure_object(y)
|
result = libops.vec_binop(x.ravel(), y.ravel(), op)
|
else:
|
# let null fall thru
|
assert lib.is_scalar(y)
|
if not isna(y):
|
y = bool(y)
|
try:
|
result = libops.scalar_binop(x, y, op)
|
except (
|
TypeError,
|
ValueError,
|
AttributeError,
|
OverflowError,
|
NotImplementedError,
|
) as err:
|
typ = type(y).__name__
|
raise TypeError(
|
f"Cannot perform '{op.__name__}' with a dtyped [{x.dtype}] array "
|
f"and scalar of type [{typ}]"
|
) from err
|
|
return result.reshape(x.shape)
|
|
|
def logical_op(left: ArrayLike, right: Any, op) -> ArrayLike:
|
"""
|
Evaluate a logical operation `|`, `&`, or `^`.
|
|
Parameters
|
----------
|
left : np.ndarray or ExtensionArray
|
right : object
|
Cannot be a DataFrame, Series, or Index.
|
op : {operator.and_, operator.or_, operator.xor}
|
Or one of the reversed variants from roperator.
|
|
Returns
|
-------
|
ndarray or ExtensionArray
|
"""
|
fill_int = lambda x: x
|
|
def fill_bool(x, left=None):
|
# if `left` is specifically not-boolean, we do not cast to bool
|
if x.dtype.kind in ["c", "f", "O"]:
|
# dtypes that can hold NA
|
mask = isna(x)
|
if mask.any():
|
x = x.astype(object)
|
x[mask] = False
|
|
if left is None or is_bool_dtype(left.dtype):
|
x = x.astype(bool)
|
return x
|
|
is_self_int_dtype = is_integer_dtype(left.dtype)
|
|
right = lib.item_from_zerodim(right)
|
if is_list_like(right) and not hasattr(right, "dtype"):
|
# e.g. list, tuple
|
right = construct_1d_object_array_from_listlike(right)
|
|
# NB: We assume extract_array has already been called on left and right
|
lvalues = ensure_wrapped_if_datetimelike(left)
|
rvalues = right
|
|
if should_extension_dispatch(lvalues, rvalues):
|
# Call the method on lvalues
|
res_values = op(lvalues, rvalues)
|
|
else:
|
if isinstance(rvalues, np.ndarray):
|
is_other_int_dtype = is_integer_dtype(rvalues.dtype)
|
rvalues = rvalues if is_other_int_dtype else fill_bool(rvalues, lvalues)
|
|
else:
|
# i.e. scalar
|
is_other_int_dtype = lib.is_integer(rvalues)
|
|
# For int vs int `^`, `|`, `&` are bitwise operators and return
|
# integer dtypes. Otherwise these are boolean ops
|
filler = fill_int if is_self_int_dtype and is_other_int_dtype else fill_bool
|
|
res_values = na_logical_op(lvalues, rvalues, op)
|
# error: Cannot call function of unknown type
|
res_values = filler(res_values) # type: ignore[operator]
|
|
return res_values
|
|
|
def get_array_op(op):
|
"""
|
Return a binary array operation corresponding to the given operator op.
|
|
Parameters
|
----------
|
op : function
|
Binary operator from operator or roperator module.
|
|
Returns
|
-------
|
functools.partial
|
"""
|
if isinstance(op, partial):
|
# We get here via dispatch_to_series in DataFrame case
|
# e.g. test_rolling_consistency_var_debiasing_factors
|
return op
|
|
op_name = op.__name__.strip("_").lstrip("r")
|
if op_name == "arith_op":
|
# Reached via DataFrame._combine_frame i.e. flex methods
|
# e.g. test_df_add_flex_filled_mixed_dtypes
|
return op
|
|
if op_name in {"eq", "ne", "lt", "le", "gt", "ge"}:
|
return partial(comparison_op, op=op)
|
elif op_name in {"and", "or", "xor", "rand", "ror", "rxor"}:
|
return partial(logical_op, op=op)
|
elif op_name in {
|
"add",
|
"sub",
|
"mul",
|
"truediv",
|
"floordiv",
|
"mod",
|
"divmod",
|
"pow",
|
}:
|
return partial(arithmetic_op, op=op)
|
else:
|
raise NotImplementedError(op_name)
|
|
|
def maybe_prepare_scalar_for_op(obj, shape: Shape):
|
"""
|
Cast non-pandas objects to pandas types to unify behavior of arithmetic
|
and comparison operations.
|
|
Parameters
|
----------
|
obj: object
|
shape : tuple[int]
|
|
Returns
|
-------
|
out : object
|
|
Notes
|
-----
|
Be careful to call this *after* determining the `name` attribute to be
|
attached to the result of the arithmetic operation.
|
"""
|
if type(obj) is datetime.timedelta:
|
# GH#22390 cast up to Timedelta to rely on Timedelta
|
# implementation; otherwise operation against numeric-dtype
|
# raises TypeError
|
return Timedelta(obj)
|
elif type(obj) is datetime.datetime:
|
# cast up to Timestamp to rely on Timestamp implementation, see Timedelta above
|
return Timestamp(obj)
|
elif isinstance(obj, np.datetime64):
|
# GH#28080 numpy casts integer-dtype to datetime64 when doing
|
# array[int] + datetime64, which we do not allow
|
if isna(obj):
|
from pandas.core.arrays import DatetimeArray
|
|
# Avoid possible ambiguities with pd.NaT
|
# GH 52295
|
if is_unitless(obj.dtype):
|
obj = obj.astype("datetime64[ns]")
|
elif not is_supported_unit(get_unit_from_dtype(obj.dtype)):
|
unit = get_unit_from_dtype(obj.dtype)
|
closest_unit = npy_unit_to_abbrev(get_supported_reso(unit))
|
obj = obj.astype(f"datetime64[{closest_unit}]")
|
right = np.broadcast_to(obj, shape)
|
return DatetimeArray(right)
|
|
return Timestamp(obj)
|
|
elif isinstance(obj, np.timedelta64):
|
if isna(obj):
|
from pandas.core.arrays import TimedeltaArray
|
|
# wrapping timedelta64("NaT") in Timedelta returns NaT,
|
# which would incorrectly be treated as a datetime-NaT, so
|
# we broadcast and wrap in a TimedeltaArray
|
# GH 52295
|
if is_unitless(obj.dtype):
|
obj = obj.astype("timedelta64[ns]")
|
elif not is_supported_unit(get_unit_from_dtype(obj.dtype)):
|
unit = get_unit_from_dtype(obj.dtype)
|
closest_unit = npy_unit_to_abbrev(get_supported_reso(unit))
|
obj = obj.astype(f"timedelta64[{closest_unit}]")
|
right = np.broadcast_to(obj, shape)
|
return TimedeltaArray(right)
|
|
# In particular non-nanosecond timedelta64 needs to be cast to
|
# nanoseconds, or else we get undesired behavior like
|
# np.timedelta64(3, 'D') / 2 == np.timedelta64(1, 'D')
|
return Timedelta(obj)
|
|
return obj
|
|
|
_BOOL_OP_NOT_ALLOWED = {
|
operator.truediv,
|
roperator.rtruediv,
|
operator.floordiv,
|
roperator.rfloordiv,
|
operator.pow,
|
roperator.rpow,
|
}
|
|
|
def _bool_arith_check(op, a, b):
|
"""
|
In contrast to numpy, pandas raises an error for certain operations
|
with booleans.
|
"""
|
if op in _BOOL_OP_NOT_ALLOWED:
|
if is_bool_dtype(a.dtype) and (
|
is_bool_dtype(b) or isinstance(b, (bool, np.bool_))
|
):
|
op_name = op.__name__.strip("_").lstrip("r")
|
raise NotImplementedError(
|
f"operator '{op_name}' not implemented for bool dtypes"
|
)
|