BUG: No kwargs in df.apply(raw=True, engine="numba")

[auderson]

Pandas version checks

• I have checked that this issue has not already been reported.

• I have confirmed this bug exists on the latest version of pandas.

• I have confirmed this bug exists on the main branch of pandas.

Reproducible Example

import pandas as pd

df = pd.DataFrame(np.nan, index=[1,2,3], columns=list("abc"))
df.apply(np.nanquantile, q=0.5, raw=True, engine="numba")

Issue Description

related: #55753
the above code raises missing a required argument: 'q':

TypingError: Failed in nopython mode pipeline (step: nopython frontend)
No implementation of function Function(<function nanquantile at 0x7fbde03c2b00>) found for signature:
 
 >>> nanquantile(array(float64, 1d, A))
 
There are 6 candidate implementations:
    - Of which 4 did not match due to:
    Overload in function 'register_quantile_function.<locals>.np_quantile': File: ../../../../../../Techfin/PublicPackages/TechFinDataRAD/rd/_deps/commons/numba_overloads/overloads.py: Line 152.
      With argument(s): '(array(float64, 1d, A))':
     Rejected as the implementation raised a specific error:
       TypingError: missing a required argument: 'q'
  raised from /home/auderson/mambaforge/envs/py3.10/lib/python3.10/site-packages/numba/core/typing/templates.py:784
    - Of which 2 did not match due to:
    Overload in function 'register_jitable.<locals>.wrap.<locals>.ov_wrap': File: numba/core/extending.py: Line 161.
      With argument(s): '(array(float64, 1d, A))':
     Rejected as the implementation raised a specific error:
       AssertionError: Implementation function returned by `@overload` has an unexpected type.  Got <function nanquantile at 0x7fbde03c2b00>
  raised from /home/auderson/mambaforge/envs/py3.10/lib/python3.10/site-packages/numba/core/typing/templates.py:807

During: resolving callee type: Function(<function nanquantile at 0x7fbde03c2b00>)
During: typing of call at /home/auderson/mambaforge/envs/py3.10/lib/python3.10/site-packages/pandas/core/_numba/executor.py (36)


File "../../../../../mambaforge/envs/py3.10/lib/python3.10/site-packages/pandas/core/_numba/executor.py", line 36:
    def nb_looper(values, axis):
        <source elided>
            dim0 = values.shape[0]
        res0 = nb_compat_func(first_elem)
        ^

Expected Behavior

kwargs should be passed to njit func

Installed Versions

/home/auderson/mambaforge/envs/py3.10/lib/python3.10/site-packages/_distutils_hack/init.py:26: UserWarning: Setuptools is replacing distutils.
warnings.warn("Setuptools is replacing distutils.")

INSTALLED VERSIONS

commit : d9cdd2e
python : 3.10.13.final.0
python-bits : 64
OS : Linux
OS-release : 5.15.0-101-generic
Version : #111-Ubuntu SMP Tue Mar 5 20:16:58 UTC 2024
machine : x86_64
processor : x86_64
byteorder : little
LC_ALL : None
LANG : en_US.UTF-8
LOCALE : en_US.UTF-8

pandas : 2.2.2
numpy : 1.26.4
pytz : 2024.1
dateutil : 2.8.2
setuptools : 69.1.0
pip : 24.0
Cython : 3.0.7
pytest : 8.0.2
hypothesis : None
sphinx : 7.2.6
blosc : None
feather : None
xlsxwriter : None
lxml.etree : None
html5lib : None
pymysql : 1.4.6
psycopg2 : 2.9.9
jinja2 : 3.1.3
IPython : 8.22.1
pandas_datareader : None
adbc-driver-postgresql: None
adbc-driver-sqlite : None
bs4 : 4.12.3
bottleneck : None
dataframe-api-compat : None
fastparquet : None
fsspec : 2024.3.1
gcsfs : None
matplotlib : 3.8.3
numba : 0.59.0
numexpr : 2.9.0
odfpy : None
openpyxl : None
pandas_gbq : None
pyarrow : 15.0.0
pyreadstat : None
python-calamine : None
pyxlsb : None
s3fs : None
scipy : 1.12.0
sqlalchemy : 2.0.27
tables : 3.9.2
tabulate : None
xarray : None
xlrd : None
zstandard : 0.22.0
tzdata : 2024.1
qtpy : 2.4.1
pyqt5 : None

[lithomas1]

Hm, this doesn't look too hard to fix, looking at the error (the nb_compat_func line).

It's been a while since I've looked at any numba stuff, though.

Are you interested in potentially contributing a PR?

[auderson]

Are you interested in potentially contributing a PR?

Sure, but I find the solution may not be that trivial:

Since numba generally requires the argument is known when function is defined, so any kind of varargs is not well supported.
It dosen't support **kwargs, and IIRC it also had some bugs with *args.

This means we can't simply add **kwargs to nb_looper:

pandas/pandas/core/_numba/executor.py

Lines 26 to 52 in 58461fe

	@numba.jit(nopython=nopython, nogil=nogil, parallel=parallel)
	def nb_looper(values, axis):
	# Operate on the first row/col in order to get
	# the output shape
	if axis == 0:
	first_elem = values[:, 0]
	dim0 = values.shape[1]
	else:
	first_elem = values[0]
	dim0 = values.shape[0]
	res0 = nb_compat_func(first_elem)
	# Use np.asarray to get shape for
	# https://github.com/numba/numba/issues/4202#issuecomment-1185981507
	buf_shape = (dim0,) + np.atleast_1d(np.asarray(res0)).shape
	if axis == 0:
	buf_shape = buf_shape[::-1]
	buff = np.empty(buf_shape)
	if axis == 1:
	buff[0] = res0
	for i in numba.prange(1, values.shape[0]):
	buff[i] = nb_compat_func(values[i])
	else:
	buff[:, 0] = res0
	for j in numba.prange(1, values.shape[1]):
	buff[:, j] = nb_compat_func(values[:, j])
	return buff

A possible solution is to dynamically create an intermediate function for each case. Use the example above:
df.apply(np.nanquantile, q=0.5, raw=True, engine="numba")

In this case, we dynamically defines:

@register_jitable
def nb_compat_func(values, q=0.5):
    return user_func(values, q=q)

which holds the **kwargs and the only required argument is values and this nb_compat_func will be used by nb_looper just as before.

However, I felt this solution is somewhat too hacky and not clean. Do you have any other ideas?

[lithomas1]

I think *args might be the best way to deal with this, if it's possible to fix this.
It's how we pass the kwargs to the other numba stuff like the numba groupby stuff.

Do you remember the issue you had with *args?

[auderson]

numba/numba#8108
Basically if *args exists then other args that follows becomes kwd-only which is not yet supported. (Maybe there are also some cache related issues from my recollection but those are irrelevant)

I think *args might be the best way to deal with this, if it's possible to fix this.
It's how we pass the kwargs to the other numba stuff like the numba groupby stuff.

I just checked out the code of numba groupby, it should be viable if we disallow **kwargs and only use *args in numba apply as well.

[auderson]

take