如何在Polars DataFrame行间高效分配数值（非迭代方式）

我们需要用Polars实现一个金额分配逻辑：将df1中每个Name对应的总金额，按规则分配到df2的对应分组行中，要求不能遍历行，保证高效性。

初始数据定义 #

import polars as pl

df1 = pl.DataFrame({"Name": ["A", "B", "C"], "Amount": [100, 300, 250]})
df2 = pl.DataFrame({"Name": ["A", "A", "B", "B", "B", "C"], "Item": ["x1", "x2", "y1", "y2", "y3", "z"],"Price": [40, 60, 50, 150, 200, 400]})

df1（总金额表） #

shape: (3, 2)
┌──────┬────────┐
│ Name ┆ Amount │
│ ---  ┆ ---    │
│ str  ┆ i64    │
╞══════╪════════╡
│ A    ┆ 100    │
│ B    ┆ 300    │
│ C    ┆ 250    │
└──────┴────────┘

df2（待分配明细表） #

shape: (6, 3)
┌──────┬──────┬───────┐
│ Name ┆ Item ┆ Price │
│ ---  ┆ ---  ┆ ---   │
│ str  ┆ str  ┆ i64   │
╞══════╪══════╪═══════╡
│ A    ┆ x1   ┆ 40    │
│ A    ┆ x2   ┆ 60    │
│ B    ┆ y1   ┆ 50    │
│ B    ┆ y2   ┆ 150   │
│ B    ┆ y3   ┆ 200   │
│ C    ┆ z    ┆ 400   │
└──────┴──────┴───────┘

分配规则 #

• 同一Name下的行已按顺序排列，优先按每行的Price金额分配
• 若剩余金额不足以覆盖当前行的Price，仅分配剩余金额
• 最终需为df2新增Assigned amount列，记录每行实际分配的金额

期望结果 #

shape: (6, 5)
┌──────┬──────┬───────┬────────┬─────────────────┐
│ Name ┆ Item ┆ Price ┆ Amount ┆ Assigned amount │
│ ---  ┆ ---  ┆ ---   ┆ ---    ┆ ---             │
│ str  ┆ str  ┆ i64   ┆ i64    ┆ i64             │
╞══════╪══════╪═══════╪════════╪═════════════════╡
│ A    ┆ x1   ┆ 40    ┆ 100    ┆ 40              │
│ A    ┆ x2   ┆ 60    ┆ 100    ┆ 60              │
│ B    ┆ y1   ┆ 50    ┆ 300    ┆ 50              │
│ B    ┆ y2   ┆ 150    ┆ 300    ┆ 150             │
│ B    ┆ y3   ┆ 200    ┆ 300    ┆ 100             │
│ C    ┆ z    ┆ 400    ┆ 250    ┆ 250             │
└──────┴──────┴───────┴────────┴─────────────────┘

高效实现方案 #

直接使用Polars的窗口函数和分组计算，无需遍历行，代码如下：

result = (
    df2
    # 关联总金额表，将每个Name的总金额带到明细行
    .join(df1, on="Name", how="left")
    # 按Name分组，计算累计金额相关中间值
    .with_columns(
        # 当前行及之前的Price累计和
        cumulative_price=pl.col("Price").cumsum().over("Name"),
        # 前一行的累计和，初始填充为0
        prev_cumulative=pl.col("Price").cumsum().over("Name").shift(1).fill_null(0)
    )
    # 计算实际分配金额：取当前Price、剩余金额的较小值，且不小于0
    .with_columns(
        assigned_amount=pl.min_horizontal(
            pl.col("Price"),
            pl.col("Amount") - pl.col("prev_cumulative")
        ).clip(lower=0)
    )
    # 调整列顺序并命名，匹配期望结果
    .select(["Name", "Item", "Price", "Amount", "assigned_amount"])
    .rename({"assigned_amount": "Assigned amount"})
)

print(result)

代码说明 #

1. 关联总金额：通过join将df1的总金额关联到df2的每一行，确保每个明细行知晓所属分组的总预算
2. 计算累计值：用cumsum().over("Name")计算同组内的Price累计和，再通过shift(1)得到前一行的累计值，用于计算当前行可分配的剩余金额
3. 计算分配金额：用min_horizontal取当前Price和剩余金额的较小值，再用clip(lower=0)避免出现负数分配
4. 整理结果：调整列顺序和名称，得到最终目标输出

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