[Python] Howto efficiently apply a gufunc to a 2D region of a Polars DataFrame

Both Polars and Numba are fantastic libraries that complement each other pretty well. There are some limitations when using Numba-compiled functions in Polars:

• Arrow columns must be converted to NumPy arrays (and then converted back).
• NumPy/Numba does not support missing data.

I'm trying to apply a Numba-compiled gufunc to a 2D region of a Polars DataFrame with two approaches:

1. Using map_batches with a custom lambda function to apply the gufunc.
2. Converting the entire DataFrame to a NumPy array, applying the gufunc on the whole array, and then concatenating the result back with the original DataFrame. A nice feature here is that you can utilize the axis keyword to do horizontal or vertical computations.

Polars map_batches is designed to apply a custom python function to a sequence of Series, it doesn't natively support applying a function to a 2D region out-of-the-box (afaik).

Is there a more efficient or graceful way to apply a gufunc over a 2D region of a Polars DataFrame?

Here is an example:

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

@nb.guvectorize([f"void(i8[:], i8, i8[:])"], "(n),()->(n)")
def foo(arr: np.ndarray, value: float, out: np.ndarray = None) -> None:
total = arr.dtype.type(0)
for i, ai in enumerate(arr):
total += ai + value
out = total

def apply_foo_batches(df: pl.DataFrame) -> pl.DataFrame:
"""Apply the gufunc via map_batches."""
return (
df.with_columns(
pl.struct(
pl.col("*").map_batches(lambda a: foo(a, 1, axis=0), return_dtype=pl.Int64))
.alias("foo"))
.with_columns(
pl.col("foo").name.prefix_fields("foo.")).unnest("foo")
)

def apply_foo_concat(df: pl.DataFrame) -> pl.DataFrame:
"""Apply the gufunc by converting to a NumPy array and concatenating the results."""
return (
pl.concat([
df,
pl.DataFrame(data=foo(df.to_numpy(), 1, axis=0),
schema=[f"foo.{c}" for c in df.columns])],
how='horizontal')
)


Example usage:

df = pl.DataFrame({"val_1": [1, 1, 1], "val_2": [10, 10, 10]})
print(df.pipe(apply_foo_batches))
print(df.pipe(apply_foo_concat))
# ┌───────┬───────┬───────────┬───────────┐
# │ val_1 ┆ val_2 ┆ foo.val_1 ┆ foo.val_2 │
# │ --- ┆ --- ┆ --- ┆ --- │
# │ i64 ┆ i64 ┆ i64 ┆ i64 │
# ╞═══════╪═══════╪═══════════╪═══════════╡
# │ 1 ┆ 10 ┆ 2 ┆ 11 │
# │ 1 ┆ 10 ┆ 4 ┆ 22 │
# │ 1 ┆ 10 ┆ 6 ┆ 33 │
# └───────┴───────┴───────────┴───────────┘


And some benchmarks for different DataFrame sizes (i7-5500U, Linux):

shape = (1_000, 50)
data = np.random.randint(1, 9, size=shape)
df = pl.DataFrame(data, schema=[f'c{i}' for i in range(shape[1])])
%timeit df.pipe(apply_foo_batches)
%timeit df.pipe(apply_foo_concat)
# 10.7 ms ± 40.9 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)
# 486 μs ± 2.74 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

shape = (10_000, 50)
data = np.random.randint(1, 9, size=shape)
df = pl.DataFrame(data, schema=[f'c{i}' for i in range(shape[1])])
%timeit df.pipe(apply_foo_batches)
%timeit df.pipe(apply_foo_concat)
# 11.3 ms ± 57.5 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)
# 1.85 ms ± 80.6 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

Continue reading...