51 lines
1.7 KiB
Python
51 lines
1.7 KiB
Python
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
|
||
|