"""
|
Arithmetic operations for PandasObjects
|
|
This is not a public API.
|
"""
|
from __future__ import annotations
|
|
import operator
|
from typing import (
|
TYPE_CHECKING,
|
cast,
|
)
|
|
import numpy as np
|
|
from pandas._libs.ops_dispatch import maybe_dispatch_ufunc_to_dunder_op
|
from pandas._typing import (
|
Axis,
|
AxisInt,
|
Level,
|
)
|
from pandas.util._decorators import Appender
|
|
from pandas.core.dtypes.common import (
|
is_array_like,
|
is_list_like,
|
)
|
from pandas.core.dtypes.generic import (
|
ABCDataFrame,
|
ABCSeries,
|
)
|
from pandas.core.dtypes.missing import isna
|
|
from pandas.core import (
|
algorithms,
|
roperator,
|
)
|
from pandas.core.ops.array_ops import (
|
arithmetic_op,
|
comp_method_OBJECT_ARRAY,
|
comparison_op,
|
get_array_op,
|
logical_op,
|
maybe_prepare_scalar_for_op,
|
)
|
from pandas.core.ops.common import (
|
get_op_result_name,
|
unpack_zerodim_and_defer,
|
)
|
from pandas.core.ops.docstrings import (
|
_flex_comp_doc_FRAME,
|
_op_descriptions,
|
make_flex_doc,
|
)
|
from pandas.core.ops.invalid import invalid_comparison
|
from pandas.core.ops.mask_ops import (
|
kleene_and,
|
kleene_or,
|
kleene_xor,
|
)
|
from pandas.core.ops.methods import add_flex_arithmetic_methods
|
from pandas.core.roperator import (
|
radd,
|
rand_,
|
rdiv,
|
rdivmod,
|
rfloordiv,
|
rmod,
|
rmul,
|
ror_,
|
rpow,
|
rsub,
|
rtruediv,
|
rxor,
|
)
|
|
if TYPE_CHECKING:
|
from pandas import (
|
DataFrame,
|
Series,
|
)
|
|
# -----------------------------------------------------------------------------
|
# constants
|
ARITHMETIC_BINOPS: set[str] = {
|
"add",
|
"sub",
|
"mul",
|
"pow",
|
"mod",
|
"floordiv",
|
"truediv",
|
"divmod",
|
"radd",
|
"rsub",
|
"rmul",
|
"rpow",
|
"rmod",
|
"rfloordiv",
|
"rtruediv",
|
"rdivmod",
|
}
|
|
|
COMPARISON_BINOPS: set[str] = {"eq", "ne", "lt", "gt", "le", "ge"}
|
|
|
# -----------------------------------------------------------------------------
|
# Masking NA values and fallbacks for operations numpy does not support
|
|
|
def fill_binop(left, right, fill_value):
|
"""
|
If a non-None fill_value is given, replace null entries in left and right
|
with this value, but only in positions where _one_ of left/right is null,
|
not both.
|
|
Parameters
|
----------
|
left : array-like
|
right : array-like
|
fill_value : object
|
|
Returns
|
-------
|
left : array-like
|
right : array-like
|
|
Notes
|
-----
|
Makes copies if fill_value is not None and NAs are present.
|
"""
|
if fill_value is not None:
|
left_mask = isna(left)
|
right_mask = isna(right)
|
|
# one but not both
|
mask = left_mask ^ right_mask
|
|
if left_mask.any():
|
# Avoid making a copy if we can
|
left = left.copy()
|
left[left_mask & mask] = fill_value
|
|
if right_mask.any():
|
# Avoid making a copy if we can
|
right = right.copy()
|
right[right_mask & mask] = fill_value
|
|
return left, right
|
|
|
# -----------------------------------------------------------------------------
|
# Series
|
|
|
def align_method_SERIES(left: Series, right, align_asobject: bool = False):
|
"""align lhs and rhs Series"""
|
# ToDo: Different from align_method_FRAME, list, tuple and ndarray
|
# are not coerced here
|
# because Series has inconsistencies described in #13637
|
|
if isinstance(right, ABCSeries):
|
# avoid repeated alignment
|
if not left.index.equals(right.index):
|
if align_asobject:
|
# to keep original value's dtype for bool ops
|
left = left.astype(object)
|
right = right.astype(object)
|
|
left, right = left.align(right, copy=False)
|
|
return left, right
|
|
|
def flex_method_SERIES(op):
|
name = op.__name__.strip("_")
|
doc = make_flex_doc(name, "series")
|
|
@Appender(doc)
|
def flex_wrapper(self, other, level=None, fill_value=None, axis: Axis = 0):
|
# validate axis
|
if axis is not None:
|
self._get_axis_number(axis)
|
|
res_name = get_op_result_name(self, other)
|
|
if isinstance(other, ABCSeries):
|
return self._binop(other, op, level=level, fill_value=fill_value)
|
elif isinstance(other, (np.ndarray, list, tuple)):
|
if len(other) != len(self):
|
raise ValueError("Lengths must be equal")
|
other = self._constructor(other, self.index)
|
result = self._binop(other, op, level=level, fill_value=fill_value)
|
result.name = res_name
|
return result
|
else:
|
if fill_value is not None:
|
self = self.fillna(fill_value)
|
|
return op(self, other)
|
|
flex_wrapper.__name__ = name
|
return flex_wrapper
|
|
|
# -----------------------------------------------------------------------------
|
# DataFrame
|
|
|
def align_method_FRAME(
|
left, right, axis, flex: bool | None = False, level: Level = None
|
):
|
"""
|
Convert rhs to meet lhs dims if input is list, tuple or np.ndarray.
|
|
Parameters
|
----------
|
left : DataFrame
|
right : Any
|
axis : int, str, or None
|
flex : bool or None, default False
|
Whether this is a flex op, in which case we reindex.
|
None indicates not to check for alignment.
|
level : int or level name, default None
|
|
Returns
|
-------
|
left : DataFrame
|
right : Any
|
"""
|
|
def to_series(right):
|
msg = "Unable to coerce to Series, length must be {req_len}: given {given_len}"
|
|
# pass dtype to avoid doing inference, which would break consistency
|
# with Index/Series ops
|
dtype = None
|
if getattr(right, "dtype", None) == object:
|
# can't pass right.dtype unconditionally as that would break on e.g.
|
# datetime64[h] ndarray
|
dtype = object
|
|
if axis is not None and left._get_axis_name(axis) == "index":
|
if len(left.index) != len(right):
|
raise ValueError(
|
msg.format(req_len=len(left.index), given_len=len(right))
|
)
|
right = left._constructor_sliced(right, index=left.index, dtype=dtype)
|
else:
|
if len(left.columns) != len(right):
|
raise ValueError(
|
msg.format(req_len=len(left.columns), given_len=len(right))
|
)
|
right = left._constructor_sliced(right, index=left.columns, dtype=dtype)
|
return right
|
|
if isinstance(right, np.ndarray):
|
if right.ndim == 1:
|
right = to_series(right)
|
|
elif right.ndim == 2:
|
# We need to pass dtype=right.dtype to retain object dtype
|
# otherwise we lose consistency with Index and array ops
|
dtype = None
|
if getattr(right, "dtype", None) == object:
|
# can't pass right.dtype unconditionally as that would break on e.g.
|
# datetime64[h] ndarray
|
dtype = object
|
|
if right.shape == left.shape:
|
right = left._constructor(
|
right, index=left.index, columns=left.columns, dtype=dtype
|
)
|
|
elif right.shape[0] == left.shape[0] and right.shape[1] == 1:
|
# Broadcast across columns
|
right = np.broadcast_to(right, left.shape)
|
right = left._constructor(
|
right, index=left.index, columns=left.columns, dtype=dtype
|
)
|
|
elif right.shape[1] == left.shape[1] and right.shape[0] == 1:
|
# Broadcast along rows
|
right = to_series(right[0, :])
|
|
else:
|
raise ValueError(
|
"Unable to coerce to DataFrame, shape "
|
f"must be {left.shape}: given {right.shape}"
|
)
|
|
elif right.ndim > 2:
|
raise ValueError(
|
"Unable to coerce to Series/DataFrame, "
|
f"dimension must be <= 2: {right.shape}"
|
)
|
|
elif is_list_like(right) and not isinstance(right, (ABCSeries, ABCDataFrame)):
|
# GH 36702. Raise when attempting arithmetic with list of array-like.
|
if any(is_array_like(el) for el in right):
|
raise ValueError(
|
f"Unable to coerce list of {type(right[0])} to Series/DataFrame"
|
)
|
# GH17901
|
right = to_series(right)
|
|
if flex is not None and isinstance(right, ABCDataFrame):
|
if not left._indexed_same(right):
|
if flex:
|
left, right = left.align(right, join="outer", level=level, copy=False)
|
else:
|
raise ValueError(
|
"Can only compare identically-labeled (both index and columns) "
|
"DataFrame objects"
|
)
|
elif isinstance(right, ABCSeries):
|
# axis=1 is default for DataFrame-with-Series op
|
axis = left._get_axis_number(axis) if axis is not None else 1
|
|
if not flex:
|
if not left.axes[axis].equals(right.index):
|
raise ValueError(
|
"Operands are not aligned. Do "
|
"`left, right = left.align(right, axis=1, copy=False)` "
|
"before operating."
|
)
|
|
left, right = left.align(
|
right, join="outer", axis=axis, level=level, copy=False
|
)
|
right = _maybe_align_series_as_frame(left, right, axis)
|
|
return left, right
|
|
|
def should_reindex_frame_op(
|
left: DataFrame, right, op, axis: int, fill_value, level
|
) -> bool:
|
"""
|
Check if this is an operation between DataFrames that will need to reindex.
|
"""
|
assert isinstance(left, ABCDataFrame)
|
|
if op is operator.pow or op is roperator.rpow:
|
# GH#32685 pow has special semantics for operating with null values
|
return False
|
|
if not isinstance(right, ABCDataFrame):
|
return False
|
|
if fill_value is None and level is None and axis == 1:
|
# TODO: any other cases we should handle here?
|
|
# Intersection is always unique so we have to check the unique columns
|
left_uniques = left.columns.unique()
|
right_uniques = right.columns.unique()
|
cols = left_uniques.intersection(right_uniques)
|
if len(cols) and not (
|
len(cols) == len(left_uniques) and len(cols) == len(right_uniques)
|
):
|
# TODO: is there a shortcut available when len(cols) == 0?
|
return True
|
|
return False
|
|
|
def frame_arith_method_with_reindex(left: DataFrame, right: DataFrame, op) -> DataFrame:
|
"""
|
For DataFrame-with-DataFrame operations that require reindexing,
|
operate only on shared columns, then reindex.
|
|
Parameters
|
----------
|
left : DataFrame
|
right : DataFrame
|
op : binary operator
|
|
Returns
|
-------
|
DataFrame
|
"""
|
# GH#31623, only operate on shared columns
|
cols, lcols, rcols = left.columns.join(
|
right.columns, how="inner", level=None, return_indexers=True
|
)
|
|
new_left = left.iloc[:, lcols]
|
new_right = right.iloc[:, rcols]
|
result = op(new_left, new_right)
|
|
# Do the join on the columns instead of using align_method_FRAME
|
# to avoid constructing two potentially large/sparse DataFrames
|
join_columns, _, _ = left.columns.join(
|
right.columns, how="outer", level=None, return_indexers=True
|
)
|
|
if result.columns.has_duplicates:
|
# Avoid reindexing with a duplicate axis.
|
# https://github.com/pandas-dev/pandas/issues/35194
|
indexer, _ = result.columns.get_indexer_non_unique(join_columns)
|
indexer = algorithms.unique1d(indexer)
|
result = result._reindex_with_indexers(
|
{1: [join_columns, indexer]}, allow_dups=True
|
)
|
else:
|
result = result.reindex(join_columns, axis=1)
|
|
return result
|
|
|
def _maybe_align_series_as_frame(frame: DataFrame, series: Series, axis: AxisInt):
|
"""
|
If the Series operand is not EA-dtype, we can broadcast to 2D and operate
|
blockwise.
|
"""
|
rvalues = series._values
|
if not isinstance(rvalues, np.ndarray):
|
# TODO(EA2D): no need to special-case with 2D EAs
|
if rvalues.dtype in ("datetime64[ns]", "timedelta64[ns]"):
|
# We can losslessly+cheaply cast to ndarray
|
rvalues = np.asarray(rvalues)
|
else:
|
return series
|
|
if axis == 0:
|
rvalues = rvalues.reshape(-1, 1)
|
else:
|
rvalues = rvalues.reshape(1, -1)
|
|
rvalues = np.broadcast_to(rvalues, frame.shape)
|
# pass dtype to avoid doing inference
|
return type(frame)(
|
rvalues, index=frame.index, columns=frame.columns, dtype=rvalues.dtype
|
)
|
|
|
def flex_arith_method_FRAME(op):
|
op_name = op.__name__.strip("_")
|
|
na_op = get_array_op(op)
|
doc = make_flex_doc(op_name, "dataframe")
|
|
@Appender(doc)
|
def f(self, other, axis: Axis = "columns", level=None, fill_value=None):
|
axis = self._get_axis_number(axis) if axis is not None else 1
|
axis = cast(int, axis)
|
|
if should_reindex_frame_op(self, other, op, axis, fill_value, level):
|
return frame_arith_method_with_reindex(self, other, op)
|
|
if isinstance(other, ABCSeries) and fill_value is not None:
|
# TODO: We could allow this in cases where we end up going
|
# through the DataFrame path
|
raise NotImplementedError(f"fill_value {fill_value} not supported.")
|
|
other = maybe_prepare_scalar_for_op(other, self.shape)
|
self, other = align_method_FRAME(self, other, axis, flex=True, level=level)
|
|
if isinstance(other, ABCDataFrame):
|
# Another DataFrame
|
new_data = self._combine_frame(other, na_op, fill_value)
|
|
elif isinstance(other, ABCSeries):
|
new_data = self._dispatch_frame_op(other, op, axis=axis)
|
else:
|
# in this case we always have `np.ndim(other) == 0`
|
if fill_value is not None:
|
self = self.fillna(fill_value)
|
|
new_data = self._dispatch_frame_op(other, op)
|
|
return self._construct_result(new_data)
|
|
f.__name__ = op_name
|
|
return f
|
|
|
def flex_comp_method_FRAME(op):
|
op_name = op.__name__.strip("_")
|
|
doc = _flex_comp_doc_FRAME.format(
|
op_name=op_name, desc=_op_descriptions[op_name]["desc"]
|
)
|
|
@Appender(doc)
|
def f(self, other, axis: Axis = "columns", level=None):
|
axis = self._get_axis_number(axis) if axis is not None else 1
|
|
self, other = align_method_FRAME(self, other, axis, flex=True, level=level)
|
|
new_data = self._dispatch_frame_op(other, op, axis=axis)
|
return self._construct_result(new_data)
|
|
f.__name__ = op_name
|
|
return f
|
|
|
__all__ = [
|
"add_flex_arithmetic_methods",
|
"align_method_FRAME",
|
"align_method_SERIES",
|
"ARITHMETIC_BINOPS",
|
"arithmetic_op",
|
"COMPARISON_BINOPS",
|
"comparison_op",
|
"comp_method_OBJECT_ARRAY",
|
"fill_binop",
|
"flex_arith_method_FRAME",
|
"flex_comp_method_FRAME",
|
"flex_method_SERIES",
|
"frame_arith_method_with_reindex",
|
"invalid_comparison",
|
"kleene_and",
|
"kleene_or",
|
"kleene_xor",
|
"logical_op",
|
"maybe_dispatch_ufunc_to_dunder_op",
|
"radd",
|
"rand_",
|
"rdiv",
|
"rdivmod",
|
"rfloordiv",
|
"rmod",
|
"rmul",
|
"ror_",
|
"rpow",
|
"rsub",
|
"rtruediv",
|
"rxor",
|
"should_reindex_frame_op",
|
"unpack_zerodim_and_defer",
|
]
|
