Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/pandas/core/_numba/executor.py: 18%

1from __future__ import annotations 

2 

3import functools 

4from typing import ( 

5 TYPE_CHECKING, 

6 Any, 

7) 

8 

9if TYPE_CHECKING: 

10 from collections.abc import Callable 

11 from pandas._typing import Scalar 

12 

13import numpy as np 

14 

15from pandas.compat._optional import import_optional_dependency 

16 

17from pandas.core.util.numba_ import jit_user_function 

18 

19 

20@functools.cache 

21def generate_apply_looper(func, nopython=True, nogil=True, parallel=False): 

22 if TYPE_CHECKING: 

23 import numba 

24 else: 

25 numba = import_optional_dependency("numba") 

26 nb_compat_func = jit_user_function(func) 

27 

28 @numba.jit(nopython=nopython, nogil=nogil, parallel=parallel) 

29 def nb_looper(values, axis, *args): 

30 # Operate on the first row/col in order to get 

31 # the output shape 

32 if axis == 0: 

33 first_elem = values[:, 0] 

34 dim0 = values.shape[1] 

35 else: 

36 first_elem = values[0] 

37 dim0 = values.shape[0] 

38 res0 = nb_compat_func(first_elem, *args) 

39 # Use np.asarray to get shape for 

40 # https://github.com/numba/numba/issues/4202#issuecomment-1185981507 

41 # Use tuple concatenation; numba doesn't support tuple unpacking syntax 

42 buf_shape = (dim0,) + np.atleast_1d(np.asarray(res0)).shape # noqa: RUF005 

43 if axis == 0: 

44 buf_shape = buf_shape[::-1] 

45 buff = np.empty(buf_shape) 

46 

47 if axis == 1: 

48 buff[0] = res0 

49 for i in numba.prange(1, values.shape[0]): 

50 buff[i] = nb_compat_func(values[i], *args) 

51 else: 

52 buff[:, 0] = res0 

53 for j in numba.prange(1, values.shape[1]): 

54 buff[:, j] = nb_compat_func(values[:, j], *args) 

55 return buff 

56 

57 return nb_looper 

58 

59 

60@functools.cache 

61def make_looper(func, result_dtype, is_grouped_kernel, nopython, nogil, parallel): 

62 if TYPE_CHECKING: 

63 import numba 

64 else: 

65 numba = import_optional_dependency("numba") 

66 

67 if is_grouped_kernel: 

68 

69 @numba.jit(nopython=nopython, nogil=nogil, parallel=parallel) 

70 def column_looper( 

71 values: np.ndarray, 

72 labels: np.ndarray, 

73 ngroups: int, 

74 min_periods: int, 

75 *args, 

76 ): 

77 result = np.empty((values.shape[0], ngroups), dtype=result_dtype) 

78 na_positions = {} 

79 for i in numba.prange(values.shape[0]): 

80 output, na_pos = func( 

81 values[i], result_dtype, labels, ngroups, min_periods, *args 

82 ) 

83 result[i] = output 

84 if len(na_pos) > 0: 

85 na_positions[i] = np.array(na_pos) 

86 return result, na_positions 

87 

88 else: 

89 

90 @numba.jit(nopython=nopython, nogil=nogil, parallel=parallel) 

91 # error: Incompatible redefinition (redefinition with type 

92 # "Callable[[ndarray[Any, Any], ndarray[Any, Any], ndarray[Any, Any], 

93 # int, VarArg(Any)], Any]", original type "Callable[[ndarray[Any, Any], 

94 # ndarray[Any, Any], int, int, VarArg(Any)], Any]") 

95 def column_looper( # type: ignore[misc] 

96 values: np.ndarray, 

97 start: np.ndarray, 

98 end: np.ndarray, 

99 min_periods: int, 

100 *args, 

101 ): 

102 result = np.empty((values.shape[0], len(start)), dtype=result_dtype) 

103 na_positions = {} 

104 for i in numba.prange(values.shape[0]): 

105 output, na_pos = func( 

106 values[i], result_dtype, start, end, min_periods, *args 

107 ) 

108 result[i] = output 

109 if len(na_pos) > 0: 

110 na_positions[i] = np.array(na_pos) 

111 return result, na_positions 

112 

113 return column_looper 

114 

115 

116default_dtype_mapping: dict[np.dtype, Any] = { 

117 np.dtype("int8"): np.int64, 

118 np.dtype("int16"): np.int64, 

119 np.dtype("int32"): np.int64, 

120 np.dtype("int64"): np.int64, 

121 np.dtype("uint8"): np.uint64, 

122 np.dtype("uint16"): np.uint64, 

123 np.dtype("uint32"): np.uint64, 

124 np.dtype("uint64"): np.uint64, 

125 np.dtype("float32"): np.float64, 

126 np.dtype("float64"): np.float64, 

127 np.dtype("complex64"): np.complex128, 

128 np.dtype("complex128"): np.complex128, 

129} 

130 

131 

132# TODO: Preserve complex dtypes 

133 

134float_dtype_mapping: dict[np.dtype, Any] = { 

135 np.dtype("int8"): np.float64, 

136 np.dtype("int16"): np.float64, 

137 np.dtype("int32"): np.float64, 

138 np.dtype("int64"): np.float64, 

139 np.dtype("uint8"): np.float64, 

140 np.dtype("uint16"): np.float64, 

141 np.dtype("uint32"): np.float64, 

142 np.dtype("uint64"): np.float64, 

143 np.dtype("float32"): np.float64, 

144 np.dtype("float64"): np.float64, 

145 np.dtype("complex64"): np.float64, 

146 np.dtype("complex128"): np.float64, 

147} 

148 

149identity_dtype_mapping: dict[np.dtype, Any] = { 

150 np.dtype("int8"): np.int8, 

151 np.dtype("int16"): np.int16, 

152 np.dtype("int32"): np.int32, 

153 np.dtype("int64"): np.int64, 

154 np.dtype("uint8"): np.uint8, 

155 np.dtype("uint16"): np.uint16, 

156 np.dtype("uint32"): np.uint32, 

157 np.dtype("uint64"): np.uint64, 

158 np.dtype("float32"): np.float32, 

159 np.dtype("float64"): np.float64, 

160 np.dtype("complex64"): np.complex64, 

161 np.dtype("complex128"): np.complex128, 

162} 

163 

164 

165def generate_shared_aggregator( 

166 func: Callable[..., Scalar], 

167 dtype_mapping: dict[np.dtype, np.dtype], 

168 is_grouped_kernel: bool, 

169 nopython: bool, 

170 nogil: bool, 

171 parallel: bool, 

172): 

173 """ 

174 Generate a Numba function that loops over the columns 2D object and applies 

175 a 1D numba kernel over each column. 

176 

177 Parameters 

178 ---------- 

179 func : function 

180 aggregation function to be applied to each column 

181 dtype_mapping: dict or None 

182 If not None, maps a dtype to a result dtype. 

183 Otherwise, will fall back to default mapping. 

184 is_grouped_kernel: bool, default False 

185 Whether func operates using the group labels (True) 

186 or using starts/ends arrays 

187 

188 If true, you also need to pass the number of groups to this function 

189 nopython : bool 

190 nopython to be passed into numba.jit 

191 nogil : bool 

192 nogil to be passed into numba.jit 

193 parallel : bool 

194 parallel to be passed into numba.jit 

195 

196 Returns 

197 ------- 

198 Numba function 

199 """ 

200 

201 # A wrapper around the looper function, 

202 # to dispatch based on dtype since numba is unable to do that in nopython mode 

203 

204 # It also post-processes the values by inserting nans where number of observations 

205 # is less than min_periods 

206 # Cannot do this in numba nopython mode 

207 # (you'll run into type-unification error when you cast int -> float) 

208 def looper_wrapper( 

209 values, 

210 start=None, 

211 end=None, 

212 labels=None, 

213 ngroups=None, 

214 min_periods: int = 0, 

215 **kwargs, 

216 ): 

217 result_dtype = dtype_mapping[values.dtype] 

218 column_looper = make_looper( 

219 func, result_dtype, is_grouped_kernel, nopython, nogil, parallel 

220 ) 

221 # Need to unpack kwargs since numba only supports *args 

222 if is_grouped_kernel: 

223 result, na_positions = column_looper( 

224 values, labels, ngroups, min_periods, *kwargs.values() 

225 ) 

226 else: 

227 result, na_positions = column_looper( 

228 values, start, end, min_periods, *kwargs.values() 

229 ) 

230 if result.dtype.kind == "i": 

231 # Look if na_positions is not empty 

232 # If so, convert the whole block 

233 # This is OK since int dtype cannot hold nan, 

234 # so if min_periods not satisfied for 1 col, it is not satisfied for 

235 # all columns at that index 

236 for na_pos in na_positions.values(): 

237 if len(na_pos) > 0: 

238 result = result.astype("float64") 

239 break 

240 # TODO: Optimize this 

241 for i, na_pos in na_positions.items(): 

242 if len(na_pos) > 0: 

243 result[i, na_pos] = np.nan 

244 return result 

245 

246 return looper_wrapper