使用Python将一个Parquet文件分成3个Parquet文件

Question

有没有一种方法可以将一个巨大的parquet文件分成较小的文件（使用Python）？保留所有列并划分行？

谢谢

- mery

Answer 1

你可以使用Dask来完成此操作。

import dask.dataframe as dd

ddf = dd.read_parquet('my_file.parquet')
ddf.repartition(3).to_parquet('my_files/')

编辑：您需要安装fastparquet或pyarrow之一。

- moshevi

---

npartitions 不是一个有效的参数，你只需要输入 3。 - rjurney

请注意，在这种情况下，整个文件必须适合您计算机的RAM。 - Victor Shytiuk

Answer 2

由于使用Dask的答案仅适用于文件大小适合计算机RAM的情况，因此我将分享使用Pyarrow的脚本，并逐页读取文件：

import os
import pyarrow as pa
import pyarrow.parquet as pq
from pyarrow import Schema


class ParquetSplitter:

    def __init__(self,
                 src_parquet_path: str,
                 target_dir: str,
                 num_chunks: int = 25
                 ):
        self._src_parquet_path = src_parquet_path
        self._target_dir = target_dir
        self._num_chunks = num_chunks

        self._src_parquet = pq.ParquetFile(
            self._src_parquet_path,
            memory_map=True,
        )

        self._total_group_num = self._src_parquet.num_row_groups
        self._schema = self._src_parquet.schema

    @property
    def num_row_groups(self):
        print(f'Total num of groups found: {self._total_group_num}')
        return self._src_parquet.num_row_groups

    @property
    def schema(self):
        return self._schema

    def read_rows(self):
        for elem in self._src_parquet.iter_batches(
                columns=['player_id', 'played_at']):
            elem: pa.RecordBatch
            print(elem.to_pydict())

    def split(self):
        for chunk_num, chunk_range in self._next_chunk_range():
            table = self._src_parquet.read_row_groups(row_groups=chunk_range)
            file_name = f'chunk_{chunk_num}.parquet'
            path = os.path.join(self._target_dir, file_name)
            print(f'Writing chunk #{chunk_num}...')
            pq.write_table(
                table=table,
                where=path,
            )

    def _next_chunk_range(self):
        upper_bound = self.num_row_groups
        
        chunk_size = upper_bound // self._num_chunks

        chunk_num = 0
        low, high = 0, chunk_size
        while low < upper_bound:
            group_range = list(range(low, high))
            
            yield chunk_num, group_range
            chunk_num += 1
            low, high = low + chunk_size, high + chunk_size
            if high > upper_bound:
                high = upper_bound

    @staticmethod
    def _get_row_hour(row: pa.RecordBatch):
        return row.to_pydict()['played_at'][0].hour


if __name__ == '__main__':
    splitter = BngParquetSplitter(
        src_parquet_path="path/to/Parquet",
        target_dir="path/to/result/dir",
        num_chunks=100,
    )
    splitter.split()

此外，您可以使用Pyspark或Apache Beam Python SDK来实现此目的。它们允许您以更高效的方式拆分文件，因为它们可以在多节点群集上运行。上述示例使用低级别的Pyarrow库，并利用一台机器上的一个进程，因此执行时间可能很长。

- Victor Shytiuk