如何在Polars中实现类似Pandas.apply的输出DataFrame多列生成方式？

核心思路 #

要实现和 Pandas 中 out_df['COL'] = df.apply(...) 一致的逐列赋值逻辑，同时高效处理field空值时用field2替代的需求，Polars 可以通过向量化空值处理+类Pandas赋值语法+自定义函数/表达式结合来完成，同时保证性能优于 Pandas 的apply。

代码实现示例 #

1. 基础场景：自定义函数 + 类Pandas赋值

假设你有输入input_df（含field、field2列），需要生成PIN、USER_CODE两个列，且自定义逻辑需要用Python函数实现：

import polars as pl

# 自定义生成PIN的函数
def create_pin(val):
    return val[:4].upper() if val else ""

# 自定义生成用户码的函数
def create_user_code(val):
    return val[-4:].lower() if val else ""

# 初始化空输出DataFrame（和Pandas逻辑一致）
out_df = pl.DataFrame()

# 生成PIN列：先通过coalesce处理空值替代，再映射自定义函数
out_df["PIN"] = input_df.select(
    pl.coalesce("field", "field2").map_elements(create_pin, return_dtype=pl.Utf8)
).to_series()

# 生成USER_CODE列：复用相同的空值处理逻辑
out_df["USER_CODE"] = input_df.select(
    pl.coalesce("field", "field2").map_elements(create_user_code, return_dtype=pl.Utf8)
).to_series()

2. 高性能优化：用Polars向量化表达式替代自定义函数

如果你的自定义逻辑可以用Polars内置表达式实现，完全不要用map_elements，向量化操作的速度会比逐行apply/map_elements快几个数量级：

import polars as pl

# 直接用Polars字符串表达式实现逻辑，无需自定义函数
out_df = pl.DataFrame()

# 生成PIN列：coalesce处理空值 + 向量化字符串操作
out_df["PIN"] = input_df.select(
    pl.coalesce("field", "field2").str.slice(0, 4).str.to_upper()
).to_series()

# 生成USER_CODE列
out_df["USER_CODE"] = input_df.select(
    pl.coalesce("field", "field2").str.slice(-4, 4).str.to_lower()
).to_series()

3. 更贴合Polars风格的链式写法（可选）

如果习惯Polars的链式操作，也可以一次性生成所有列，性能更优：

out_df = (
    input_df
    # 先生成处理好空值的中间列
    .with_columns(pl.coalesce("field", "field2").alias("filled_field"))
    # 一次性生成所有目标列
    .with_columns(
        PIN=pl.col("filled_field").str.slice(0, 4).str.to_upper(),
        USER_CODE=pl.col("filled_field").str.slice(-4, 4).str.to_lower()
    )
    # 丢弃中间列
    .drop("filled_field")
)

关键知识点说明 #

• pl.coalesce：高效空值替代：Polars内置的coalesce函数会自动取传入列中第一个非空的值，完美替代field为空时用field2填充的逻辑，比Python函数判断空值快得多。
• 类Pandas赋值语法：Polars支持out_df["COL"] = series的赋值方式，和Pandas完全一致，满足你的使用习惯。
• 优先用向量化表达式：Polars的内置表达式（如str.slice、str.to_upper）都是向量化实现，避免了逐行处理的开销，这也是Polars比Pandas快的核心原因。如果必须用自定义逻辑，再用map_elements。

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