如何用Python Polars构建产品替代转换表

解决方案代码 #

import polars as pl

def get_final_supersession(df: pl.DataFrame) -> pl.DataFrame:
    # 初始化Last_new_code列为原始New_code的值
    df = df.with_columns(pl.col("New_code").alias("Last_new_code"))
    # 创建旧编码到新编码的映射字典
    supersession_map = dict(zip(df["Old_code"], df["New_code"]))

    while True:
        # 对Last_new_code进行一轮替换：如果当前值属于旧编码，就替换为对应的新编码
        updated_df = df.with_columns(
            pl.col("Last_new_code")
            .replace(supersession_map, default=pl.col("Last_new_code"))
            .alias("Updated_last")
        )
        # 检查是否所有值都已到达最终节点（不再能被替换）
        if updated_df["Last_new_code"].series_equal(updated_df["Updated_last"]):
            break
        # 更新Last_new_code为最新值，继续循环
        df = updated_df.drop("Last_new_code").rename({"Updated_last": "Last_new_code"})
    
    return df

# 测试示例数据
df = pl.DataFrame(
    {
        'Old_code': ['A', 'B', 'C', 'K'],
        'New_code': ['B', 'C', 'D', 'C']
    }
)

result = get_final_supersession(df)
print(result)

方案优势 #

1. 时间复杂度优化：替代原O(n²)的双层循环，时间复杂度仅与最长替代链的长度相关（比如A→B→C→D的链长度为3，只需3轮迭代），处理大规模数据时性能提升显著。
2. 符合Polars最佳实践：全程使用Polars向量化操作，避免直接遍历DataFrame/Series，也未使用scatter()这类反模式方法，不会阻碍Polars的查询优化。
3. 兼容性强：支持任意长度的链式替代、多商品指向同一新品的场景，逻辑简洁易维护。

输出结果 #

shape: (4, 3)
┌──────────┬──────────┬───────────────┐
│ Old_code ┆ New_code ┆ Last_new_code │
│ ---      ┆ ---      ┆ ---           │
│ str      ┆ str      ┆ str           │
╞══════════╪══════════╪═══════════════╡
│ A        ┆ B        ┆ D             │
├╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ B        ┆ C        ┆ D             │
├╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ C        ┆ D        ┆ D             │
├╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ K        ┆ C        ┆ D             │
└──────────┴──────────┴───────────────┘

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