Polars 中高效替换 DataFrame 部分数据：基于条件的跨行值更新

本文介绍如何在 Polars 中高效地将满足某条件的行（如 distribution_id == 2）中指定列（如 "var_1" 和 "var_N"）的值，替换为同一数据集内另一组匹配行（如 distribution_id == 1）对应列的值，核心方法是结合 filter、select 与 update 实现精准对齐更新。

这篇东西想聊的是，如何在Polars里高效地完成一个看起来简单、实际坑不少的操作——把满足某条件的行（比如 distribution_id == 2）中指定列（比如 "var_1" 和 "var_N"）的值，替换成同一数据集内另一组匹配行（比如 distribution_id == 1）的对应值。核心思路说到底就一句话：构造一个结构一致、键对齐的更新源表，然后用 .update() 一把搞定。

如果你在Polars里直接拿 when/then/otherwise 对整列做“跨行引用”，十有八九会撞上 ShapeError。原因很简单：df.filter(...) 返回的是一个子集DataFrame，它没法直接作为标量或对齐序列参与 zip_with。那该怎么办？正确做法是——先造一个更新源，再让 update 方法原子化覆盖。

✅ 推荐方案：filter → select → update

第一步，从原始 DataFrame 里把 distribution_id == 1 的行拎出来，把它的 distribution_id 改成 2，同时只保留用于对齐的键列（"year", "district_id"）和要更新的目标列（"var_1", "var_N"）：

import polars as pldf = 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],})# 构建更新源：取 distribution_id == 1 的行，映射为 distribution_id == 2update_source = (    df.filter(pl.col("distribution_id") == 1)    .select(        "year",        "district_id",        pl.col("^var_.+$"),  # 正则匹配所有以 "var_" 开头的列        distribution_id=pl.lit(2, dtype=pl.Int64)  # 强制设为 2    ))# 执行更新：基于 ["year", "district_id", "distribution_id"] 三元组精确对齐result = df.update(update_source, on=["year", "district_id", "distribution_id"])print(result)

输出结果完全符合预期：

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           │└──────┴─────────────┴─────────────────┴───────┴───────┴───────────────┘

⚠️ 关键注意事项

• on 参数必须唯一标识目标行：update() 要求 on 列在原始 DataFrame 和 update_source 中均构成唯一键（或至少能一对一匹配）。本例中 ["year", "district_id", "distribution_id"] 满足该条件；若存在重复组合，Polars 将报错或产生非预期覆盖。
• 列名与类型需严格一致：update_source 中的列名（包括重命名后的 distribution_id）和数据类型必须与目标 DataFrame 完全匹配，否则 update() 会静默忽略或报类型错误。
• 非目标列不受影响：如 "unrelated_var" 不在 update_source 中，其值保持不变，体现了 .update() 的安全性和局部性。
• 避免链式赋值陷阱：df.update(...) 返回新 DataFrame（默认 in_place=False），务必显式赋值（如 df = df.update(...)）或使用 in_place=True（仅限可变操作场景）。

? 扩展建议

• 若需动态指定列，可用列表推导式生成 select 参数：

  cols_to_copy = ["var_1", "var_N"].select("year", "district_id", *cols_to_copy, distribution_id=pl.lit(2))

• 对于更复杂的映射逻辑（如多对一、加权替换），可先用 join 构造映射表，再 update；但本场景下 filter + update 是最简洁、高效且语义清晰的解法。

通过这一模式，你就能在 Polars 中实现类似 SQL UPDATE ... JOIN 或 SQLAlchemy 更新语句的表达能力，兼具性能、可读性与工程健壮性。