REF: make operate_blockwise into a BlockManager method (#34281)

Author: jbrockmendel

pandas/core/internals/managers.py

@@ -47,6 +47,7 @@
     get_block_type,
     make_block,
 )
+from pandas.core.internals.ops import operate_blockwise
 
 # TODO: flexible with index=None and/or items=None
 
@@ -352,6 +353,12 @@ def reduce(self, func, *args, **kwargs):
 
         return res
 
+    def operate_blockwise(self, other: "BlockManager", array_op) -> "BlockManager":
+        """
+        Apply array_op blockwise with another (aligned) BlockManager.
+        """
+        return operate_blockwise(self, other, array_op)
+
     def apply(self: T, f, align_keys=None, **kwargs) -> T:
         """
         Iterate over the blocks, collect and create a new BlockManager.

pandas/core/ops/blockwise.py renamed to pandas/core/internals/ops.py

@@ -5,22 +5,24 @@
 from pandas._typing import ArrayLike
 
 if TYPE_CHECKING:
+    from pandas.core.internals.managers import BlockManager  # noqa:F401
     from pandas.core.internals.blocks import Block  # noqa:F401
 
 
-def operate_blockwise(left, right, array_op):
-    # At this point we have already checked
-    #  assert right._indexed_same(left)
+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"] = []
-    rmgr = right._mgr
-    for n, blk in enumerate(left._mgr.blocks):
+    for n, blk in enumerate(left.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)
+        rblks = right._slice_take_blocks_ax0(locs.indexer, only_slice=True)
 
         # Assertions are disabled for performance, but should hold:
         # if left_ea:
@@ -51,11 +53,11 @@ def operate_blockwise(left, right, array_op):
     # 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 nlocs == len(left.items), (nlocs, len(left.items))
     #  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)
+    new_mgr = type(right)(res_blks, axes=right.axes, do_integrity_check=False)
     return new_mgr
 
 

pandas/core/ops/__init__.py

@@ -26,7 +26,6 @@
     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 (
@@ -327,7 +326,7 @@ def dispatch_to_series(left, right, func, str_rep=None, axis=None):
         assert right._indexed_same(left)
 
         array_op = get_array_op(func, str_rep=str_rep)
-        bm = operate_blockwise(left, right, array_op)
+        bm = left._mgr.operate_blockwise(right._mgr, array_op)
         return type(left)(bm)
 
     elif isinstance(right, ABCSeries) and axis == "columns":