如何使用同一字典映射Polars数据框的两列节点至索引

数据准备 #

我有一个描述网络连接的Polars数据框，记录节点间的连接权重：

import polars as pl
df_network = pl.DataFrame(
    {"node1": ["a", "c", "c"], "node2": ["b", "d", "a"], "weight": [1, 2, 3]}
)
# shape: (3, 3)
# ┌───────┬───────┬────────┐
# │ node1 ┆ node2 ┆ weight │
# │ ---   ┆ ---   ┆ ---    │
# │ str   ┆ str   ┆ i64    │
# ╞═══════╪═══════╪════════╡
# │ a     ┆ b     ┆ 1      │
# │ c     ┆ d     ┆ 2      │
# │ c     ┆ a     ┆ 3      │
# └───────┴───────┴────────┘

另有一个节点到整数索引的映射数据框：

df_node_index = pl.DataFrame({"node": ["d", "b", "c", "a"], "index": [0, 1, 2, 3]})
# shape: (4, 2)
# ┌──────┬───────┐
# │ node ┆ index │
# │ ---  ┆ ---   │
# │ str  ┆ i64   │
# ╞══════╪═══════╡
# │ d    ┆ 0     │
# │ b    ┆ 1     │
# │ c    ┆ 2     │
# │ a    ┆ 3     │
# └──────┴───────┘

需要将node1和node2替换为对应的索引值，得到目标数据框：

df_desirable = pl.DataFrame({"index1": [3, 2, 2], "index2": [1, 0, 3], "weight":[1,2,3]})
# shape: (3, 3)
# ┌────────┬────────┬────────┐
# │ index1 ┆ index2 ┆ weight │
# │ ---    ┆ ---    ┆ ---    │
# │ i64    ┆ i64    ┆ i64    │
# ╞════════╪════════╪════════╡
# │ 3      ┆ 1      ┆ 1      │
# │ 2      ┆ 0      ┆ 2      │
# │ 2      ┆ 3      ┆ 3      │
# └────────┴────────┴────────┘

当前实现方案 #

我目前通过两次join实现需求，虽然可行但不够简洁：

df_network.join(
    df_node_index.rename({"index": "index1"}),
    how="left",
    left_on="node1",
    right_on="node",
).join(
    df_node_index.rename({"index": "index2"}),
    how="left",
    left_on="node2",
    right_on="node",
).select(
    ["index1", "index2", "weight"]
).equals(df_desirable) # True

更优雅的Polars实现方案 #

方案1：使用lookup方法（原生高效） #

Polars的lookup方法可以直接在列上执行映射查询，无需多次join，代码更简洁且性能更优：

result = df_network.with_columns(
    index1=pl.col("node1").lookup(df_node_index, "node", "index"),
    index2=pl.col("node2").lookup(df_node_index, "node", "index")
).drop("node1", "node2")

print(result.equals(df_desirable)) # True

方案2：使用字典映射（简洁直观） #

先将映射数据框转换为Python字典，再用map_dict方法替换列值，适合小数据集场景：

# 构建节点到索引的字典
node_map = df_node_index.select(pl.col("node"), pl.col("index")).to_dict(as_series=False)
node_map = dict(zip(node_map["node"], node_map["index"]))

result = df_network.with_columns(
    index1=pl.col("node1").map_dict(node_map),
    index2=pl.col("node2").map_dict(node_map)
).drop("node1", "node2")

print(result.equals(df_desirable)) # True

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