PySpark实现多列Join通用化紧凑写法的方法问询

解决方案 #

完全可以通过遍历阈值列名列表生成等值条件，再批量拼接关联条件实现通用封装，这里以Spark DataFrame场景为例给出实现代码：
首先导入所需依赖：

from functools import reduce
from pyspark.sql import DataFrame
from typing import List

函数实现 #

def join_dfs(df1: DataFrame, df2: DataFrame, thr_cols: List[str]) -> DataFrame:
    # 固定关联条件
    base_cond = (df1.event_date < df2.risk_date) & (df1.client_id == df2.client_id_risk)
    # 无阈值列时直接用固定条件关联
    if not thr_cols:
        return df1.join(df2, on=base_cond, how="left")
    # 生成所有阈值列的等值判断条件
    thr_conditions = [df1[col] == df2[col] for col in thr_cols]
    # 批量拼接所有阈值条件
    all_thr_cond = reduce(lambda a, b: a & b, thr_conditions)
    # 拼接完整关联条件
    total_cond = base_cond & all_thr_cond
    return df1.join(df2, on=total_cond, how="left")

调用示例 #

当你有阈值列表thr = ["col_thr_param_1", "col_thr_param_2", "col_thr_param_3", "col_thr_param_4"]时，直接调用即可：

joined_df = join_dfs(df1, df2, thr_cols=thr)

扩展说明 #

• 如果阈值列在两个DataFrame中名称不一致，可以把入参改为字典格式的映射关系，比如thr_cols={"df1_col1": "df2_col1", "df1_col2": "df2_col2"}，遍历的时候取对应列名生成条件即可
• 如果是Pandas场景，不等值关联可以用pandas.merge配合条件过滤实现，条件生成逻辑和上面完全一致，仅需要调整API调用写法
• 你也可以把固定条件里的字段、关联方式也改成入参，进一步提升函数的通用性

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