如何用Polars扩展DataFrame，使每个id对应4周完整行数据？

我有一个包含id、week、num1、num2字段的DataFrame，希望扩展该DataFrame，让每个id都对应4周的完整行数据（缺失周数的num1、num2字段填充NaN）。

示例数据及初始化代码 #

首先用Polars创建示例数据：

import polars as pl

data = {
    'id': ['a', 'a', 'b', 'c', 'c', 'c'],
    'week': ['1', '2', '3', '4', '3', '1'],
    'num1': [1, 3, 5, 4, 3, 6],
    'num2': [4, 5, 3, 4, 6, 6]
}
df = pl.DataFrame(data)

原始DataFrame内容：

shape: (6, 4)
┌────┬──────┬──────┬──────┐
│ id ┆ week ┆ num1 ┆ num2 │
│ ---┆ ---  ┆ ---  ┆ ---  │
│ str┆ str  ┆ i64  ┆ i64  │
╞════╪══════╪══════╪══════╡
│ a  ┆ 1    ┆ 1    ┆ 4    │
│ a  ┆ 2    ┆ 3    ┆ 5    │
│ b  ┆ 3    ┆ 5    ┆ 3    │
│ c  ┆ 4    ┆ 4    ┆ 4    │
│ c  ┆ 3    ┆ 3    ┆ 6    │
│ c  ┆ 1    ┆ 6    ┆ 6    │
└────┴──────┴──────┴──────┘

Polars实现方式 #

通过构建完整的id-week笛卡尔积，再与原始数据做左连接即可实现需求：

# 生成所有id和1-4周的笛卡尔积
full_combos = df.select('id').unique().join(
    pl.DataFrame({'week': ['1', '2', '3', '4']}),
    how='cross'
)

# 左连接原始数据，缺失字段自动填充为null（对应Pandas的NaN）
result_df = full_combos.join(df, on=['id', 'week'], how='left')

也可以用链式写法简化代码：

result_df = (
    df.select('id').unique()
      .join(pl.DataFrame({'week': ['1', '2', '3', '4']}), how='cross')
      .join(df, on=['id', 'week'], how='left')
)

扩展后的DataFrame #

执行后得到的结果如下：

shape: (12, 4)
┌────┬──────┬──────┬──────┐
│ id ┆ week ┆ num1 ┆ num2 │
│ ---┆ ---  ┆ ---  ┆ ---  │
│ str┆ str  ┆ i64  ┆ i64  │
╞════╪══════╪══════╪══════╡
│ a  ┆ 1    ┆ 1    ┆ 4    │
│ a  ┆ 2    ┆ 3    ┆ 5    │
│ a  ┆ 3    ┆ null ┆ null │
│ a  ┆ 4    ┆ null ┆ null │
│ b  ┆ 1    ┆ null ┆ null │
│ b  ┆ 2    ┆ null ┆ null │
│ b  ┆ 3    ┆ 5    ┆ 3    │
│ b  ┆ 4    ┆ null ┆ null │
│ c  ┆ 1    ┆ 6    ┆ 6    │
│ c  ┆ 2    ┆ null ┆ null │
│ c  ┆ 3    ┆ 3    ┆ 6    │
│ c  ┆ 4    ┆ 4    ┆ 4    │
└────┴──────┴──────┴──────┘

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