计算 Polars 数据框中 int_range 列的组平均值

Question

更新： 此问题已在 Polars 中解决。代码现在运行没有错误。

---

我在极坐标中有以下代码：

import datetime
import polars as pl

df = pl.DataFrame(
    {
        "id": [1, 2, 1, 2, 1, 2, 3],
        "date": [
            datetime.date(2022, 1, 1),
            datetime.date(2022, 1, 1),
            datetime.date(2022, 1, 11),
            datetime.date(2022, 1, 11),
            datetime.date(2022, 2, 1),
            datetime.date(2022, 2, 1),
            datetime.date(2022, 2, 1),
        ],
        "value": [1, 2, 3, None, 5, 6, None],
    }
)

(df.group_by_dynamic("date", group_by="id", every="1mo", period="1mo", closed="both")
   .agg(
       pl.int_range(1, pl.len() + 1) 
       - pl.int_range(1, pl.len() + 1).filter(pl.col("value").is_not_null()).mean(),
   )
)

但是，当我运行它时，我收到以下错误，我不太明白。

pyo3_runtime.PanicException: index out of bounds: the len is 1 but the index is 1

我想要实现的行为是：对于每个组，创建一个从 1 到该组中的行数的自然序列，并从中减去该组中“值”列中非空的平均值。 （如果该组中的所有“值”均为 null，则返回 null）。

更具体地说，我想要的结果是

shape: (5, 3)
┌─────┬────────────┬──────────────────┐
│ id  ┆ date       ┆ arange           │
│ --- ┆ ---        ┆ ---              │
│ i64 ┆ date       ┆ list[f64]        │
╞═════╪════════════╪══════════════════╡
│ 1   ┆ 2022-01-01 ┆ [-1.0, 0.0, 1.0] │
│ 1   ┆ 2022-02-01 ┆ [0.0]            │
│ 2   ┆ 2022-01-01 ┆ [-1.0, 2.0, 1.0] │
│ 2   ┆ 2022-02-01 ┆ [0.0]            │
│ 3   ┆ 2022-02-01 ┆ [null]           │
└─────┴────────────┴──────────────────┘

我怎样才能实现这个目标？

Answer 1

看起来也许你需要“爆炸”，然后实现你的逻辑。

编辑：与

"value": [1, None, 3, 4, 5, 6, None]

(
   df
   .groupby_dynamic(index_column="date", by="id", every="1mo", period="1mo", closed="both").agg(
      pl.exclude("date"))
   .with_row_count("group")
   .explode("value")
   .with_columns(
      (pl.arange(1, pl.count() + 1) 
       - pl.arange(1, pl.count() + 1)
           .filter(pl.col("value").is_not_null())
           .mean())
      .over("group"))
)

shape: (9, 5)
┌───────┬─────┬────────────┬───────┬────────┐
│ group | id  | date       | value | arange │
│ ---   | --- | ---        | ---   | ---    │
│ u32   | i64 | date       | i64   | f64    │
╞═══════╪═════╪════════════╪═══════╪════════╡
│ 0     | 1   | 2022-01-01 | 1     | -1.0   │
│ 0     | 1   | 2022-01-01 | 3     | 0.0    │
│ 0     | 1   | 2022-01-01 | 5     | 1.0    │
│ 1     | 1   | 2022-02-01 | 5     | 0.0    │
│ 2     | 2   | 2022-01-01 | null  | -1.5   │
│ 2     | 2   | 2022-01-01 | 4     | -0.5   │
│ 2     | 2   | 2022-01-01 | 6     | 0.5    │
│ 3     | 2   | 2022-02-01 | 6     | 0.0    │
│ 4     | 3   | 2022-02-01 | null  | null   │
└───────┴─────┴────────────┴───────┴────────┘

之后您可以使用

group

列将它们重新组合在一起。