即使有块，Dask也会耗尽内存

Posted 2023-03-11

【中文标题】即使有块，Dask也会耗尽内存

【英文标题】：Dask running out of memory even with chunks

【发布时间】：2020-05-08 23:00:25

【问题描述】：

我正在处理大型 CSV 文件，我需要制作笛卡尔积（合并操作）。由于内存错误，我尝试解决 Pandas 的问题（您可以检查 Panda 的代码和数据格式示例同样的问题，here）但没有成功。现在，我正在尝试使用 Dask，它应该可以管理大量数据集，即使它的大小大于可用 RAM。

首先我阅读了两个 CSV：

from dask import dataframe as dd

BLOCKSIZE = 64000000  # = 64 Mb chunks


df1_file_path = './mRNA_TCGA_breast.csv'
df2_file_path = './miRNA_TCGA_breast.csv'

# Gets Dataframes
df1 = dd.read_csv(
    df1_file_path,
    delimiter='\t',
    blocksize=BLOCKSIZE
)
first_column = df1.columns.values[0]
df1.set_index(first_column)
df2 = dd.read_csv(
    df2_file_path,
    delimiter='\t',
    blocksize=BLOCKSIZE
)
first_column = df2.columns.values[0]
df2.set_index(first_column)

# Filter common columns
common_columns = df1.columns.intersection(df2.columns)
df1 = df1[common_columns]
df2 = df2[common_columns]

然后，我将操作存储在磁盘上以防止内存错误：

# Computes a Cartesian product
df1['_tmpkey'] = 1
df2['_tmpkey'] = 1

# Neither of these two options work
# df1.merge(df2, on='_tmpkey').drop('_tmpkey', axis=1).to_hdf('/tmp/merge.*.hdf', key='/merge_data')
# df1.merge(df2, on='_tmpkey').drop('_tmpkey', axis=1).to_parquet('/tmp/')

我制作了a repo to try with exactly the same CSV files that I'm using。我尝试使用较小的 blocksize 值，但我得到了同样的错误。我错过了什么吗？任何形式的帮助都将不胜感激。

【问题讨论】：

你也可以分享一下简单的 Pandas 代码吗？你有一个输出应该是什么样子的例子吗？

感谢您的关注！我当然可以！我已编辑问题以添加指向另一个 Stack Overflow 问题的链接，您可以在其中欣赏代码和数据格式示例

【参考方案1】：

我使用以下方法成功运行了您的代码，内存限制为 32GB。

我已经摆脱了 BLOCKSIZE 的论点，并在 df1 和 df2 上使用了 repartition。

df1 = df1.repartition(npartitions=50)
df2 = df2.repartition(npartitions=1)

请注意，与 df1 相比，df2 的大小非常小（2.5 MB vs 23.75 MB），这就是为什么我只为 df2 保留一个分区并进行剪切df1 分成 50 个分区。

这样做应该会使代码为您工作。 对我来说，使用的内存保持在 12GB 以下。

为了检查，我计算了结果的 len：

len(df) # 3001995

按照上述内容创建一个包含 50 个分区的 parquet 文件。 你可以再次使用repartition 来获得你想要的partition_size。

注意：

添加这个应该可以加快你的代码：

from dask.distributed import Client
client = Client()

就我而言，由于我的运行环境，我不得不使用参数Client(processes=False)。

【讨论】：

非常感谢！！我试试看告诉你

不客气。也许你可以尝试少于 50 个。我一开始尝试了 100 个。

当然！因为它可以并行运行多次，所以我可能会选择一个消耗更少内存的重新分区大小以避免错误。再次比你！

嗨！有用！不幸的是，它非常慢，因为我必须设置更多的分区（我只有 8 GB 的 RAM）。我会按照Pandas' docs 的建议尝试使用 Cython 和 Numba，看看会发生什么

我不知道您是否有可能，但您可以使用colab.research.google.com 使用超过 8GB 的​​内存