如何在Python Polars中链式调用表达式列表中的when表达式？

问题场景 #

需要实现以下逻辑：

• 基于正则规则字典对data列进行匹配替换，生成new_data列
• 若data列未匹配任何正则规则，则使用ISO_codes列的值通过映射字典填充new_data
• 要求避免逐行循环，改用Polars原生的链式pl.when表达式实现

解决方案代码 #

import polars as pl

# 1. 定义示例数据、正则规则、ISO映射
df = pl.DataFrame({
    "data": ["US-123", "CA-456", "DE-789", "FR-012", "UNKNOWN-999"],
    "ISO_codes": ["USA", "CAN", "DEU", "FRA", "ESP"]
})

regex_rules = {
    r"^US-(.*)": "USA_$1",
    r"^CA-(.*)": "CAN_$1",
    r"^DE-(.*)": "DEU_$1"
}

iso_mapping = {
    "USA": "USA_DEFAULT",
    "CAN": "CAN_DEFAULT",
    "DEU": "DEU_DEFAULT",
    "FRA": "FRA_DEFAULT",
    "ESP": "ESP_DEFAULT"
}

# 2. 构建链式when表达式
# 先设置兜底逻辑：用ISO_codes映射值填充
expr = pl.col("ISO_codes").replace(iso_mapping).alias("new_data")

# 链式添加正则匹配规则：按顺序匹配，匹配成功则替换
for pattern, replacement in regex_rules.items():
    expr = pl.when(pl.col("data").str.contains(pattern)) \
             .then(pl.col("data").str.replace(pattern, replacement)) \
             .otherwise(expr)

# 3. 应用表达式生成结果
result_df = df.with_columns(expr)
print(result_df)

输出结果 #

shape: (5, 3)
┌──────────────┬───────────┬─────────────┐
│ data         ┆ ISO_codes ┆ new_data    │
│ ---          ┆ ---       ┆ ---         │
│ str          ┆ str       ┆ str         │
╞══════════════╪═══════════╪═════════════╡
│ US-123       ┆ USA       ┆ USA_123     │
│ CA-456       ┆ CAN       ┆ CAN_456     │
│ DE-789       ┆ DEU       ┆ DEU_789     │
│ FR-012       ┆ FRA       ┆ FRA_DEFAULT │
│ UNKNOWN-999  ┆ ESP       ┆ ESP_DEFAULT │
└──────────────┴───────────┴─────────────┘

逻辑说明 #

• 兜底逻辑优先初始化：先把expr设置为ISO映射后的结果，确保所有未匹配正则的行都能得到兜底值
• 链式正则匹配：遍历正则规则字典，每次用pl.when检查data列是否匹配当前正则，匹配成功则执行替换，否则保留之前的表达式结果（可能是之前匹配的正则结果，也可能是兜底值）
• 向量化效率：这种方式完全利用Polars的向量化操作，比逐行循环处理效率高几个数量级，尤其适合大数据集

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