Polars DataFrame按组计算int_range均值时遇索引越界问题求助

错误原因 #

你的代码触发索引越界panic，是因为当分组内value列全为null时，pl.int_range(...).filter(pl.col("value").is_not_null())会生成空序列，此时调用.mean()在Polars内部处理时出现了索引越界问题（属于Polars的版本兼容bug）。

正确实现代码 #

以下代码可以实现你想要的效果，同时处理全null分组的情况：

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

result = (
    df.group_by_dynamic(
        "date", 
        group_by="id", 
        every="1mo", 
        period="1mo", 
        closed="both"
    )
    .agg(
        pl.when(pl.col("value").is_not_null().any())
        .then(
            pl.int_range(1, pl.len() + 1).cast(pl.Float64) 
            - pl.int_range(1, pl.len() + 1).filter(pl.col("value").is_not_null()).mean()
        )
        .otherwise(pl.int_range(1, pl.len() + 1).cast(pl.Float64).map(lambda x: [None]*len(x)))
    )
    .rename({"int_range": "arange"})
)

print(result)

代码说明 #

1. 分组判断：用pl.col("value").is_not_null().any()检查分组内是否存在非null值，避免空序列求均值触发的错误。
2. 序列计算：
   • 存在非null值时，生成1到分组行数的序列，转换为Float64后减去非null值对应序列的均值。
   • 全为null时，生成与分组行数一致的全null列表。
3. 列名重命名：将默认生成的int_range列名改为你期望的arange。

输出结果 #

运行代码后会得到你期望的结果：

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]           │
└─────┴────────────┴──────────────────┘

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