如何在Polars中简洁实现多when条件判断？

你可以通过循环或者functools.reduce简洁遍历字典的键值对，自动生成对应的pl.when().then()链式表达式，无需手动逐个编写分支。

方法一：循环实现（直观易读） #

适合大多数场景，逻辑清晰，便于调试：

import polars as pl

# 示例字典：键是要返回的字符串值，值是对应的Polars条件表达式
condition_dict = {
    "high": pl.col("score") > 80,
    "medium": pl.col("score").is_between(60, 80),
    "low": pl.col("score") < 60
}

# 初始化表达式
case_expr = None
for label, condition in condition_dict.items():
    if case_expr is None:
        # 第一个条件，创建初始when表达式
        case_expr = pl.when(condition).then(pl.lit(label))
    else:
        # 后续条件，链式拼接when-then
        case_expr = case_expr.when(condition).then(pl.lit(label))

# 添加otherwise分支并命名新列
case_expr = case_expr.otherwise(None).alias("level")

# 应用到DataFrame
df = df.with_columns(case_expr)

方法二：reduce实现（更简洁） #

借助functools.reduce将字典的键值对迭代拼接成完整表达式：

import polars as pl
from functools import reduce

condition_dict = {
    "high": pl.col("score") > 80,
    "medium": pl.col("score").is_between(60, 80),
    "low": pl.col("score") < 60
}

if not condition_dict:
    # 处理空字典的边界情况
    case_expr = pl.lit(None).alias("level")
else:
    # 取出第一个键值对作为初始表达式
    first_label, first_cond = next(iter(condition_dict.items()))
    # 用reduce拼接剩余的when-then分支
    case_expr = reduce(
        lambda acc, (label, cond): acc.when(cond).then(pl.lit(label)),
        list(condition_dict.items())[1:],
        pl.when(first_cond).then(pl.lit(first_label))
    )
    case_expr = case_expr.otherwise(None).alias("level")

df = df.with_columns(case_expr)

注意事项 #

• 条件顺序：Polars的when是按顺序匹配的，第一个满足的条件会返回对应值，因此要确保字典的键值对顺序符合你的优先级要求（Python 3.7+字典默认保留插入顺序）。
• 空字典处理：如果字典为空，直接返回None避免表达式构建报错。

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