使用pyarrow读取保存为镶木地板的数据帧，将文件名保存在列中

Question

我想读取一个包含 pandas DataFrames 的 parquet 文件的文件夹。除了我正在读取的数据之外，我还想将读取数据的文件名存储在“file_origin”列中。在 pandas 我可以这样做：

import pandas as pd
from pathlib import Path

data_dir = Path("path_of_folder_with_files")
df = pd.concat(
                pd.read_parquet(parquet_file).assign(file_origin=parquet_file.name)
                for parquet_file in data_dir.glob("*")
            )

不幸的是，这很慢。有没有类似的方法可以用 pyarrow（或任何其他有效的包）来做到这一点？

import pyarrow.parquet as pq

table = pq.read_table(data_dir, use_threads=True)
df = table.to_pandas()

Answer 1

您可以使用箭头代替 pandas:

batches = []
for file_name in data_dir.glob("*"):
    table = pq.read_table(file_name)
    table = table.append_column("file_name", pa.array([file_name]*len(table), pa.string()))
    batches.extend(table.to_batches())
return pa.Table.from_batches(batches)

我不希望它明显更快，除非你的 table 中有很多字符串和对象（在 pandas 中很慢）。

使用pyarrow读取保存为镶木地板的数据帧，将文件名保存在列中

Reading DataFrames saved as parquet with pyarrow, save filenames in columns

python

pandas

pyarrow