如何在Polars DataFrame中补全缺失因子并填充0值？

原DataFrame定义 #

import polars as pl

df = pl.DataFrame({
    'date':['date1','date1','date1','date2','date3','date3'], 
    'factor':['A','B','C','B','B','C'], 
    'val':[1,2,3,3,1,5]
})

原DataFrame结构：

shape: (6, 3)
┌───────┬────────┬─────┐
│ date  ┆ factor ┆ val │
│ ---   ┆ ---    ┆ --- │
│ str   ┆ str    ┆ i64 │
╞═══════╪════════╪═════╡
│ date1 ┆ A      ┆ 1   │
│ date1 ┆ B      ┆ 2   │
│ date1 ┆ C      ┆ 3   │
│ date2 ┆ B      ┆ 3   │
│ date3 ┆ B      ┆ 1   │
│ date3 ┆ C      ┆ 5   │
└───────┴────────┴─────┘

解决方案代码 #

# 获取所有唯一的日期和因子
unique_dates = df.select(pl.col('date').unique())
unique_factors = df.select(pl.col('factor').unique())

# 生成日期和因子的笛卡尔积（所有可能组合）
all_combinations = unique_dates.join(unique_factors, how='cross')

# 左连接原表，填充缺失值为0，并重命名列
result = all_combinations.join(df, on=['date', 'factor'], how='left').fill_null(0).rename({'val': 'value'})

print(result)

代码解释 #

• 获取唯一值：提取date和factor的所有唯一值，确保覆盖所有可能的日期和因子类型。
• 笛卡尔积：通过cross连接生成所有日期-因子的组合，补全缺失的配对。
• 左连接与填充：将笛卡尔积结果和原表左连接，保留所有组合，原表中不存在的配对会产生空值，用fill_null(0)将空值替换为0。
• 重命名列：将val列重命名为value，与目标结果格式一致。

输出结果 #

shape: (9, 3)
┌───────┬────────┬───────┐
│ date  ┆ factor ┆ value │
│ ---   ┆ ---    ┆ ---   │
│ str   ┆ str    ┆ i64   │
╞═══════╪════════╪═══════╡
│ date1 ┆ A      ┆ 1     │
│ date1 ┆ B      ┆ 2     │
│ date1 ┆ C      ┆ 3     │
│ date2 ┆ A      ┆ 0     │
│ date2 ┆ B      ┆ 3     │
│ date2 ┆ C      ┆ 0     │
│ date3 ┆ A      ┆ 0     │
│ date3 ┆ B      ┆ 1     │
│ date3 ┆ C      ┆ 5     │
└───────┴────────┴───────┘

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