根据另一个数据帧对一个数据帧的列进行求和

Question

我有两个看起来像这样的数据框：

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

我想对第二个 df (df2) 的每一年列进行求和，仅考虑 df1 中对应年份为同一年或之后的名称。所需输出：

┌──────┬──────┐
│ Year ┆  Sum │
╞══════╪══════╡
│ 2001 ┆  111 │  
│ 2002 ┆  221 │
│ 2003 ┆  996 │  (= 331 + 332 + 333)
│ 2004 ┆ 1770 │  (= 441 + 442 + 443 + 444)
└──────┴──────┘

我是 Polars 新手（来自 Pandas），我不知道该怎么做。任何帮助将不胜感激。

Answer 1

您可以

unpivot

df2

将其放在与

df1

更“兼容”的形式中。它将指定的列标题“逆透视”到行中。此时，

Year

应从

int

:

 投射到 

str

列中

df2.unpivot(index='Name', variable_name='Year').with_columns(
    pl.col('Year').cast(pl.Int64)
)

shape: (16, 3)
┌──────┬──────┬───────┐
│ Name ┆ Year ┆ value │
│ ---  ┆ ---  ┆ ---   │
│ str  ┆ i64  ┆ i64   │
╞══════╪══════╪═══════╡
│ A    ┆ 2001 ┆ 111   │
│ B    ┆ 2001 ┆ 112   │
│ C    ┆ 2001 ┆ 113   │
│ D    ┆ 2001 ┆ 114   │
│ A    ┆ 2002 ┆ 221   │
│ B    ┆ 2002 ┆ 222   │
│ C    ┆ 2002 ┆ 223   │
│ D    ┆ 2002 ┆ 224   │
│ A    ┆ 2003 ┆ 331   │
│ B    ┆ 2003 ┆ 332   │
│ C    ┆ 2003 ┆ 333   │
│ D    ┆ 2003 ┆ 334   │
│ A    ┆ 2004 ┆ 441   │
│ B    ┆ 2004 ┆ 442   │
│ C    ┆ 2004 ┆ 443   │
│ D    ┆ 2004 ┆ 444   │
└──────┴──────┴───────┘

然后，您可以使用

join

上的

df1

Name

，并按您指定的年份条件过滤结果行：

(...).join(df1, on='Name').filter(pl.col('Year') >= pl.col('Year_right'))

从那里开始，在

group_by

上有一个简单的

agg

和

Year

区域。大家一起：

df2.unpivot(index='Name', variable_name='Year').with_columns(
    pl.col('Year').cast(pl.Int64)
).join(df1, on='Name').filter(pl.col('Year') >= pl.col('Year_right')).group_by(
    'Year', maintain_order=True
).agg(
    Sum=pl.col('value').sum()
)

shape: (4, 2)
┌──────┬──────┐
│ Year ┆ Sum  │
│ ---  ┆ ---  │
│ i64  ┆ i64  │
╞══════╪══════╡
│ 2001 ┆ 111  │
│ 2002 ┆ 221  │
│ 2003 ┆ 996  │
│ 2004 ┆ 1770 │
└──────┴──────┘

（

maintain_order

是可选的，但速度较慢，如果确实需要，可以在

.sort('Year')

之后使用

agg

达到相同的效果。）

Answer 2

我希望我正确理解了你的问题：

df2 = df2.join(df1, on='Name')

out = {}
for c in df2.select(pl.col('^\d+$')):
    m = int(c.name) >= df2['Year']
    out[c.name] = c.filter(m).sum()

out = pl.DataFrame(out)
print(out)

打印：

shape: (1, 4)
┌──────┬──────┬──────┬──────┐
│ 2001 ┆ 2002 ┆ 2003 ┆ 2004 │
│ ---  ┆ ---  ┆ ---  ┆ ---  │
│ i64  ┆ i64  ┆ i64  ┆ i64  │
╞══════╪══════╪══════╪══════╡
│ 111  ┆ 221  ┆ 996  ┆ 1770 │
└──────┴──────┴──────┴──────┘

---

或转置输出：

out = out.transpose(
    include_header=True, header_name="Year", column_names=["Sum"]
)
print(out)

打印：

shape: (4, 2)
┌──────┬──────┐
│ Year ┆ Sum  │
│ ---  ┆ ---  │
│ str  ┆ i64  │
╞══════╪══════╡
│ 2001 ┆ 111  │
│ 2002 ┆ 221  │
│ 2003 ┆ 996  │
│ 2004 ┆ 1770 │
└──────┴──────┘