使用PyArrow从多个文件中读取分区的parquet数据集，并根据文件名添加分区键

Question

我有一堆parquet文件，每个文件包含我的数据集的一个子集。假设文件命名为data-N.parquet，其中N为整数。

我可以读取它们并随后转换为pandas数据帧:

files = glob.glob("data-**.parquet")
ds = pq.ParquetDataset(
files,
metadata_nthreads=64,
).read_table(use_threads=True)
df = ds.to_pandas()

这很好。它希望在最后的数据帧中有一个额外的列，指示数据来自哪个文件。

据我所知，ds数据是分区的，每个文件一个分区。因此，在数据帧中包含分区键将是一个问题。

这可行吗?

Answer 1

分区键目前包含在数据框中。但是，所有现有的分区方案都使用目录名作为键。因此，如果您的数据是/N/data.parquet或/batch=N/data.parquet，则会发生这种情况(在读取数据集时需要提供分区对象)。

目前(pyarrow中)无法从返回的结果中获取文件名。