如何用pandas.read_parquet自动分组不同Job ID的分块Parquet文件

问题描述 #

现有带Job ID的分块Parquet数据集，初始目录下为同一Job ID（如abc）的分块文件，可通过vaex或pandas直接读取。但ETL管道会向该目录追加其他Job ID（如xyz）的分块文件，导致目录包含多组不同Job ID的文件。已知Job ID唯一但未知具体值，请问是否可通过pandas.read_parquet自动按Job ID分组读取，返回以Job ID为键、对应DataFrame为值的字典？示例预期输出如下：

{
 'abc': pd.DataFrame, # 读取part*.abc.parquet
 'xyz': pd.DataFrame  # 读取part*.xyz.parquet
} 

目前已尝试使用glob读取。

解决方案 #

pandas.read_parquet本身没有内置按文件名规则分组读取的功能，但可以结合文件遍历、分组逻辑实现需求，具体步骤和代码如下：

实现步骤 #

• 遍历目标目录下所有Parquet文件，提取每个文件名中的Job ID
• 按Job ID对文件路径进行分组
• 对每组文件调用pandas.read_parquet读取并合并为单个DataFrame

代码示例 #

import pandas as pd
import os
from collections import defaultdict

# 替换为你的Parquet文件所在目录
data_directory = "/your/target/directory"

# 用于按Job ID分组存储文件路径
job_file_groups = defaultdict(list)

# 遍历目录下的Parquet文件
for file_name in os.listdir(data_directory):
    if file_name.endswith(".parquet"):
        # 从文件名中提取Job ID（假设格式为part*.{job_id}.parquet）
        job_id = file_name.split(".")[-2]
        full_path = os.path.join(data_directory, file_name)
        job_file_groups[job_id].append(full_path)

# 生成Job ID对应DataFrame的字典
job_dataframes = {}
for job_id, file_paths in job_file_groups.items():
    # 读取同Job ID的所有分块文件并合并
    job_dataframes[job_id] = pd.read_parquet(file_paths)

# 输出结果
print(job_dataframes)

注意事项 #

• 如果文件名中Job ID的位置或格式不同，只需调整job_id = file_name.split(".")[-2]这一行的提取逻辑即可（比如用正则表达式匹配）
• pd.read_parquet支持直接传入文件路径列表，会自动合并所有分块数据，无需手动拼接
• 使用collections.defaultdict可以简化分组过程，无需手动判断Job ID是否已存在于字典中

内容的提问来源于stack exchange，提问作者alvas