如何用Polars.Expr实现将其他组值替换为指定组同时间戳的值

原始数据 #

import polars as pl

df = pl.DataFrame(
    {
        "timestamp": [1, 2, 3, 1, 2, 3],
        "var1": [1, 2, 3, 3, 4, 5], 
        "group": ["a", "a", "a", "b", "b", "b"],
    }
)
print(df)

输出：

shape: (6, 3)
┌───────────┬──────┬───────┐
│ timestamp ┆ var1 ┆ group │
│ ---       ┆ ---  ┆ ---   │
│ i64       ┆ i64  ┆ str   │
╞═══════════╪══════╪═══════╡
│ 1         ┆ 1    ┆ a     │
│ 2         ┆ 2    ┆ a     │
│ 3         ┆ 3    ┆ a     │
│ 1         ┆ 3    ┆ b     │
│ 2         ┆ 4    ┆ b     │
│ 3         ┆ 5    ┆ b     │
└───────────┴──────┴───────┘

需求说明 #

将group为'b'的var1值，替换为相同timestamp下group为'a'的var1值，期望输出如下：

shape: (6, 3)
┌───────────┬──────┬───────┐
│ timestamp ┆ var1 ┆ group │
│ ---       ┆ ---  ┆ ---   │
│ i64       ┆ i64  ┆ str   │
╞═══════════╪══════╪═══════╡
│ 1         ┆ 1    ┆ a     │
│ 2         ┆ 2    ┆ a     │
│ 3         ┆ 3    ┆ a     │
│ 1         ┆ 1    ┆ b     │
│ 2         ┆ 2    ┆ b     │
│ 3         ┆ 3    ┆ b     │
└───────────┴──────┴───────┘

原辅助DataFrame解法 #

def group_value_replacer(
    df: pl.DataFrame, 
    target_group_col: str,
    target_var: str,
    target_group: str,
):
    helper_df = df.filter(pl.col(target_group_col) == target_group)
    df = df.drop(target_var).join(
        helper_df.drop(target_group_col),
        on=["timestamp"],
        how="left",
    )
    return df

result = group_value_replacer(df, "group", "var1", "a")
print(result)

输出：

shape: (6, 3)
┌───────────┬───────┬──────┐
│ timestamp ┆ group ┆ var1 │
│ ---       ┆ ---   ┆ ---  │
│ i64       ┆ str   ┆ i64  │
╞═══════════╪═══════╪══════╡
│ 1         ┆ a     ┆ 1    │
│ 2         ┆ a     ┆ 2    │
│ 3         ┆ a     ┆ 3    │
│ 1         ┆ b     ┆ 1    │
│ 2         ┆ b     ┆ 2    │
│ 3         ┆ b     ┆ 3    │
└───────────┴───────┴──────┘

使用Polars表达式优化实现 #

无需创建辅助DataFrame，直接通过with_columns结合分组表达式和条件判断完成替换：

通用函数实现 #

def group_value_replacer_expr(
    df: pl.DataFrame,
    target_group_col: str,
    target_var: str,
    target_group: str,
):
    # 按timestamp分组，提取目标组的对应变量值
    target_values = pl.col(target_var).filter(pl.col(target_group_col) == target_group).first().over("timestamp")
    # 条件替换非目标组的变量值
    return df.with_columns(
        pl.when(pl.col(target_group_col) != target_group)
        .then(target_values)
        .otherwise(pl.col(target_var))
        .alias(target_var)
    )

result_expr = group_value_replacer_expr(df, "group", "var1", "a")
print(result_expr)

简化直接实现 #

如果不需要通用函数，可直接写表达式：

result = df.with_columns(
    pl.when(pl.col("group") != "a")
    .then(pl.col("var1").filter(pl.col("group") == "a").first().over("timestamp"))
    .otherwise(pl.col("var1"))
    .alias("var1")
)
print(result)

输出结果 #

两种方式输出均与期望一致：

shape: (6, 3)
┌───────────┬──────┬───────┐
│ timestamp ┆ var1 ┆ group │
│ ---       ┆ ---  ┆ ---   │
│ i64       ┆ i64  ┆ str   │
╞═══════════╪══════╪═══════╡
│ 1         ┆ 1    ┆ a     │
│ 2         ┆ 2    ┆ a     │
│ 3         ┆ 3    ┆ a     │
│ 1         ┆ 1    ┆ b     │
│ 2         ┆ 2    ┆ b     │
│ 3         ┆ 3    ┆ b     │
└───────────┴──────┴───────┘

逻辑说明 #

• pl.col("var1").filter(pl.col("group") == "a").first().over("timestamp")：按timestamp分组，提取每个组内group=a的var1值（因每个timestamp对应唯一的a组值，first可确保取到正确值）
• when/then/otherwise：判断当前行的group是否为目标组，若非则替换为提取到的目标组值，否则保留原数值

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