基于配置文件复制Polars DataFrame行值的代码问题排查

我有一个包含ID、Local、Entity、Field和Global列的Polars DataFrame，数据如下：

# 创建数据字典
data = {'ID': [4, 4, 4, 4, 4],
        'Local': ['A', 'B', 'C', 'D', 'E'],
        'Field': ['P', 'Q', 'R', 'S', 'T'],   
        'Entity': ['K', 'L', 'M', 'N', 'O'],
        'Global': ['F', 'G', 'H', 'I', 'J'],}

# 创建DataFrame
table = pl.DataFrame(data)
print(table)

输出结果：

shape: (5, 5)
┌─────┬───────┬───────┬──────────┬────────┐
│ ID  ┆ Local ┆ Field ┆ Entity   ┆ Global │
│ --- ┆ ---   ┆ ---   ┆ ---      ┆ ---    │
│ i64 ┆ str   ┆ str   ┆ str      ┆ str    │
╞═════╪═══════╪═══════╪══════════╪════════╡
│ 4   ┆ A     ┆ P     ┆ K        ┆ F      │
│ 4   ┆ B     ┆ Q     ┆ L        ┆ G      │
│ 4   ┆ C     ┆ R     ┆ M        ┆ H      │
│ 4   ┆ D     ┆ S     ┆ N        ┆ I      │
│ 4   ┆ E     ┆ T     ┆ O        ┆ J      │
└─────┴───────┴───────┴──────────┴────────┘

需要根据配置文件复制数据集中的指定行，配置文件信息如下：

copying:
- column_name: P
  source_table: K
  destination_table: X
- column_name: S
  source_table: N
  destination_table: W

配置文件中column_name对应DataFrame的Field列，source_table对应Entity列，destination_table为新行的Entity列值，目标是基于现有行丰富数据。预期结果如下：

shape: (7, 5)
    ┌─────┬───────┬───────┬──────────┬────────┐
    │ ID  ┆ Local ┆ Field ┆ Entity   ┆ Global │
    │ --- ┆ ---   ┆ ---   ┆ ---      ┆ ---    │
    │ i64 ┆ str   ┆ str   ┆ str      ┆ str    │
    ╞═════╪═══════╪═══════╪══════════╪════════╡
    │ 4   ┆ A     ┆ P     ┆ K        ┆ F      │
    │ 4   ┆ B     ┆ Q     ┆ L        ┆ G      │
    │ 4   ┆ C     ┆ R     ┆ M        ┆ H      │
    │ 4   ┆ D     ┆ S     ┆ N        ┆ I      │
    │ 4   ┆ E     ┆ T     ┆ O        ┆ J      │
    │ 4   ┆ A     ┆ P     ┆ X        ┆ F      │
    │ 4   ┆ D     ┆ S     ┆ W        ┆ I      │
    └─────┴───────┴───────┴──────────┴────────┘

数据集为Polars DataFrame，配置文件通过omegaconf加载。我尝试了以下代码：

conf.copying = [
    {"column_name": "P", "source_table": "K", "destination_table": "X"},
    {"column_name": "S", "source_table": "N", "destination_table": "W"},
]

# 遍历配置文件
for i in range(len(conf.copying)):
    # 筛选匹配配置中column_name和source_table的行
    match_rows = table.filter(
        (pl.col("Field") == conf.copying[i]["column_name"])
        & (pl.col("Entity") == conf.copying[i]["source_table"])
    )
    # 选择指定列
    match_rows = match_rows.select(
        [
            "ID",
            "Local",
            "Field",
            "Global",
        ]
    )
    # 添加新的Entity列
    match_rows = match_rows.with_columns(
        pl.lit(conf.copying[i]["destination_table"]).alias("Entity")
    )
    # 调整列顺序
    match_rows = match_rows[
        [
            "ID",
            "Local",
            "Field",
            "Entity",
            "Global",
        ]
    ]
    # 将新行追加到原DataFrame
df_copy = match_rows.vstack(table)

但数据未按预期复制并添加到现有数据集中，请问我哪里出错了？

你的代码存在两个关键问题：

1. 循环内未累积新行：每次循环生成的match_rows都会覆盖之前的结果，最后只保留了最后一次循环生成的行，再和原表拼接，导致只添加了最后一条配置对应的新行。
2. 拼接顺序错误：match_rows.vstack(table)是把原表追加到新行后面，和预期的「原行在前、新行在后」的顺序相反。

修正后的代码如下：

import polars as pl
from omegaconf import OmegaConf

# 初始化配置
conf = OmegaConf.create({
    "copying": [
        {"column_name": "P", "source_table": "K", "destination_table": "X"},
        {"column_name": "S", "source_table": "N", "destination_table": "W"},
    ]
})

# 创建原始DataFrame
data = {'ID': [4, 4, 4, 4, 4],
        'Local': ['A', 'B', 'C', 'D', 'E'],
        'Field': ['P', 'Q', 'R', 'S', 'T'],   
        'Entity': ['K', 'L', 'M', 'N', 'O'],
        'Global': ['F', 'G', 'H', 'I', 'J'],}
table = pl.DataFrame(data)

# 初始化空的DataFrame用于存储所有新行
new_rows = pl.DataFrame()

# 遍历配置生成新行
for item in conf.copying:
    # 筛选匹配行
    match_row = table.filter(
        (pl.col("Field") == item["column_name"])
        & (pl.col("Entity") == item["source_table"])
    )
    # 修改Entity列并添加到新行集合
    updated_row = match_row.with_columns(
        pl.lit(item["destination_table"]).alias("Entity")
    )
    new_rows = new_rows.vstack(updated_row)

# 将新行追加到原表后面
df_copy = table.vstack(new_rows)
print(df_copy)

关键优化点 #

• 新增new_rows变量累积每次循环生成的新行，避免覆盖。
• 使用table.vstack(new_rows)确保原行在前，新行在后，符合预期顺序。
• 简化了列选择和顺序调整的步骤：原始行的列顺序已经正确，直接修改Entity列即可，无需重新选择列。

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