@@ -1119,98 +1119,6 @@ def rank(
1119
1119
return ranks
1120
1120
1121
1121
1122
- def checked_add_with_arr (
1123
- arr : npt .NDArray [np .int64 ],
1124
- b : int | npt .NDArray [np .int64 ],
1125
- arr_mask : npt .NDArray [np .bool_ ] | None = None ,
1126
- b_mask : npt .NDArray [np .bool_ ] | None = None ,
1127
- ) -> npt .NDArray [np .int64 ]:
1128
- """
1129
- Perform array addition that checks for underflow and overflow.
1130
-
1131
- Performs the addition of an int64 array and an int64 integer (or array)
1132
- but checks that they do not result in overflow first. For elements that
1133
- are indicated to be NaN, whether or not there is overflow for that element
1134
- is automatically ignored.
1135
-
1136
- Parameters
1137
- ----------
1138
- arr : np.ndarray[int64] addend.
1139
- b : array or scalar addend.
1140
- arr_mask : np.ndarray[bool] or None, default None
1141
- array indicating which elements to exclude from checking
1142
- b_mask : np.ndarray[bool] or None, default None
1143
- array or scalar indicating which element(s) to exclude from checking
1144
-
1145
- Returns
1146
- -------
1147
- sum : An array for elements x + b for each element x in arr if b is
1148
- a scalar or an array for elements x + y for each element pair
1149
- (x, y) in (arr, b).
1150
-
1151
- Raises
1152
- ------
1153
- OverflowError if any x + y exceeds the maximum or minimum int64 value.
1154
- """
1155
- # For performance reasons, we broadcast 'b' to the new array 'b2'
1156
- # so that it has the same size as 'arr'.
1157
- b2 = np .broadcast_to (b , arr .shape )
1158
- if b_mask is not None :
1159
- # We do the same broadcasting for b_mask as well.
1160
- b2_mask = np .broadcast_to (b_mask , arr .shape )
1161
- else :
1162
- b2_mask = None
1163
-
1164
- # For elements that are NaN, regardless of their value, we should
1165
- # ignore whether they overflow or not when doing the checked add.
1166
- if arr_mask is not None and b2_mask is not None :
1167
- not_nan = np .logical_not (arr_mask | b2_mask )
1168
- elif arr_mask is not None :
1169
- not_nan = np .logical_not (arr_mask )
1170
- elif b_mask is not None :
1171
- # error: Argument 1 to "__call__" of "_UFunc_Nin1_Nout1" has
1172
- # incompatible type "Optional[ndarray[Any, dtype[bool_]]]";
1173
- # expected "Union[_SupportsArray[dtype[Any]], _NestedSequence
1174
- # [_SupportsArray[dtype[Any]]], bool, int, float, complex, str
1175
- # , bytes, _NestedSequence[Union[bool, int, float, complex, str
1176
- # , bytes]]]"
1177
- not_nan = np .logical_not (b2_mask ) # type: ignore[arg-type]
1178
- else :
1179
- not_nan = np .empty (arr .shape , dtype = bool )
1180
- not_nan .fill (True )
1181
-
1182
- # gh-14324: For each element in 'arr' and its corresponding element
1183
- # in 'b2', we check the sign of the element in 'b2'. If it is positive,
1184
- # we then check whether its sum with the element in 'arr' exceeds
1185
- # np.iinfo(np.int64).max. If so, we have an overflow error. If it
1186
- # it is negative, we then check whether its sum with the element in
1187
- # 'arr' exceeds np.iinfo(np.int64).min. If so, we have an overflow
1188
- # error as well.
1189
- i8max = lib .i8max
1190
- i8min = iNaT
1191
-
1192
- mask1 = b2 > 0
1193
- mask2 = b2 < 0
1194
-
1195
- if not mask1 .any ():
1196
- to_raise = ((i8min - b2 > arr ) & not_nan ).any ()
1197
- elif not mask2 .any ():
1198
- to_raise = ((i8max - b2 < arr ) & not_nan ).any ()
1199
- else :
1200
- to_raise = ((i8max - b2 [mask1 ] < arr [mask1 ]) & not_nan [mask1 ]).any () or (
1201
- (i8min - b2 [mask2 ] > arr [mask2 ]) & not_nan [mask2 ]
1202
- ).any ()
1203
-
1204
- if to_raise :
1205
- raise OverflowError ("Overflow in int64 addition" )
1206
-
1207
- result = arr + b
1208
- if arr_mask is not None or b2_mask is not None :
1209
- np .putmask (result , ~ not_nan , iNaT )
1210
-
1211
- return result
1212
-
1213
-
1214
1122
# ---- #
1215
1123
# take #
1216
1124
# ---- #
0 commit comments