如何将Polars DataFrame转换为指定结构的嵌套字典？

我有如下结构的Polars DataFrame：

import polars as pl

df = pl.DataFrame(
    {
        "file": ["A", "A", "A", "B", "B", "B"],
        "user": ["u1", "u2", "u3", "u1", "u2", "u3"],
        "data1": [1, 2, 3, 4, 5, 6],
        "data2": [7, 8, 9, 10, 11, 12],
        "data3": [13, 14, 15, 16, 17, 18],
    }
)

输出的DataFrame样式：

shape: (6, 5)
┌──────┬──────┬───────┬───────┬───────┐
│ file ┆ user ┆ data1 ┆ data2 ┆ data3 │
│ ---  ┆ ---  ┆ ---   ┆ ---   ┆ ---   │
│ str  ┆ str  ┆ i64   ┆ i64   ┆ i64   │
╞══════╪══════╪═══════╪═══════╪═══════╡
│ A    ┆ u1   ┆ 1     ┆ 7     ┆ 13    │
│ A    ┆ u2   ┆ 2     ┆ 8     ┆ 14    │
│ A    ┆ u3   ┆ 3     ┆ 9     ┆ 15    │
│ B    ┆ u1   ┆ 4     ┆ 10    ┆ 16    │
│ B    ┆ u2   ┆ 5     ┆ 11    ┆ 17    │
│ B    ┆ u3   ┆ 6     ┆ 12    ┆ 18    │
└──────┴──────┴───────┴───────┴───────┘

希望转换为以下嵌套字典格式：

{
    'file': {
        'A': {
            'user': {
                'u1': {'data1': 1, 'data2': 7, 'data3': 13},
                'u2': {'data1': 2, 'data2': 8, 'data3': 14},
                'u3': {'data1': 3, 'data2': 9, 'data3': 15}
            }
        },
        'B': {
            'user': {
                'u1': {'data1': 4, 'data2': 10, 'data3': 16},
                'u2': {'data1': 5, 'data2': 11, 'data3': 17},
                'u3': {'data1': 6, 'data2': 12, 'data3': 18}
            }
        }
    }
}

我尝试过两种方法：

1. 使用rows_by_key：

df.rows_by_key(key=["file", "user"], unique=True, named=True)

得到的是元组为键的字典，不符合需求：

{('A', 'u1'): {'data1': 1, 'data2': 7, 'data3': 13},
 ('A', 'u2'): {'data1': 2, 'data2': 8, 'data3': 14},
 ('A', 'u3'): {'data1': 3, 'data2': 9, 'data3': 15},
 ('B', 'u1'): {'data1': 4, 'data2': 10, 'data3': 16},
 ('B', 'u2'): {'data1': 5, 'data2': 11, 'data3': 17},
 ('B', 'u3'): {'data1': 6, 'data2': 12, 'data3': 18}}

2. 手动循环分组：

d = {}
for file, df_file in df.group_by("file", maintain_order=True):
    d_user = {}
    for user, df_file_user in df_file.group_by("user", maintain_order=True):
        d_user[user[0]] = df_file_user.drop("file", "user").to_dicts()[0]
    d[file[0]] = {"user": d_user}
print({"file": d})

这个方法能得到正确结果，但流程繁琐，想知道有没有更直接的实现方式。

方法1：Polars原生嵌套聚合+字典转换 #

通过两层分组聚合，将数据打包成结构体后直接转换为目标格式：

final_dict = (
    df
    .group_by("file", "user", maintain_order=True)
    .agg(pl.col("data1", "data2", "data3").first())
    .group_by("file", maintain_order=True)
    .agg(
        pl.struct(pl.col("user"), pl.struct(pl.col("data1", "data2", "data3")).alias("data"))
        .map_elements(lambda x: {u: d for u, d in zip(x["user"], x["data"])}, return_dtype=pl.Object)
        .alias("user")
    )
    .to_dict(as_series=False)
)

该方法全程通过Polars的聚合和映射操作完成，避免手动循环，代码更简洁高效。

方法2：基于rows_by_key的二次处理 #

如果已经通过rows_by_key得到元组键的字典，可利用defaultdict快速嵌套：

from collections import defaultdict

raw_dict = df.rows_by_key(key=["file", "user"], unique=True, named=True)

nested_dict = defaultdict(dict)
for (file, user), data in raw_dict.items():
    if file not in nested_dict:
        nested_dict[file] = {"user": {}}
    nested_dict[file]["user"][user] = data

final_dict = {"file": nested_dict}

这种方式在已有rows_by_key结果的基础上，用简单的循环和字典操作完成结构转换，代码量极少。

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