如何在Polars中根据条件动态选择分组列实现指定列的向后填充？

嘿，我完全get到你的需求了——你需要让call_id列的null值，根据它能匹配到的id1/id2/id3列来做向后填充（也就是用后面的非null值补前面的空），而且得自动判断用哪一列分组。这种动态选分组的场景确实有点绕，不过用Polars的条件判断+窗口函数就能搞定。

先理清楚核心逻辑：你的示例里所有非null的call_id值，刚好都和某一个id列的值对应（比如call_id=2对应id1=2，call_id=7对应id2=7），所以我们可以基于这个匹配关系，先给每个可能的分组列计算好填充结果，再筛选出符合匹配条件的填充值，最后合并到原列里。

下面是具体的实现代码，我一步步给你解释：

import polars as pl

# 你的原始数据
df = pl.from_repr("""
┌─────┬─────┬─────┬─────────┐
│ id1 ┆ id2 ┆ id3 ┆ call_id │
│ --- ┆ --- ┆ --- ┆ ---     │
│ i64 ┆ i64 ┆ i64 ┆ i64     │
╞═════╪═════╪═════╪═════════╡
│ 1   ┆ 4   ┆ 9   ┆ null    │
│ 1   ┆ 5   ┆ 9   ┆ null    │
│ 1   ┆ 5   ┆ 9   ┆ null    │
│ 2   ┆ 5   ┆ 9   ┆ null    │
│ 2   ┆ 6   ┆ 9   ┆ 2       │
│ 2   ┆ 7   ┆ 10  ┆ null    │
│ 3   ┆ 7   ┆ 11  ┆ null    │
│ 3   ┆ 7   ┆ 12  ┆ null    │
│ 3   ┆ 7   ┆ 13  ┆ 7       │
│ 3   ┆ 8   ┆ 13  ┆ null    │
└─────┴─────┴─────┴─────────┘
""")

# 第一步：为每个分组列计算向后填充结果，只保留和id列值匹配的填充值
df = df.with_columns(
    # 针对id1分组：只有当id1的值等于填充后的call_id时，才保留这个填充值
    fill_id1=pl.when(pl.col('id1') == pl.col('call_id').backward_fill().over('id1'))
               .then(pl.col('call_id').backward_fill().over('id1'))
               .otherwise(None),
    # 同理处理id2分组
    fill_id2=pl.when(pl.col('id2') == pl.col('call_id').backward_fill().over('id2'))
               .then(pl.col('call_id').backward_fill().over('id2'))
               .otherwise(None),
    # 同理处理id3分组
    fill_id3=pl.when(pl.col('id3') == pl.col('call_id').backward_fill().over('id3'))
               .then(pl.col('call_id').backward_fill().over('id3'))
               .otherwise(None)
)

# 第二步：合并填充结果，优先保留原call_id的非null值，再取填充列的有效值
df = df.with_columns(
    call_id=pl.coalesce('call_id', 'fill_id1', 'fill_id2', 'fill_id3')
)

# 清理临时的填充列
df = df.drop(['fill_id1', 'fill_id2', 'fill_id3'])

# 查看最终结果
print(df)

运行这段代码后，你会得到和预期完全一致的输出：

┌─────┬─────┬─────┬─────────┐
│ id1 ┆ id2 ┆ id3 ┆ call_id │
│ --- ┆ --- ┆ --- ┆ ---     │
│ i64 ┆ i64 ┆ i64 ┆ i64     │
╞═════╪═════╪═════╪═════════╡
│ 1   ┆ 4   ┆ 9   ┆ null    │
│ 1   ┆ 5   ┆ 9   ┆ null    │
│ 1   ┆ 5   ┆ 9   ┆ null    │
│ 2   ┆ 5   ┆ 9   ┆ 2       │
│ 2   ┆ 6   ┆ 9   ┆ 2       │
│ 2   ┆ 7   ┆ 10  ┆ 7       │
│ 3   ┆ 7   ┆ 11  ┆ 7       │
│ 3   ┆ 7   ┆ 12  ┆ 7       │
│ 3   ┆ 7   ┆ 13  ┆ 7       │
│ 3   ┆ 8   ┆ 13  ┆ null    │
└─────┴─────┴─────┴─────────┘

简单拆解下逻辑：

1. 先分别对id1/id2/id3做分组，计算call_id的向后填充结果；
2. 通过when条件过滤，只保留填充结果和对应id列值相等的情况（这一步确保我们只取到和call_id匹配的分组填充值）；
3. 最后用coalesce函数，优先保留原call_id的非null值，再依次取三个填充列的有效值，完美实现动态分组的向后填充。

备注：内容来源于stack exchange，提问作者epistemetrica