快速加载常用数据集

通常情况下，我们会针对不同的数据格式采用不同的数据读取方式，例如使用 json 来读取 .json 文件，使用 jsonlines 来读取 .jsonl 文件，使用 pandas 或者 Python 的文件读取方式来获取 .csv 文件。但这些方法的读取时间各不相同，有快有慢，具体可参考《Python 数据读取方式以及时间比对》。

本文主要介绍如何快速加载常用的数据集。

jsonlines 读取数据集

此时，我们有一个 13G jsonlines 格式的数据集，我们使用 jsonlines 包进行读取。

  • 读取代码：
1
2
3
4
%%time
file = "/data/xueyou/data/taobao_note/redbook/0516/data.json"
with open(file, "r", encoding="utf-8") as file:
    datasets_redbook = [line for line in jsonlines.Reader(file)]
  • 所需时间：
1
2
CPU times: user 2min 7s, sys: 10.2 s, total: 2min 17s
Wall time: 2min 20s

需要花费较长的时间（2分20秒）才能够将数据全部读入，如果需要经常读取该数据，但又不想等这么久时间该怎么办呢？

转换为 HDF5

我们借助 vaex 包将数据转换为 HDF5 格式，关于 vaex 包的内容可以参考 GitHub。

  • 首先，将数据组织成 pandas.DataFrame 格式。
1
df_dataset = pd.DataFrame(data=datasets_redbook)
  • 然后，调用 vaex 的 from_pandas()，将其转换为 vaex 的 DataFrame 对象。
1
df_vaex = vaex.from_pandas(df_dataset)
  • 导出成 HDF5 格式的数据集。
1
2
3
4
5
%%time
df_vaex.export_hdf5(os.path.join(project_dir, "data/detect/redbook/redbook.hdf5"))

CPU times: user 21.3 s, sys: 6.05 s, total: 27.3 s
Wall time: 2min 21s

可以看到，光导出成 HDF5 格式的数据集就需要 2 分 21 秒了。虽然数据转换工作也需要大量的时间，但为了后续快速读取数据，这都是值得的，马上就是见证奇迹的时刻了！

读取 HDF5 格式数据集

使用 vaex 的 open() 函数，我们可以加载 HDF5 格式的数据集。

  • 读取代码：
1
2
%%time
df_test_vaex = vaex.open(os.path.join(project_dir, "data/detect/redbook/redbook.hdf5"))
  • 所需时间：
1
2
CPU times: user 20 ms, sys: 728 ms, total: 748 ms
Wall time: 885 ms

仅需要 885 毫秒！我们再把它转换为 Python 的列表。

1
2
3
4
5
%%time
content_list = df_test_vaex["content"].tolist()

CPU times: user 8.26 s, sys: 1.77 s, total: 10 s
Wall time: 13.8 s

也就是说，后续每次打开 + 转换为列表操作总共不超过 14 秒！这快了近乎 10 倍！对于频繁使用的大规模数据集，我们都可以尝试将其转换为 HDF5 格式存储。