如何用Polars递归查询员工的多级管理者ID？

问题背景 #

有个带employee_id和manager_1_id列的Polars DataFrame，需要递归生成每位员工的manager_2_id到manager_12_id列。之前用Pandas的实现代码是这样的：

id_index = df.set_index("employee_id")["manager_1_id"]
for i in range(1, 12):
    df[f"manager_{str(i + 1)}_id"] = df[f"manager_{str(i)}_id"].map(id_index)

已知每个管理者ID本身也是员工ID，要得到包含所有层级管理者ID的结果。试过循环左连接但效果不好，求Polars的高效实现方式。

示例数据 #

employee_id | manager_1_id
    1       |      3
    2       |      5
    3       |      4
    4       |      5
    5       |

目标结果 #

employee_id | manager_1_id | manager_2_id | manager_3_id | ...
     1              3             4              5
     2              5       
     3              4             5
     5              5

说明：员工1汇报给3，3汇报给4，4汇报给5；5是CEO，无上级，所以自身的上级设为自己。

解决方案 #

不用循环左连接，直接用Polars的replace方法结合字典映射，逻辑和Pandas版本对齐，性能更优：

完整代码 #

import polars as pl

# 假设原始DataFrame是df
df = pl.DataFrame({
    "employee_id": [1,2,3,4,5],
    "manager_1_id": [3,5,4,5,None]
})

# 第一步：构建员工ID到直接上级的映射字典，处理CEO的空值
manager_map = (
    df.select("employee_id", "manager_1_id")
    .with_columns(pl.col("manager_1_id").fill_null(pl.col("employee_id")))
    .to_dict(as_series=False)
)
manager_map = dict(zip(manager_map["employee_id"], manager_map["manager_1_id"]))

# 第二步：初始化结果DataFrame，先处理manager_1_id的空值
df_result = df.with_columns(
    pl.col("manager_1_id").fill_null(pl.col("employee_id")).alias("manager_1_id")
)

# 第三步：循环生成manager_2到manager_12列
for i in range(1, 12):
    current_col = f"manager_{i}_id"
    next_col = f"manager_{i+1}_id"
    df_result = df_result.with_columns(
        pl.col(current_col).replace(manager_map, default=None).alias(next_col)
    )

print(df_result)

方案优势 #

• 避免了循环左连接的性能损耗，字典映射是O(n)级别的操作，远快于重复表连接
• 逻辑和你之前的Pandas代码完全一致，上手成本低
• Polars的懒执行特性会自动优化整个计算流程，比手动循环连接高效得多

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