如何按组将V2列对应行值替换为V1列首个非null值？

问题背景 #

现有Polars DataFrame定义如下：

fr = pl.DataFrame({'Cat': ['A']*4 + ['B']*4, 'V1': [None, 1, 2, None]*2, 'V2': [None, None, None, 555]*2})

其结构为：

shape: (8, 3)
┌─────┬──────┬──────┐
│ Cat ┆ V1   ┆ V2   │
│ --- ┆ ---  ┆ ---  │
│ str ┆ i64  ┆ i64  │
╞═════╪══════╪══════╡
│ A   ┆ null ┆ null │
│ A   ┆ 1    ┆ null │
│ A   ┆ 2    ┆ null │
│ A   ┆ null ┆ 555  │
│ B   ┆ null ┆ null │
│ B   ┆ 1    ┆ null │
│ B   ┆ 2    ┆ null │
│ B   ┆ null ┆ 555  │
└─────┴──────┴──────┘

需求：按Cat分组，把每组中V1列第一个非null值所在行的V2值替换为该非null的V1值，预期结果如下：

┌─────┬──────┬──────┐
│ Cat ┆ V1   ┆ V2   │
│ --- ┆ ---  ┆ ---  │
│ str ┆ i64  ┆ i64  │
╞═════╪══════╪══════╡
│ A   ┆ null ┆ null │
│ A   ┆ 1    ┆ 1    │
│ A   ┆ 2    ┆ null │
│ A   ┆ null ┆ 555  │
│ B   ┆ null ┆ null │
│ B   ┆ 1    ┆ 1    │
│ B   ┆ 2    ┆ null │
│ B   ┆ null ┆ 555  │
└─────┴──────┴──────┘

之前尝试用when/then结合is_not_null().first()时，会把整个V2列都替换掉，达不到预期效果，下面是正确的实现方案。

正确实现方案 #

核心思路是在分组内精准定位到第一个V1不为null的行，只对该行的V2值进行替换，其他行保持原样。以下是几种可行的写法：

方法1：用累积求和标记目标行 #

result = fr.with_row_index().group_by("Cat", maintain_order=True).map_groups(
    lambda df: df.with_columns(
        V2=pl.when(
            (pl.col("V1").is_not_null()) & (pl.col("V1").is_not_null().cum_sum() == 1)
        ).then(pl.col("V1")).otherwise(pl.col("V2"))
    )
).drop("index")

方法2：通过first_valid_index定位行位置 #

result = fr.group_by("Cat", maintain_order=True).map_groups(
    lambda df: df.with_columns(
        V2=pl.when(pl.int_range(0, len(df)) == df["V1"].first_valid_index())
        .then(df["V1"].first())
        .otherwise(df["V2"])
    )
)

方法3：更简洁的分组内逻辑 #

result = fr.group_by("Cat", maintain_order=True).with_columns(
    V2=pl.when(
        pl.col("V1").is_not_null()
        & (pl.col("V1").is_not_null().cum_sum() == 1)
    ).then(pl.col("V1")).otherwise(pl.col("V2"))
)

说明 #

• 以上方法都会在分组内只修改第一个V1非null的行的V2值，其他行的V2保持原有数据不变。
• maintain_order=True参数用于保留原DataFrame的行顺序，避免分组后行序被打乱。

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