1
1
"""Common utilities for Numba operations"""
2
2
from __future__ import annotations
3
3
4
- import types
5
4
from typing import (
6
5
TYPE_CHECKING ,
7
6
Callable ,
8
7
)
9
8
10
- import numpy as np
11
-
12
9
from pandas .compat ._optional import import_optional_dependency
13
10
from pandas .errors import NumbaUtilError
14
11
@@ -63,27 +60,20 @@ def get_jit_arguments(
63
60
return {"nopython" : nopython , "nogil" : nogil , "parallel" : parallel }
64
61
65
62
66
- def jit_user_function (
67
- func : Callable , nopython : bool , nogil : bool , parallel : bool
68
- ) -> Callable :
63
+ def jit_user_function (func : Callable ) -> Callable :
69
64
"""
70
- JIT the user's function given the configurable arguments.
65
+ If user function is not jitted already, mark the user's function
66
+ as jitable.
71
67
72
68
Parameters
73
69
----------
74
70
func : function
75
71
user defined function
76
- nopython : bool
77
- nopython parameter for numba.JIT
78
- nogil : bool
79
- nogil parameter for numba.JIT
80
- parallel : bool
81
- parallel parameter for numba.JIT
82
72
83
73
Returns
84
74
-------
85
75
function
86
- Numba JITed function
76
+ Numba JITed function, or function marked as JITable by numba
87
77
"""
88
78
if TYPE_CHECKING :
89
79
import numba
@@ -94,19 +84,6 @@ def jit_user_function(
94
84
# Don't jit a user passed jitted function
95
85
numba_func = func
96
86
else :
97
-
98
- @numba .generated_jit (nopython = nopython , nogil = nogil , parallel = parallel )
99
- def numba_func (data , * _args ):
100
- if getattr (np , func .__name__ , False ) is func or isinstance (
101
- func , types .BuiltinFunctionType
102
- ):
103
- jf = func
104
- else :
105
- jf = numba .jit (func , nopython = nopython , nogil = nogil )
106
-
107
- def impl (data , * _args ):
108
- return jf (data , * _args )
109
-
110
- return impl
87
+ numba_func = numba .extending .register_jitable (func )
111
88
112
89
return numba_func
0 commit comments