PERF: BlockManager.equals blockwise (#35357)

Author: jbrockmendel

pandas/core/dtypes/missing.py

@@ -10,7 +10,7 @@
 from pandas._libs import lib
 import pandas._libs.missing as libmissing
 from pandas._libs.tslibs import NaT, iNaT
-from pandas._typing import DtypeObj
+from pandas._typing import ArrayLike, DtypeObj
 
 from pandas.core.dtypes.common import (
     DT64NS_DTYPE,
@@ -484,6 +484,18 @@ def _array_equivalent_object(left, right, strict_nan):
     return True
 
 
+def array_equals(left: ArrayLike, right: ArrayLike) -> bool:
+    """
+    ExtensionArray-compatible implementation of array_equivalent.
+    """
+    if not is_dtype_equal(left.dtype, right.dtype):
+        return False
+    elif isinstance(left, ABCExtensionArray):
+        return left.equals(right)
+    else:
+        return array_equivalent(left, right, dtype_equal=True)
+
+
 def _infer_fill_value(val):
     """
     infer the fill value for the nan/NaT from the provided

pandas/core/internals/managers.py

@@ -19,7 +19,6 @@
 from pandas.core.dtypes.common import (
     DT64NS_DTYPE,
     is_datetimelike_v_numeric,
-    is_dtype_equal,
     is_extension_array_dtype,
     is_list_like,
     is_numeric_v_string_like,
@@ -28,10 +27,9 @@
 from pandas.core.dtypes.concat import concat_compat
 from pandas.core.dtypes.dtypes import ExtensionDtype
 from pandas.core.dtypes.generic import ABCDataFrame, ABCSeries
-from pandas.core.dtypes.missing import array_equivalent, isna
+from pandas.core.dtypes.missing import array_equals, isna
 
 import pandas.core.algorithms as algos
-from pandas.core.arrays import ExtensionArray
 from pandas.core.arrays.sparse import SparseDtype
 from pandas.core.base import PandasObject
 import pandas.core.common as com
@@ -49,7 +47,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
 
@@ -1449,26 +1447,9 @@ def equals(self, other: "BlockManager") -> bool:
                 return False
             left = self.blocks[0].values
             right = other.blocks[0].values
-            if not is_dtype_equal(left.dtype, right.dtype):
-                return False
-            elif isinstance(left, ExtensionArray):
-                return left.equals(right)
-            else:
-                return array_equivalent(left, right)
+            return array_equals(left, right)
 
-        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):
-                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
+        return blockwise_all(self, other, array_equals)
 
     def unstack(self, unstacker, fill_value) -> "BlockManager":
         """

pandas/core/internals/ops.py

@@ -1,4 +1,5 @@
-from typing import TYPE_CHECKING, List, Tuple
+from collections import namedtuple
+from typing import TYPE_CHECKING, Iterator, List, Tuple
 
 import numpy as np
 
@@ -9,13 +10,17 @@
     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"
+) -> Iterator[BlockPairInfo]:
     # 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 +39,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 +101,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 +118,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