Polars实现类pandas .loc分组首行赋值的简洁方案

不需要手动创建伪索引、做关联匹配，Polars 原生窗口函数就可以非常简洁地实现需求，性能也比 join 方案更好。

核心实现代码 #

import polars as pl

df = (
    df
    # 先按Time全局排序，保证分组内顺序正确
    .sort("Time")
    .with_columns(
        # 按DayOfWeek分区生成分组内从0开始的序号，序号为0即分组首行
        pl.when(pl.int_range(pl.len()).over("DayOfWeek") == 0)
        .then(pl.lit("High"))
        .otherwise(pl.lit("Low"))
        .alias("Risk")
    )
)

逻辑说明 #

• 先对全表按Time字段排序，确保每个DayOfWeek分组内的记录是按时间从小到大排列的
• pl.int_range(pl.len()).over("DayOfWeek") 是Polars原生的窗口计算逻辑，会在每个分组内独立生成从0开始的连续行号，完全不依赖全局索引
• 直接判断行号是否为0，对应给首行赋值High，其余行赋值Low，全程不需要额外建中间列、做表连接、空值填充操作，同时支持懒计算，执行效率更高。

注意：如果使用0.19.0之前的Polars版本，把代码里的pl.len()替换为pl.count()即可，逻辑完全一致。

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