如何在 polars 中调用带有两个位置参数的 numpy ufunc？

Question

我想调用一个在 polars 中有两个位置参数的 numpy 通用函数 (ufunc)。

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

以上尝试导致以下错误消息

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

还有其他方法可以执行此计算，例如，

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

...但我正在尝试将其与 polars.LazyFrame.

一起使用

我正在使用 polars 0.13.13 和 Python 3.8。

Answer 1

编辑：从 Polars 0.13.19 开始，apply 方法无需 Numpy item 方法即可将 Numpy 数据类型转换为 Polars 数据类型。

当您只需要将一列从 polars 传递到 ufunc 时（如您的示例所示），最简单的方法是在特定列上使用 apply 函数。

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   │
└──────┴────────┘

如果您需要将多列从 Polars 传递到 ufunc，请使用 struct 表达式和 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     │
└──────┴───────┴────────┘

一个注意事项：在未来的 Polars 版本中可能不再需要使用 numpy item 方法。（目前，apply 方法并不总是自动在 numpy dtypes 和 Polars dtypes 之间转换。）

这有帮助吗？

如何在 polars 中调用带有两个位置参数的 numpy ufunc？

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

python-polars