如何将Polars中pl.corr结果转为指定格式的相关系数DataFrame？

现有Polars DataFrame：

import polars as pl

df = pl.DataFrame(
    {
        "group": [1, 1, 1, 1, 2, 2, 2, 2],
        "a": [1, 2, 3, 4, 1, 2, 3, 4],
        "b": [5, 1, 7, 9, 2, 4, 9, 7],
        "c": [2, 6, 3, 9, 1, 5, 3, 6],
    }
)

希望生成按group分组的如下结构的相关系数矩阵DataFrame：

┌───────┬──────┬──────────┬──────────┬──────────┐
│ group ┆ name ┆ a        ┆ b        ┆ c        │
│ ---   ┆ ---  ┆ ---      ┆ ---      ┆ ---      │
│ i64   ┆ str  ┆ f64      ┆ f64      ┆ f64      │
╞═══════╪══════╪══════════╪══════════╪══════════╡
│ 1     ┆ a    ┆ 1.0      ┆ 0.680336 ┆ 0.734847 │
│ 1     ┆ b    ┆ 0.680336 ┆ 1.0      ┆ 0.246885 │
│ 1     ┆ c    ┆ 0.734847 ┆ 0.246885 ┆ 1.0      │
│ 2     ┆ a    ┆ 1.0      ┆ 0.830455 ┆ 0.756889 │
│ 2     ┆ b    ┆ 0.830455 ┆ 1.0      ┆ 0.410983 │
│ 2     ┆ c    ┆ 0.756889 ┆ 0.410983 ┆ 1.0      │
└───────┴──────┴──────────┴──────────┴──────────┘

目前尝试的代码生成了扁平格式结果：

df.group_by("group").agg(
    pl.corr(col1, col2).alias(f"{col1}_{col2}")
    for col1 in ["a", "b", "c"]
    for col2 in ["a", "b", "c"]
)

得到的结果：

shape: (2, 10)
┌───────┬─────┬──────────┬──────────┬───┬──────────┬──────────┬──────────┬─────┐
│ group ┆ a_a ┆ a_b      ┆ a_c      ┆ … ┆ b_c      ┆ c_a      ┆ c_b      ┆ c_c │
│ ---   ┆ --- ┆ ---      ┆ ---      ┆   ┆ ---      ┆ ---      ┆ ---      ┆ --- │
│ i64   ┆ f64 ┆ f64      ┆ f64      ┆   ┆ f64      ┆ f64      ┆ f64      ┆ f64 │
╞═══════╪═════╪══════════╪══════════╪═══╪══════════╪══════════╪══════════╪═════╡
│ 2     ┆ 1.0 ┆ 0.830455 ┆ 0.756889 ┆ … ┆ 0.410983 ┆ 0.756889 ┆ 0.410983 ┆ 1.0 │
│ 1     ┆ 1.0 ┆ 0.680336 ┆ 0.734847 ┆ … ┆ 0.246885 ┆ 0.734847 ┆ 0.246885 ┆ 1.0 │
└───────┴─────┴──────────┴──────────┴───┴──────────┴──────────┴──────────┴─────┘

需要将结果转换为目标格式，或者找到更直接的生成方法。

方法一：从扁平结果转换 #

通过拆分列名、重塑数据结构，将现有扁平结果转换为目标格式：

# 先获取扁平结果
flat_df = df.group_by("group").agg(
    pl.corr(col1, col2).alias(f"{col1}_{col2}")
    for col1 in ["a", "b", "c"]
    for col2 in ["a", "b", "c"]
)

# 转换为目标格式
result = (
    flat_df
    # 将非group列转为长格式：列名作为配对标识，值为相关系数
    .melt(id_vars="group", variable_name="pair", value_name="corr")
    # 拆分配对标识，得到行名和列名
    .with_columns(
        pl.col("pair").str.split("_").list.get(0).alias("name"),
        pl.col("pair").str.split("_").list.get(1).alias("col")
    )
    # 透视数据，将列名转为表头
    .pivot(index=["group", "name"], columns="col", values="corr")
    # 调整列顺序，匹配目标格式
    .select("group", "name", "a", "b", "c")
)

print(result)

方法二：直接生成目标格式（更高效） #

无需生成扁平结果，直接在分组聚合时构造结构化数据，一步得到目标格式：

cols = ["a", "b", "c"]

result = (
    df
    .group_by("group")
    # 对每个分组，构造嵌套结构的相关矩阵
    .agg(
        pl.struct(
            {
                col1: pl.struct(
                    {col2: pl.corr(col1, col2) for col2 in cols}
                )
                for col1 in cols
            }
        ).alias("corr_matrix")
    )
    # 展开嵌套结构，得到行名和对应列的相关系数
    .melt(id_vars="group", value_name="corr_data")
    .rename({"variable": "name"})
    # 将嵌套字段展开为单独列
    .unnest("corr_data")
)

print(result)

这个方法避免了冗余的扁平列生成，逻辑更直接，性能更优。

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