如何用Polars基于系数、数值和嵌套列表计算加权值？

需求说明 #

需要将包含列表列的Polars DataFrame按列表元素拆分，同时根据给定的系数字典，结合每行对应元素的系数总和重新计算value列的值，全程保留Decimal类型精度，避免浮点数误差。

解决方案代码 #

import polars as pl
from decimal import Decimal

# 原DataFrame
df = pl.from_dicts(
    [
        {"id": 1, "value": Decimal("100.0"), "items": ["A"]},
        {"id": 2, "value": Decimal("150.000"), "items": ["A", "B"]},
        {"id": 3, "value": Decimal("70.0000"), "items": ["A", "B", "C"]},
    ]
)

# 系数字典
coef = {"A": Decimal("0.2"), "B": Decimal("0.35"), "C": Decimal("0.45")}

# 转换系数字典为Polars DataFrame
coef_df = pl.from_dicts([{"item": k, "coef": v} for k, v in coef.items()])

# 核心计算流程
result = (
    df
    # 展开items列表为单行
    .explode("items")
    # 关联系数表，获取每个item对应的系数
    .join(coef_df, left_on="items", right_on="item")
    # 按id分组，计算组内系数总和
    .with_columns(
        pl.col("coef").sum().over("id").alias("coef_sum")
    )
    # 计算新value：(系数/组内系数总和) * 原value
    .with_columns(
        ((pl.col("coef") / pl.col("coef_sum")) * pl.col("value")).alias("value")
    )
    # 保留目标列并调整顺序
    .select(["id", "value", "item"])
)

print(result)

代码解释 #

1. 展开列表列：用explode("items")将每个id对应的items列表拆分为多行，每个元素单独占一行。
2. 关联系数表：将原DataFrame与系数DataFrame通过item字段关联，获取每个元素对应的系数值。
3. 计算组内系数总和：借助over("id")窗口函数，按id分组计算该组内所有元素的系数总和。
4. 计算新value值：按照公式(系数/组内系数总和)*原value计算新的value，全程使用Decimal类型保证精度。
5. 整理结果列：选择并调整id、value、item列的顺序，得到目标结构。

输出结果 #

┌─────┬────────────────┬──────┐
│ id  ┆ value          ┆ item │
│ --- ┆ ---            ┆ ---  │
│ i64 ┆ decimal[*,10]  ┆ str  │
╞═════╪════════════════╪══════╡
│ 1   ┆ 100.0000000000 ┆ A    │
│ 2   ┆ 54.5454000000  ┆ A    │
│ 2   ┆ 95.4544500000  ┆ B    │
│ 3   ┆ 14.0000000000  ┆ A    │
│ 3   ┆ 24.5000000000  ┆ B    │
│ 3   ┆ 31.5000000000  ┆ C    │
└─────┴────────────────┴──────┘

内容的提问来源于stack exchange，提问作者scūriolus