PERF: BlockManager.equals blockwise

pandas/core/internals/managers.py

@@ -49,7 +49,7 @@
     get_block_type,
     make_block,
 )
-from pandas.core.internals.ops import operate_blockwise
+from pandas.core.internals.ops import blockwise_all, operate_blockwise
 
 # TODO: flexible with index=None and/or items=None
 
@@ -1422,32 +1422,25 @@ def equals(self, other: "BlockManager") -> bool:
         if not all(ax1.equals(ax2) for ax1, ax2 in zip(self_axes, other_axes)):
             return False
 
-        if self.ndim == 1:
-            # For SingleBlockManager (i.e.Series)
-            if other.ndim != 1:
-                return False
-            left = self.blocks[0].values
-            right = other.blocks[0].values
+        def blk_func(left, right):
             if not is_dtype_equal(left.dtype, right.dtype):
                 return False
             elif isinstance(left, ExtensionArray):
                 return left.equals(right)
+            elif isinstance(right, ExtensionArray):
+                return False
             else:
-                return array_equivalent(left, right)
+                return array_equivalent(left, right, dtype_equal=True)
 
-        for i in range(len(self.items)):
-            # Check column-wise, return False if any column doesn't match
-            left = self.iget_values(i)
-            right = other.iget_values(i)
-            if not is_dtype_equal(left.dtype, right.dtype):
+        if self.ndim == 1:
+            # For SingleBlockManager (i.e.Series)
+            if other.ndim != 1:
                 return False
-            elif isinstance(left, ExtensionArray):
-                if not left.equals(right):
-                    return False
-            else:
-                if not array_equivalent(left, right, dtype_equal=True):
-                    return False
-        return True
+            left = self.blocks[0].values
+            right = other.blocks[0].values
+            return blk_func(left, right)
+
+        return blockwise_all(self, other, blk_func)
 
     def unstack(self, unstacker, fill_value) -> "BlockManager":
         """

pandas/core/internals/ops.py

@@ -1,3 +1,4 @@
+from collections import namedtuple
 from typing import TYPE_CHECKING, List, Tuple
 
 import numpy as np
@@ -9,13 +10,15 @@
     from pandas.core.internals.managers import BlockManager  # noqa:F401
 
 
-def operate_blockwise(
-    left: "BlockManager", right: "BlockManager", array_op
-) -> "BlockManager":
+BlockPairInfo = namedtuple(
+    "BlockPairInfo", ["lvals", "rvals", "locs", "left_ea", "right_ea", "rblk"],
+)
+
+
+def _iter_block_pairs(left: "BlockManager", right: "BlockManager"):
     # At this point we have already checked the parent DataFrames for
     #  assert rframe._indexed_same(lframe)
 
-    res_blks: List["Block"] = []
     for n, blk in enumerate(left.blocks):
         locs = blk.mgr_locs
         blk_vals = blk.values
@@ -34,21 +37,32 @@ def operate_blockwise(
             right_ea = not isinstance(rblk.values, np.ndarray)
 
             lvals, rvals = _get_same_shape_values(blk, rblk, left_ea, right_ea)
+            info = BlockPairInfo(lvals, rvals, locs, left_ea, right_ea, rblk)
+            yield info
 
-            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)
+def operate_blockwise(
+    left: "BlockManager", right: "BlockManager", array_op
+) -> "BlockManager":
+    # At this point we have already checked the parent DataFrames for
+    #  assert rframe._indexed_same(lframe)
+
+    res_blks: List["Block"] = []
+    for lvals, rvals, locs, left_ea, right_ea, rblk in _iter_block_pairs(left, right):
+        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)
 
-            _reset_block_mgr_locs(nbs, locs)
+        _reset_block_mgr_locs(nbs, locs)
 
-            res_blks.extend(nbs)
+        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}
@@ -85,7 +99,7 @@ def _get_same_shape_values(
     # Require that the indexing into lvals be slice-like
     assert rblk.mgr_locs.is_slice_like, rblk.mgr_locs
 
-    # TODO(EA2D): with 2D EAs pnly this first clause would be needed
+    # TODO(EA2D): with 2D EAs only 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)
@@ -102,3 +116,14 @@ def _get_same_shape_values(
         rvals = rvals[0, :]
 
     return lvals, rvals
+
+
+def blockwise_all(left: "BlockManager", right: "BlockManager", op) -> bool:
+    """
+    Blockwise `all` reduction.
+    """
+    for info in _iter_block_pairs(left, right):
+        res = op(info.lvals, info.rvals)
+        if not res:
+            return False
+    return True