用字典中的值替换极坐标中的空值（使用其他列作为键）

Question

我有这个代码

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']}

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

基本上，我想使用

null

作为键，将

month

列中的所有

dict_months

条目替换为

item_id

中的值。我可以假设该字典包含所有缺失 id 的键（通过构造），但不包含其他 id。

当我运行上面的代码时，我收到错误

PanicException: python function failed KeyError: 2

，这似乎意味着 Polars 正在尝试查找 id

2

的值，但它不应该这样做，因为 id 2 有一个月。

如何解决这个问题？

Answer 1

这是使用

replace_strict

和

default

参数的一种方法。我还将

when/then

替换为

coalesce

:

print(
    df.with_columns(
        pl.coalesce(
            "month", pl.col("item_id").replace_strict(dict_months, default=None)
        )
    )
)

输出：

shape: (4, 2)
┌─────────┬─────────┐
│ item_id ┆ month   │
│ ---     ┆ ---     │
│ i64     ┆ str     │
╞═════════╪═════════╡
│ 1       ┆ 2024-01 │
│ 2       ┆ 2023-07 │
│ 3       ┆ 2024-03 │
│ 4       ┆ 2023-08 │
└─────────┴─────────┘

完整代码：

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"]
}

print(
    df.with_columns(
        pl.coalesce(
            "month", pl.col("item_id").replace_strict(dict_months, default=None)
        )
    )
)