Polars DataFrame中slice表达式索引异常问题排查与解决

问题描述 #

我有一个Polars DataFrame，结构如下：

shape: (2_655_541, 4)
┌────────────┬────────────┬─────────────────┬─────────────────────┐
│ streamflow ┆ sm_surface ┆ basin_id        ┆ time                │
│ ---        ┆ ---        ┆ ---             ┆ ---                 │
│ f32        ┆ f32        ┆ str             ┆ datetime[μs]        │
╞════════════╪════════════╪═════════════════╪═════════════════════╡
│ null       ┆ null       ┆ camels_01022500 ┆ 2015-01-01 03:00:00 │
│ null       ┆ null       ┆ camels_01022500 ┆ 2015-01-01 06:00:00 │
│ null       ┆ null       ┆ camels_01022500 ┆ 2015-01-01 09:00:00 │
│ null       ┆ null       ┆ camels_01022500 ┆ 2015-01-01 12:00:00 │
│ null       ┆ null       ┆ camels_01022500 ┆ 2015-01-01 15:00:00 │
│ …          ┆ …          ┆ …               ┆ …                   │
│ 0.718293   ┆ 0.40595    ┆ HML_LOBO3       ┆ 2016-12-30 18:00:00 │
│ null       ┆ 0.40601    ┆ HML_LOBO3       ┆ 2016-12-30 21:00:00 │
│ null       ┆ 0.406075   ┆ HML_LOBO3       ┆ 2016-12-31 00:00:00 │
│ null       ┆ 0.406177   ┆ HML_LOBO3       ┆ 2016-12-31 03:00:00 │
│ null       ┆ 0.406333   ┆ HML_LOBO3       ┆ 2016-12-31 06:00:00 │
└────────────┴────────────┴─────────────────┴─────────────────────┘

我需要按basin_id分组后对每个分组切片，执行以下代码：

df1 = (valid_ds.y_origin.group_by('basin_id', maintain_order=True).agg(pl.all().slice(0, 2865)).explode(pl.exclude('basin_id')))

得到结果行数为2_604_285。但出现以下异常：

• 将slice参数改为(1,2865)时，结果起始时间变化，但总行数未改变；
• 改为slice(100,2865)时，结果行数变为2_564_641，计算得每个basin平均行数为2821，并非预期的2865-100=2765。

问题原因 #

Polars中的slice(offset, length)方法，第二个参数是要提取的行数，而非结束索引。具体来说：

• slice(0,2865)：从第0行开始，提取2865行；
• slice(1,2865)：从第1行开始，提取2865行；
• slice(100,2865)：从第100行开始，提取2865行。

你的异常源于对参数的误解：当你用slice(100,2865)时，你以为是取第100到2865行（共2765行），但实际是从第100行开始取2865行。如果某个分组的总行数≥100+2865，就会取满2865行；如果分组总行数不足，则只能取到分组末尾的所有行，这就导致最终平均行数低于预期的2765。

而slice(0,2865)和slice(1,2865)总行数不变，是因为大部分分组的总行数≥2865（slice(1,2865)仍能取到2865行），只有少数分组行数不足，但对总行数影响极小。

解决方法 #

根据你的需求（从第offset行开始，取2865 - offset行，即每个分组保留前2865行中的后2865-offset行），可以用以下两种方式调整：

方法1：先取前2865行，再从偏移量位置切片 #

先确保每个分组只保留前2865行，再从指定偏移量开始提取剩余行，这样能保证每个分组最多取2865 - offset行：

offset = 100
df1 = (
    valid_ds.y_origin
    .group_by('basin_id', maintain_order=True)
    .agg(pl.all().head(2865).slice(offset))  # 先取前2865行，再从第100行取到末尾
    .explode(pl.exclude('basin_id'))
)

方法2：明确指定切片长度为目标行数 #

直接将slice的第二个参数设置为2865 - offset，这样每个分组从偏移量开始取固定长度的行：

offset = 100
target_length = 2865 - offset
df1 = (
    valid_ds.y_origin
    .group_by('basin_id', maintain_order=True)
    .agg(pl.all().slice(offset, target_length))
    .explode(pl.exclude('basin_id'))
)

如果分组总行数不足offset + target_length，这种方式会取到分组末尾的所有行，符合多数水文数据处理的预期（避免补无效空值）。

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