Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/pandas/core/internals/ops.py: 25%

1from __future__ import annotations 

2 

3from typing import ( 

4 TYPE_CHECKING, 

5 NamedTuple, 

6) 

7 

8from pandas.core.dtypes.common import is_1d_only_ea_dtype 

9 

10if TYPE_CHECKING: 

11 from collections.abc import ( 

12 Callable, 

13 Iterator, 

14 ) 

15 

16 from pandas._libs.internals import BlockPlacement 

17 from pandas._typing import ArrayLike 

18 

19 from pandas.core.internals.blocks import Block 

20 from pandas.core.internals.managers import BlockManager 

21 

22 

23class BlockPairInfo(NamedTuple): 

24 lvals: ArrayLike 

25 rvals: ArrayLike 

26 locs: BlockPlacement 

27 left_ea: bool 

28 right_ea: bool 

29 rblk: Block 

30 

31 

32def _iter_block_pairs( 

33 left: BlockManager, right: BlockManager 

34) -> Iterator[BlockPairInfo]: 

35 # At this point we have already checked the parent DataFrames for 

36 # assert rframe._indexed_same(lframe) 

37 

38 for blk in left.blocks: 

39 locs = blk.mgr_locs 

40 blk_vals = blk.values 

41 

42 left_ea = blk_vals.ndim == 1 

43 

44 rblks = right._slice_take_blocks_ax0(locs.indexer, only_slice=True) 

45 

46 # Assertions are disabled for performance, but should hold: 

47 # if left_ea: 

48 # assert len(locs) == 1, locs 

49 # assert len(rblks) == 1, rblks 

50 # assert rblks[0].shape[0] == 1, rblks[0].shape 

51 

52 for rblk in rblks: 

53 right_ea = rblk.values.ndim == 1 

54 

55 lvals, rvals = _get_same_shape_values(blk, rblk, left_ea, right_ea) 

56 info = BlockPairInfo(lvals, rvals, locs, left_ea, right_ea, rblk) 

57 yield info 

58 

59 

60def operate_blockwise( 

61 left: BlockManager, 

62 right: BlockManager, 

63 array_op: Callable[[ArrayLike, ArrayLike], ArrayLike], 

64) -> BlockManager: 

65 # At this point we have already checked the parent DataFrames for 

66 # assert rframe._indexed_same(lframe) 

67 

68 res_blks: list[Block] = [] 

69 for lvals, rvals, locs, left_ea, right_ea, rblk in _iter_block_pairs(left, right): 

70 res_values = array_op(lvals, rvals) 

71 if ( 

72 left_ea 

73 and not right_ea 

74 and hasattr(res_values, "reshape") 

75 and not is_1d_only_ea_dtype(res_values.dtype) 

76 ): 

77 res_values = res_values.reshape(1, -1) # pyright: ignore[reportAttributeAccessIssue] 

78 nbs = rblk._split_op_result(res_values) 

79 

80 # Assertions are disabled for performance, but should hold: 

81 # if right_ea or left_ea: 

82 # assert len(nbs) == 1 

83 # else: 

84 # assert res_values.shape == lvals.shape, (res_values.shape, lvals.shape) 

85 

86 _reset_block_mgr_locs(nbs, locs) 

87 

88 res_blks.extend(nbs) 

89 

90 # Assertions are disabled for performance, but should hold: 

91 # slocs = {y for nb in res_blks for y in nb.mgr_locs.as_array} 

92 # nlocs = sum(len(nb.mgr_locs.as_array) for nb in res_blks) 

93 # assert nlocs == len(left.items), (nlocs, len(left.items)) 

94 # assert len(slocs) == nlocs, (len(slocs), nlocs) 

95 # assert slocs == set(range(nlocs)), slocs 

96 

97 # TODO shallow copy axes? 

98 new_mgr = type(right)(tuple(res_blks), axes=right.axes, verify_integrity=False) 

99 return new_mgr 

100 

101 

102def _reset_block_mgr_locs(nbs: list[Block], locs: BlockPlacement) -> None: 

103 """ 

104 Reset mgr_locs to correspond to our original DataFrame. 

105 """ 

106 for nb in nbs: 

107 nblocs = locs[nb.mgr_locs.indexer] 

108 nb.mgr_locs = nblocs 

109 # Assertions are disabled for performance, but should hold: 

110 # assert len(nblocs) == nb.shape[0], (len(nblocs), nb.shape) 

111 # assert all(x in locs.as_array for x in nb.mgr_locs.as_array) 

112 

113 

114def _get_same_shape_values( 

115 lblk: Block, rblk: Block, left_ea: bool, right_ea: bool 

116) -> tuple[ArrayLike, ArrayLike]: 

117 """ 

118 Slice lblk.values to align with rblk. Squeeze if we have EAs. 

119 """ 

120 lvals = lblk.values 

121 rvals = rblk.values 

122 

123 # Require that the indexing into lvals be slice-like 

124 assert rblk.mgr_locs.is_slice_like, rblk.mgr_locs 

125 

126 # TODO(EA2D): with 2D EAs only this first clause would be needed 

127 if not (left_ea or right_ea): 

128 # error: No overload variant of "__getitem__" of "ExtensionArray" matches 

129 # argument type "Tuple[Union[ndarray, slice], slice]" 

130 lvals = lvals[rblk.mgr_locs.indexer, :] # type: ignore[call-overload] 

131 assert lvals.shape == rvals.shape, (lvals.shape, rvals.shape) 

132 elif left_ea and right_ea: 

133 assert lvals.shape == rvals.shape, (lvals.shape, rvals.shape) 

134 elif right_ea: 

135 # lvals are 2D, rvals are 1D 

136 

137 # error: No overload variant of "__getitem__" of "ExtensionArray" matches 

138 # argument type "Tuple[Union[ndarray, slice], slice]" 

139 lvals = lvals[rblk.mgr_locs.indexer, :] # type: ignore[call-overload] 

140 assert lvals.shape[0] == 1, lvals.shape 

141 lvals = lvals[0, :] 

142 else: 

143 # lvals are 1D, rvals are 2D 

144 assert rvals.shape[0] == 1, rvals.shape 

145 # error: No overload variant of "__getitem__" of "ExtensionArray" matches 

146 # argument type "Tuple[int, slice]" 

147 rvals = rvals[0, :] # type: ignore[call-overload] 

148 

149 return lvals, rvals 

150 

151 

152def blockwise_all( 

153 left: BlockManager, 

154 right: BlockManager, 

155 op: Callable[[ArrayLike, ArrayLike], bool], 

156) -> bool: 

157 """ 

158 Blockwise `all` reduction. 

159 """ 

160 for info in _iter_block_pairs(left, right): 

161 res = op(info.lvals, info.rvals) 

162 if not res: 

163 return False 

164 return True