ENH add *args support for numba apply

[auderson]

• closes BUG: No kwargs in df.apply(raw=True, engine="numba") #58712
• Tests added and passed if fixing a bug or adding a new feature
• All code checks passed.
• Added type annotations to new arguments/methods/functions.
• Added an entry in the latest doc/source/whatsnew/vX.X.X.rst file if fixing a bug or adding a new feature.

added *args to nb_looper().
use get_jit_arguments() to disallow **kwargs in raw + numba engine.

[auderson]

@lithomas1
I notice that *args support for raw=False + numba engine can be also added, should I open a another PR or just put them together?

[lithomas1]

@lithomas1 I notice that *args support for raw=False + numba engine can be also added, should I open a another PR or just put them together?

If it's not too involved to add support there, you're welcome to put it in the same PR.

Thanks for the quick PR btw.

[auderson]

It's all done. Could you take a look when you have time?

[lithomas1]

Does this also handle **kwargs by e.g. passing them as *args?

It might be worth exploring that.

This LGTM otherwise.

[auderson]

That's a good point, I'll have a try later when I have time.

[auderson]

BTW, since this PR was labeled as enhancement, should I move it to the enhancement section in whatsnew?

[auderson]

@lithomas1
Hi, I added an util function as you suggested. It would be nice if you spare some time to have a look at it.

[lithomas1]

Does this work if you pass nopython=False in the engine_kwargs?

[auderson]

The pandas test raises missing a required argument error when nopython=False.

FAILED
pandas/tests/apply/test_frame_apply.py:63 (test_apply_args[numba-False-True-0])
float_frame =                A         B         C         D
foo_0   0.189053 -0.522748 -0.413064 -2.441467
foo_1   1.799707  1.1441...47218  0.968478 -0.955145
foo_28  0.354112 -1.968397  0.899274 -0.158248
foo_29 -0.967681  1.678419  0.765355  0.045808
axis = 0, raw = True, engine = 'numba', nopython = False

    @pytest.mark.parametrize("axis", [0, 1])
    @pytest.mark.parametrize("raw", [True, False])
    @pytest.mark.parametrize("nopython", [True, False])
    def test_apply_args(float_frame, axis, raw, engine, nopython):
        engine_kwargs = {"nopython": nopython}
        result = float_frame.apply(
            lambda x, y: x + y, axis, args=(1,), raw=raw, engine=engine, engine_kwargs=engine_kwargs
        )
        expected = float_frame + 1
        tm.assert_frame_equal(result, expected)
    
        # GH:58712
        result = float_frame.apply(
            lambda x, a, b: x + a + b, args=(1,), b=2, raw=raw, engine=engine, engine_kwargs=engine_kwargs
        )
        expected = float_frame + 3
        tm.assert_frame_equal(result, expected)
    
        if engine == "numba":
            # keyword-only arguments are not supported in numba
            with pytest.raises(
                pd.errors.NumbaUtilError,
                match="numba does not support kwargs with nopython=True",
            ):
>               float_frame.apply(
                    lambda x, a, *, b: x + a + b, args=(1,), b=2, raw=raw, engine=engine, engine_kwargs=engine_kwargs
                )

test_frame_apply.py:88: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
../../core/frame.py:10353: in apply
    return op.apply().__finalize__(self, method="apply")
../../core/apply.py:886: in apply
    return self.apply_raw(engine=self.engine, engine_kwargs=self.engine_kwargs)
../../core/apply.py:991: in apply_raw
    result = nb_looper(self.values, self.axis, *args)
/home/auderson/miniconda3/envs/pandas-dev/lib/python3.10/site-packages/numba/core/dispatcher.py:468: in _compile_for_args
    error_rewrite(e, 'typing')
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

e = TypingError('Failed in nopython mode pipeline (step: nopython frontend)\n\x1b[1m\x1b[1m\x1b[1mNo implementation of fun...0 = values.shape[0]\n\x1b[1m        res0 = nb_compat_func(first_elem, *args)\n\x1b[0m        \x1b[1m^\x1b[0m\x1b[0m\n')
issue_type = 'typing'

    def error_rewrite(e, issue_type):
        """
        Rewrite and raise Exception `e` with help supplied based on the
        specified issue_type.
        """
        if config.SHOW_HELP:
            help_msg = errors.error_extras[issue_type]
            e.patch_message('\n'.join((str(e).rstrip(), help_msg)))
        if config.FULL_TRACEBACKS:
            raise e
        else:
>           raise e.with_traceback(None)
E           numba.core.errors.TypingError: Failed in nopython mode pipeline (step: nopython frontend)
E           No implementation of function Function(<function test_apply_args.<locals>.<lambda> at 0x7f0d2f6081f0>) found for signature:
E            
E            >>> <lambda>(readonly array(float64, 1d, A), int64)
E            
E           There are 2 candidate implementations:
E               - Of which 2 did not match due to:
E               Overload in function 'register_jitable.<locals>.wrap.<locals>.ov_wrap': File: numba/core/extending.py: Line 161.
E                 With argument(s): '(readonly array(float64, 1d, A), int64)':
E                Rejected as the implementation raised a specific error:
E                  TypeError: missing a required argument: 'a'
E             raised from /home/auderson/miniconda3/envs/pandas-dev/lib/python3.10/inspect.py:3101
E           
E           During: resolving callee type: Function(<function test_apply_args.<locals>.<lambda> at 0x7f0d2f6081f0>)
E           During: typing of call at /mnt/c/Users/auderson/Documents/Works/OpenSourcePackages/pandas/pandas/core/_numba/executor.py (38)
E           
E           
E           File "../../core/_numba/executor.py", line 38:
E               def nb_looper(values, axis, *args):
E                   <source elided>
E                       dim0 = values.shape[0]
E                   res0 = nb_compat_func(first_elem, *args)
E                   ^

/home/auderson/miniconda3/envs/pandas-dev/lib/python3.10/site-packages/numba/core/dispatcher.py:409: TypingError

An equivalent reproducer:

from numba import jit

@jit(nopython=False)
def foo(*args, kwarg=None):
    print(args)
    print(kwarg)

@jit(nopython=False)
def bar(a, *args):
    foo(a, *args, kwarg='foobar')

bar(1, 2, 3)

UnsupportedError: Failed in object mode pipeline (step: analyzing bytecode)

CALL_FUNCTION_EX with **kwargs not supported.
If you are not using **kwargs this may indicate that
you have a large number of kwargs and are using inlined control
flow. You can resolve this issue by moving the control flow out of
the function call. For example, if you have

    f(a=1 if flag else 0, ...)

Replace that with:

    a_val = 1 if flag else 0
    f(a=a_val, ...)

[auderson]

Maybe here need to be modified:

pandas/pandas/core/util/numba_.py

Lines 57 to 61 in 3b48b17

	if kwargs and nopython:
	raise NumbaUtilError(
	"numba does not support kwargs with nopython=True: "
	"https://github.com/numba/numba/issues/2916"
	)

Regardless nopython is True or False, the "numba does not support kwargs" should be raised:

 if kwargs: 
     raise NumbaUtilError( 
         "numba does not support keyword-only arguments" 
         "https://github.com/numba/numba/issues/2916" 
     ) 

[lithomas1]

Good catch, updating the error message is probably the right move here.

[auderson]

We can also extend prepare_function_arguments by adding a option: num_required_args, so that all the numba functions' args & kwargs stuff can be handled by this function, instead of in get_jit_arguments .
For example in groupby, we have numba_func(group, group_index, *args), in this case num_required_args = 2:

pandas/pandas/core/groupby/numba_.py

Lines 102 to 120 in 3b48b17

	@numba.jit(nopython=nopython, nogil=nogil, parallel=parallel)
	def group_agg(
	values: np.ndarray,
	index: np.ndarray,
	begin: np.ndarray,
	end: np.ndarray,
	num_columns: int,
	*args: Any,
	) -> np.ndarray:
	assert len(begin) == len(end)
	num_groups = len(begin)
	result = np.empty((num_groups, num_columns))
	for i in numba.prange(num_groups):
	group_index = index[begin[i] : end[i]]
	for j in numba.prange(num_columns):
	group = values[begin[i] : end[i], j]
	result[i, j] = numba_func(group, group_index, *args)
	return result

[lithomas1]

We can also extend prepare_function_arguments by adding a option: num_required_args, so that all the numba functions' args & kwargs stuff can be handled by this function, instead of in get_jit_arguments . For example in groupby, we have numba_func(group, group_index, *args), in this case num_required_args = 2:

pandas/pandas/core/groupby/numba_.py

Lines 102 to 120 in 3b48b17

	@numba.jit(nopython=nopython, nogil=nogil, parallel=parallel)
	def group_agg(
	values: np.ndarray,
	index: np.ndarray,
	begin: np.ndarray,
	end: np.ndarray,
	num_columns: int,
	*args: Any,
	) -> np.ndarray:
	assert len(begin) == len(end)
	num_groups = len(begin)
	result = np.empty((num_groups, num_columns))
	for i in numba.prange(num_groups):
	group_index = index[begin[i] : end[i]]
	for j in numba.prange(num_columns):
	group = values[begin[i] : end[i], j]
	result[i, j] = numba_func(group, group_index, *args)
	return result

Good idea.

Do you mind doing this in a separate PR after this?

[auderson]

Good idea.

Do you mind doing this in a separate PR after this?

Sure, I'll do that after I finish this PR.

[auderson]

Looks all green :) @lithomas1

[auderson]

@lithomas1 Hi, if you may have time, please take a look so that I will start the other PR which enhances prepare_function_arguments.