Polars中对两列应用行级自定义函数的语法是什么？

Polars的设计核心是列式向量化计算，因此优先推荐向量化实现（性能远高于逐行处理），仅当自定义函数无法适配向量化逻辑时，才考虑行级逐行操作。

一、向量化实现（优先选择） #

如果你的自定义函数可以接收Polars列对象（而非单个标量值），直接在with_columns中调用即可：

示例代码：

import polars as pl

# 自定义向量化函数：接收两个列对象，返回计算后的列
def vectorized_calc(col1, col2):
    return col1 * 2 + col2  # 替换为你的业务逻辑

# 构造示例DataFrame
df = pl.DataFrame({"col_1": [1, 2, 3], "col_2": [4, 5, 6]})

# 应用函数并生成新列
df = df.with_columns(
    col_3=vectorized_calc(pl.col("col_1"), pl.col("col_2"))
)

如果函数原本是基于标量的，可通过map_batches适配为向量化操作：

def scalar_func(a, b):
    return a + b * 3

df = df.with_columns(
    col_3=pl.struct(["col_1", "col_2"]).map_batches(
        lambda s: scalar_func(s.struct.field("col_1"), s.struct.field("col_2")),
        return_dtype=pl.Int64
    )
)

二、行级逐行处理（仅在必要时使用） #

当无法实现向量化时，可通过以下两种方式逐行处理：

方法1：使用map_elements #

def row_level_func(a, b):
    return a * b + 10  # 你的自定义行级逻辑

df = df.with_columns(
    col_3=pl.struct(["col_1", "col_2"]).map_elements(
        lambda row: row_level_func(row["col_1"], row["col_2"]),
        return_dtype=pl.Int64
    )
)

方法2：使用apply（行级） #

df = df.with_columns(
    col_3=df.apply(
        lambda row: row_level_func(row["col_1"], row["col_2"]),
        axis=1
    ).alias("col_3")
)

注意：逐行处理的性能远低于向量化操作，处理大数据集时强烈建议优先优化为向量化实现。

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