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

1"""Common utilities for Numba operations""" 

2 

3from __future__ import annotations 

4 

5import inspect 

6import types 

7from typing import TYPE_CHECKING 

8 

9import numpy as np 

10 

11from pandas.compat._optional import import_optional_dependency 

12from pandas.errors import NumbaUtilError 

13 

14if TYPE_CHECKING: 

15 from collections.abc import Callable 

16 

17GLOBAL_USE_NUMBA: bool = False 

18 

19 

20def maybe_use_numba(engine: str | None) -> bool: 

21 """Signal whether to use numba routines.""" 

22 return engine == "numba" or (engine is None and GLOBAL_USE_NUMBA) 

23 

24 

25def set_use_numba(enable: bool = False) -> None: 

26 global GLOBAL_USE_NUMBA 

27 if enable: 

28 import_optional_dependency("numba") 

29 GLOBAL_USE_NUMBA = enable 

30 

31 

32def get_jit_arguments(engine_kwargs: dict[str, bool] | None = None) -> dict[str, bool]: 

33 """ 

34 Return arguments to pass to numba.JIT, falling back on pandas default JIT settings. 

35 

36 Parameters 

37 ---------- 

38 engine_kwargs : dict, default None 

39 user passed keyword arguments for numba.JIT 

40 

41 Returns 

42 ------- 

43 dict[str, bool] 

44 nopython, nogil, parallel 

45 

46 Raises 

47 ------ 

48 NumbaUtilError 

49 """ 

50 if engine_kwargs is None: 

51 engine_kwargs = {} 

52 

53 nopython = engine_kwargs.get("nopython", True) 

54 nogil = engine_kwargs.get("nogil", False) 

55 parallel = engine_kwargs.get("parallel", False) 

56 return {"nopython": nopython, "nogil": nogil, "parallel": parallel} 

57 

58 

59def jit_user_function(func: Callable) -> Callable: 

60 """ 

61 If user function is not jitted already, mark the user's function 

62 as jitable. 

63 

64 Parameters 

65 ---------- 

66 func : function 

67 user defined function 

68 

69 Returns 

70 ------- 

71 function 

72 Numba JITed function, or function marked as JITable by numba 

73 """ 

74 if TYPE_CHECKING: 

75 import numba 

76 else: 

77 numba = import_optional_dependency("numba") 

78 

79 if numba.extending.is_jitted(func): 

80 # Don't jit a user passed jitted function 

81 numba_func = func 

82 elif getattr(np, func.__name__, False) is func or isinstance( 

83 func, types.BuiltinFunctionType 

84 ): 

85 # Not necessary to jit builtins or np functions 

86 # This will mess up register_jitable 

87 numba_func = func # type: ignore[assignment] 

88 else: 

89 numba_func = numba.extending.register_jitable(func) # type: ignore[arg-type] 

90 

91 return numba_func 

92 

93 

94_sentinel = object() 

95 

96 

97def prepare_function_arguments( 

98 func: Callable, args: tuple, kwargs: dict, *, num_required_args: int 

99) -> tuple[tuple, dict]: 

100 """ 

101 Prepare arguments for jitted function. As numba functions do not support kwargs, 

102 we try to move kwargs into args if possible. 

103 

104 Parameters 

105 ---------- 

106 func : function 

107 User defined function 

108 args : tuple 

109 User input positional arguments 

110 kwargs : dict 

111 User input keyword arguments 

112 num_required_args : int 

113 The number of leading positional arguments we will pass to udf. 

114 These are not supplied by the user. 

115 e.g. for groupby we require "values", "index" as the first two arguments: 

116 `numba_func(group, group_index, *args)`, in this case num_required_args=2. 

117 See :func:`pandas.core.groupby.numba_.generate_numba_agg_func` 

118 

119 Returns 

120 ------- 

121 tuple[tuple, dict] 

122 args, kwargs 

123 

124 """ 

125 if not kwargs: 

126 return args, kwargs 

127 

128 # the udf should have this pattern: def udf(arg1, arg2, ..., *args, **kwargs):... 

129 signature = inspect.signature(func) 

130 arguments = signature.bind(*[_sentinel] * num_required_args, *args, **kwargs) 

131 arguments.apply_defaults() 

132 # Ref: https://peps.python.org/pep-0362/ 

133 # Arguments which could be passed as part of either *args or **kwargs 

134 # will be included only in the BoundArguments.args attribute. 

135 args = arguments.args 

136 kwargs = arguments.kwargs 

137 

138 if kwargs: 

139 # Note: in case numba supports keyword-only arguments in 

140 # a future version, we should remove this check. But this 

141 # seems unlikely to happen soon. 

142 

143 raise NumbaUtilError( 

144 "numba does not support keyword-only arguments" 

145 "https://github.com/numba/numba/issues/2916, " 

146 "https://github.com/numba/numba/issues/6846" 

147 ) 

148 

149 args = args[num_required_args:] 

150 return args, kwargs