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[..., Any]) -> 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 nogil: bool, 

70 parallel: bool, 

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

72 """ 

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

74 

75 1. jit the user's function 

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

77 

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

79 function _AND_ the groupby evaluation loop. 

80 

81 Parameters 

82 ---------- 

83 func : function 

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

85 nogil : bool 

86 nogil to be passed into numba.jit 

87 parallel : bool 

88 parallel to be passed into numba.jit 

89 

90 Returns 

91 ------- 

92 Numba function 

93 """ 

94 numba_func = jit_user_function(func) 

95 if TYPE_CHECKING: 

96 import numba 

97 else: 

98 numba = import_optional_dependency("numba") 

99 

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

101 def group_agg( 

102 values: np.ndarray, 

103 index: np.ndarray, 

104 begin: np.ndarray, 

105 end: np.ndarray, 

106 num_columns: int, 

107 *args: Any, 

108 ) -> np.ndarray: 

109 assert len(begin) == len(end) 

110 num_groups = len(begin) 

111 

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

113 for i in numba.prange(num_groups): 

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

115 for j in numba.prange(num_columns): 

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

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

118 return result 

119 

120 return group_agg 

121 

122 

123@functools.cache 

124def generate_numba_transform_func( 

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

126 nogil: bool, 

127 parallel: bool, 

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

129 """ 

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

131 

132 1. jit the user's function 

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

134 

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

136 function _AND_ the groupby evaluation loop. 

137 

138 Parameters 

139 ---------- 

140 func : function 

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

142 nogil : bool 

143 nogil to be passed into numba.jit 

144 parallel : bool 

145 parallel to be passed into numba.jit 

146 

147 Returns 

148 ------- 

149 Numba function 

150 """ 

151 numba_func = jit_user_function(func) 

152 if TYPE_CHECKING: 

153 import numba 

154 else: 

155 numba = import_optional_dependency("numba") 

156 

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

158 def group_transform( 

159 values: np.ndarray, 

160 index: np.ndarray, 

161 begin: np.ndarray, 

162 end: np.ndarray, 

163 num_columns: int, 

164 *args: Any, 

165 ) -> np.ndarray: 

166 assert len(begin) == len(end) 

167 num_groups = len(begin) 

168 

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

170 for i in numba.prange(num_groups): 

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

172 for j in numba.prange(num_columns): 

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

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

175 return result 

176 

177 return group_transform