import pandas as pd import pandas as pd def read_and_merge_h5_data(h5_filename, key, columns, df=None): """ 读取 HDF5 文件中的数据,根据指定的 columns 筛选数据, 如果传入 df 参数,则将其与读取的数据根据 ts_code 和 trade_date 合并。 参数: - h5_filename: HDF5 文件名 - key: 数据存储在 HDF5 文件中的 key - columns: 要读取的列名列表 - df: 需要合并的 DataFrame(如果为空,则不进行合并) 返回: - 合并后的 DataFrame """ # 处理 _ 开头的列名 processed_columns = [] for col in columns: if col.startswith('_'): processed_columns.append(col[1:]) # 去掉下划线 else: processed_columns.append(col) # 从 HDF5 文件读取数据,选择需要的列 data = pd.read_hdf(h5_filename, key=key, columns=processed_columns) # 修改列名,如果列名以前有 _,加上 _ for col in data.columns: if col not in columns: # 只有不在 columns 中的列才需要加下划线 new_col = f'_{col}' data.rename(columns={col: new_col}, inplace=True) # 如果传入的 df 不为空,则进行合并 if df is not None and not df.empty: # 确保两个 DataFrame 都有 ts_code 和 trade_date 列 df['trade_date'] = pd.to_datetime(df['trade_date'], format='%Y%m%d') data['trade_date'] = pd.to_datetime(data['trade_date'], format='%Y%m%d') # 根据 ts_code 和 trade_date 合并 merged_df = pd.merge(df, data, on=['ts_code', 'trade_date'], how='left') else: # 如果 df 为空,则直接返回读取的数据 merged_df = data return merged_df