如何在Polars DataFrame中使用numpy arctan2函数？

Polars 0.16.x版本中，直接在with_columns里调用numpy函数会因为Polars列数据类型与numpy数组的适配问题报错，以下是两种可行的解决方法：

方法1：使用map_batches适配numpy函数 #

通过map_batches将Polars列转换为numpy数组完成计算，再转回Polars列：

import polars as pl
import numpy as np

# 示例DataFrame
df = pl.DataFrame({
    "u": [1.0, 0.0, -1.0, 0.0],
    "v": [0.0, 1.0, 0.0, -1.0]
})

# 计算wd列
df = df.with_columns(
    pl.map_batches(
        ["u", "v"],
        lambda cols: np.arctan2(cols[1], cols[0]),
        return_dtype=pl.Float64
    ).alias("wd")
)

print(df)

方法2：使用Polars原生arctan2函数 #

Polars内置了原生arctan2函数，无需依赖numpy，写法更简洁且性能更优：

import polars as pl

# 示例DataFrame
df = pl.DataFrame({
    "u": [1.0, 0.0, -1.0, 0.0],
    "v": [0.0, 1.0, 0.0, -1.0]
})

# 计算wd列（参数顺序与numpy一致：arctan2(y, x)）
df = df.with_columns(
    pl.arctan2(pl.col("v"), pl.col("u")).alias("wd")
)

print(df)

常见失败写法的原因 #

• 直接在Polars表达式上调用numpy函数：np.arctan2(pl.col("v"), pl.col("u"))，numpy无法识别Polars的表达式对象，只能处理实际数值数组，因此报错。
• 用apply逐元素处理：pl.col("v").apply(lambda x: np.arctan2(x, pl.col("u")))，apply仅能单元素处理，无法同时获取另一列的对应元素，逻辑不成立。

内容的提问来源于stack exchange，提问作者bgk