PERF: block-wise arithmetic for frame-with-frame

pandas/core/arrays/datetimelike.py

@@ -1299,7 +1299,7 @@ def _addsub_object_array(self, other: np.ndarray, op):
         result : same class as self
         """
         assert op in [operator.add, operator.sub]
-        if len(other) == 1:
+        if len(other) == 1 and self.ndim == other.ndim == 1:
             return op(self, other[0])
 
         warnings.warn(

pandas/core/ops/__init__.py

@@ -5,7 +5,7 @@
 """
 import datetime
 import operator
-from typing import TYPE_CHECKING, Optional, Set, Tuple
+from typing import TYPE_CHECKING, List, Optional, Set, Tuple
 
 import numpy as np
 
@@ -58,6 +58,7 @@
 
 if TYPE_CHECKING:
     from pandas import DataFrame  # noqa:F401
+    from pandas.core.internals.blocks import Block  # noqa: F401
 
 # -----------------------------------------------------------------------------
 # constants
@@ -353,6 +354,70 @@ def fill_binop(left, right, fill_value):
 # Dispatch logic
 
 
+def operate_blockwise(left, right, array_op):
+    assert right._indexed_same(left)
+
+    res_blks: List["Block"] = []
+    rmgr = right._data
+    for n, blk in enumerate(left._data.blocks):
+        locs = blk.mgr_locs
+
+        blk_vals = blk.values
+
+        if not isinstance(blk_vals, np.ndarray):
+            # 1D EA
+            assert len(locs) == 1, locs
+            rser = right.iloc[:, locs[0]]
+            rvals = extract_array(rser, extract_numpy=True)
+            res_values = array_op(blk_vals, rvals)
+            nbs = blk._split_op_result(res_values)
+            res_blks.extend(nbs)
+            continue
+
+        rblks = rmgr._slice_take_blocks_ax0(locs.indexer)
+
+        for k, rblk in enumerate(rblks):
+            lvals = blk_vals[rblk.mgr_locs.indexer, :]
+            rvals = rblk.values
+
+            if not isinstance(rvals, np.ndarray):
+                # 1D EA
+                assert lvals.shape[0] == 1, lvals.shape
+                lvals = lvals[0, :]
+                res_values = array_op(lvals, rvals)
+                nbs = rblk._split_op_result(res_values)
+                assert len(nbs) == 1
+                nb = nbs[0]
+                nb.mgr_locs = locs.as_array[nb.mgr_locs]
+                res_blks.append(nb)
+                continue
+
+            assert lvals.shape == rvals.shape, (lvals.shape, rvals.shape)
+
+            res_values = array_op(lvals, rvals)
+            assert res_values.shape == lvals.shape, (res_values.shape, lvals.shape)
+            nbs = rblk._split_op_result(res_values)
+            for nb in nbs:
+                # TODO: maybe optimize by sticking with slices?
+                nb_mgr_locs = nb.mgr_locs
+                nblocs = locs.as_array[nb_mgr_locs.indexer]
+                nb.mgr_locs = nblocs
+                assert len(nblocs) == nb.shape[0], (len(nblocs), nb.shape)
+                assert all(x in locs.as_array for x in nb.mgr_locs.as_array)
+
+            res_blks.extend(nbs)
+
+    slocs = {y for nb in res_blks for y in nb.mgr_locs.as_array}
+    nlocs = sum(len(nb.mgr_locs.as_array) for nb in res_blks)
+    assert nlocs == len(left.columns), (nlocs, len(left.columns))
+    assert len(slocs) == nlocs, (len(slocs), nlocs)
+    assert slocs == set(range(nlocs)), slocs
+
+    # TODO: once this is working, pass do_integrity_check=False
+    new_mgr = type(rmgr)(res_blks, axes=rmgr.axes)
+    return new_mgr
+
+
 def dispatch_to_series(left, right, func, str_rep=None, axis=None):
     """
     Evaluate the frame operation func(left, right) by evaluating
@@ -385,8 +450,9 @@ def dispatch_to_series(left, right, func, str_rep=None, axis=None):
     elif isinstance(right, ABCDataFrame):
         assert right._indexed_same(left)
 
-        def column_op(a, b):
-            return {i: func(a.iloc[:, i], b.iloc[:, i]) for i in range(len(a.columns))}
+        array_op = get_array_op(func, str_rep=str_rep)
+        bm = operate_blockwise(left, right, array_op)
+        return type(left)(bm)
 
     elif isinstance(right, ABCSeries) and axis == "columns":
         # We only get here if called via _combine_series_frame,

pandas/core/ops/array_ops.py

@@ -5,6 +5,7 @@
 from functools import partial
 import operator
 from typing import Any, Optional
+import warnings
 
 import numpy as np
 
@@ -132,7 +133,7 @@ def masked_arith_op(x: np.ndarray, y, op):
     return result
 
 
-def define_na_arithmetic_op(op, str_rep: str):
+def define_na_arithmetic_op(op, str_rep: Optional[str]):
     def na_op(x, y):
         return na_arithmetic_op(x, y, op, str_rep)
 
@@ -163,15 +164,18 @@ def na_arithmetic_op(left, right, op, str_rep: Optional[str], is_cmp: bool = Fal
     """
     import pandas.core.computation.expressions as expressions
 
-    try:
-        result = expressions.evaluate(op, str_rep, left, right)
-    except TypeError:
-        if is_cmp:
-            # numexpr failed on comparison op, e.g. ndarray[float] > datetime
-            #  In this case we do not fall back to the masked op, as that
-            #  will handle complex numbers incorrectly, see GH#32047
-            raise
-        result = masked_arith_op(left, right, op)
+    with warnings.catch_warnings():
+        # suppress warnings from numpy about element-wise comparison
+        warnings.simplefilter("ignore", DeprecationWarning)
+        try:
+            result = expressions.evaluate(op, str_rep, left, right)
+        except TypeError:
+            if is_cmp:
+                # numexpr failed on comparison op, e.g. ndarray[float] > datetime
+                #  In this case we do not fall back to the masked op, as that
+                #  will handle complex numbers incorrectly, see GH#32047
+                raise
+            result = masked_arith_op(left, right, op)
 
     if is_cmp and (is_scalar(result) or result is NotImplemented):
         # numpy returned a scalar instead of operating element-wise

pandas/tests/arithmetic/common.py

@@ -70,7 +70,14 @@ def assert_invalid_comparison(left, right, box):
     result = right != left
     tm.assert_equal(result, ~expected)
 
-    msg = "Invalid comparison between|Cannot compare type|not supported between"
+    msg = "|".join(
+        [
+            "Invalid comparison between",
+            "Cannot compare type",
+            "not supported between",
+            "invalid type promotion",
+        ]
+    )
     with pytest.raises(TypeError, match=msg):
         left < right
     with pytest.raises(TypeError, match=msg):

pandas/tests/arithmetic/test_datetime64.py

@@ -964,7 +964,9 @@ def test_dt64arr_sub_dt64object_array(self, box_with_array, tz_naive_fixture):
         obj = tm.box_expected(dti, box_with_array)
         expected = tm.box_expected(expected, box_with_array)
 
-        warn = PerformanceWarning if box_with_array is not pd.DataFrame else None
+        warn = None
+        if box_with_array is not pd.DataFrame or tz_naive_fixture is None:
+            warn = PerformanceWarning
         with tm.assert_produces_warning(warn):
             result = obj - obj.astype(object)
         tm.assert_equal(result, expected)
@@ -1388,8 +1390,7 @@ def test_dt64arr_add_mixed_offset_array(self, box_with_array):
         s = DatetimeIndex([Timestamp("2000-1-1"), Timestamp("2000-2-1")])
         s = tm.box_expected(s, box_with_array)
 
-        warn = None if box_with_array is pd.DataFrame else PerformanceWarning
-        with tm.assert_produces_warning(warn):
+        with tm.assert_produces_warning(PerformanceWarning):
             other = pd.Index([pd.offsets.DateOffset(years=1), pd.offsets.MonthEnd()])
             other = tm.box_expected(other, box_with_array)
             result = s + other