Pandas DataFrame 大数据分批处理与外部API调用优化指南

import pandas as pd
from sklearn.datasets import load_diabetes # 用于生成示例数据
import time
import os

# 模拟一个大型DataFrame
# 在实际应用中，这里会加载您真实的50万行数据
data = load_diabetes().data
columns = load_diabetes().feature_names
df = pd.DataFrame(data, columns=columns)

# 模拟一些需要处理的额外列
df['dummy_col_1'] = df['age'] * 10
df['dummy_col_2'] = df['bmi'] / 2

# 定义批次大小，例如每批处理100行
batch_size = 100

# 为DataFrame添加一个批次编号列
# df.index // batch_size 会根据索引值自动生成批次号
df['batch_num'] = df.index // batch_size

print(f"原始DataFrame总行数: {len(df)}")
print(f"总批次数量: {df['batch_num'].nunique()}")
print(f"示例批次分配:\n{df[['age', 'batch_num']].head(batch_size + 5)}")

output_csv_path = 'processed_data_batched.csv'

# 确保输出文件是干净的，以便重新运行示例
if os.path.exists(output_csv_path):
    os.remove(output_csv_path)
    print(f"已删除现有文件: {output_csv_path}")

# 存储处理结果的列表（如果选择先收集再合并）
# processed_batches = []

# 遍历所有唯一的批次编号
for i, batch_id in enumerate(df['batch_num'].unique()):
    # 获取当前批次的数据子集
    # 使用 .copy() 避免 SettingWithCopyWarning
    current_batch_df = df[df['batch_num'] == batch_id].copy()

    print(f"\n正在处理第 {i+1}/{df['batch_num'].nunique()} 批次 (批次ID: {batch_id})，包含 {len(current_batch_df)} 行数据...")

    # --- 在此模拟批次内的操作 ---

    # 1. 模拟 df.merge 操作：
    # 例如，根据现有列创建新列，模拟合并外部数据
    current_batch_df['merged_data_sim'] = current_batch_df['s1'] + current_batch_df['s2']

    # 2. 模拟 df.apply 操作，特别是涉及外部API调用的场景：
    def custom_api_call_sim(row):
        # 模拟一个耗时的API调用，例如Google Maps API请求
        # 在实际应用中，这里会是您真实的API调用逻辑
        # time.sleep(0.01) # 模拟每行数据的网络延迟，或在批次结束后统一延迟
        return f"Processed_{row['age']}_{row['bmi']}_via_API"

    # 对当前批次的数据应用模拟的API调用函数
    current_batch_df['api_result'] = current_batch_df.apply(custom_api_call_sim, axis=1)

    # 3. 模拟其他 df.apply 或数据转换
    current_batch_df['transformed_data'] = current_batch_df['bmi'] * 100

    # --- 结果持久化：写入CSV文件 ---
    # 选择需要输出的列
    output_columns = ['age', 'sex', 'bmi', 'bp', 'merged_data_sim', 'api_result', 'transformed_data']

    if i == 0: # 对于第一个批次，写入时包含CSV头
        current_batch_df[output_columns].to_csv(output_csv_path, mode='w', header=True, index=False)
        print(f"已创建文件 {output_csv_path} 并写入首批数据。")
    else: # 对于后续批次，以追加模式写入，不包含CSV头
        current_batch_df[output_columns].to_csv(output_csv_path, mode='a', header=False, index=False)
        print(f"已将批次 {batch_id} 数据追加到 {output_csv_path}。")

    # 可选：在批次之间引入延迟，以遵守API速率限制
    # time.sleep(1) # 每处理完一个批次暂停1秒

print(f"\n所有批次处理完毕。结果已保存到 {output_csv_path}")

# 验证最终输出文件（可选）
final_df_check = pd.read_csv(output_csv_path)
print(f"\n最终CSV文件 '{output_csv_path}' 总行数: {len(final_df_check)}")
print("最终CSV文件前5行数据:\n", final_df_check.head())