如何在Polars的with_columns中用Lambda引用多列并优雅处理空值

先明确场景：我们需要对DataFrame的多列执行含空值判断的复杂逻辑，先给出示例数据、错误实现、预期结果，再说明优化方案。

示例DataFrame #

import polars as pl

df = pl.DataFrame({
    "a": [1, None, 3],
    "b": [None, 2, 4]
})

原错误实现及问题 #

假设我们想实现：当a和b都非空时返回两列之和，否则返回-1。错误的with_columns写法要么无法正确引用多列，要么代码冗余易出错，比如：

# 错误写法：直接用map_elements处理多列，依赖列顺序且可读性差
df_bad = df.with_columns(
    pl.col(["a", "b"]).map_elements(lambda x: x[0] + x[1] if x[0] is not None and x[1] is not None else -1).alias("sum_ab")
)

预期结果 #

shape: (3, 3)
┌──────┬──────┬────────┐
│ a    ┆ b    ┆ sum_ab │
│ ---  ┆ ---  ┆ ---    │
│ i64  ┆ i64  ┆ i64    │
╞══════╪══════╪════════╡
│ 1    ┆ null ┆ -1     │
│ null ┆ 2    ┆ -1     │
│ 3    ┆ 4    ┆ 7      │
└──────┴──────┴────────┘

自行实现的解法 #

通常会用struct包装多列后处理：

df_self_solved = df.with_columns(
    pl.struct(["a", "b"]).map_elements(lambda s: s["a"] + s["b"] if s["a"] is not None and s["b"] is not None else -1).alias("sum_ab")
)

更优雅的实现方式 #

1. 用pl.lambda_.map（Polars 0.19+） #

直接将列名列表传入，lambda参数按顺序对应列，无需包装struct，代码更简洁直观：

df_elegant = df.with_columns(
    pl.lambda_.map(lambda a, b: a + b if a is not None and b is not None else -1, ["a", "b"]).alias("sum_ab")
)

2. 用pl.apply #

pl.apply支持直接接收多列，lambda可通过参数名或索引引用列值，写法和pandas的apply逻辑对齐：

# 通过参数名引用（可读性更强）
df_apply1 = df.with_columns(
    pl.apply(["a", "b"], lambda a, b: a + b if a is not None and b is not None else -1).alias("sum_ab")
)

# 通过索引引用
df_apply2 = df.with_columns(
    pl.apply(["a", "b"], lambda row: row[0] + row[1] if row[0] is not None and row[1] is not None else -1).alias("sum_ab")
)

3. 优先用矢量化表达式（性能最优） #

如果逻辑能拆解为Polars内置操作，完全不需要lambda，性能远高于逐行处理：

df_vectorized = df.with_columns(
    pl.when((pl.col("a").is_not_null()) & (pl.col("b").is_not_null()))
      .then(pl.col("a") + pl.col("b"))
      .otherwise(-1)
      .alias("sum_ab")
)

在with_columns中像apply那样引用多列的方法 #

核心有三种可靠方式：

• pl.lambda_.map：将目标列名列表作为第二个参数传入，lambda的参数按顺序对应列，直接用参数名引用（如lambda a, b: ...中a对应列"a"）。
• pl.struct+map_elements：把多列包装成struct，lambda里通过键名引用列（如s["a"]），列名较多时更清晰。
• pl.apply：直接传入列名列表，lambda可通过参数名或索引引用每一列的值，写法贴近传统apply逻辑。

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