如何用Polars实现分组下不同日期对的交集相关性计算？

核心思路 #

通过透视生成宽表（避免转置），利用Polars原生的corr函数自动过滤缺失值（对应仅保留id交集的需求），再按class分组计算所有日期对的相关性，最终整理为目标格式。

实现代码 #

import polars as pl

# 示例数据
df = pl.DataFrame(
    {
        "class": [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
        "day": [1, 1, 1, 2, 2, 2, 3, 3, 3, 4, 4, 4, 4],
        "id": [1, 2, 3, 2, 3, 4, 1, 2, 5, 2, 1, 3, 4],
        "value": [1, 2, 2, 3, 5, 2, 1, 2, 7, 3, 5, 3, 4],
    }
)

# 1. 透视生成宽表：每个class+id对应一行，各day的value为列
pivoted_df = df.pivot(
    index=["class", "id"],
    columns="day",
    values="value",
    aggregate_function=pl.first  # 确保每个class+id+day仅保留一个值
)

# 2. 获取所有日期列，并生成无序日期对（避免重复计算）
day_columns = [col for col in pivoted_df.columns if col not in ["class", "id"]]
day_pairs = [(d1, d2) for idx, d1 in enumerate(day_columns) for d2 in day_columns[idx+1:]]

# 3. 按class分组，计算每对日期的相关性
result_df = pivoted_df.group_by("class").agg(
    [
        pl.corr(pl.col(d1), pl.col(d2)).alias(f"cor_day_{d1}_{d2}")
        for d1, d2 in day_pairs
    ]
)

print(result_df)

关键细节说明 #

1. 透视操作：以class和id为索引，将day转为列，直接得到每个id在不同日期的value值，无需转置，Polars对列操作的效率远高于行操作。
2. 相关性计算：pl.corr会自动忽略任一日期值为null的行（即该id仅在其中一个日期有数据的情况），恰好实现"仅考虑id交集"的需求。
3. 日期对生成：生成无序对（如仅计算day1-day2而非重复计算day2-day1），减少冗余计算。

输出结果 #

运行代码后，输出格式完全匹配需求：

shape: (1, 7)
┌───────┬─────────────┬─────────────┬─────────────┬─────────────┬─────────────┬─────────────┐
│ class ┆ cor_day_1_2 ┆ cor_day_1_3 ┆ cor_day_1_4 ┆ cor_day_2_3 ┆ cor_day_2_4 ┆ cor_day_3_4 │
│ ---   ┆ ---         ┆ ---         ┆ ---         ┆ ---         ┆ ---         ┆ ---         │
│ i64   ┆ f64         ┆ f64         ┆ f64         ┆ f64         ┆ f64         ┆ f64         │
╞═══════╪═════════════╪═════════════╪═════════════╪═════════════╪═════════════╪═════════════╡
│ 1     ┆ 0.981981    ┆ 1.0         ┆ -0.981981   ┆ 0.981981    ┆ 0.981981    ┆ -0.981981   │
└───────┴─────────────┴─────────────┴─────────────┴─────────────┴─────────────┴─────────────┘

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