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

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 nogil, parallel 

45 

46 Raises 

47 ------ 

48 NumbaUtilError 

49 """ 

50 if engine_kwargs is None: 

51 engine_kwargs = {} 

52 

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

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

55 return {"nogil": nogil, "parallel": parallel} 

56 

57 

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

59 """ 

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

61 as jitable. 

62 

63 Parameters 

64 ---------- 

65 func : function 

66 user defined function 

67 

68 Returns 

69 ------- 

70 function 

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

72 """ 

73 if TYPE_CHECKING: 

74 import numba 

75 else: 

76 numba = import_optional_dependency("numba") 

77 

78 if numba.extending.is_jitted(func): # pyright: ignore[reportAttributeAccessIssue] 

79 # Don't jit a user passed jitted function 

80 numba_func = func 

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

82 func, types.BuiltinFunctionType 

83 ): 

84 # Not necessary to jit builtins or np functions 

85 # This will mess up register_jitable 

86 numba_func = func # type: ignore[assignment] 

87 else: 

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

89 

90 return numba_func 

91 

92 

93_sentinel = object() 

94 

95 

96def prepare_function_arguments( 

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

98) -> tuple[tuple, dict]: 

99 """ 

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

101 we try to move kwargs into args if possible. 

102 

103 Parameters 

104 ---------- 

105 func : function 

106 User defined function 

107 args : tuple 

108 User input positional arguments 

109 kwargs : dict 

110 User input keyword arguments 

111 num_required_args : int 

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

113 These are not supplied by the user. 

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

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

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

117 

118 Returns 

119 ------- 

120 tuple[tuple, dict] 

121 args, kwargs 

122 

123 """ 

124 if not kwargs: 

125 return args, kwargs 

126 

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

128 signature = inspect.signature(func) 

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

130 arguments.apply_defaults() 

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

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

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

134 args = arguments.args 

135 kwargs = arguments.kwargs 

136 

137 if kwargs: 

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

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

140 # seems unlikely to happen soon. 

141 

142 raise NumbaUtilError( 

143 "numba does not support keyword-only arguments" 

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

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

146 ) 

147 

148 args = args[num_required_args:] 

149 return args, kwargs