The first step would be to check if your task can be solved natively using Polars Expressions.

If a custom function is neccessary, .map_elements() can be used to apply one on a row by row basis.

To pass in values from multiple columns, you can utilize the Struct data type.

e.g. with pl.struct()

Copy>>> df.select(pl.struct(pl.all())) # all columns
shape: (3, 1)
┌───────────┐
│ foo       │
│ ---       │
│ struct[3] │
╞═══════════╡
│ {1,4,7}   │
│ {2,5,8}   │
│ {3,6,9}   │
└───────────┘

Using pl.struct(...).map_elements will pass the values to the custom function as a dict argument.

Copydef my_complicated_function(row: dict) -> int:
    """
    A function that cannot utilize polars expressions.
    This should be avoided.
    """

    # a dict with column names as keys
    print(f"[DEBUG]: {row=}")
    
    # do some work
    return row["foo"] + row["bar"] + row["baz"]


df = pl.DataFrame({
    "foo": [1, 2, 3], 
    "bar": [4, 5, 6], 
    "baz": [7, 8, 9]
})

df = df.with_columns(
    pl.struct(pl.all())
      .map_elements(my_complicated_function, return_dtype=pl.Int64)
      .alias("foo + bar + baz")
)

Copy# [DEBUG]: row={'foo': 1, 'bar': 4, 'baz': 7}
# [DEBUG]: row={'foo': 2, 'bar': 5, 'baz': 8}
# [DEBUG]: row={'foo': 3, 'bar': 6, 'baz': 9}

Copyshape: (3, 4)
┌─────┬─────┬─────┬─────────────────┐
│ foo ┆ bar ┆ baz ┆ foo + bar + baz │
│ --- ┆ --- ┆ --- ┆ ---             │
│ i64 ┆ i64 ┆ i64 ┆ i64             │
╞═════╪═════╪═════╪═════════════════╡
│ 1   ┆ 4   ┆ 7   ┆ 12              │
│ 2   ┆ 5   ┆ 8   ┆ 15              │
│ 3   ┆ 6   ┆ 9   ┆ 18              │
└─────┴─────┴─────┴─────────────────┘

You can use the expression syntax to select all columns with pl.all() and then map_batches the numpy np.log2(..) function over the columns.

df.select(
    pl.all().map_batches(np.log2)
)

Note that we choose map_batches here as map_elements would call the function upon each value.

map_elements = pl.Series(np.log2(value) for value in pl.Series([1, 2, 3]))

But np.log2 can be called once with multiple values, which would be faster.

map_batches = np.log2(pl.Series([1, 2, 3]))

See the User guide for more.

• map_elements: Call a function separately on each value in the Series.
• map_batches: Always passes the full Series to the function.

Numpy

Polars expressions also support numpy universal functions.

That means you can pass a polars expression to a numpy ufunc:

df.select(
    np.log2(pl.all())
)