dask read_parquet 内存不足

Question

我正在尝试读取一个大（不适合内存）镶木地板数据集，然后从中采样。数据集的每个分区都非常适合内存。

数据集是磁盘上大约 20Gb 的数据，分为 104 个分区，每个分区大约 200Mb。我不想在任何时候使用超过 40Gb 的内存，所以我相应地设置了 n_workers 和 memory_limit。

我的假设是 Dask 会加载它可以处理的尽可能多的分区，从中采样，从内存中删除它们，然后继续加载下一个。或者类似的东西。

相反，从执行图来看（104 个并行加载操作，在每个样本之后），看起来它试图同时加载所有分区，因此工作人员不断因内存不足而被杀死。

我错过了什么吗？

这是我的代码：

from datetime import datetime

from dask.distributed import Client
client = Client(n_workers=4, memory_limit=10e9) #Gb per worker

import dask.dataframe as dd

df = dd.read_parquet('/path/to/dataset/')
df = df.sample(frac=0.01)

df = df.compute()

要重现错误，您可以创建一个大小为我尝试使用此代码加载的数据集大小的 1/10 的模拟数据集，并尝试使用 1GB memory_limit=1e9 的代码进行补偿。

from dask.distributed import Client
client = Client() #add restrictions depending on your system here

from dask import datasets

df = datasets.timeseries(end='2002-12-31')
df = df.repartition(npartitions=104)
df.to_parquet('./mock_dataset')

Answer 1

Parquet 是一种高效的二进制格式，具有编码和压缩功能。很有可能在内存中，它占用的空间比你想象的要多得多。

为了以 1% 的采样率对数据进行采样，每个分区都被加载并整体扩展到内存中，然后再进行子选择。这伴随着缓冲区副本的大量内存开销。每个工作线程都需要容纳当前处理的块，以及迄今为止在该工作人员上累积的结果，然后一个任务将复制所有这些以进行最终的 concat 操作（这也涉及副本和开销）。

一般建议是每个工作人员都应该能够“多次”访问每个分区的内存大小，在您的情况下，这些内存大小约为 2GB 磁盘和更大的内存。