能否用Polars的.replace()替代map_elements处理嵌套账单字典映射？

问题背景 #

现有一个存储账单数据的嵌套字典billing_data，需要将其中的usage字段映射到Polars DataFrame的billing_cycle列对应行中。当前使用map_elements实现，但该方法性能较低，希望用更高效的.replace()替代。

原实现代码：

import polars as pl
import datetime

billing_data = {
    datetime.date(2022, 3, 11): {'usage': 896, 'amount': 119.7}, 
    datetime.date(2022, 4, 12): {'usage': 775, 'amount': 229.49}, 
    datetime.date(2022, 5, 12): {'usage': 1046, 'amount': 135.8}, 
    datetime.date(2022, 6, 13): {'usage': 1280, 'amount': 162.6}
}

df = pl.DataFrame({
    "billing_cycle": [datetime.date(2022, 3, 11), datetime.date(2022, 3, 12), datetime.date(2022, 3, 13), datetime.date(2022, 3, 14)]
})

# 性能较低的map_elements实现
df.with_columns(pl.col("billing_cycle").map_elements(lambda x: billing_data[x]["usage"]).alias("billing_usage"))

解决方案：使用.replace()替代 #

完全可以用.replace()实现，且性能会大幅提升——map_elements是逐行处理的Python级循环，而.replace()是Polars的矢量化操作，效率更高。

步骤如下：

1. 将嵌套字典转换为键为日期、值为对应usage的扁平字典
2. 调用.replace()完成映射，可通过default参数指定无匹配值时的默认结果

# 转换嵌套字典为扁平映射
usage_map = {date: data['usage'] for date, data in billing_data.items()}

# 矢量化replace实现映射
df = df.with_columns(
    pl.col("billing_cycle").replace(usage_map, default=None).alias("billing_usage")
)

扩展：批量映射多字段 #

如果需要同时映射usage和amount，可以分别构建映射字典批量处理：

# 构建两个字段的映射字典
usage_map = {date: data['usage'] for date, data in billing_data.items()}
amount_map = {date: data['amount'] for date, data in billing_data.items()}

df = df.with_columns(
    pl.col("billing_cycle").replace(usage_map, default=None).alias("billing_usage"),
    pl.col("billing_cycle").replace(amount_map, default=None).alias("billing_amount")
)

若数据量较大或字段较多，更推荐将字典转为Polars DataFrame后做左连接，性能更优：

# 将账单字典转为DataFrame
billing_df = pl.DataFrame(
    [(date, data['usage'], data['amount']) for date, data in billing_data.items()],
    schema=["billing_cycle", "billing_usage", "billing_amount"]
)

# 左连接完成映射
df = df.join(billing_df, on="billing_cycle", how="left")

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