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 

12 from pandas._typing import Scalar 

13 

14import numpy as np 

15 

16from pandas.compat._optional import import_optional_dependency 

17 

18from pandas.core.util.numba_ import jit_user_function 

19 

20 

21@functools.cache 

22def generate_apply_looper( 

23 func: Callable[..., Any], nogil: bool = True, parallel: bool = False 

24) -> Callable[..., np.ndarray]: 

25 if TYPE_CHECKING: 

26 import numba 

27 else: 

28 numba = import_optional_dependency("numba") 

29 nb_compat_func = jit_user_function(func) 

30 

31 @numba.jit(nogil=nogil, parallel=parallel) 

32 def nb_looper(values: np.ndarray, axis: int, *args: Any) -> np.ndarray: 

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

34 # the output shape 

35 if axis == 0: 

36 first_elem = values[:, 0] 

37 dim0 = values.shape[1] 

38 else: 

39 first_elem = values[0] 

40 dim0 = values.shape[0] 

41 res0 = nb_compat_func(first_elem, *args) 

42 # Use np.asarray to get shape for 

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

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

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

46 if axis == 0: 

47 buf_shape = buf_shape[::-1] 

48 buff = np.empty(buf_shape) 

49 

50 if axis == 1: 

51 buff[0] = res0 

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

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

54 else: 

55 buff[:, 0] = res0 

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

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

58 return buff 

59 

60 return nb_looper 

61 

62 

63@functools.cache 

64def make_looper( 

65 func: Callable[..., Any], 

66 result_dtype: np.dtype, 

67 is_grouped_kernel: bool, 

68 nogil: bool, 

69 parallel: bool, 

70) -> Callable[..., tuple[np.ndarray, dict[int, np.ndarray]]]: 

71 if TYPE_CHECKING: 

72 import numba 

73 else: 

74 numba = import_optional_dependency("numba") 

75 

76 if is_grouped_kernel: 

77 

78 @numba.jit(nogil=nogil, parallel=parallel) 

79 def column_looper( # pyright: ignore[reportRedeclaration] 

80 values: np.ndarray, 

81 labels: np.ndarray, 

82 ngroups: int, 

83 min_periods: int, 

84 *args: Any, 

85 ) -> tuple[np.ndarray, dict[int, np.ndarray]]: 

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

87 na_positions = {} 

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

89 output, na_pos = func( 

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

91 ) 

92 result[i] = output 

93 if len(na_pos) > 0: 

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

95 return result, na_positions 

96 

97 else: 

98 

99 @numba.jit(nogil=nogil, parallel=parallel) 

100 def column_looper( 

101 values: np.ndarray, 

102 start: np.ndarray, 

103 end: np.ndarray, 

104 min_periods: int, 

105 *args: Any, 

106 ) -> tuple[np.ndarray, dict[int, np.ndarray]]: 

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

108 na_positions = {} 

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

110 output, na_pos = func( 

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

112 ) 

113 result[i] = output 

114 if len(na_pos) > 0: 

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

116 return result, na_positions 

117 

118 return column_looper 

119 

120 

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

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

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

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

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

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

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

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

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

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

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

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

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

134} 

135 

136 

137# TODO: Preserve complex dtypes 

138 

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

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

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

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

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

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

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

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

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

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

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

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

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

152} 

153 

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

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

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

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

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

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

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

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

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

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

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

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

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

167} 

168 

169 

170def generate_shared_aggregator( 

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

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

173 is_grouped_kernel: bool, 

174 nogil: bool, 

175 parallel: bool, 

176) -> Callable[..., np.ndarray]: 

177 """ 

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

179 a 1D numba kernel over each column. 

180 

181 Parameters 

182 ---------- 

183 func : function 

184 aggregation function to be applied to each column 

185 dtype_mapping: dict or None 

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

187 Otherwise, will fall back to default mapping. 

188 is_grouped_kernel: bool, default False 

189 Whether func operates using the group labels (True) 

190 or using starts/ends arrays 

191 

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

193 nogil : bool 

194 nogil to be passed into numba.jit 

195 parallel : bool 

196 parallel to be passed into numba.jit 

197 

198 Returns 

199 ------- 

200 Numba function 

201 """ 

202 

203 # A wrapper around the looper function, 

204 # to dispatch based on dtype 

205 

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

207 # is less than min_periods 

208 def looper_wrapper( 

209 values: np.ndarray, 

210 start: np.ndarray | None = None, 

211 end: np.ndarray | None = None, 

212 labels: np.ndarray | None = None, 

213 ngroups: int | None = None, 

214 min_periods: int = 0, 

215 **kwargs: Any, 

216 ) -> np.ndarray: 

217 result_dtype = dtype_mapping[values.dtype] 

218 column_looper = make_looper( 

219 func, result_dtype, is_grouped_kernel, 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, 

225 labels, # pyright: ignore[reportArgumentType] 

226 ngroups, # pyright: ignore[reportArgumentType] 

227 min_periods, 

228 *kwargs.values(), 

229 ) 

230 else: 

231 result, na_positions = column_looper( 

232 values, 

233 start, # pyright: ignore[reportArgumentType] 

234 end, # pyright: ignore[reportArgumentType] 

235 min_periods, 

236 *kwargs.values(), 

237 ) 

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

239 # Look if na_positions is not empty 

240 # If so, convert the whole block 

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

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

243 # all columns at that index 

244 for na_pos in na_positions.values(): 

245 if len(na_pos) > 0: 

246 result = result.astype("float64") 

247 break 

248 # TODO: Optimize this 

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

250 if len(na_pos) > 0: 

251 result[i, na_pos] = np.nan 

252 return result 

253 

254 return looper_wrapper