如何在Polars链式调用中内联调用含两个pl.Series参数的UDF

要把你的prefix_cat_based_on_num_order函数直接塞进with_columns实现链式代码，核心是把接收pl.Series的UDF包装成Polars可识别的表达式（pl.Expr）——因为with_columns只接受表达式对象，无法直接传入普通的Series处理函数。

具体实现步骤 #

假设你的UDF逻辑如下（示例逻辑，可替换为你的实际代码）：

import polars as pl

def prefix_cat_based_on_num_order(s1: pl.Series, s2: pl.Series) -> pl.Series:
    # 示例逻辑：根据s2的排序给s1添加递增前缀，最后恢复原顺序
    sorted_indices = s2.argsort()
    prefixed_vals = s1[sorted_indices].str.concat("_", pl.Series(range(1, len(s1)+1)))
    return prefixed_vals.sort_by(sorted_indices)

原来的拆分写法

# 先构建DataFrame
df = pl.DataFrame({"col1": ["a", "b", "c"], "col2": [3, 1, 2]})
# 单独调用UDF生成新列
new_col = prefix_cat_based_on_num_order(df["col1"], df["col2"])
# 再合并到DataFrame
df = df.with_columns(new_col.alias("prefixed_col"))

改进后的链式内联写法

先写一个包装函数，把你的UDF转换成Polars表达式：

def prefix_cat_expr(col1_name: str, col2_name: str) -> pl.Expr:
    # 用pl.struct打包目标列，通过map_batches批量传入UDF
    return pl.struct([col1_name, col2_name]).map_batches(
        lambda batch_struct: prefix_cat_based_on_num_order(
            batch_struct[col1_name], 
            batch_struct[col2_name]
        )
    )

之后就能直接在链式调用里内联使用了：

df = (
    pl.DataFrame({"col1": ["a", "b", "c"], "col2": [3, 1, 2]})
    .with_columns(prefix_cat_expr("col1", "col2").alias("prefixed_col"))
)

# 输出结果
print(df)
# shape: (3, 3)
# ┌──────┬──────┬──────────────┐
# │ col1 ┆ col2 ┆ prefixed_col │
# │ ---  ┆ ---  ┆ ---          │
# │ str  ┆ i64  ┆ str          │
# ╞══════╪══════╪══════════════╡
# │ a    ┆ 3    ┆ c_3          │
# │ b    ┆ 1    ┆ b_1          │
# │ c    ┆ 2    ┆ a_2          │
# └──────┴──────┴──────────────┘

原理说明 #

• pl.struct([col1_name, col2_name])将指定两列打包为结构体列，实现一次性批量传入map_batches
• map_batches是对整列批量处理，而非逐行操作，不会损失原UDF的性能优势
• 包装后的prefix_cat_expr返回pl.Expr对象，完全符合with_columns的参数要求

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