如何基于另一个PySpark DataFrame更新目标PySpark DataFrame？

问题描述 #

给定两个PySpark DataFrame：

• df1（主数据集）包含完整业务字段：

+------------------------------------------------------
|ID|  NAME|ADDRESS|DELETE_FLAG|INSERT_DATE|UPDATE_DATE|
+------------------------------------------------------
| 1|sravan|delhi  |false      |25/01/2023 |25/01/2023|
| 2|ojasvi|patna  |false      |25/01/2023 |25/01/2023|
| 3|rohith|jaipur |false      |25/01/2023 |25/01/2023|

• df2（匹配条件集）仅包含匹配键字段：

+----------
|ID|  NAME|
+----------
| 1|sravan|
| 2|ojasvi|

需求：基于指定匹配键（此处为ID和NAME），将df1中与df2匹配的行的DELETE_FLAG设为true，UPDATE_DATE更新为指定日期（示例为02/02/2023），未匹配行保持不变，得到目标df3。要求实现通用方案，支持通过字符串或列表指定匹配键。

通用解决方案 #

实现逻辑 #

1. 统一匹配键格式：将输入的单个键（字符串）或多个键（列表）统一转为列表，便于拼接连接条件。
2. 左连接标记匹配行：通过左连接将主表与匹配表关联，生成匹配标记列区分是否需要更新。
3. 条件更新字段：根据匹配标记更新DELETE_FLAG和UPDATE_DATE，未匹配行保留原字段值。
4. 清理临时列：移除连接生成的冗余列，返回结构与主表一致的结果。

代码实现 #

from pyspark.sql import functions as F

def batch_update(main_df, match_df, match_keys, update_date=None):
    # 统一匹配键为列表格式
    match_keys = [match_keys] if isinstance(match_keys, str) else match_keys
    
    # 构建左连接条件
    join_conditions = [main_df[k] == match_df[k] for k in match_keys]
    
    # 左连接并标记匹配行
    joined = main_df.join(
        match_df.select(match_keys),
        on=join_conditions,
        how="left"
    ).withColumn(
        "is_matched",
        F.when(F.col(match_keys[0]).isNotNull(), F.lit(True)).otherwise(F.lit(False))
    )
    
    # 处理更新日期：自定义日期或当前日期
    update_date_col = F.lit(update_date) if update_date else F.current_date().cast("string")
    
    # 更新目标字段
    updated = joined.withColumn(
        "DELETE_FLAG",
        F.when(F.col("is_matched"), F.lit(True)).otherwise(F.col("DELETE_FLAG"))
    ).withColumn(
        "UPDATE_DATE",
        F.when(F.col("is_matched"), update_date_col).otherwise(F.col("UPDATE_DATE"))
    )
    
    # 清理临时列并返回
    return updated.drop(*match_keys[1:], "is_matched")

# 示例调用
if __name__ == "__main__":
    # 假设df1和df2已通过Spark创建
    match_keys = ["ID", "NAME"]
    target_update_date = "02/02/2023"
    
    df3 = batch_update(df1, df2, match_keys, target_update_date)
    df3.show()

关键特性 #

• 匹配键灵活：支持单个键（如"ID"）或复合键（如["ID", "NAME"]）输入。
• 日期可控：可传入自定义更新日期，也可默认使用系统当前日期（自动转为字符串格式）。
• 性能优化：仅选择匹配表的必要键进行连接，避免冗余数据加载。
• 结构保留：返回结果的字段结构与主表完全一致，无需额外调整。

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