如何在Polars中实现类似Pandas的Multi-Index reindex功能

当然可以直接在Polars中实现相同的结果，无需依赖Pandas。我们可以通过生成分组与日期的笛卡尔积，再与原表左连接来补全每个分组的完整日期范围，具体实现如下：

Polars 实现代码 #

import polars as pl

# 构造原始数据
df = pl.from_repr("""
┌────────────┬───────┬───────┐
│ date       ┆ group ┆ value │
│ ---        ┆ ---   ┆ ---   │
│ date       ┆ str   ┆ i64   │
╞════════════╪═══════╪═══════╡
│ 2022-01-01 ┆ A     ┆ 10    │
│ 2022-02-01 ┆ A     ┆ 20    │
│ 2022-03-01 ┆ A     ┆ 30    │
│ 2022-04-01 ┆ B     ┆ 40    │
│ 2022-05-01 ┆ B     ┆ 50    │
│ 2022-06-01 ┆ B     ┆ 60    │
└────────────┴───────┴───────┘
""")

# 获取所有唯一分组和日期
unique_groups = df.select(pl.col("group").unique())
unique_dates = df.select(pl.col("date").unique())

# 生成分组与日期的笛卡尔积
full_index = unique_groups.cross_join(unique_dates)

# 左连接原表，补全缺失值
result = full_index.join(df, on=["group", "date"], how="left")

# 查看结果
with pl.Config(tbl_rows=13):
    print(result)

运行结果 #

shape: (12, 3)
┌───────┬────────────┬───────┐
│ group ┆ date       ┆ value │
│ ---   ┆ ---        ┆ ---   │
│ str   ┆ date       ┆ i64   │
╞═══════╪════════════╪═══════╡
│ A     ┆ 2022-01-01 ┆ 10    │
│ A     ┆ 2022-02-01 ┆ 20    │
│ A     ┆ 2022-03-01 ┆ 30    │
│ A     ┆ 2022-04-01 ┆ null  │
│ A     ┆ 2022-05-01 ┆ null  │
│ A     ┆ 2022-06-01 ┆ null  │
│ B     ┆ 2022-01-01 ┆ null  │
│ B     ┆ 2022-02-01 ┆ null  │
│ B     ┆ 2022-03-01 ┆ null  │
│ B     ┆ 2022-04-01 ┆ 40    │
│ B     ┆ 2022-05-01 ┆ 50    │
│ B     ┆ 2022-06-01 ┆ 60    │
└───────┴────────────┴───────┘

实现说明 #

• 先用unique()获取所有不重复的分组和日期
• 通过cross_join()生成两者的笛卡尔积，得到每个分组对应的完整日期范围
• 最后用join(how="left")将原表数据补全到全量索引中，缺失的value会自动填充为null

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