File size: 3,584 Bytes
60b5883 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 | #!/usr/bin/env python3
"""
合并 res2.csv 和 dataset_all.csv 数据集
res2.csv 结构: 编号, JSON列表(函数信息)
dataset_all.csv 结构: 编号(row[0]), 文件内容(row[1]), 其他内容...
目标: 找到 res2.csv 中存在的所有编号,从 dataset_all.csv 中提取对应的记录,整合成新数据集
"""
import pandas as pd
import json
from tqdm import tqdm
import os
def main():
print("开始处理数据集合并...")
# 1. 读取 res2.csv,获取所有存在的编号
print("\n步骤 1: 读取 res2.csv 获取编号列表...")
res2_path = '/home/weifengsun/tangou1/step2/res2.csv'
# 读取编号列(第一列)
res2_ids = set()
chunk_size = 100000
for chunk in tqdm(pd.read_csv(res2_path, chunksize=chunk_size, header=None, usecols=[0]),
desc="读取res2.csv"):
res2_ids.update(chunk[0].tolist())
print(f"从 res2.csv 中找到 {len(res2_ids)} 个唯一编号")
# 2. 读取 dataset_all.csv 并筛选匹配的记录
print("\n步骤 2: 从 dataset_all.csv 中筛选匹配的记录...")
dataset_all_path = '/home/weifengsun/tangou1/domain_code/src/datasets/data_merged/dataset_all.csv'
output_path = '/home/weifengsun/tangou1/step2/merged_dataset.csv'
# 统计信息
total_rows = 0
matched_rows = 0
# 分块读取和处理
first_chunk = True
for chunk in tqdm(pd.read_csv(dataset_all_path, chunksize=chunk_size, low_memory=False),
desc="处理dataset_all.csv"):
total_rows += len(chunk)
# 筛选编号在 res2_ids 中的行
# 假设第一列是编号
matched_chunk = chunk[chunk.iloc[:, 0].isin(res2_ids)]
matched_rows += len(matched_chunk)
# 写入输出文件
if len(matched_chunk) > 0:
if first_chunk:
matched_chunk.to_csv(output_path, index=True, mode='w')
first_chunk = False
else:
matched_chunk.to_csv(output_path, index=True, mode='a', header=False)
print(f"\n处理完成!")
print(f"总共处理行数: {total_rows}")
print(f"匹配的行数: {matched_rows}")
print(f"匹配率: {matched_rows/total_rows*100:.2f}%")
print(f"输出文件: {output_path}")
# 3. 可选: 创建一个包含函数信息的增强版本
print("\n步骤 3: 创建包含函数信息的增强数据集...")
enhanced_output_path = '/home/weifengsun/tangou1/step2/enhanced_dataset.csv'
# 读取res2.csv到字典
print("加载res2.csv函数信息...")
res2_dict = {}
for chunk in tqdm(pd.read_csv(res2_path, chunksize=chunk_size, header=None),
desc="加载res2"):
for idx, row in chunk.iterrows():
res2_dict[row[0]] = row[1]
# 读取刚生成的merged_dataset.csv并添加函数信息
print("合并函数信息...")
merged_df = pd.read_csv(output_path, low_memory=False)
# 添加函数信息列
merged_df['function_info'] = merged_df.iloc[:, 0].map(res2_dict)
# 保存增强数据集
merged_df.to_csv(enhanced_output_path, index=True)
print(f"\n增强数据集已保存到: {enhanced_output_path}")
print(f"增强数据集行数: {len(merged_df)}")
print(f"增强数据集列数: {len(merged_df.columns)}")
# 显示示例
print("\n数据集前5行预览:")
print(merged_df.head())
print("\n增强数据集列名:")
print(merged_df.columns.tolist())
if __name__ == "__main__":
main()
|