基于另一DataFrame对Polars DataFrame年份列求和的方法

步骤拆解 #

要实现需求，核心是先把宽格式的df2转为长格式，再结合df1的条件过滤后分组求和，具体步骤如下：

1. 将df2宽表转长表
   把df2中以年份命名的列转为行数据，同时将年份列名转为整数类型（方便后续数值比较）：

   df2_long = df2.melt(
       id_vars="Name",
       variable_name="Year",
       value_name="Value"
   ).with_columns(pl.col("Year").cast(pl.Int32))
   

2. 合并两个DataFrame
   按Name列将df1和转长后的df2关联：

   merged = df1.join(df2_long, on="Name", how="inner")
   

3. 过滤符合条件的行
   只保留df1的Year ≤ 当前行的Year的记录：

   filtered = merged.filter(pl.col("Year_right") >= pl.col("Year_left"))
   

4. 分组求和并整理结果
   按目标年份（Year_right）分组，对Value求和，最后重命名列并排序：

   result = filtered.group_by("Year_right").agg(
       pl.col("Value").sum().alias("Sum")
   ).rename({"Year_right": "Year"}).sort("Year")
   

完整代码 #

import polars as pl

df1 = pl.DataFrame(
    {
        "Name": ["A", "B", "C", "D"],
        "Year": [2001, 2003, 2003, 2004]
    }
)

df2 = pl.DataFrame(
    {
        "Name": ["A", "B", "C", "D"],
        "2001": [111, 112, 113, 114],
        "2002": [221, 222, 223, 224],
        "2003": [331, 332, 333, 334],
        "2004": [441, 442, 443, 444]
    }
)

# 宽转长并转换年份类型
df2_long = df2.melt(id_vars="Name", variable_name="Year", value_name="Value").with_columns(pl.col("Year").cast(pl.Int32))

# 合并+过滤+分组求和
result = (
    df1.join(df2_long, on="Name")
    .filter(pl.col("Year_right") >= pl.col("Year_left"))
    .group_by("Year_right")
    .agg(pl.col("Value").sum().alias("Sum"))
    .rename({"Year_right": "Year"})
    .sort("Year")
)

print(result)

输出验证 #

运行后会得到和期望一致的结果：

┌──────┬──────┐
│ Year ┆ Sum  │
╞══════╪══════╡
│ 2001 ┆ 111  │
│ 2002 ┆ 221  │
│ 2003 ┆ 996  │
│ 2004 ┆ 1770 │
└──────┴──────┘

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