PySpark 加入多列

Question

我有 2 个数据框，我想知道是否可以以更通用和更紧凑的方式跨多个列连接。例如，这是一种非常明确的方式，很难在函数中泛化：

df = df1.join(df2,
         on=[
              (df1.event_date < df2.risk_date)
                 &
              (df1.client_id == df2.client_id_risk)
                 &
              (df1.col_thr_param_1  == df2.col_thr_param_1)
                 &
              (df1.col_thr_param_2  == df2.col_thr_param_2)
                 &
              (df1.col_thr_param_3  == df2.col_thr_param_3)
                 &
              (df1.col_thr_param_4  == df2.col_thr_param_4)
    ], how="left"
)

如果我有一个包含我想要加入的阈值列名称的列表：

thr = [“col_thr_param_1”、“col_thr_param_2”、“col_thr_param_3”、“col_thr_param_4”]

是否可以在函数中传递它并泛化连接？或者我总是需要求助于调用 df1 和 df2？像这样。

def join_dfs(df1: DataFrame, df2: DataFrame, thr_cols: List[str]):
   df = df1.join(df2,
           on=[
                (df1.event_date < df2.risk_date)
                   &
                (df1.client_id == df2.client_id_risk)
                   &
                **df1.thr_cols == **df2.thr_cols
           ], how="left"
   ) 

Answer 1

理想情况下，您可以使用带有col() 的列表的别名来加入。您可以尝试以下方法：

from pyspark.sql import functions as F

def join_dfs(df1, df2, thr_cols):
  df = df1.alias("df1").join(df2.alias("df2"),
           on=[
                [(F.col("df1.event_date") < F.col("df2.risk_date"))
                   ,
                (F.col("df1.client_id") == F.col("df2.client_id_risk")) ]+
                   
                [F.col(f"df1.{col}")==F.col(f"df2.{col}") for col in thr_cols]
           ], how="left"
        ) 
  return df