python - Dask DataFrame.to_parquet 读取 - 重新分区 - 写入操作失败

Question

我有以下工作流程。

def read_file(path, indx):
    df = pd.read_parquet(path)
    df.index = [indx] * len(df)
    return df

files_list = get_all_files()  # list of 10k parquet files, each about 1MB
df = dask.dataframe.from_delayed([dask.delayed(read_file)(x, indx) for (indx, x) in enumerate(files_list)])
df.divisions = list(range(10000)) + [9999] # each divisions include 1 file
new_divisions = [0, 10, 23, 45, ...., 9999] # new_divisions that reduces number of partitions by putting a bunch of files into same partitions. 
df = df.repartition(divisions = new_divisions)
df.to_parquet("fewer_files") # This causes dask to essentially freeze and no files get written

选择新分区时，每个分区中文件的总内存不超过 1000 MB。然而，最终的 to_parquet 调用永远挂起。 dask 仪表板上没有任何事件。所有工作人员消耗的内存仍然非常小(55MB)，至少在仪表板中是如此；但我怀疑它可能只是没有更新，因为一切都变得 super 慢。运行代码的python进程不断增加内存消耗(Mac中的虚拟内存不断增加；我让它达到30GB)。

如果 files_list 中只有大约 200 个文件，则代码可以正常工作。以下是当 files_list 中有 236 个文件被重新分区为 41 个分区时 df.visualize() 的样子:

当有 10k 个文件时，您知道什么可能导致 df.to_parquet 卡住吗？当我在计算之前打印 df 时，它显示以下内容:

npartitions=65, Dask Name: repartition-merge, 26417 tasks

此外，我可以让 df.get_partition(0).to_parquet 或其他分区相当快地工作。但是，整个数据集上的 df.to_parquet 失败。对于我的笔记本电脑中的 4 个工作人员来说，26K 任务是否太多了？

Answer 1

使用dask.dataframe.read_parquet或其他 dask I/O 实现，尽可能不使用 dask.delayed 包装 pandas I/O 操作。让 dask 直接访问文件对象或文件路径允许调度程序快速评估作业中的步骤并准确估计作业大小和要求，而无需执行完整的工作流程。

说明

通过将 dask.delayed 与 pandas read_parquet reader 一起使用，您实际上剥夺了 dask 窥视文件结构的能力，以帮助安排作业，以及在运行时多次打开和关闭文件完整的工作(你甚至还没有解决的问题)。

当所有内容都整齐地装入内存时，使用 dask.dataframe.read_parquet 和您使用的延迟方法非常相似。当最佳策略不是简单地“读入所有数据，然后弄清楚如何处理它”时，就会出现差异。具体来说，您正在执行许多重新索引和排序操作，所有这些操作都需要 dask 在安排索引操作操作之前了解有关文件内容的大量信息。

本质上，将某些内容包装在 dask.delayed 中告诉 dask“这是一个未知的代码块。将其作为纯 python 黑盒运行很多次。dask.dataframe与 pandas 和 numpy 接口(interface)相比，code> 和 dask.array 接口(interface)具有更小的 API 和更少的互操作性，但你得到的是 dask 实际上知道幕后发生了什么，并且可以为你优化它. 当您使用 dask.delayed 时，您将获得灵活性，但代价是 dask 无法为您调整操作。

示例

作为一个例子，我将创建大量的小文件:

In [9]: tinydf = pd.DataFrame({"col1": [11, 21], "col2": [12, 22]})
   ...: for i in range(1000):
   ...:     tinydf.to_parquet(f"myfile_{i}.parquet")

dask.dataframe.read_parquet

现在，让我们用 dask.dataframe.read_parquet 来阅读此内容:

In [10]: df = dask.dataframe.read_parquet([f"myfile_{i}.parquet" for i in range(1000)])

请注意，这速度快如闪电。我们可以通过检查 dask 属性来查看高级任务图:

In [13]: df.dask
Out[13]:
HighLevelGraph with 1 layers.
<dask.highlevelgraph.HighLevelGraph object at 0x15f79e2f0>
 0. read-parquet-e38709bfe39c7f8dfb5c4abf2fd08b50

请注意，dask.dataframe.read_parquet 对于 dask 来说是一个单一的概念。它可以根据任务的需要进行调整和优化。这包括“查看”文件以了解其列结构、查看元数据文件/属性等，而无需读取所有数据。

In [30]: df.divisions = list(range(0, 2001, 2))

In [31]: df = df.repartition(divisions=list(range(0, 2001, 500)))

In [33]: df.dask
Out[33]:
HighLevelGraph with 2 layers.
<dask.highlevelgraph.HighLevelGraph object at 0x168b5fcd0>
 0. read-parquet-e38709bfe39c7f8dfb5c4abf2fd08b50
 1. repartition-merge-bc42fb2f09234f7656901995bf3b29fa

完整工作流程的高级图表有两个步骤! Dask 了解文件 I/O 和重新分区方面的操作。它可以决定如何拆分这些任务，以保持在内存限制内并将工作负载分散到工作人员之间，而不会导致调度程序陷入困境。

dask.delayed(pd.read_parquet)

另一方面，如果我们使用 dask.delayed 执行此操作会发生什么？

In [14]: def read_file(path, indx):
    ...:     df = pd.read_parquet(path)
    ...:     df.index = [indx] * len(df)
    ...:     return df
    ...:
    ...:
    ...: files_list = [f"myfile_{i}.parquet" for i in range(1000)]
    ...: df = dask.dataframe.from_delayed(
    ...:     [dask.delayed(read_file)(x, indx) for (indx, x) in enumerate(files_list)]
    ...: )

数据帧预览最终看起来很相似，但如果我们在高级任务图的引擎盖下查看，我们可以看到 dask 需要读入所有数据，然后才能知道索引是什么样子!

In [16]: df.dask
Out[16]:
HighLevelGraph with 1001 layers.
<dask.highlevelgraph.HighLevelGraph object at 0x168bf6230>
 0. read_file-b7aed020-1dc7-4872-a37d-d514b407a7d8
 1. read_file-a0462606-999b-4af1-9977-acb562edab67
 2. read_file-286df439-df34-4a5a-baf9-75dd0a5ae09b
 3. read_file-4db8c178-a67e-4775-b117-228ac607f02f
 4. read_file-a19d6144-5560-4da7-a1f5-8dc92b3ccf1c
 
 # yeah... really there are 1000 of these...

 998. read_file-d0cbd4a4-c255-4a77-a905-199bc289a0b5
 999. read_file-45a80080-426a-48fd-8dcb-9ba7565307f1
 1000. from-delayed-833eff6e232da1e10ca7221b961c21c1

更糟糕的是，每个 pd.read_parquet 使用默认的 pandas 读取行为，即假设数据可以装入内存并立即读取整个文件。 Pandas 不返回文件对象 - 它会加载所有数据并在 dask 看到它之前返回 DataFrame。

因此，在所有读取完成之前，dask 基本上无法进入调度位，并且在工作负载平衡、内存管理等方面几乎没有什么可做的。它可以尝试通过执行第一个任务来了解工作负载，但这仍然是对整个第一个文件的读取。

当我们开始尝试打乱索引时，情况只会变得更糟。我不会在这里详细讨论，但你明白了......