求助：修改Parquet读取代码以实现DataFrame过滤与指定列保留

原始代码问题 #

你原本的代码实现了从ParquetDataset读取数据并存入列表，但需要补充两个处理需求，同时代码里存在一些可优化的点（比如不必要的vars()使用、无意义的dfList[0].head(3)调用）。

修改后的完整代码 #

import pandas as pd
import pyarrow.parquet as pq

def readfile() -> list[pd.DataFrame]:
    df_list = []
    val1 = [1, 2, 4]
    # 修正语法错误后的列映射字典（原df2/df3的PATH_ID缺少键值对，这里按示例补全，可按需调整）
    col_mapping = {
        "df0": {"id": "P_ID", "result_date": "DATE", "test_name": "CODE", "unit": "T", "result": "val"},
        "df1": {"id": "T_ID", "DIA_DATE": "DATE", "DI_CODE": "CODE", "D_CODE_TYPE": "TYPE"},
        "df2": {"id": "P_id", "PATH_ID": "PATH", "PR_DATE": "DATE", "PR_CODE": "CODE", "PR_CODE_TYPE": "TYPE"},
        "df3": {"id": "P_ID", "PATH_ID": "PATH", "birth": "DATE"},
    }
    # 需要执行空值填充的目标数据集名称
    fill_target_dfs = {"df1", "df2"}

    for df_name in ["df0", "df1", "df2", "df3"]:
        # 读取并过滤原始数据
        df = pq.ParquetDataset(
            f"{path}/{df_name}.parquet",
            filesystem=s3
        ).read_pandas().to_pandas().query("id in @val1")

        # 需求1：对指定数据集执行空值填充
        if df_name in fill_target_dfs:
            if df_name == "df1":
                df[["TYPE"]] = ""
            elif df_name == "df2":
                df[["val"]] = ""
        # 注：原示例中的df4不存在，故忽略；df0无需填充，原代码中按列表索引操作易出错，改为按名称匹配更可靠

        # 需求2：保留指定列并按字典重命名
        current_map = col_mapping[df_name]
        # 筛选数据中实际存在的列，避免KeyError
        valid_cols = [col for col in current_map.keys() if col in df.columns]
        # 筛选列并重命名
        processed_df = df[valid_cols].rename(columns=current_map)

        df_list.append(processed_df)
        print(f"已完成处理：{df_name}")

    return df_list

关键优化点说明 #

• 移除冗余操作：删掉了vars()的使用，直接用临时变量存储每个数据集，代码更简洁易维护。
• 需求1优化：不再依赖列表索引（如dfList[0]）匹配数据集，改为按名称判断，避免因循环顺序变化导致的错误。
• 需求2鲁棒性增强：增加了valid_cols筛选步骤，确保只保留数据中实际存在的列，避免因原数据集缺少字典中定义的列而抛出异常。
• 修正字典语法：原列映射字典存在语法错误（PATH_ID未定义映射值），已补全，可根据实际业务需求调整映射关系。

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