How can I call a numpy ufunc with two positional arguments in polars?

user_18696658 0
Viewed 192

I would like to call a numpy universal function (ufunc) that has two positional arguments in polars.

df.with_column(
    numpy.left_shift(pl.col('col1'), 8)
)

Above attempt results in the following error message

Traceback (most recent call last):
  File "<stdin>", line 2, in <module>
  File "/usr/local/lib/python3.8/dist-packages/polars/internals/expr.py", line 181, in __array_ufunc__
    out_type = ufunc(np.array([1])).dtype
TypeError: left_shift() takes from 2 to 3 positional arguments but 1 were given

There are other ways to perform this computation, e.g.,

df['col1'] = numpy.left_shift(df['col1'], 8)

... but I'm trying to use this with a polars.LazyFrame.

I'm using polars 0.13.13 and Python 3.8.

python-polars
1 Answers

Edit: as of Polars 0.13.19, the apply method converts Numpy datatypes to Polars datatypes without requiring the Numpy item method.

When you need to pass only one column from polars to the ufunc (as in your example), the easist method is to use the apply function on the particular column.

import numpy as np
import polars as pl

df = pl.DataFrame({"col1": [2, 4, 8, 16]}).lazy()

df.with_column(
    pl.col("col1").apply(lambda x: np.left_shift(x, 8).item()).alias("result")
).collect()

shape: (4, 2)
┌──────┬────────┐
│ col1 ┆ result │
│ ---  ┆ ---    │
│ i64  ┆ i64    │
╞══════╪════════╡
│ 2    ┆ 512    │
├╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┤
│ 4    ┆ 1024   │
├╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┤
│ 8    ┆ 2048   │
├╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┤
│ 16   ┆ 4096   │
└──────┴────────┘

If you need to pass multiple columns from Polars to the ufunc, then use the struct expression with apply.

df = pl.DataFrame({"col1": [2, 4, 8, 16], "shift": [1, 1, 2, 2]}).lazy()

df.with_column(
    pl.struct(["col1", "shift"])
    .apply(lambda cols: np.left_shift(cols["col1"], cols["shift"]).item())
    .alias("result")
).collect()

shape: (4, 3)
┌──────┬───────┬────────┐
│ col1 ┆ shift ┆ result │
│ ---  ┆ ---   ┆ ---    │
│ i64  ┆ i64   ┆ i64    │
╞══════╪═══════╪════════╡
│ 2    ┆ 1     ┆ 4      │
├╌╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┤
│ 4    ┆ 1     ┆ 8      │
├╌╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┤
│ 8    ┆ 2     ┆ 32     │
├╌╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┤
│ 16   ┆ 2     ┆ 64     │
└──────┴───────┴────────┘

One Note: the use of the numpy item method may no longer be needed in future releases of Polars. (Presently, the apply method does not always automatically translate between numpy dtypes and Polars dtypes.)

Does this help?

user_185598750
Related