如何通过字典映射将Polars DataFrame的姓名列替换为用户ID？

我有一份包含用户薪资记录的Excel/CSV文件，已在数据库中为每位用户创建账户，希望把Polars DataFrame里的first_name、middle_name、last_name列替换为数据库对应的用户ID。现有数据如下：

原Polars DataFrame #

import polars as pl

# 创建包含重复用户的DataFrame
df = pl.DataFrame({
    "first_name": ["John", "Jane", "Alice", "Bob", "John", "Jane", "John", "Jane"],
    "middle_name": ["A.", "B.", "C.", "D.", "A.", "B.", "A.", "B."],
    "last_name": ["Doe", "Smith", "Johnson", "Brown", "Doe", "Smith", "Doe", "Smith"],
    "salary": [50000, 60000, 55000, 62000, 50000, 60000, 50000, 60000],
    "date": ["2023-01-15", "2023-02-20", "2023-03-05", "2023-04-10", "2023-05-15", "2023-06-20", "2023-07-15", "2023-08-20"]
})
print(df)

输出结果：

shape: (8, 5)
┌────────────┬────────────┬────────────┬────────┬────────────┐
│ first_name │ middle_name│ last_name  │ salary │ date       │
│ ---        │ ---        │ ---        │ ---    │ ---        │
│ str        │ str        │ str        │ i64    │ date       │
├────────────┼────────────┼────────────┼────────┼────────────┤
│ John       │ A.         │ Doe        │ 50000  │ 2023-01-15 │
│ Jane       │ B.         │ Smith      │ 60000  │ 2023-02-20 │
│ Alice      │ C.         │ Johnson    │ 55000  │ 2023-03-05 │
│ Bob        │ D.         │ Brown      │ 62000  │ 2023-04-10 │
│ John       │ A.         │ Doe        │ 50000  │ 2023-05-15 │
│ Jane       │ B.         │ Smith      │ 60000  │ 2023-06-20 │
│ John       │ A.         │ Doe        │ 50000  │ 2023-07-15 │
│ Jane       │ B.         │ Smith      │ 60000  │ 2023-08-20 │
└────────────┴────────────┴────────────┴────────┴────────────┘

包含对应ID的用户列表 #

users = [
    {'id': 1, 'first_name': 'John', 'middle_name': 'A.', 'last_name': 'Doe'},
    {'id': 2, 'first_name': 'Jane', 'middle_name': 'B.', 'last_name': 'Smith'},
    {'id': 3, 'first_name': 'Alice', 'middle_name': 'C.', 'last_name': 'Johnson'},
    {'id': 4, 'first_name': 'Bob', 'middle_name': 'D.', 'last_name': 'Brown'}
]

解决方案 #

方法一：使用Polars原生Join操作（推荐，适合大数据量） #

Polars的Join操作是优化过的原生方法，性能远优于循环或映射，适合处理大规模数据：

import polars as pl

# 原DataFrame
df = pl.DataFrame({
    "first_name": ["John", "Jane", "Alice", "Bob", "John", "Jane", "John", "Jane"],
    "middle_name": ["A.", "B.", "C.", "D.", "A.", "B.", "A.", "B."],
    "last_name": ["Doe", "Smith", "Johnson", "Brown", "Doe", "Smith", "Doe", "Smith"],
    "salary": [50000, 60000, 55000, 62000, 50000, 60000, 50000, 60000],
    "date": ["2023-01-15", "2023-02-20", "2023-03-05", "2023-04-10", "2023-05-15", "2023-06-20", "2023-07-15", "2023-08-20"]
})

# 将用户列表转换为Polars DataFrame
users_df = pl.DataFrame(users)

# 按姓名三列执行左连接，确保所有原记录都保留
result_df = df.join(users_df, on=["first_name", "middle_name", "last_name"], how="left")

# 移除原姓名列，保留需要的字段
result_df = result_df.drop(["first_name", "middle_name", "last_name"])

print(result_df)

输出结果：

shape: (8, 3)
┌──────┬────────┬────────────┐
│ id   │ salary │ date       │
│ ---  │ ---    │ ---        │
│ i64  │ i64    │ str        │
├──────┼────────┼────────────┤
│ 1    │ 50000  │ 2023-01-15 │
│ 2    │ 60000  │ 2023-02-20 │
│ 3    │ 55000  │ 2023-03-05 │
│ 4    │ 62000  │ 2023-04-10 │
│ 1    │ 50000  │ 2023-05-15 │
│ 2    │ 60000  │ 2023-06-20 │
│ 1    │ 50000  │ 2023-07-15 │
│ 2    │ 60000  │ 2023-08-20 │
└──────┴────────┴────────────┘

方法二：创建映射字典（适合小数据集） #

如果数据量较小，也可以用字典映射的方式实现，代码更直观：

import polars as pl

# 原DataFrame
df = pl.DataFrame({
    "first_name": ["John", "Jane", "Alice", "Bob", "John", "Jane", "John", "Jane"],
    "middle_name": ["A.", "B.", "C.", "D.", "A.", "B.", "A.", "B."],
    "last_name": ["Doe", "Smith", "Johnson", "Brown", "Doe", "Smith", "Doe", "Smith"],
    "salary": [50000, 60000, 55000, 62000, 50000, 60000, 50000, 60000],
    "date": ["2023-01-15", "2023-02-20", "2023-03-05", "2023-04-10", "2023-05-15", "2023-06-20", "2023-07-15", "2023-08-20"]
})

# 创建姓名组合到ID的映射（用元组作为字典键）
user_map = {(u['first_name'], u['middle_name'], u['last_name']): u['id'] for u in users}

# 添加id列，通过映射获取对应值
df = df.with_columns(
    pl.struct(["first_name", "middle_name", "last_name"])
    .map_elements(lambda x: user_map[(x['first_name'], x['middle_name'], x['last_name'])])
    .alias("id")
)

# 移除原姓名列
df = df.drop(["first_name", "middle_name", "last_name"])

print(df)

输出结果和方法一完全一致。

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