Polars：如何高效用同DataFrame的部分数据替换另一部分？

你需要将distribution_id=2的var_1和var_N列值，替换为同一year和district_id分组下distribution_id=1的对应值。先看示例数据：

import polars as pl

df = pl.DataFrame({
  "year": [2020,2021,2020,2021],
  "district_id": [1,2,1,2],
  "distribution_id": [1, 1, 2, 2],
  "var_1": [1,2,0.1,0.3],
  "var_N": [1,2,0.3,0.5],
  "unrelated_var": [0.2,0.5,0.3,0.7],
})

你之前代码报错的原因 #

你用then(df.filter(...))返回的是行数更少的DataFrame，而pl.when要求then的表达式长度必须和原DataFrame一致，因此触发ShapeError。

解决方案一：窗口函数（推荐，无需额外Join） #

利用over窗口按year和district_id分组，提取distribution_id=1的对应列值并广播到同分组所有行，再通过when完成替换：

columns_to_copy = ["var_1", "var_N"]

df_updated = df.with_columns(
    [
        pl.when(pl.col("distribution_id") == 2)
        .then(pl.col(col).filter(pl.col("distribution_id") == 1).over(["year", "district_id"]))
        .otherwise(pl.col(col))
        .alias(col)
        for col in columns_to_copy
    ]
)

print(df_updated)

解决方案二：Join映射表 #

先提取distribution_id=1的行作为映射表，再和原表Join后替换目标列：

columns_to_copy = ["var_1", "var_N"]

# 构建仅包含有效映射的表
lookup_df = df.filter(pl.col("distribution_id") == 1).select(["year", "district_id"] + columns_to_copy)

# 左连接后替换，最后清理临时列
df_updated = df.join(lookup_df, on=["year", "district_id"], suffix="_lookup").with_columns(
    [
        pl.when(pl.col("distribution_id") == 2)
        .then(pl.col(f"{col}_lookup"))
        .otherwise(pl.col(col))
        .alias(col)
        for col in columns_to_copy
    ]
).drop([f"{col}_lookup" for col in columns_to_copy])

print(df_updated)

两种方法都会得到你期望的结果：

shape: (4, 6)
┌──────┬─────────────┬─────────────────┬───────┬───────┬──────────────┐
│ year ┆ district_id ┆ distribution_id ┆ var_1 ┆ var_N ┆ unrelated_var│
│ ---  ┆ ---         ┆ ---             ┆ ---   ┆ ---   ┆ ---          │
│ i64  ┆ i64         ┆ i64             ┆ f64   ┆ f64   ┆ f64          │
╞══════╪═════════════╪═════════════════╪═══════╪═══════╪══════════════╡
│ 2020 ┆ 1           ┆ 1               ┆ 1.0   ┆ 1.0   ┆ 0.2          │
│ 2021 ┆ 2           ┆ 1               ┆ 2.0   ┆ 2.0   ┆ 0.5          │
│ 2020 ┆ 1           ┆ 2               ┆ 1.0   ┆ 1.0   ┆ 0.3          │
│ 2021 ┆ 2           ┆ 2               ┆ 2.0   ┆ 2.0   ┆ 0.7          │
└──────┴─────────────┴─────────────────┴───────┴───────┴──────────────┘

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