From 3086c7b01c9b598e08bc42eeb8f42258c0ce34fd Mon Sep 17 00:00:00 2001 From: Joris Van den Bossche Date: Wed, 17 Mar 2021 10:45:05 +0100 Subject: [PATCH 1/3] REF: prepare (upcast) scalar before dispatching to arithmetic array ops --- pandas/core/arrays/numpy_.py | 1 + pandas/core/frame.py | 1 + pandas/core/ops/__init__.py | 2 ++ pandas/core/ops/array_ops.py | 6 +++--- pandas/core/series.py | 1 + 5 files changed, 8 insertions(+), 3 deletions(-) diff --git a/pandas/core/arrays/numpy_.py b/pandas/core/arrays/numpy_.py index 89988349132e6..7465e3b4d94e5 100644 --- a/pandas/core/arrays/numpy_.py +++ b/pandas/core/arrays/numpy_.py @@ -375,6 +375,7 @@ def _cmp_method(self, other, op): if isinstance(other, PandasArray): other = other._ndarray + other = ops.prepare_scalar_for_op(other, (len(self),)) pd_op = ops.get_array_op(op) result = pd_op(self._ndarray, other) diff --git a/pandas/core/frame.py b/pandas/core/frame.py index 9fdd979ce8eca..69f9d7996e88b 100644 --- a/pandas/core/frame.py +++ b/pandas/core/frame.py @@ -6617,6 +6617,7 @@ def _arith_method(self, other, op): return ops.frame_arith_method_with_reindex(self, other, op) axis = 1 # only relevant for Series other case + other = ops.prepare_scalar_for_op(other, (self.shape[axis],)) self, other = ops.align_method_FRAME(self, other, axis, flex=True, level=None) diff --git a/pandas/core/ops/__init__.py b/pandas/core/ops/__init__.py index 8ace64fedacb9..82aa879ec2a63 100644 --- a/pandas/core/ops/__init__.py +++ b/pandas/core/ops/__init__.py @@ -36,6 +36,7 @@ comparison_op, get_array_op, logical_op, + prepare_scalar_for_op, ) from pandas.core.ops.common import ( # noqa:F401 get_op_result_name, @@ -429,6 +430,7 @@ def f(self, other, axis=default_axis, level=None, fill_value=None): axis = self._get_axis_number(axis) if axis is not None else 1 + other = prepare_scalar_for_op(other, (self.shape[axis],)) self, other = align_method_FRAME(self, other, axis, flex=True, level=level) if isinstance(other, ABCDataFrame): diff --git a/pandas/core/ops/array_ops.py b/pandas/core/ops/array_ops.py index 6f6972c34f0a9..a3043b9e8b346 100644 --- a/pandas/core/ops/array_ops.py +++ b/pandas/core/ops/array_ops.py @@ -194,12 +194,12 @@ def arithmetic_op(left: ArrayLike, right: Any, op): """ # NB: We assume that extract_array has already been called - # on `left` and `right`. + # on `left` and `right`, + # and `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) 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): # Timedelta is included because numexpr will fail on it, see GH#31457 @@ -422,7 +422,7 @@ def get_array_op(op): raise NotImplementedError(op_name) -def _maybe_upcast_for_op(obj, shape: Shape): +def prepare_scalar_for_op(obj, shape: Shape): """ Cast non-pandas objects to pandas types to unify behavior of arithmetic and comparison operations. diff --git a/pandas/core/series.py b/pandas/core/series.py index 662c7abb33e33..3fa7a1cd0d0ba 100644 --- a/pandas/core/series.py +++ b/pandas/core/series.py @@ -5107,6 +5107,7 @@ def _arith_method(self, other, op): lvalues = self._values rvalues = extract_array(other, extract_numpy=True) + rvalues = ops.prepare_scalar_for_op(rvalues, lvalues.shape) result = ops.arithmetic_op(lvalues, rvalues, op) return self._construct_result(result, name=res_name) From a0dfcf1cc122e2e9e0c0c0b58f790530ebd3be8a Mon Sep 17 00:00:00 2001 From: Joris Van den Bossche Date: Fri, 23 Apr 2021 13:59:00 +0200 Subject: [PATCH 2/3] prepare_scalar_for_op -> maybe_prepare_scalar_for_op --- pandas/core/arrays/numpy_.py | 2 +- pandas/core/frame.py | 2 +- pandas/core/ops/__init__.py | 4 ++-- pandas/core/ops/array_ops.py | 4 ++-- pandas/core/series.py | 2 +- 5 files changed, 7 insertions(+), 7 deletions(-) diff --git a/pandas/core/arrays/numpy_.py b/pandas/core/arrays/numpy_.py index 385a8aa178c5e..52900d9b62dc2 100644 --- a/pandas/core/arrays/numpy_.py +++ b/pandas/core/arrays/numpy_.py @@ -395,7 +395,7 @@ def _cmp_method(self, other, op): if isinstance(other, PandasArray): other = other._ndarray - other = ops.prepare_scalar_for_op(other, (len(self),)) + other = ops.maybe_prepare_scalar_for_op(other, (len(self),)) pd_op = ops.get_array_op(op) other = ensure_wrapped_if_datetimelike(other) with np.errstate(all="ignore"): diff --git a/pandas/core/frame.py b/pandas/core/frame.py index 3090b50251372..8ac94111cca56 100644 --- a/pandas/core/frame.py +++ b/pandas/core/frame.py @@ -6804,7 +6804,7 @@ def _arith_method(self, other, op): return ops.frame_arith_method_with_reindex(self, other, op) axis = 1 # only relevant for Series other case - other = ops.prepare_scalar_for_op(other, (self.shape[axis],)) + other = ops.maybe_prepare_scalar_for_op(other, (self.shape[axis],)) self, other = ops.align_method_FRAME(self, other, axis, flex=True, level=None) diff --git a/pandas/core/ops/__init__.py b/pandas/core/ops/__init__.py index 5bcf493d5845a..1e8bb13f055e3 100644 --- a/pandas/core/ops/__init__.py +++ b/pandas/core/ops/__init__.py @@ -35,7 +35,7 @@ comparison_op, get_array_op, logical_op, - prepare_scalar_for_op, + maybe_prepare_scalar_for_op, ) from pandas.core.ops.common import ( # noqa:F401 get_op_result_name, @@ -429,7 +429,7 @@ def f(self, other, axis=default_axis, level=None, fill_value=None): axis = self._get_axis_number(axis) if axis is not None else 1 - other = prepare_scalar_for_op(other, (self.shape[axis],)) + other = maybe_prepare_scalar_for_op(other, (self.shape[axis],)) self, other = align_method_FRAME(self, other, axis, flex=True, level=level) if isinstance(other, ABCDataFrame): diff --git a/pandas/core/ops/array_ops.py b/pandas/core/ops/array_ops.py index 7b7de374ec1ec..01549d1e31c49 100644 --- a/pandas/core/ops/array_ops.py +++ b/pandas/core/ops/array_ops.py @@ -196,7 +196,7 @@ def arithmetic_op(left: ArrayLike, right: Any, op): """ # NB: We assume that extract_array and ensure_wrapped_if_datetimelike # have already been called on `left` and `right`, - # and `prepare_scalar_for_op` has already been called on `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) @@ -418,7 +418,7 @@ def get_array_op(op): raise NotImplementedError(op_name) -def prepare_scalar_for_op(obj, shape: Shape): +def maybe_prepare_scalar_for_op(obj, shape: Shape): """ Cast non-pandas objects to pandas types to unify behavior of arithmetic and comparison operations. diff --git a/pandas/core/series.py b/pandas/core/series.py index 5a7ee19b6c38c..23359532a4cbb 100644 --- a/pandas/core/series.py +++ b/pandas/core/series.py @@ -5314,7 +5314,7 @@ def _arith_method(self, other, op): lvalues = self._values rvalues = extract_array(other, extract_numpy=True, extract_range=True) - rvalues = ops.prepare_scalar_for_op(rvalues, lvalues.shape) + rvalues = ops.maybe_prepare_scalar_for_op(rvalues, lvalues.shape) rvalues = ensure_wrapped_if_datetimelike(rvalues) with np.errstate(all="ignore"): From 489df401ffeace500bd78d091004e02d1b266b93 Mon Sep 17 00:00:00 2001 From: Joris Van den Bossche Date: Wed, 28 Apr 2021 15:49:29 +0200 Subject: [PATCH 3/3] broadcast to full frame --- pandas/core/ops/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pandas/core/ops/__init__.py b/pandas/core/ops/__init__.py index 1e8bb13f055e3..9cccf1cff60a1 100644 --- a/pandas/core/ops/__init__.py +++ b/pandas/core/ops/__init__.py @@ -429,7 +429,7 @@ def f(self, other, axis=default_axis, level=None, fill_value=None): axis = self._get_axis_number(axis) if axis is not None else 1 - other = maybe_prepare_scalar_for_op(other, (self.shape[axis],)) + 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):