使用iter_rows修改大型Polars DataFrame时遇段错误求助

针对你的需求，完全可以避免iter_rows的低效遍历，改用Polars原生的向量化操作来实现，既提升速度又避免段错误。以下提供两种高效实现方式：

方法一：宽表转长表（Melt）+ 关联替换 + 转回宽表（Pivot） #

这种方法适合个体列（ind1/ind2...）数量较多的场景，无需手动遍历列，完全依赖Polars的批量处理能力：

代码示例 #

import polars as pl

# 示例数据（替换为你的实际数据）
df_large = pl.LazyFrame({
    "x": ["a", "a", "b", "b"],
    "y": [1, 2, 1, 2],
    "ind1": ["val1", "val2", "val3", "val4"],
    "ind2": ["val5", "val6", "val7", "val8"]
})

df_rep = pl.DataFrame({
    "x": ["a", "b"],
    "y": [1, 2],
    "ind": ["ind1", "ind2"]
})

# 核心处理逻辑
result = (
    df_large
    # 将宽表转为长表，保留坐标列作为标识符
    .melt(id_vars=["x", "y"], variable_name="ind", value_name="value")
    # 关联需要替换的记录，标记待替换行
    .join(
        df_rep.with_columns(pl.lit(True).alias("to_replace")),
        on=["x", "y", "ind"],
        how="left"
    )
    # 执行替换：标记为待替换的行改为"./."，其余保留原值
    .with_columns(
        pl.when(pl.col("to_replace"))
          .then("./.")
          .otherwise(pl.col("value"))
          .alias("value")
    )
    .drop("to_replace")
    # 转回宽表结构，恢复原数据格式
    .pivot(index=["x", "y"], columns="ind", values="value")
    .collect()
)

print(result)

优势 #

• 全程向量化操作，Polars会自动优化执行计划，内存占用远低于逐行遍历
• 无需手动处理每一列，适配任意数量的个体列
• 避免多次修改DataFrame导致的执行计划膨胀，彻底解决段错误问题

方法二：按个体列分组生成替换条件 #

如果个体列数量不多，可以先对df_rep按个体列分组，生成每个列的替换条件，批量应用到df_large：

代码示例 #

import polars as pl

# 预处理df_rep：按个体列分组，收集对应的坐标对
rep_groups = df_rep.group_by("ind").agg(
    pl.struct(["x", "y"]).alias("coords")
)

# 生成每个个体列的替换表达式
update_exprs = []
for group in rep_groups.iter_rows(named=True):
    ind_col = group["ind"]
    # 把坐标对转为集合，用于快速匹配
    target_coords = set((coord["x"], coord["y"]) for coord in group["coords"])
    # 构造匹配条件：当前行的(x,y)在目标集合中
    replace_condition = pl.struct(["x", "y"]).is_in(target_coords)
    # 生成替换表达式
    update_expr = pl.when(replace_condition)
                    .then("./.")
                    .otherwise(pl.col(ind_col))
                    .alias(ind_col)
    update_exprs.append(update_expr)

# 批量应用替换规则
result = df_large.with_columns(update_exprs).collect()

优势 #

• 仅遍历分组后的df_rep（行数远小于原df_rep），比逐行遍历高效得多
• 每个个体列仅执行一次替换操作，执行计划简洁

为什么原方法会出错？ #

你之前用iter_rows逐行调用with_columns，每次调用都会生成新的LazyFrame执行计划，多次叠加后会导致执行计划异常复杂，Polars在处理时内存占用飙升，最终触发段错误。而上述两种方法都是一次性生成完整的执行计划，Polars可以高效优化内存和执行流程，完美适配大型数据集。

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