Shortcuts on this page
r m x toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1from __future__ import annotations
3from typing import (
4 TYPE_CHECKING,
5 Iterator,
6 NamedTuple,
7)
9from pandas._typing import ArrayLike
11if TYPE_CHECKING:
12 from pandas._libs.internals import BlockPlacement
14 from pandas.core.internals.blocks import Block
15 from pandas.core.internals.managers import BlockManager
18class BlockPairInfo(NamedTuple):
19 lvals: ArrayLike
20 rvals: ArrayLike
21 locs: BlockPlacement
22 left_ea: bool
23 right_ea: bool
24 rblk: Block
27def _iter_block_pairs(
28 left: BlockManager, right: BlockManager
29) -> Iterator[BlockPairInfo]:
30 # At this point we have already checked the parent DataFrames for
31 # assert rframe._indexed_same(lframe)
33 for blk in left.blocks:
34 locs = blk.mgr_locs
35 blk_vals = blk.values
37 left_ea = blk_vals.ndim == 1
39 rblks = right._slice_take_blocks_ax0(locs.indexer, only_slice=True)
41 # Assertions are disabled for performance, but should hold:
42 # if left_ea:
43 # assert len(locs) == 1, locs
44 # assert len(rblks) == 1, rblks
45 # assert rblks[0].shape[0] == 1, rblks[0].shape
47 for rblk in rblks:
48 right_ea = rblk.values.ndim == 1
50 lvals, rvals = _get_same_shape_values(blk, rblk, left_ea, right_ea)
51 info = BlockPairInfo(lvals, rvals, locs, left_ea, right_ea, rblk)
52 yield info
55def operate_blockwise(
56 left: BlockManager, right: BlockManager, array_op
57) -> BlockManager:
58 # At this point we have already checked the parent DataFrames for
59 # assert rframe._indexed_same(lframe)
61 res_blks: list[Block] = []
62 for lvals, rvals, locs, left_ea, right_ea, rblk in _iter_block_pairs(left, right):
63 res_values = array_op(lvals, rvals)
64 if left_ea and not right_ea and hasattr(res_values, "reshape"):
65 res_values = res_values.reshape(1, -1)
66 nbs = rblk._split_op_result(res_values)
68 # Assertions are disabled for performance, but should hold:
69 # if right_ea or left_ea:
70 # assert len(nbs) == 1
71 # else:
72 # assert res_values.shape == lvals.shape, (res_values.shape, lvals.shape)
74 _reset_block_mgr_locs(nbs, locs)
76 res_blks.extend(nbs)
78 # Assertions are disabled for performance, but should hold:
79 # slocs = {y for nb in res_blks for y in nb.mgr_locs.as_array}
80 # nlocs = sum(len(nb.mgr_locs.as_array) for nb in res_blks)
81 # assert nlocs == len(left.items), (nlocs, len(left.items))
82 # assert len(slocs) == nlocs, (len(slocs), nlocs)
83 # assert slocs == set(range(nlocs)), slocs
85 new_mgr = type(right)(tuple(res_blks), axes=right.axes, verify_integrity=False)
86 return new_mgr
89def _reset_block_mgr_locs(nbs: list[Block], locs) -> None:
90 """
91 Reset mgr_locs to correspond to our original DataFrame.
92 """
93 for nb in nbs:
94 nblocs = locs[nb.mgr_locs.indexer]
95 nb.mgr_locs = nblocs
96 # Assertions are disabled for performance, but should hold:
97 # assert len(nblocs) == nb.shape[0], (len(nblocs), nb.shape)
98 # assert all(x in locs.as_array for x in nb.mgr_locs.as_array)
101def _get_same_shape_values(
102 lblk: Block, rblk: Block, left_ea: bool, right_ea: bool
103) -> tuple[ArrayLike, ArrayLike]:
104 """
105 Slice lblk.values to align with rblk. Squeeze if we have EAs.
106 """
107 lvals = lblk.values
108 rvals = rblk.values
110 # Require that the indexing into lvals be slice-like
111 assert rblk.mgr_locs.is_slice_like, rblk.mgr_locs
113 # TODO(EA2D): with 2D EAs only this first clause would be needed
114 if not (left_ea or right_ea):
115 # error: No overload variant of "__getitem__" of "ExtensionArray" matches
116 # argument type "Tuple[Union[ndarray, slice], slice]"
117 lvals = lvals[rblk.mgr_locs.indexer, :] # type: ignore[call-overload]
118 assert lvals.shape == rvals.shape, (lvals.shape, rvals.shape)
119 elif left_ea and right_ea:
120 assert lvals.shape == rvals.shape, (lvals.shape, rvals.shape)
121 elif right_ea:
122 # lvals are 2D, rvals are 1D
124 # error: No overload variant of "__getitem__" of "ExtensionArray" matches
125 # argument type "Tuple[Union[ndarray, slice], slice]"
126 lvals = lvals[rblk.mgr_locs.indexer, :] # type: ignore[call-overload]
127 assert lvals.shape[0] == 1, lvals.shape
128 lvals = lvals[0, :]
129 else:
130 # lvals are 1D, rvals are 2D
131 assert rvals.shape[0] == 1, rvals.shape
132 # error: No overload variant of "__getitem__" of "ExtensionArray" matches
133 # argument type "Tuple[int, slice]"
134 rvals = rvals[0, :] # type: ignore[call-overload]
136 return lvals, rvals
139def blockwise_all(left: BlockManager, right: BlockManager, op) -> bool:
140 """
141 Blockwise `all` reduction.
142 """
143 for info in _iter_block_pairs(left, right):
144 res = op(info.lvals, info.rvals)
145 if not res:
146 return False
147 return True