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

1"""Common utilities for Numba operations with groupby ops""" 

2 

3from __future__ import annotations 

4 

5import functools 

6import inspect 

7from typing import ( 

8 TYPE_CHECKING, 

9 Any, 

10) 

11 

12import numpy as np 

13 

14from pandas.compat._optional import import_optional_dependency 

15 

16from pandas.core.util.numba_ import ( 

17 NumbaUtilError, 

18 jit_user_function, 

19) 

20 

21if TYPE_CHECKING: 

22 from collections.abc import Callable 

23 

24 from pandas._typing import Scalar 

25 

26 

27def validate_udf(func: Callable) -> None: 

28 """ 

29 Validate user defined function for ops when using Numba with groupby ops. 

30 

31 The first signature arguments should include: 

32 

33 def f(values, index, ...): 

34 ... 

35 

36 Parameters 

37 ---------- 

38 func : function, default False 

39 user defined function 

40 

41 Returns 

42 ------- 

43 None 

44 

45 Raises 

46 ------ 

47 NumbaUtilError 

48 """ 

49 if not callable(func): 

50 raise NotImplementedError( 

51 "Numba engine can only be used with a single function." 

52 ) 

53 udf_signature = list(inspect.signature(func).parameters.keys()) 

54 expected_args = ["values", "index"] 

55 min_number_args = len(expected_args) 

56 if ( 

57 len(udf_signature) < min_number_args 

58 or udf_signature[:min_number_args] != expected_args 

59 ): 

60 raise NumbaUtilError( 

61 f"The first {min_number_args} arguments to {func.__name__} must be " 

62 f"{expected_args}" 

63 ) 

64 

65 

66@functools.cache 

67def generate_numba_agg_func( 

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

69 nopython: bool, 

70 nogil: bool, 

71 parallel: bool, 

72) -> Callable[[np.ndarray, np.ndarray, np.ndarray, np.ndarray, int, Any], np.ndarray]: 

73 """ 

74 Generate a numba jitted agg function specified by values from engine_kwargs. 

75 

76 1. jit the user's function 

77 2. Return a groupby agg function with the jitted function inline 

78 

79 Configurations specified in engine_kwargs apply to both the user's 

80 function _AND_ the groupby evaluation loop. 

81 

82 Parameters 

83 ---------- 

84 func : function 

85 function to be applied to each group and will be JITed 

86 nopython : bool 

87 nopython to be passed into numba.jit 

88 nogil : bool 

89 nogil to be passed into numba.jit 

90 parallel : bool 

91 parallel to be passed into numba.jit 

92 

93 Returns 

94 ------- 

95 Numba function 

96 """ 

97 numba_func = jit_user_function(func) 

98 if TYPE_CHECKING: 

99 import numba 

100 else: 

101 numba = import_optional_dependency("numba") 

102 

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

104 def group_agg( 

105 values: np.ndarray, 

106 index: np.ndarray, 

107 begin: np.ndarray, 

108 end: np.ndarray, 

109 num_columns: int, 

110 *args: Any, 

111 ) -> np.ndarray: 

112 assert len(begin) == len(end) 

113 num_groups = len(begin) 

114 

115 result = np.empty((num_groups, num_columns)) 

116 for i in numba.prange(num_groups): 

117 group_index = index[begin[i] : end[i]] 

118 for j in numba.prange(num_columns): 

119 group = values[begin[i] : end[i], j] 

120 result[i, j] = numba_func(group, group_index, *args) 

121 return result 

122 

123 return group_agg 

124 

125 

126@functools.cache 

127def generate_numba_transform_func( 

128 func: Callable[..., np.ndarray], 

129 nopython: bool, 

130 nogil: bool, 

131 parallel: bool, 

132) -> Callable[[np.ndarray, np.ndarray, np.ndarray, np.ndarray, int, Any], np.ndarray]: 

133 """ 

134 Generate a numba jitted transform function specified by values from engine_kwargs. 

135 

136 1. jit the user's function 

137 2. Return a groupby transform function with the jitted function inline 

138 

139 Configurations specified in engine_kwargs apply to both the user's 

140 function _AND_ the groupby evaluation loop. 

141 

142 Parameters 

143 ---------- 

144 func : function 

145 function to be applied to each window and will be JITed 

146 nopython : bool 

147 nopython to be passed into numba.jit 

148 nogil : bool 

149 nogil to be passed into numba.jit 

150 parallel : bool 

151 parallel to be passed into numba.jit 

152 

153 Returns 

154 ------- 

155 Numba function 

156 """ 

157 numba_func = jit_user_function(func) 

158 if TYPE_CHECKING: 

159 import numba 

160 else: 

161 numba = import_optional_dependency("numba") 

162 

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

164 def group_transform( 

165 values: np.ndarray, 

166 index: np.ndarray, 

167 begin: np.ndarray, 

168 end: np.ndarray, 

169 num_columns: int, 

170 *args: Any, 

171 ) -> np.ndarray: 

172 assert len(begin) == len(end) 

173 num_groups = len(begin) 

174 

175 result = np.empty((len(values), num_columns)) 

176 for i in numba.prange(num_groups): 

177 group_index = index[begin[i] : end[i]] 

178 for j in numba.prange(num_columns): 

179 group = values[begin[i] : end[i], j] 

180 result[begin[i] : end[i], j] = numba_func(group, group_index, *args) 

181 return result 

182 

183 return group_transform