diff --git a/pandas/core/internals/managers.py b/pandas/core/internals/managers.py index aa21c0f2bca5e..6a7df1c8d88a1 100644 --- a/pandas/core/internals/managers.py +++ b/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. diff --git a/pandas/core/ops/blockwise.py b/pandas/core/internals/ops.py similarity index 84% rename from pandas/core/ops/blockwise.py rename to pandas/core/internals/ops.py index f41a30b136637..fd9a9a5ef6c93 100644 --- a/pandas/core/ops/blockwise.py +++ b/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 diff --git a/pandas/core/ops/__init__.py b/pandas/core/ops/__init__.py index 585e6d0eb0811..3cbfa280a4e30 100644 --- a/pandas/core/ops/__init__.py +++ b/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":