[Python] Turn polars dataframe into nested dictionary

I have a polars dataframe like such:

print(
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],
}
)
)



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 │
└──────┴──────┴───────┴───────┴───────┘


And I would like to turn it into a bested dictionary like this

{
'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}
}
}
}
}


I tried my luck with grouping and selecting structs, but I'm unable to do the conversion using only polars functions. What's the right approach here?

EDIT:

I tried two things:

1. Using rows_by_key:

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


but that gives the keys as tuples, i.e.,

{('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}}

1. I also tried to manually loop over the groups:

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})


This yields the right output:

{
'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}
}
}
}
}


But this feels cumbersome and I wonder whether there's a more direct way of doing this.

Continue reading...