在分组Polars DataFrame中获取列首个非空值的行索引

需求说明 #

需要在分组后的Polars DataFrame中，获取trade列首个非空值对应的原始行号，实现类似Pandas中first_valid_index()的功能。

示例数据 #

import polars as pl

df = pl.DataFrame(
    {
        "symbol": ["s1", "s1", "s2", "s2"],
        "trade": [None, 1, -1, None],
    }
)

实现方案 #

Polars本身没有内置行索引，需先手动添加行号列，再通过分组聚合筛选首非空值对应的行号：

# 添加行索引列
df_with_idx = df.with_row_index("row_idx")

# 分组聚合获取首非空值的行号
result = df_with_idx.group_by("symbol").agg(
    pl.col("row_idx").filter(pl.col("trade").is_not_null()).first().alias("first-non-null")
)

print(result)

结果输出 #

shape: (2, 2)
┌────────┬────────────────┐
│ symbol ┆ first-non-null │
│ ---    ┆ ---            │
│ str    ┆ i64            │
╞════════╪════════════════╡
│ s1     ┆ 1              │
│ s2     ┆ 0              │
└────────┴────────────────┘

代码解释 #

• with_row_index("row_idx")：给DataFrame添加从0开始的行号列row_idx，对应原始数据的行位置。
• group_by("symbol")：按symbol列分组。
• pl.col("row_idx").filter(pl.col("trade").is_not_null()).first()：在每个分组内，筛选出trade非空的行，取第一个行号。
• alias("first-non-null")：将聚合结果列重命名为目标列名。

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