使用Polars表达式填充年龄并补全行，构建完整客户年度数据集

现有包含缺失年份的Polars DataFrame，部分客户的1999、2000、2001年数据不全，需要补全所有年份，并按每年增长1岁的规则推算年龄。

原始数据 #

import polars as pl

df = pl.DataFrame({
    "cust_id": [1, 2 ,2, 2, 3, 3],
    "year": [2000,1999,2000,2001,1999,2001],
    "cust_age": [21,31,32,33,44,46]
})

目标结果 #

target_df = pl.DataFrame({
    "cust_id": [1, 1, 1, 2 ,2, 2, 3, 3, 3 ],
    "year": [1999,2000,2001,1999,2000,2001,1999,2000,2001],
    "cust_age": [20,21,22,31,32,33,44,45,46]
})

Polars 地道实现方案 #

这里提供两种Polars风格的简洁实现，避免迭代拼接，充分利用其向量化和窗口函数能力：

方法一：交叉连接+窗口函数推算 #

# 生成所有客户与目标年份的完整组合
full_combinations = (
    df.select("cust_id").unique()
    .join(pl.DataFrame({"year": [1999, 2000, 2001]}), how="cross")
)

# 合并原始数据并推算年龄
result = (
    full_combinations.join(df, on=["cust_id", "year"], how="left")
    .with_columns(
        # 对每个客户，基于已知的年龄和对应年份，计算所有年份的年龄
        pl.col("cust_age").fill_null(
            pl.col("cust_age").filter(pl.col("cust_age").is_not_null()).first()
            + (pl.col("year") - pl.col("year").filter(pl.col("cust_age").is_not_null()).first())
        ).over("cust_id")
    )
    .sort(["cust_id", "year"])
)

方法二：分组聚合+序列生成（更简洁） #

result = (
    df
    .group_by("cust_id")
    .agg(
        # 生成1999-2001的年份序列
        pl.int_range(start=1999, end=2002).alias("year"),
        # 基于客户的已知年龄和对应年份，直接推算所有年份的年龄
        (
            pl.col("cust_age").filter(pl.col("cust_age").is_not_null()).first()
            + (pl.int_range(1999, 2002) - pl.col("year").filter(pl.col("cust_age").is_not_null()).first())
        ).alias("cust_age")
    )
    .explode(["year", "cust_age"])
    .sort(["cust_id", "year"])
)

实现说明 #

• 两种方法都基于Polars的向量化操作，避免了低效的循环拼接。
• 核心逻辑：对每个客户，找到其任意一条已知的(year, cust_age)数据，通过当前年份与基准年份的差值推算年龄，确保每年增长1岁的规则准确生效。
• 方法二更紧凑，直接在分组聚合阶段生成完整的年份和年龄序列，再通过explode展开成行。

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