如何为 tensorflow/pytorch 加载大型多文件 parquet 文件答案

【问题标题】：How to load large multi file parquet files for tensorflow/pytorch如何为 tensorflow/pytorch 加载大型多文件 parquet 文件
【发布时间】：2020-06-26 08:49:56
【问题描述】：

我正在尝试将一些 parquet 文件从目录加载到 Python 中以用于 tensorflow/pytorch。

文件太大，无法通过 pyarrow.parquet 函数加载

import pyarrow.parquet as pq
dataset = pq.ParquetDataset('dir')
table = dataset.read()

这给出了out of memory error。

我也尝试过使用petastorm，但这不适用于make_reader()，因为它不是petastorm 类型。

with make_batch_reader('dir') as reader:
  dataset = make_petastorm_dataset(reader)

当我使用make_batch_reader() 和make_petastorm_dataset(reader) 时，它再次给出了zip not iterable error 或类似的东西。

我不确定如何将文件加载到 Python 中以进行 ML 训练。一些快速帮助将不胜感激。

谢谢扎什

【问题讨论】：

标签： python-3.x tensorflow machine-learning pytorch parquet

【解决方案1】：

对于pyarrow，你可以用Python列出目录，遍历*.parquet文件，以pq.ParquetFile打开每个文件，然后读取one row group at a time。这将减轻内存压力，但如果没有并行化就不会超快。

对于petastorm，您可以使用make_batch_reader()。事实上，错误消息并不总是有用的。但您可以检查堆栈跟踪并调查它在 petastorm code 中的来源。

【讨论】：

【解决方案2】：

您可以使用以下代码使用 dask 加载整个数据。您还可以通过使用索引仅计算那些行来在需要时仅加载数据块。 [假设您有不同的索引]。

import dask.dataframe as dd
from dask import delayed
from fastparquet import ParquetFile
import glob

@delayed
def load_chunk(pth):
    x = ParquetFile(pth).to_pandas()
    x = x.drop('[unwanted_columns_to_save_space]',axis=1)
    return x

files = glob.glob('./your_path/*.parquet')

ddf = dd.from_delayed([load_chunk(f) for f in files])
df = ddf.compute()

【讨论】：