如何在Polars DataFrame中为字典未匹配项设置新列默认值

环境 #

• Python 3.12.3
• Polars 1.8.2
• Polars-lts-cpu: 1.10.0
• OS: Linux-lite 24.04 虚拟机

问题描述 #

现有如下数据集和字典：

import polars as pl

countries = ['usa', 'france', 'japan', 'brazil', 'new_zealand']
calling_codes = [1, 33, 81, 55, 64]

df = pl.DataFrame({'country': countries, 'calling_code': calling_codes })

capitals_dict = {'usa':'washington_dc', 'france': 'paris', 'brazil': 'brasilia'}

需要新增capital列：当country值存在于capitals_dict的键中时，填入对应首都；不存在则填入默认值[default]。使用replace方法时，未匹配的行保留了原国家名称，不符合需求。

解决方案 #

方法1：使用replace的default参数 #

replace方法支持通过default参数指定未匹配项的默认值，直接修改原有代码即可：

df.with_columns(capital = pl.col('country').replace(capitals_dict, default="[default]"))

输出结果：

shape: (5, 3)
┌─────────────┬──────────────┬───────────────┐
│ country     ┆ calling_code ┆ capital       │
│ ---         ┆ ---          ┆ ---           │
│ str         ┆ i64          ┆ str           │
╞═════════════╪══════════════╪═══════════════╡
│ usa         ┆ 1            ┆ washington_dc │
│ france      ┆ 33           ┆ paris         │
│ japan       ┆ 81           ┆ [default]     │
│ brazil      ┆ 55           ┆ brasilia      │
│ new_zealand ┆ 64           ┆ [default]     │
└─────────────┴──────────────┴───────────────┘

方法2：使用map_dict方法 #

map_dict是专门用于字典映射的方法，同样支持设置默认值：

df.with_columns(capital = pl.col('country').map_dict(capitals_dict, default="[default]"))

方法3：使用条件表达式when/then/otherwise #

通过条件判断实现映射逻辑，适合更复杂的场景：

df.with_columns(
    capital = pl.when(pl.col('country').is_in(capitals_dict.keys()))
               .then(pl.col('country').replace(capitals_dict))
               .otherwise("[default]")
)

说明 #

• 方法1和方法2更简洁，适合简单的字典映射场景；
• 方法3灵活性更高，可扩展添加更多条件判断逻辑。

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