Polars: 在固定点扩展窗口

Question

我有一个叫做 Polars 的数据表，里面有三列：组别、日期和数值。我的目标是计算每个组在每年第一次出现的日期之前（包括这一天）的数值累加和。

比如，对于下面这个示例数据表：

import polars as pl

df = pl.DataFrame(
    {
        "date": [
            "2020-03-01",
            "2020-05-01",
            "2020-11-01",
            "2021-01-01",
            "2021-02-03",
            "2021-06-08",
            "2022-01-05",
            "2020-07-01",
            "2020-09-01",
            "2022-01-05",
            "2023-02-04",
        ],
        "group": [1, 1, 1, 1, 1, 1, 1, 2, 2, 2, 2],
        "value": [1, 2, 3, 4, 5, 6, 7, 1, 2, 3, 4],
    },
).with_columns(pl.col("date").str.strptime(pl.Date))

我想要的结果是：

┌────────────┬───────┬───────┐
│ date       ┆ group ┆ value │
│ ---        ┆ ---   ┆ ---   │
│ date       ┆ i64   ┆ i64   │
╞════════════╪═══════╪═══════╡
│ 2020-03-01 ┆ 1     ┆ 1     │
│ 2021-01-01 ┆ 1     ┆ 10    │
│ 2022-01-05 ┆ 1     ┆ 28    │
│ 2020-07-01 ┆ 2     ┆ 1     │
│ 2022-01-05 ┆ 2     ┆ 6     │
│ 2023-02-04 ┆ 2     ┆ 10    │
└────────────┴───────┴───────┘

简单来说，就是在每年的第一天，计算从开始到（包括）这个特定日期的数值累加和，分别针对每个组。

我试过 group_by_dynamic 和 rolling，但还是找不到一个简单明了的方法来解决这个问题。

欢迎任何想法。谢谢！

Answer 1

你可能需要把这个问题分成两个步骤来解决：

1. 计算每个条目的累计值和年份
2. 找出每年中的第一个日期（放到一个单独的数据框中）
3. 把这两个结果根据组和年份合并成一个综合结果。

下面是一种实现这个过程的方法：

import polars as pl

df = pl.DataFrame(
    {
        "date": [
            "2020-03-01",
            "2020-05-01",
            "2020-11-01",
            "2021-01-01",
            "2021-02-03",
            "2021-06-08",
            "2022-01-05",
            "2020-07-01",
            "2020-09-01",
            "2022-01-05",
            "2023-02-04",
        ],
        "group": [1, 1, 1, 1, 1, 1, 1, 2, 2, 2, 2],
        "value": [1, 2, 3, 4, 5, 6, 7, 1, 2, 3, 4],
    },
).with_columns([
    pl.col("date").str.strptime(pl.Date)
])

# Calculate the cumulative sum of "value" for each group
df = df.with_columns(
    [
        pl.col("value").cum_sum().over("group").alias("cumulative_value"),
        pl.col("date").dt.year().alias("year")
    ]
)


# Identify the first date of each year for each group
first_date_each_year = df.group_by(["group", "year"]).agg(pl.col("date").min().alias("first_date_of_year"))

print(first_date_each_year)

# Join to get the cumulative sum at the first date of each year for each group
result_df = df.join(
    first_date_each_year,
    on=["group", "year"],
    how="inner"
).filter(pl.col("date") == pl.col("first_date_of_year"))

# Select and rename columns as needed
result_df = result_df.select(
    [
        pl.col("first_date_of_year").alias("date"),
        "group",
        pl.col("cumulative_value").alias("value")
    ]
)

print(result_df)

这样就能得到你想要的结果：

┌────────────┬───────┬───────┐
│ date       ┆ group ┆ value │
│ ---        ┆ ---   ┆ ---   │
│ date       ┆ i64   ┆ i64   │
╞════════════╪═══════╪═══════╡
│ 2020-03-01 ┆ 1     ┆ 1     │
│ 2021-01-01 ┆ 1     ┆ 10    │
│ 2022-01-05 ┆ 1     ┆ 28    │
│ 2020-07-01 ┆ 2     ┆ 1     │
│ 2022-01-05 ┆ 2     ┆ 6     │
│ 2023-02-04 ┆ 2     ┆ 10    │
└────────────┴───────┴───────┘

Answer 2

一种“简单”的方法是先计算出完整的 cum_sum，然后再保留每个组的第一行年度数据。

df.with_columns(
   cum_sum = pl.col("value").cum_sum().over("group")
).filter(
   pl.col("date").dt.year().is_first_distinct().over("group")
)

shape: (6, 4)
┌────────────┬───────┬───────┬─────────┐
│ date       ┆ group ┆ value ┆ cum_sum │
│ ---        ┆ ---   ┆ ---   ┆ ---     │
│ date       ┆ i64   ┆ i64   ┆ i64     │
╞════════════╪═══════╪═══════╪═════════╡
│ 2020-03-01 ┆ 1     ┆ 1     ┆ 1       │
│ 2021-01-01 ┆ 1     ┆ 4     ┆ 10      │
│ 2022-01-05 ┆ 1     ┆ 7     ┆ 28      │
│ 2020-07-01 ┆ 2     ┆ 1     ┆ 1       │
│ 2022-01-05 ┆ 2     ┆ 3     ┆ 6       │
│ 2023-02-04 ┆ 2     ┆ 4     ┆ 10      │
└────────────┴───────┴───────┴─────────┘

另一种更“直接”（高效）的方法是先计算每年的总和，然后再计算 cum_sum：

(df.group_by(
      "group", 
      pl.col("date").dt.year().alias("year"),
      maintain_order=True
   )
   .agg(
      pl.col("date").first(),
      pl.col("value").first(), 
      pl.col("value").sum().alias("sum"),
   )
   .with_columns(
      pl.when(pl.col("group").is_first_distinct())
        .then(pl.col("value"))
        .otherwise(
           pl.col("value") + pl.col("sum").cum_sum().shift().over("group")
        )
        .alias("cum_sum")
   )
)

shape: (6, 6)
┌───────┬──────┬────────────┬───────┬─────┬─────────┐
│ group ┆ year ┆ date       ┆ value ┆ sum ┆ cum_sum │
│ ---   ┆ ---  ┆ ---        ┆ ---   ┆ --- ┆ ---     │
│ i64   ┆ i32  ┆ date       ┆ i64   ┆ i64 ┆ i64     │
╞═══════╪══════╪════════════╪═══════╪═════╪═════════╡
│ 1     ┆ 2020 ┆ 2020-03-01 ┆ 1     ┆ 6   ┆ 1       │
│ 1     ┆ 2021 ┆ 2021-01-01 ┆ 4     ┆ 15  ┆ 10      │
│ 1     ┆ 2022 ┆ 2022-01-05 ┆ 7     ┆ 7   ┆ 28      │
│ 2     ┆ 2020 ┆ 2020-07-01 ┆ 1     ┆ 3   ┆ 1       │
│ 2     ┆ 2022 ┆ 2022-01-05 ┆ 3     ┆ 3   ┆ 6       │
│ 2     ┆ 2023 ┆ 2023-02-04 ┆ 4     ┆ 4   ┆ 10      │
└───────┴──────┴────────────┴───────┴─────┴─────────┘