[Python] Polars and Mixed Type guvectorize Numba Function

I am attempting to write a fast numba guvectorized function to use with in Polars (see here). I've simplified things down to examples below. Immediately below is an example that works fine:

import numpy as np
import polars as pl
import numba as nb

data = pl.DataFrame({'a': np.random.random(10), 'b': np.random.random(10)})
lazy = pl.LazyFrame(data)

# Note output type matches input types
@nb.guvectorize([(nb.float64[:], nb.float64[:], nb.float64[:])], '(n),(n)->(n)')
def to_float(x,y, res):
for i in range(x.shape[0]):
res=x+y

lazy = lazy.with_columns(c=to_float(pl.col('a'), pl.col('b')))
print(lazy.collect())


If I change the output type of the result array to not match the input array types, it doesn't work, see below. Trying print(to_int(np.random.random(10), np.random.random(10))) works fine, so it appears that the mismatch is how Polars is calling the to_int function. I've done a fair amount of searching, but I can't seem to find the way to tell Polars that the result array type is different than the input array types. I can't keep the result array as a float64 and do a .cast(pl.Int64) because my ints (in my non-simplified function) could have more than 15 digits.

Any ideas on how to fix this would be appreciated. Thanks!

I'm on Ubuntu, numba version = 0.61.0, numpy = 2.1.3, polars = 1.27.1, and python 3.12.8.

import numpy as np
import polars as pl
import numba as nb

data = pl.DataFrame({'a': np.random.random(10), 'b': np.random.random(10)})
lazy = pl.LazyFrame(data)

# Note output type mismatches input types
@nb.guvectorize([(nb.float64[:], nb.float64[:], nb.int64[:])], '(n),(n)->(n)')
def to_int(x,y, res):
for i in range(x.shape[0]):
res=int(x+y)

lazy = lazy.with_columns(c=to_int(pl.col('a'), pl.col('b')))
print(lazy.collect())


Traceback (most recent call last):
File "/tmp/fails.py", line 16, in <module>
print(lazy.collect())
^^^^^^^^^^^^^^
File "/home/ubuntu/miniforge3/lib/python3.12/site-packages/polars/_utils/deprecation.py", line 93, in wrapper
return function(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/ubuntu/miniforge3/lib/python3.12/site-packages/polars/lazyframe/frame.py", line 2206, in collect
return wrap_df(ldf.collect(engine, callback))
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
polars.exceptions.ComputeError: TypeError: No loop matching the specified signature and casting was found for ufunc to_int

Continue reading...