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

asv_bench/benchmarks/arithmetic.py

@@ -101,6 +101,56 @@ def time_frame_op_with_series_axis1(self, opname):
         getattr(operator, opname)(self.df, self.ser)
 
 
+class FrameWithFrameWide:
+    # Many-columns, mixed dtypes
+
+    params = [
+        [
+            operator.add,
+            operator.sub,
+            operator.mul,
+            operator.truediv,
+            operator.floordiv,
+            operator.pow,
+            operator.mod,
+            operator.eq,
+            operator.ne,
+            operator.gt,
+            operator.ge,
+            operator.lt,
+            operator.le,
+        ]
+    ]
+    param_names = ["op"]
+
+    def setup(self, op):
+        # we choose dtypes so as to make the blocks
+        #  a) not perfectly match between right and left
+        #  b) appreciably bigger than single columns
+        arr = np.random.randn(10 ** 6).reshape(500, 2000).astype(np.float64)
+        df = pd.DataFrame(arr)
+        df[1000] = df[1000].astype(np.float32)
+        df.iloc[:, 1000:] = df.iloc[:, 1000:].astype(np.float32)
+        df._consolidate_inplace()
+
+        # TODO: GH#33198 the setting here shoudlnt need two steps
+        df2 = pd.DataFrame(arr)
+        df2[1000] = df2[1000].astype(np.int64)
+        df2.iloc[:, 500:1500] = df2.iloc[:, 500:1500].astype(np.int64)
+        df2._consolidate_inplace()
+
+        self.left = df
+        self.right = df
+
+    def time_op_different_blocks(self, op):
+        # blocks (and dtypes) are not aligned
+        op(self.left, self.right)
+
+    def time_op_same_blocks(self, op):
+        # blocks (and dtypes) are aligned
+        op(self.left, self.left)
+
+
 class Ops:
 
     params = [[True, False], ["default", 1]]

doc/source/whatsnew/v1.1.0.rst

@@ -459,7 +459,7 @@ Performance improvements
   :meth:`DataFrame.sparse.from_spmatrix` constructor (:issue:`32821`,
   :issue:`32825`,  :issue:`32826`, :issue:`32856`, :issue:`32858`).
 - Performance improvement in reductions (sum, prod, min, max) for nullable (integer and boolean) dtypes (:issue:`30982`, :issue:`33261`, :issue:`33442`).
-
+- Performance improvement in arithmetic operations between two :class:`DataFrame` objects (:issue:`32779`)
 
 .. ---------------------------------------------------------------------------
 

pandas/core/arrays/datetimelike.py

@@ -1360,18 +1360,20 @@ def _addsub_object_array(self, other: np.ndarray, op):
         """
         assert op in [operator.add, operator.sub]
         if len(other) == 1:
+            # If both 1D then broadcasting is unambiguous
+            # TODO(EA2D): require self.ndim == other.ndim here
             return op(self, other[0])
 
         warnings.warn(
-            "Adding/subtracting array of DateOffsets to "
+            "Adding/subtracting object-dtype array to "
             f"{type(self).__name__} not vectorized",
             PerformanceWarning,
         )
 
         # Caller is responsible for broadcasting if necessary
         assert self.shape == other.shape, (self.shape, other.shape)
 
-        res_values = op(self.astype("O"), np.array(other))
+        res_values = op(self.astype("O"), np.asarray(other))
         result = array(res_values.ravel())
         result = extract_array(result, extract_numpy=True).reshape(self.shape)
         return result

pandas/core/frame.py

@@ -457,6 +457,7 @@ def __init__(
             mgr = self._init_mgr(
                 data, axes=dict(index=index, columns=columns), dtype=dtype, copy=copy
             )
+
         elif isinstance(data, dict):
             mgr = init_dict(data, index, columns, dtype=dtype)
         elif isinstance(data, ma.MaskedArray):
@@ -5654,10 +5655,11 @@ def _construct_result(self, result) -> "DataFrame":
         -------
         DataFrame
         """
-        out = self._constructor(result, index=self.index, copy=False)
+        out = self._constructor(result, copy=False)
         # Pin columns instead of passing to constructor for compat with
         #  non-unique columns case
         out.columns = self.columns
+        out.index = self.index
         return out
 
     def combine(

pandas/core/internals/managers.py

@@ -1233,12 +1233,21 @@ def reindex_indexer(
 
         return type(self).from_blocks(new_blocks, new_axes)
 
-    def _slice_take_blocks_ax0(self, slice_or_indexer, fill_value=lib.no_default):
+    def _slice_take_blocks_ax0(
+        self, slice_or_indexer, fill_value=lib.no_default, only_slice: bool = False
+    ):
         """
         Slice/take blocks along axis=0.
 
         Overloaded for SingleBlock
 
+        Parameters
+        ----------
+        slice_or_indexer : slice, ndarray[bool], or list-like of ints
+        fill_value : scalar, default lib.no_default
+        only_slice : bool, default False
+            If True, we always return views on existing arrays, never copies.
+
         Returns
         -------
         new_blocks : list of Block
@@ -1309,11 +1318,25 @@ def _slice_take_blocks_ax0(self, slice_or_indexer, fill_value=lib.no_default):
                         blocks.append(newblk)
 
                 else:
-                    blocks.append(
-                        blk.take_nd(
-                            blklocs[mgr_locs.indexer], axis=0, new_mgr_locs=mgr_locs,
-                        )
-                    )
+                    # GH#32779 to avoid the performance penalty of copying,
+                    #  we may try to only slice
+                    taker = blklocs[mgr_locs.indexer]
+                    max_len = max(len(mgr_locs), taker.max() + 1)
+                    if only_slice:
+                        taker = lib.maybe_indices_to_slice(taker, max_len)
+
+                    if isinstance(taker, slice):
+                        nb = blk.getitem_block(taker, new_mgr_locs=mgr_locs)
+                        blocks.append(nb)
+                    elif only_slice:
+                        # GH#33597 slice instead of take, so we get
+                        #  views instead of copies
+                        for i, ml in zip(taker, mgr_locs):
+                            nb = blk.getitem_block([i], new_mgr_locs=ml)
+                            blocks.append(nb)
+                    else:
+                        nb = blk.take_nd(taker, axis=0, new_mgr_locs=mgr_locs)
+                        blocks.append(nb)
 
         return blocks
 

pandas/core/ops/__init__.py

@@ -26,6 +26,7 @@
     logical_op,
 )
 from pandas.core.ops.array_ops import comp_method_OBJECT_ARRAY  # noqa:F401
+from pandas.core.ops.blockwise import operate_blockwise
 from pandas.core.ops.common import unpack_zerodim_and_defer
 from pandas.core.ops.dispatch import should_series_dispatch
 from pandas.core.ops.docstrings import (
@@ -325,8 +326,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

@@ -6,6 +6,7 @@
 from functools import partial
 import operator
 from typing import Any, Optional, Tuple
+import warnings
 
 import numpy as np
 
@@ -120,7 +121,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)
 
@@ -254,8 +255,13 @@ def comparison_op(
         res_values = comp_method_OBJECT_ARRAY(op, lvalues, rvalues)
 
     else:
-        with np.errstate(all="ignore"):
-            res_values = na_arithmetic_op(lvalues, rvalues, op, str_rep, is_cmp=True)
+        with warnings.catch_warnings():
+            # suppress warnings from numpy about element-wise comparison
+            warnings.simplefilter("ignore", DeprecationWarning)
+            with np.errstate(all="ignore"):
+                res_values = na_arithmetic_op(
+                    lvalues, rvalues, op, str_rep, is_cmp=True
+                )
 
     return res_values
 

pandas/core/ops/blockwise.py

@@ -0,0 +1,91 @@
+from typing import TYPE_CHECKING, List, Tuple
+
+import numpy as np
+
+from pandas._typing import ArrayLike
+
+if TYPE_CHECKING:
+    from pandas.core.internals.blocks import Block  # noqa:F401
+
+
+def operate_blockwise(left, right, array_op):
+    assert right._indexed_same(left)
+
+    def get_same_shape_values(
+        lblk: "Block", rblk: "Block", left_ea: bool, right_ea: bool
+    ) -> Tuple[ArrayLike, ArrayLike]:
+        """
+        Slice lblk.values to align with rblk.  Squeeze if we have EAs.
+        """
+        lvals = lblk.values
+        rvals = rblk.values
+
+        # TODO(EA2D): with 2D EAs pnly this first clause would be needed
+        if not (left_ea or right_ea):
+            lvals = lvals[rblk.mgr_locs.indexer, :]
+            assert lvals.shape == rvals.shape, (lvals.shape, rvals.shape)
+        elif left_ea and right_ea:
+            assert lvals.shape == rvals.shape, (lvals.shape, rvals.shape)
+        elif right_ea:
+            # lvals are 2D, rvals are 1D
+            lvals = lvals[rblk.mgr_locs.indexer, :]
+            assert lvals.shape[0] == 1, lvals.shape
+            lvals = lvals[0, :]
+        else:
+            # lvals are 1D, rvals are 2D
+            assert rvals.shape[0] == 1, rvals.shape
+            rvals = rvals[0, :]
+
+        return lvals, rvals
+
+    res_blks: List["Block"] = []
+    rmgr = right._mgr
+    for n, blk in enumerate(left._mgr.blocks):
+        locs = blk.mgr_locs
+        blk_vals = blk.values
+
+        left_ea = not isinstance(blk_vals, np.ndarray)
+
+        rblks = rmgr._slice_take_blocks_ax0(locs.indexer, only_slice=True)
+
+        # Assertions are disabled for performance, but should hold:
+        # if left_ea:
+        #    assert len(locs) == 1, locs
+        #    assert len(rblks) == 1, rblks
+        #    assert rblks[0].shape[0] == 1, rblks[0].shape
+
+        for k, rblk in enumerate(rblks):
+            right_ea = not isinstance(rblk.values, np.ndarray)
+
+            lvals, rvals = get_same_shape_values(blk, rblk, left_ea, right_ea)
+
+            res_values = array_op(lvals, rvals)
+            if left_ea and not right_ea and hasattr(res_values, "reshape"):
+                res_values = res_values.reshape(1, -1)
+            nbs = rblk._split_op_result(res_values)
+
+            # Assertions are disabled for performance, but should hold:
+            # if right_ea or left_ea:
+            #    assert len(nbs) == 1
+            # else:
+            #    assert res_values.shape == lvals.shape, (res_values.shape, lvals.shape)
+
+            for nb in nbs:
+                # Reset mgr_locs to correspond to our original DataFrame
+                nblocs = locs.as_array[nb.mgr_locs.indexer]
+                nb.mgr_locs = nblocs
+                # Assertions are disabled for performance, but should hold:
+                #  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)
+
+    # Assertions are disabled for performance, but should hold:
+    #  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
+
+    new_mgr = type(rmgr)(res_blks, axes=rmgr.axes, do_integrity_check=False)
+    return new_mgr

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)
@@ -1444,8 +1446,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

pandas/tests/frame/test_arithmetic.py

@@ -49,9 +49,11 @@ def check(df, df2):
                 )
                 tm.assert_frame_equal(result, expected)
 
-                msg = re.escape(
-                    "Invalid comparison between dtype=datetime64[ns] and ndarray"
-                )
+                msgs = [
+                    r"Invalid comparison between dtype=datetime64\[ns\] and ndarray",
+                    "invalid type promotion",
+                ]
+                msg = "|".join(msgs)
                 with pytest.raises(TypeError, match=msg):
                     x >= y
                 with pytest.raises(TypeError, match=msg):