如何在序列化为Pandas数据框时展平嵌套的数据类？

Question

我有一个包含其他数据类作为字段的数据类：

@dataclass
class Bar:
    abc: int
    bed: int
    asd: int


@dataclass
class Foo:
    xy: int
    yz: Bar

接着我尝试使用Pandas将其序列化为CSV文件，代码如下：

dataset = [Foo(xy=1, yz=Bar(abc=1, bed=2, asd=3))]
pd_dataset = pandas.DataFrame(vars(row) for row in dataset)
pd_dataset.to_csv('dataset_example.csv', index=False)

但是我得到的结果与我想要实现的有些不同。准确地说，我现在得到的是：

xy,yz
1,"Bar(abc=1, bed=2, asd=3)"

and I want:

xy,yz_abc,yz_bed,yz_asd
1,1,2,3

你能帮我把它弄对吗？我试着编写自己的序列化函数并执行以下操作：

pandas.DataFrame(asdict(row, dict_factory=row_to_dict) for row in dataset)

但我不知道如何正确地编写它。

- Gustaw Ohler

Answer 1

无需使用外部库，Pandas 以 pd.json_normalize 的形式为您提供所需一切:

>>> import pandas as pd
... from dataclasses import asdict, dataclass
... 
... @dataclass
... class Bar:
...     abc: int
...     bed: int
...     asd: int
... 
... @dataclass
... class Foo:
...     xy: int
...     yz: Bar
... 
... dataset = [
...     Foo(xy=1, yz=Bar(abc=1, bed=2, asd=3)),
...     Foo(xy=10, yz=Bar(abc=10, bed=20, asd=30)),
... ]

>>> dataset
[Foo(xy=1, yz=Bar(abc=1, bed=2, asd=3)),
 Foo(xy=10, yz=Bar(abc=10, bed=20, asd=30))]

>>> df = pd.json_normalize(asdict(obj) for obj in dataset)
>>> df
   xy  yz.abc  yz.bed  yz.asd
0   1       1       2       3
1  10      10      20      30

>>> print(df.to_csv(index=False))
xy,yz.abc,yz.bed,yz.asd
1,1,2,3
10,10,20,30

我个人更喜欢使用默认的"."分隔符，但如果您强烈偏爱下划线，Pandas也为您提供了支持：

>>> pd.json_normalize((asdict(obj) for obj in dataset), sep="_")
   xy  yz_abc  yz_bed  yz_asd
0   1       1       2       3
1  10      10      20      30

- Thrastylon

Answer 2

小心地从vars(Bar)中创建所需的关键字，可以实现您想要的功能。

dataset = [Foo(xy=1, yz=Bar(abc=1, bed=2, asd=3))]

res = []

for obj in dataset:
    d = {}
    for k, v in vars(obj).items():
        if isinstance(v, Bar):
            for k_, v_ in vars(vars(obj)[k]).items():
                d[f'{k}_{k_}'] = v_
        else:
            d[k] = v
    res.append(d)

print(res)
'''
[{'xy': 1, 'yz_abc': 1, 'yz_bed': 2, 'yz_asd': 3}]
'''

pd_dataset = pd.DataFrame.from_records(res)

print(pd_dataset)
'''
   xy  yz_abc  yz_bed  yz_asd
0   1       1       2       3
'''

- Ynjxsjmh

Answer 3

好的，我在发布问题后自己解决了。为了解决这个问题，我需要下载一个叫做flatten-dict的库。然后像这样使用它：

pd_dataset = pandas.DataFrame(flatten(asdict(row), reducer='underscore') for row in dataset)

如果这种方法还有改进的空间，请告诉我，但我觉得它非常简洁和简单。

- Gustaw Ohler