Airflow ETL任务优化：解决DataFrame对比错误与提速需求

1. 修复"Can only compare identically-labeled DataFrame objects"错误 #

该错误的核心原因是用于对比的两个DataFrame存在列名不匹配、列顺序不一致、索引未对齐的问题，解决步骤如下：

• 强制对齐列名与顺序：只保留业务需要的对比字段，确保两个DataFrame的列名、顺序完全一致
• 重置索引：去掉无关索引，避免因索引差异导致对比失败
• 示例修复代码：

# df_csv为CSV读取的数据集，df_db为数据库读取的数据集
# 定义需要对比的核心字段
target_columns = ["fee_id", "amount", "create_time", "user_id"]

# 清洗两个数据集：保留目标字段、重置索引
df_csv_clean = df_csv[target_columns].reset_index(drop=True)
df_db_clean = df_db[target_columns].reset_index(drop=True)

# 执行对比
diff_result = df_csv_clean.compare(df_db_clean)

2. 优化Extract & Load任务耗时 #

提取阶段优化 #

• 替换全量查询为增量查询：通过Airflow的Variable或XCom记录上次任务成功的时间戳，SQL Server查询时只拉取该时间戳之后的数据
• 分块读取大结果集：用chunksize参数分批读取，避免内存溢出同时提升读取效率

from airflow.models import Variable

# 获取上次成功提取的时间
last_extract_time = Variable.get("last_fee_extract_time", default_var="2020-01-01 00:00:00")
# 增量查询SQL
extract_sql = f"""
    SELECT fee_id, amount, create_time, user_id
    FROM source_table
    WHERE update_time > '{last_extract_time}'
"""
# 分块读取
for chunk in pd.read_sql(extract_sql, sql_server_conn, chunksize=10000):
    # 处理单块数据
    process_chunk(chunk)

加载阶段优化 #

• 批量插入：使用to_sql的method='multi'参数，减少数据库交互次数，同时设置chunksize控制单批插入量
• 临时关闭目标表索引：插入前关闭索引，插入完成后重建（适合非实时查询场景，大幅提升写入速度）

# 批量插入示例
chunk.to_sql(
    name="Fee",
    con=target_db_conn,
    if_exists="append",
    method="multi",
    chunksize=5000,
    index=False
)

3. 正确的多源数据对比逻辑（CSV vs 数据库） #

避免全量对比整个DataFrame，基于唯一标识字段（如fee_id）做精准增量判断：

1. 提取CSV与数据库中的唯一ID集合
2. 计算新增、更新、删除的ID范围
3. 针对不同范围分别执行写入、更新、删除操作

完整逻辑代码示例 #

def sync_fee_data():
    # 读取CSV数据
    df_csv = pd.read_csv("/path/to/fee_data.csv")
    # 读取数据库现有数据（仅需唯一ID与核心对比字段）
    df_db = pd.read_sql("SELECT fee_id, amount, update_time FROM Fee", target_db_conn)
    
    # 提取ID集合
    csv_ids = set(df_csv["fee_id"].tolist())
    db_ids = set(df_db["fee_id"].tolist())
    
    # 处理新增数据
    new_ids = csv_ids - db_ids
    new_data = df_csv[df_csv["fee_id"].isin(new_ids)]
    new_data.to_sql("Fee", target_db_conn, if_exists="append", method="multi", index=False)
    
    # 处理更新数据：先对比字段再执行更新
    update_ids = csv_ids & db_ids
    update_data = df_csv[df_csv["fee_id"].isin(update_ids)]
    # 合并数据对比字段差异
    merged_df = pd.merge(update_data, df_db, on="fee_id", suffixes=("_csv", "_db"))
    changed_rows = merged_df[
        (merged_df["amount_csv"] != merged_df["amount_db"]) |
        (merged_df["update_time_csv"] != merged_df["update_time_db"])
    ]
    # 批量执行更新SQL
    update_sql_template = """
        UPDATE Fee 
        SET amount = %(amount)s, update_time = %(update_time)s
        WHERE fee_id = %(fee_id)s
    """
    target_db_conn.execute(update_sql_template, changed_rows[["fee_id", "amount_csv", "update_time_csv"]].to_dict("records"))
    
    # 处理删除数据（可选，根据业务需求开启）
    delete_ids = db_ids - csv_ids
    if delete_ids:
        delete_sql = f"DELETE FROM Fee WHERE fee_id IN ({','.join(map(str, delete_ids))})"
        target_db_conn.execute(delete_sql)
    
    # 更新上次提取时间
    Variable.set("last_fee_extract_time", df_csv["update_time"].max())

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