如何在Polars中以指定列为键，用字典值替换空值？

错误原因 #

你遇到的KeyError: 2是因为：虽然when(pl.col('month').is_null())限制了替换的行范围，但map_elements是对整个item_id列执行Python函数，包括那些month不为null的行——这些行的item_id不在dict_months中，因此触发键不存在的错误。

修复方案 #

优先使用Polars的向量化操作替代逐行的map_elements，既高效又避免报错：

方案1：用replace结合fill_null（最简洁）

直接针对month列的null值，用item_id匹配字典中的值填充：

import polars as pl

def get_month(item_id: int):
    # In practice, fetch month from some DB ...
    return f'2024-{item_id:02.0f}'

df = pl.DataFrame({
    'item_id': [1, 2, 3, 4],
    'month': [None, '2023-07', None, '2023-08']
})

dict_months = {item_id: get_month(item_id) for item_id in df.filter(pl.col('month').is_null())['item_id']}

# 仅填充month列的null值，用字典匹配对应的item_id
df = df.with_columns(
    pl.col('month').fill_null(pl.col('item_id').replace(dict_months))
)

print(df)

方案2：用when/then配合replace

如果需要保留原有的条件判断逻辑，用replace替代map_elements：

df = df.with_columns(
    pl.when(pl.col('month').is_null())
    .then(pl.col('item_id').replace(dict_months))
    .otherwise(pl.col('month'))
    .alias('month')
)

replace是Polars的向量化操作，只会替换字典中存在的键，不存在的键会保留原item_id，但由于when仅选中month为null的行（这些行的item_id都在字典中），因此不会出现键不存在的问题。

方案3：给map_elements加容错（不推荐）

如果一定要用map_elements，可以在lambda中用get方法做容错处理（性能不如前两种向量化方案）：

df = df.with_columns(
    pl.when(pl.col('month').is_null())
    .then(pl.col('item_id').map_elements(lambda id: dict_months.get(id, id), return_dtype=pl.String))
    .otherwise(pl.col('month'))
    .alias('month')
)

内容的提问来源于stack exchange，提问作者Phil-ZXX